# oMLX Cache Stack Architecture: Two-Tier Caching for Vision Features > Explore the oMLX cache stack architecture. Discover its two-tier caching with in-memory LRU and SSD persistence to speed up vision feature tensor computation. - Repository: [Jun Kim/omlx](https://github.com/jundot/omlx) - Tags: architecture - Published: 2026-05-11 --- **The oMLX cache stack implements a hybrid two-tier architecture featuring an in-memory LRU cache backed by optional SSD persistence to eliminate redundant computation of vision feature tensors.** The `jundot/omlx` repository employs a sophisticated caching system designed to accelerate vision-language model inference by reusing computed image embeddings. At its core lies the `VisionFeatureSSDCache` class, which balances memory speed with disk durability through a composite key-value store indexed by image hash and model name. ## Core Components of the Cache Stack The architecture centers on a single primary class that manages two distinct storage layers working in tandem to minimize latency. ### VisionFeatureSSDCache The `VisionFeatureSSDCache` class, defined in [`omlx/cache/vision_feature_cache.py`](https://github.com/jundot/omlx/blob/main/omlx/cache/vision_feature_cache.py), serves as the main entry point for all vision feature caching operations. It abstracts the complexity of tiered storage behind a simple `put` and `get` interface while maintaining statistics on cache performance through a dedicated `stats` attribute. ### Memory Tier (LRU Cache) The first tier resides entirely in RAM as a Python dictionary implementing LRU (Least Recently Used) eviction semantics. This layer stores `mx.ndarray` objects keyed by composite tuples of `(image_hash, model_name)`, allowing simultaneous caching for multiple vision models. The memory tier enforces capacity constraints through the `max_memory_entries` parameter. When the cache exceeds this limit, the least recently accessed entry is evicted to make room for new data. Access operations refresh the LRU order, ensuring frequently requested features remain hot in memory. ### SSD Tier (Disk Persistence) The second tier provides durable storage through serialized tensor files written to the filesystem. When configured with a `cache_dir` path (typically derived from the scheduler's `paged_ssd_cache_dir`), the cache writes files named according to the composite key pattern `_.npy`. This tier acts as a spill-over for evicted memory entries and enables cache survival across process restarts. On initialization, the cache can reload previously computed features from disk, avoiding cold-start penalties for repeated image queries. ## How the Two-Tier Cache Works The system follows a strict hierarchy when storing and retrieving vision features, optimizing for speed while ensuring durability. ### Composite Key Structure All cache entries are indexed using a tuple combination of `image_hash` and `model_name`. This design prevents collisions when serving multiple vision models simultaneously and ensures that different preprocessing configurations maintain separate cache lines. ### Write Operations (Put Path) When the VLM engine computes new vision features, it invokes `cache.put(image_hash, model_name, features)`: 1. The method first acquires the `_memory_lock` to ensure thread-safe updates. 2. The tensor is stored in the memory dictionary with its access time updated. 3. If the SSD tier is enabled, the tensor serializes to disk using the composite key filename. 4. Eviction occurs automatically if `max_memory_entries` is exceeded, removing the oldest entry from memory while preserving the disk copy. ### Read Operations (Get Path) Retrieval follows a cascading lookup strategy via `cache.get(image_hash, model_name)`: 1. **Memory Check:** The system queries the in-memory dictionary first. On success, it updates the LRU order, increments `stats.hits`, and returns the tensor immediately. 2. **SSD Fallback:** If the memory check fails, the cache probes the filesystem for the corresponding `.npy` file. Successful disk reads deserialize the tensor, insert it back into the memory tier (promoting it to hot status), increment `stats.hits`, and return the result. 3. **Complete Miss:** If absent from both tiers, the method returns `None` and increments `stats.misses`, triggering the caller to compute the features fresh. ### Statistics Tracking The `stats` attribute exposes a dataclass tracking `hits`, `misses`, and `evictions`, providing observability into cache efficiency and memory pressure. ## Integration with the VLM Engine The cache integrates deeply with the vision-language model pipeline in [`omlx/engine/vlm.py`](https://github.com/jundot/omlx/blob/main/omlx/engine/vlm.py). During the `_compute_vision_features` method execution, the engine checks the cache before invoking the vision encoder. When processing new images, the engine calls `self._vision_feature_cache.put()` immediately after tensor computation. Subsequent generation steps or batch requests retrieve cached features via `get()`, eliminating redundant forward passes through the vision backbone. This integration point demonstrates the cache's role as a transparent optimization layer that requires no changes to model architecture while delivering significant latency reductions for repeated image queries. ## Thread Safety and Concurrency The implementation ensures safe concurrent access through the `_memory_lock` threading primitive. This lock protects both the memory dictionary mutations and the SSD I/O path, preventing race conditions when multiple generation threads share a single cache instance. The locking strategy is coarse-grained but effective for the typical oMLX workload, where cache operations are fast relative to model inference times. ## Configuration and Directory Structure Cache persistence requires configuring the `paged_ssd_cache_dir` parameter in the scheduler configuration. The system organizes cached vision features under subdirectories within this path, maintaining isolation from other cache types (such as token embedding caches). Users can instantiate a memory-only cache by passing `cache_dir=None`, suitable for ephemeral workloads or environments without writable persistent storage. ## Practical Implementation Example The following example demonstrates instantiating and using the cache stack: ```python from pathlib import Path from omlx.cache.vision_feature_cache import VisionFeatureSSDCache # Configure SSD-backed cache ssd_dir = Path("/tmp/omlx_vision_cache") cache = VisionFeatureSSDCache( cache_dir=ssd_dir, max_memory_entries=100 ) # Store computed features image_hash = "a1b2c3d4e5f6" model_name = "clip_vit_large" features = mx.random.normal((1, 768)) # Example vision tensor cache.put(image_hash, model_name, features) # Retrieve with automatic tier fallback retrieved = cache.get(image_hash, model_name) # Monitor performance print(f"Cache hits: {cache.stats.hits}") print(f"Cache misses: {cache.stats.misses}") ``` ## Summary - The **oMLX cache stack** centers on `VisionFeatureSSDCache` in [`omlx/cache/vision_feature_cache.py`](https://github.com/jundot/omlx/blob/main/omlx/cache/vision_feature_cache.py), implementing a two-tier storage system. - **Memory tier** uses an LRU dictionary keyed by `(image_hash, model_name)` tuples with configurable `max_memory_entries` limits. - **SSD tier** persists tensors as serialized files under the scheduler's `paged_ssd_cache_dir`, enabling survival across restarts. - **Read operations** cascade from memory to disk, promoting SSD hits back to RAM while tracking statistics via `stats.hits` and `stats.misses`. - **Thread safety** is enforced through `_memory_lock`, protecting concurrent access in multi-threaded generation scenarios. - **VLM integration** occurs in [`omlx/engine/vlm.py`](https://github.com/jundot/omlx/blob/main/omlx/engine/vlm.py) through the `_compute_vision_features` method, which transparently caches and retrieves vision tensors. ## Frequently Asked Questions ### What is the maximum size of the memory cache in oMLX? The memory cache size is controlled by the `max_memory_entries` parameter passed to `VisionFeatureSSDCache` during initialization. Once this limit is reached, the cache automatically evicts the least recently used entry according to LRU policy, though the evicted data remains available in the SSD tier if persistence is enabled. ### How does oMLX handle cache misses when SSD persistence is enabled? On a cache miss in the memory tier, the system probes the filesystem for a serialized tensor file matching the composite key. If found, it deserializes the tensor, inserts it back into the memory cache (marking it as recently used), and returns the result. Only if the data is absent from both tiers does the system return `None` and increment the miss counter. ### Is the oMLX cache stack thread-safe for concurrent requests? Yes, the implementation uses a `_memory_lock` threading lock to synchronize access to both the memory dictionary and SSD I/O operations. This ensures that multiple threads can safely call `put()` and `get()` methods simultaneously without corrupting the cache state or causing race conditions during file writes. ### Where are the cached vision features stored on disk? When SSD persistence is enabled, features are stored as `.npy` files within the directory specified by the `cache_dir` parameter, typically derived from the scheduler's `paged_ssd_cache_dir` configuration. Files follow the naming convention `_.npy`, organizing tensors by their composite key for efficient retrieval.