# How OMLX's Tiered KV Cache (Hot + Cold SSD) Works for Token Caching

> Discover how OMLX's tiered KV cache (hot RAM + cold SSD) optimizes token caching, allowing sequences to exceed memory limits for efficient language model performance.

- Repository: [Jun Kim/omlx](https://github.com/jundot/omlx)
- Tags: internals
- Published: 2026-05-11

---

**OMLX implements a two-tier KV cache architecture where the hot tier stores active KV blocks in RAM via `PrefixCache` while the cold tier persists evicted blocks to SSD using `PagedSSDCacheManager`, enabling the system to cache token sequences far exceeding memory limits.**

The `jundot/omlx` repository provides a high-performance inference engine that leverages a **tiered KV cache** to optimize large language model serving. By separating frequently accessed tokens (hot) from archived history (cold), the system maintains low-latency access for active requests while utilizing persistent storage for long-term token reuse. This architecture centers on the `TieredCacheManager` class in [`omlx/cache/tiered_manager.py`](https://github.com/jundot/omlx/blob/main/omlx/cache/tiered_manager.py), which orchestrates data movement between fast memory and SSD-backed storage according to configurable resource limits.

## Architecture Overview

The tiered system comprises two specialized managers that present a unified interface to the scheduler.

### The Hot Tier: In-Memory PrefixCache

The **hot tier** is implemented by `BlockAwarePrefixCache` in [`omlx/cache/prefix_cache.py`](https://github.com/jundot/omlx/blob/main/omlx/cache/prefix_cache.py). This component maintains a fast, RAM-resident circular buffer for KV blocks currently in use by active requests. When a new block is computed during prefill or decode, the hot tier stores the keys and values as MLX tensors, providing immediate access for subsequent attention operations.

### The Cold Tier: SSD-Backed PagedSSDCacheManager

The **cold tier** is handled by `PagedSSDCacheManager` in [`omlx/cache/paged_ssd_cache.py`](https://github.com/jundot/omlx/blob/main/omlx/cache/paged_ssd_cache.py). This manager writes evicted KV blocks to disk as safetensors files, preserving them for future reuse without recomputation. Each block is stored with a unique hash and reference count, allowing the system to validate data integrity during restoration.

## Eviction and Restoration Workflow

Movement between tiers follows a deterministic FIFO policy to optimize for autoregressive generation patterns.

### Moving Hot Blocks to Cold Storage

When the hot cache exceeds the `max_kv_cache_memory` budget, `TieredCacheManager.evict_blocks_to_cold()` triggers. The implementation walks the circular buffer from oldest to newest, serializing each selected block via `PagedSSDCacheManager.save_block()`. The function writes the tensors to the configured `ssd_dir` using the safetensors format, then removes the block from RAM. This process updates internal statistics including `blocks_cold` and `cold_tokens` tracked by `MemoryMonitor`.

### Restoring Blocks from SSD

When a request requires a block absent from the hot tier, the scheduler invokes `TieredCacheManager.restore_block_from_cold(block_id, block_hash)`. The cold manager locates the file path via `get_block_path()`, loads the tensors using `mx.load`, and reinserts them into the hot buffer via `HotCache.store_block()`. If the block file exists on disk, restoration succeeds instantly; otherwise, the system marks a cache miss and recomputes the KV values via a forward pass.

## Integration with Request Lifecycle

The tiered cache integrates directly into OMLX's request processing pipeline defined in [`omlx/scheduler.py`](https://github.com/jundot/omlx/blob/main/omlx/scheduler.py). During the initial prefill phase, KV blocks populate the hot cache until the memory threshold triggers eviction. In the decode phase, the scheduler first attempts hot lookups; on misses, it transparently restores from cold storage before generating the next token. Per-request bookkeeping maintained by the tiered manager ensures that when a request completes, its hot blocks are released for eviction while cold copies persist for subsequent sessions.

## Configuration and Monitoring

Instantiate the cache by providing memory limits and storage paths to `TieredCacheManager`:

```python
from pathlib import Path
from omlx.cache.tiered_manager import TieredCacheManager

tiered_cache = TieredCacheManager(
    max_kv_cache_memory=4 * 1024**3,  # 4 GiB hot limit

    ssd_dir=Path("/tmp/omlx_ssd"),    # Cold storage location

)

```

The `MemoryMonitor` class in [`omlx/memory_monitor.py`](https://github.com/jundot/omlx/blob/main/omlx/memory_monitor.py) exposes real-time metrics including `kv_cache_memory` (current hot usage), `blocks_cold`, and `cold_tokens`. These values are served via the administrative endpoint `/admin/cache_probe` defined in [`omlx/admin/routes.py`](https://github.com/jundot/omlx/blob/main/omlx/admin/routes.py), allowing operators to observe eviction rates and cache saturation.

## Practical Implementation Examples

### Evicting Blocks to Free Memory

To proactively reclaim RAM, call the eviction method specifying the bytes to free:

```python
bytes_freed = tiered_cache.evict_blocks_to_cold(200 * 1024**2)
print(f"Moved {bytes_freed} bytes to cold SSD storage")

```

### Restoring Historical Context

When handling a new request that continues a previous conversation:

```python
block_id = 42
block_hash = b'\x12\x34\x56...'  # Retrieved from request metadata

if tiered_cache.restore_block_from_cold(block_id, block_hash):
    print("Successfully restored KV block from SSD")
else:
    print("Cache miss: Recomputing KV tensors")

```

### Checking Cache Statistics

Access current utilization metrics through the stats interface:

```python
stats = tiered_cache.stats()
print(f"Hot tier usage: {stats['kv_cache_memory']} bytes")
print(f"Cold blocks stored: {stats['blocks_cold']} ({stats['cold_tokens']} tokens)")

```

## Summary

- **OMLX combines** a RAM-based `PrefixCache` hot tier with an SSD-backed `PagedSSDCacheManager` cold tier to extend KV cache capacity beyond physical memory limits.
- **Eviction follows** a FIFO policy moving oldest blocks from hot to cold storage as safetensors files when `max_kv_cache_memory` is exceeded.
- **Restoration is lazy**: blocks return to hot storage via `mx.load` only when actively requested, minimizing unnecessary memory pressure.
- **Configuration** requires setting the hot memory budget and SSD directory path, with metrics exposed through `/admin/cache_probe`.

## Frequently Asked Questions

### What file format does OMLX use for cold tier storage?

OMLX serializes KV blocks to disk using the **safetensors** format, as implemented in [`omlx/cache/paged_ssd_cache.py`](https://github.com/jundot/omlx/blob/main/omlx/cache/paged_ssd_cache.py). This format provides efficient tensor serialization with zero-copy loading capabilities when restoring blocks via `mx.load`.

### How does the eviction policy handle concurrent requests?

The `TieredCacheManager` maintains per-request bookkeeping in a thread-safe manner. When eviction triggers, it only targets blocks not currently locked by active requests, ensuring that in-use KV data remains in the hot tier while older, idle blocks move to cold SSD storage.

### Can the cold tier be distributed across multiple SSDs?

The `ssd_dir` parameter accepts any valid filesystem path, allowing administrators to mount multiple drives at a single mount point or specify a path on a high-speed NVMe volume. The `PagedSSDCacheManager` treats the directory as a flat namespace for block files without internal striping logic.

### What happens when a cold restore fails?

If `restore_block_from_cold()` cannot locate the block file on disk (indicating corruption or manual deletion), the function returns `False` and the scheduler treats this as a cache miss. The system then recomputes the KV tensors for that token window via a standard forward pass, maintaining correct inference results while updating metrics to reflect the miss.