# How to Set Up Paged SSD Cache for KV Blocks in oMLX

> Learn to set up paged SSD cache for KV blocks in oMLX. Configure cache directory, leverage LRU eviction, and use asynchronous background writes for efficient disk storage.

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

---

**Configure `CacheConfig.paged_ssd_cache_dir` or pass `--paged-ssd-cache-dir` to enable persistent disk storage for KV blocks, with automatic LRU eviction managing size limits while the background writer thread handles I/O asynchronously.**

oMLX implements a three-tier caching architecture for large language model inference, with the **Paged SSD cache** serving as a persistent on-disk storage layer for KV blocks. This feature in the `jundot/omlx` repository allows cache reuse across requests and process restarts by storing blocks as `safetensors` files. Setting up the SSD tier requires only directory configuration and optional size limits—the `CacheFactory` and `SchedulerConfig` automatically wire the rest of the stack.

## Architecture Overview

The SSD cache implementation spans several key components across the codebase:

- **`CacheConfig`** (`omlx/cache/factory.py#L23-L48`): Dataclass holding `paged_ssd_cache_dir`, `max_paged_ssd_cache_size`, and block size parameters.
- **`CacheFactory.create_paged_ssd_cache`** (`omlx/cache/factory.py#L99-L126`): Instantiates `PagedSSDCacheManager` when a directory is supplied.
- **`PagedSSDCacheManager`** ([`omlx/cache/paged_ssd_cache.py`](https://github.com/jundot/omlx/blob/main/omlx/cache/paged_ssd_cache.py)): Implements the `CacheManager` ABC, handling directory layout (`/a/...`, `/b/...` sharding by hex prefix), safetensors serialization, background writer threads, and LRU eviction.
- **`BlockAwarePrefixCache`** ([`omlx/cache/prefix_cache.py`](https://github.com/jundot/omlx/blob/main/omlx/cache/prefix_cache.py)): Frontend used by the scheduler to store and retrieve KV blocks, coordinating between GPU and SSD tiers.
- **`SchedulerConfig`** (`omlx/scheduler.py#L382-L526`): Runtime configuration that passes SSD settings to the engine.
- **Settings Conversion** (`omlx/settings.py#L1191-L1205`): Resolves the final SSD directory path and creates the `SchedulerConfig`.

During inference, the scheduler requests blocks through the prefix cache. If a block is not in the hot cache, the SSD manager attempts to load it from disk via `load_block`, deserializing safetensors files back into `mx.array` objects. New blocks are queued for writing via `save_block` and persisted by a background thread.

## Enabling the SSD Cache

### Via Command Line

The simplest activation method uses the CLI flags defined in [`omlx/cli.py`](https://github.com/jundot/omlx/blob/main/omlx/cli.py):

```bash
omlx run \
    --paged-ssd-cache-dir /tmp/omlx_ssd_cache \
    --paged-ssd-cache-max-size 200GB

```

The directory is required; the size limit defaults to **100 GB** if unspecified. The scheduler automatically creates a model-scoped subdirectory (e.g., `/tmp/omlx_ssd_cache/llama-3b`) to isolate different model weights.

### Via Python API

For custom inference loops, construct the cache stack manually:

```python
from pathlib import Path
from omlx.cache.factory import CacheConfig, CacheFactory

# Configure SSD settings

cfg = CacheConfig(
    block_size=64,
    paged_ssd_cache_dir=Path("/tmp/omlx_ssd_cache"),
    max_paged_ssd_cache_size=200 * 1024**3,  # 200 GB

)

# Build the three-tier stack

paged_cache = CacheFactory.create_paged_cache(cfg, num_layers=32)
ssd_cache = CacheFactory.create_paged_ssd_cache(cfg, model_name="llama-3b")
prefix_cache = CacheFactory.create_prefix_cache(
    cfg,
    model=None,
    paged_cache=paged_cache,
    paged_ssd_cache=ssd_cache
)

# Use for inference

kv_data = [...]  # Your KV tensors

block_hash = b'\x12\x34' * 16  # 32-byte hash

prefix_cache.save_block(block_hash, kv_data, token_count=64)

```

The `CacheFactory` ensures the SSD manager receives a **model-scoped path** (`cache_dir / model_name`) to prevent collisions between different architectures.

## Configuration Options

### Size Limits and LRU Eviction

The `PagedSSDCacheIndex` maintains an in-memory ordered dictionary tracking access order. Before writing new blocks, the manager calls `_enforce_size_limit_for_new_block()` to evict least-recently used files until the total size falls under `max_size_bytes`. Eviction queues special `("unlink", path)` tasks to the background writer thread, preventing race conditions between writes and deletions.

### Hot-Cache-Only Mode

To disable disk persistence and use only RAM (useful for testing or strict latency requirements), instantiate the manager directly:

```python
from omlx.cache.paged_ssd_cache import PagedSSDCacheManager

ssd_cache = PagedSSDCacheManager(
    cache_dir=None,
    max_size_bytes=0,
    hot_cache_max_bytes=4 * 1024**3,  # 4 GB

    hot_cache_only=True,
)

```

In this mode, the background writer thread is disabled, and all blocks reside in the in-memory hot cache.

## Complete Working Example

This minimal script demonstrates the full lifecycle from configuration to block retrieval:

```python
import time
from pathlib import Path
from omlx.cache.factory import CacheConfig, CacheFactory
import mlx.core as mx
import numpy as np

# 1. Configure SSD cache

cfg = CacheConfig(
    paged_ssd_cache_dir=Path("/tmp/omlx_ssd_cache"),
    max_paged_ssd_cache_size=150 * 1024**3,  # 150 GB

)

# 2. Build cache stack for model "tiny"

paged_cache = CacheFactory.create_paged_cache(cfg, num_layers=4)
ssd_cache = CacheFactory.create_paged_ssd_cache(cfg, model_name="tiny")
prefix_cache = CacheFactory.create_prefix_cache(
    cfg, 
    model=None,
    paged_cache=paged_cache,
    paged_ssd_cache=ssd_cache
)

# 3. Simulate block generation (dummy KV tensors)

kv_block = [
    (mx.array(np.random.rand(1, 64, 32)), mx.array(np.random.rand(1, 64, 32))),
    (mx.array(np.random.rand(1, 64, 32)), mx.array(np.random.rand(1, 64, 32))),
]

block_hash = b'\x00' * 32  # Simulated SHA-256 hash

# 4. Save to cache (non-blocking write to SSD)

prefix_cache.save_block(block_hash, kv_block, token_count=64)

# 5. Load back (hits SSD if evicted from hot cache)

loaded = prefix_cache.load_block(block_hash)
assert loaded is not None
print(f"Restored block with {len(loaded)} layers")

```

You can inspect persisted blocks through the internal index:

```python
for meta in ssd_cache._index.get_all_metadata():
    print(f"Hash: {meta.block_hash.hex()[:8]}... Size: {meta.file_size} bytes")

```

## Summary

- **Configuration**: Set `paged_ssd_cache_dir` in `CacheConfig` or `--paged-ssd-cache-dir` via CLI to activate the SSD tier.
- **Automatic Management**: The `PagedSSDCacheManager` handles safetensors serialization, background writing, and LRU eviction when `max_paged_ssd_cache_size` is reached.
- **Model Isolation**: The factory automatically scopes cache directories by model name to prevent key collisions.
- **Flexibility**: Enable `hot_cache_only` mode to use the same API surface without disk persistence.
- **Persistence**: Blocks survive process restarts and are reusable across inference requests, enabling effectively unlimited KV cache capacity bounded only by SSD size.

## Frequently Asked Questions

### What file format does the SSD cache use for KV blocks?

The `PagedSSDCacheManager` serializes blocks using **safetensors** format via `_write_safetensors_no_mx`, storing each block as a separate file in a sharded directory structure based on the hash's first hex digit (`/a/...`, `/b/...`). This ensures fast single-block reads and writes without loading entire checkpoints.

### How does eviction work when the cache reaches its size limit?

Before persisting a new block, the manager calls `_enforce_size_limit_for_new_block()`, which removes the least-recently used files from the `PagedSSDCacheIndex` until the total size stays under `max_size_bytes`. Eviction tasks are queued to the background writer thread to prevent race conditions with active writes.

### Can I reuse the SSD cache between different model runs?

Yes, but only within the same model architecture. The `CacheFactory.create_paged_ssd_cache` function creates model-scoped subdirectories (e.g., `cache_dir/llama-3b`). Blocks saved for one model cannot be loaded for a different architecture because the hash keys are content-dependent and specific to layer dimensions and token sequences.

### What happens if the SSD cache directory is on a slow disk?

The implementation uses a **background writer thread** to decouple inference from I/O latency. However, initial cache misses requiring disk reads will block the request while deserializing safetensors files into `mx.array` objects. For high-throughput deployments, use NVMe SSDs and consider increasing the hot cache size to reduce disk hits.