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

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): 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): 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:

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:

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:

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:

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:

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.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}