vllm_mlx.ssd_cache

SSD KV cache tiering for vllm-mlx.

View the complete module source at #L1-L1248.

API details

Each callable below includes its exact signature, type annotations, inputs, defaults, return contract, documented exceptions, implementation source, and parsed docstring sections when the source provides them.

vllm_mlx.ssd_cache

SSD KV cache tiering for vllm-mlx.

This module provides a cold-tier disk cache that sits behind MemoryAwarePrefixCache. Evicted entries spill to NVMe instead of being discarded, and cold-tier fetches reload from disk asynchronously with RAM budget reservation before the read completes.

Key design: - SQLite for atomic metadata index (no mutable JSON) - Async writer thread for non-blocking spills - Per-layer serializer interface for hybrid cache types - Atomic temp-file + rename writes for crash consistency - Metrics exposed from day one

vllm_mlx.ssd_cache.logger module-attribute

logger = logging.getLogger(__name__)

vllm_mlx.ssd_cache._BYTES_PER_MB module-attribute

_BYTES_PER_MB = 1024 * 1024

vllm_mlx.ssd_cache._BYTES_PER_GB module-attribute

_BYTES_PER_GB = 1024 * 1024 * 1024

vllm_mlx.ssd_cache._PREFIX_FILTER_TOKENS module-attribute

_PREFIX_FILTER_TOKENS = 16

vllm_mlx.ssd_cache.SERIALIZER_SUPPORT_MATRIX module-attribute

SERIALIZER_SUPPORT_MATRIX = {'KVCache': 'supported', 'RotatingKVCache': 'supported', 'ArraysCache': 'supported', 'MambaCache': 'supported', '_QuantizedCacheWrapper': 'supported_via_dequant_on_spill', 'QuantizedKVCache': 'supported_via_dequant_on_spill'}

vllm_mlx.ssd_cache.SSDCacheConfig dataclass

SSDCacheConfig(cache_dir: str | None = None, max_size_gb: float = 10.0, max_entries: int = 10000, file_permissions: int = 384, dir_permissions: int = 448, spill_queue_size: int = 64, retention_seconds: int | None = None)

Configuration for SSD cache tier.

Attributes:

• cache_dir (str | None) –

  Directory for SSD cache files. None = auto-detect (~/.cache/vllm-mlx/ssd_cache/{model}/).

• max_size_gb (float) –

  Maximum total size of SSD cache in GB.

• max_entries (int) –

  Maximum number of entries in SSD cache.

• file_permissions (int) –

  Unix permission bits for cache data files.

• dir_permissions (int) –

  Unix permission bits for cache directories.

• spill_queue_size (int) –

  Max pending spill operations before dropping.

• retention_seconds (int | None) –

  Optional max age for cache entries (None = no expiry).

vllm_mlx.ssd_cache.SSDCacheConfig.cache_dir class-attribute instance-attribute

cache_dir: str | None = None

vllm_mlx.ssd_cache.SSDCacheConfig.max_size_gb class-attribute instance-attribute

max_size_gb: float = 10.0

vllm_mlx.ssd_cache.SSDCacheConfig.max_entries class-attribute instance-attribute

max_entries: int = 10000

vllm_mlx.ssd_cache.SSDCacheConfig.file_permissions class-attribute instance-attribute

file_permissions: int = 384

vllm_mlx.ssd_cache.SSDCacheConfig.dir_permissions class-attribute instance-attribute

dir_permissions: int = 448

vllm_mlx.ssd_cache.SSDCacheConfig.spill_queue_size class-attribute instance-attribute

spill_queue_size: int = 64

vllm_mlx.ssd_cache.SSDCacheConfig.retention_seconds class-attribute instance-attribute

retention_seconds: int | None = None

vllm_mlx.ssd_cache.SSDCacheConfig.max_size_bytes property

max_size_bytes: int

Maximum cache size in bytes.

vllm_mlx.ssd_cache.SSDCacheConfig.__post_init__

__post_init__() -> None

Source code in vllm_mlx/ssd_cache.py

def __post_init__(self) -> None:
    if self.max_size_gb <= 0:
        raise ValueError(f"max_size_gb must be > 0, got {self.max_size_gb}")
    if self.max_entries < 1:
        raise ValueError(f"max_entries must be >= 1, got {self.max_entries}")
    if self.spill_queue_size < 1:
        raise ValueError(
            f"spill_queue_size must be >= 1, got {self.spill_queue_size}"
        )

vllm_mlx.ssd_cache.SSDCacheStats dataclass

SSDCacheStats(spill_count: int = 0, spill_bytes: int = 0, ssd_hits: int = 0, ssd_misses: int = 0, reload_latency_sum: float = 0.0, reload_bytes: int = 0, promotion_failures: int = 0)

Statistics for SSD cache tier — exposed from day one.

Attributes:

• spill_count (int) –

  Number of entries spilled to SSD.

• spill_bytes (int) –

  Total bytes written to SSD.

• ssd_hits (int) –

  Number of successful SSD cache lookups.

• ssd_misses (int) –

  Number of SSD cache lookup misses.

• reload_latency_sum (float) –

  Sum of reload latencies in seconds.

• reload_bytes (int) –

  Total bytes read from SSD.

• promotion_failures (int) –

  Number of failed promotions (RAM budget exhausted).

vllm_mlx.ssd_cache.SSDCacheStats.spill_count class-attribute instance-attribute

spill_count: int = 0

vllm_mlx.ssd_cache.SSDCacheStats.spill_bytes class-attribute instance-attribute

spill_bytes: int = 0

vllm_mlx.ssd_cache.SSDCacheStats.ssd_hits class-attribute instance-attribute

ssd_hits: int = 0

vllm_mlx.ssd_cache.SSDCacheStats.ssd_misses class-attribute instance-attribute

ssd_misses: int = 0

vllm_mlx.ssd_cache.SSDCacheStats.reload_latency_sum class-attribute instance-attribute

reload_latency_sum: float = 0.0

vllm_mlx.ssd_cache.SSDCacheStats.reload_bytes class-attribute instance-attribute

reload_bytes: int = 0

vllm_mlx.ssd_cache.SSDCacheStats.promotion_failures class-attribute instance-attribute

promotion_failures: int = 0

vllm_mlx.ssd_cache.SSDCacheStats.to_dict

to_dict() -> dict

Return spill, lookup, reload, and promotion statistics.

Source code in vllm_mlx/ssd_cache.py

def to_dict(self) -> dict:
    """Return spill, lookup, reload, and promotion statistics."""

    total_lookups = self.ssd_hits + self.ssd_misses
    hit_rate = self.ssd_hits / total_lookups if total_lookups > 0 else 0.0
    avg_latency_ms = (
        (self.reload_latency_sum / self.ssd_hits * 1000)
        if self.ssd_hits > 0
        else 0.0
    )
    return {
        "spill_count": self.spill_count,
        "spill_bytes": self.spill_bytes,
        "ssd_hits": self.ssd_hits,
        "ssd_misses": self.ssd_misses,
        "ssd_hit_rate": round(hit_rate, 4),
        "reload_latency_sum_s": round(self.reload_latency_sum, 4),
        "avg_reload_latency_ms": round(avg_latency_ms, 2),
        "reload_bytes": self.reload_bytes,
        "promotion_failures": self.promotion_failures,
    }

vllm_mlx.ssd_cache.SSDIndex

SSDIndex(cache_dir: str)

SQLite-backed index for SSD cache entries.

Uses SQLite for atomic metadata operations instead of mutable JSON. The token sequence is stored as a binary blob for prefix-searchable representation. The primary key is a SHA-256 hash of the token sequence.

Thread safety: All operations are serialized through a threading.Lock. The SQLite connection uses WAL mode for concurrent read/write safety.

Source code in vllm_mlx/ssd_cache.py

def __init__(self, cache_dir: str) -> None:
    self._cache_dir = cache_dir
    self._db_lock = threading.Lock()
    db_path = os.path.join(cache_dir, "index.db")
    self._conn = sqlite3.connect(db_path, check_same_thread=False)
    self._conn.execute("PRAGMA journal_mode=WAL")
    self._conn.execute("PRAGMA synchronous=NORMAL")
    self._conn.row_factory = sqlite3.Row
    self._create_tables()

vllm_mlx.ssd_cache.SSDIndex._SCHEMA_VERSION class-attribute instance-attribute

_SCHEMA_VERSION = 1

vllm_mlx.ssd_cache.SSDIndex._cache_dir instance-attribute

_cache_dir = cache_dir

vllm_mlx.ssd_cache.SSDIndex._db_lock instance-attribute

_db_lock = threading.Lock()

vllm_mlx.ssd_cache.SSDIndex._conn instance-attribute

_conn = sqlite3.connect(db_path, check_same_thread=False)

vllm_mlx.ssd_cache.SSDIndex._create_tables

_create_tables() -> None

Source code in vllm_mlx/ssd_cache.py

def _create_tables(self) -> None:
    schema_sql = """
        CREATE TABLE IF NOT EXISTS schema_version (
            version INTEGER NOT NULL
        );

        CREATE TABLE IF NOT EXISTS entries (
            token_hash   TEXT PRIMARY KEY,
            tokens_blob  BLOB NOT NULL,
            prefix_hash  TEXT,
            num_tokens   INTEGER NOT NULL,
            file_path    TEXT NOT NULL,
            memory_bytes INTEGER NOT NULL,
            created_at   REAL NOT NULL,
            accessed_at  REAL NOT NULL
        );

        """
    self._conn.executescript(schema_sql)
    self._ensure_column("entries", "prefix_hash", "TEXT")
    self._conn.executescript("""
        CREATE INDEX IF NOT EXISTS idx_entries_accessed
            ON entries(accessed_at);

        CREATE INDEX IF NOT EXISTS idx_entries_num_tokens
            ON entries(num_tokens);

        CREATE INDEX IF NOT EXISTS idx_entries_prefix_hash_num_tokens
            ON entries(prefix_hash, num_tokens);
        """)
    # Insert schema version if not present
    cur = self._conn.execute("SELECT COUNT(*) FROM schema_version")
    if cur.fetchone()[0] == 0:
        self._conn.execute(
            "INSERT INTO schema_version (version) VALUES (?)",
            (self._SCHEMA_VERSION,),
        )
    self._backfill_prefix_hashes()
    self._conn.commit()

vllm_mlx.ssd_cache.SSDIndex._ensure_column

_ensure_column(table: str, column: str, definition: str) -> None

Source code in vllm_mlx/ssd_cache.py

def _ensure_column(self, table: str, column: str, definition: str) -> None:
    cur = self._conn.execute(f"PRAGMA table_info({table})")
    if column not in {row["name"] for row in cur.fetchall()}:
        self._conn.execute(f"ALTER TABLE {table} ADD COLUMN {column} {definition}")

vllm_mlx.ssd_cache.SSDIndex._backfill_prefix_hashes

_backfill_prefix_hashes() -> None

Source code in vllm_mlx/ssd_cache.py

def _backfill_prefix_hashes(self) -> None:
    cur = self._conn.execute(
        "SELECT token_hash, tokens_blob FROM entries WHERE prefix_hash IS NULL"
    )
    rows = cur.fetchall()
    for row in rows:
        tokens = _blob_to_tokens(row["tokens_blob"])
        self._conn.execute(
            "UPDATE entries SET prefix_hash = ? WHERE token_hash = ?",
            (_prefix_hash(tokens), row["token_hash"]),
        )

vllm_mlx.ssd_cache.SSDIndex.insert_entry

insert_entry(tokens_key: tuple[int, ...], file_path: str, memory_bytes: int, num_tokens: int) -> None

Insert or replace a cache entry in the index.

Source code in vllm_mlx/ssd_cache.py

def insert_entry(
    self,
    tokens_key: tuple[int, ...],
    file_path: str,
    memory_bytes: int,
    num_tokens: int,
) -> None:
    """Insert or replace a cache entry in the index."""
    now = time.time()
    token_hash = _tokens_hash(tokens_key)
    prefix_hash = _prefix_hash(tokens_key)
    tokens_blob = _tokens_to_blob(tokens_key)
    with self._db_lock:
        self._conn.execute(
            """
            INSERT OR REPLACE INTO entries
                (token_hash, tokens_blob, prefix_hash, num_tokens, file_path,
                 memory_bytes, created_at, accessed_at)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?)
            """,
            (
                token_hash,
                tokens_blob,
                prefix_hash,
                num_tokens,
                file_path,
                memory_bytes,
                now,
                now,
            ),
        )
        self._conn.commit()

vllm_mlx.ssd_cache.SSDIndex.lookup_exact

lookup_exact(tokens_key: tuple[int, ...]) -> dict | None

Look up an exact token sequence. Returns dict or None.

Source code in vllm_mlx/ssd_cache.py

def lookup_exact(self, tokens_key: tuple[int, ...]) -> dict | None:
    """Look up an exact token sequence. Returns dict or None."""
    token_hash = _tokens_hash(tokens_key)
    with self._db_lock:
        cur = self._conn.execute(
            "SELECT file_path, memory_bytes, num_tokens FROM entries WHERE token_hash = ?",
            (token_hash,),
        )
        row = cur.fetchone()
    if row is None:
        return None
    return {
        "file_path": row["file_path"],
        "memory_bytes": row["memory_bytes"],
        "num_tokens": row["num_tokens"],
    }

vllm_mlx.ssd_cache.SSDIndex.lookup_prefix

lookup_prefix(query_tokens: tuple[int, ...]) -> list[dict]

Find entries whose token sequence is a prefix of query_tokens.

Uses a bounded token-prefix hash to avoid scanning all entries, then compares the full stored token blob against the corresponding prefix of query_tokens.

Returns list of dicts sorted by num_tokens descending (longest prefix first).

Source code in vllm_mlx/ssd_cache.py

def lookup_prefix(self, query_tokens: tuple[int, ...]) -> list[dict]:
    """Find entries whose token sequence is a prefix of query_tokens.

    Uses a bounded token-prefix hash to avoid scanning all entries, then
    compares the full stored token blob against the corresponding prefix of
    query_tokens.

    Returns list of dicts sorted by num_tokens descending (longest prefix first).
    """
    query_len = len(query_tokens)
    query_blob = _tokens_to_blob(query_tokens)
    prefix_hashes = {
        _tokens_hash(query_tokens[:n])
        for n in range(1, min(query_len, _PREFIX_FILTER_TOKENS) + 1)
    }
    if not prefix_hashes:
        return []

    with self._db_lock:
        placeholders = ",".join("?" for _ in prefix_hashes)
        cur = self._conn.execute(
            "SELECT token_hash, tokens_blob, num_tokens, file_path, memory_bytes "
            f"FROM entries WHERE num_tokens <= ? AND prefix_hash IN ({placeholders}) "
            "ORDER BY num_tokens DESC",
            (query_len, *prefix_hashes),
        )
        rows = cur.fetchall()

    results = []
    for row in rows:
        stored_blob = row["tokens_blob"]
        n = row["num_tokens"]
        prefix_blob = query_blob[: n * 4]
        if stored_blob == prefix_blob:
            results.append(
                {
                    "token_hash": row["token_hash"],
                    "file_path": row["file_path"],
                    "memory_bytes": row["memory_bytes"],
                    "num_tokens": n,
                }
            )
    return results

vllm_mlx.ssd_cache.SSDIndex.delete_entry

delete_entry(tokens_key: tuple[int, ...]) -> None

Delete an entry by token sequence.

Source code in vllm_mlx/ssd_cache.py

def delete_entry(self, tokens_key: tuple[int, ...]) -> None:
    """Delete an entry by token sequence."""
    token_hash = _tokens_hash(tokens_key)
    with self._db_lock:
        self._conn.execute(
            "DELETE FROM entries WHERE token_hash = ?", (token_hash,)
        )
        self._conn.commit()

vllm_mlx.ssd_cache.SSDIndex.get_lru

get_lru(limit: int = 10) -> list[dict]

Get the least recently used entries, ordered oldest first.

Source code in vllm_mlx/ssd_cache.py

def get_lru(self, limit: int = 10) -> list[dict]:
    """Get the least recently used entries, ordered oldest first."""
    with self._db_lock:
        cur = self._conn.execute(
            "SELECT token_hash, tokens_blob, num_tokens, file_path, memory_bytes "
            "FROM entries ORDER BY accessed_at ASC LIMIT ?",
            (limit,),
        )
        rows = cur.fetchall()
    results = []
    for row in rows:
        results.append(
            {
                "token_hash": row["token_hash"],
                "tokens_blob": row["tokens_blob"],
                "file_path": row["file_path"],
                "memory_bytes": row["memory_bytes"],
                "num_tokens": row["num_tokens"],
            }
        )
    return results

vllm_mlx.ssd_cache.SSDIndex.get_total_bytes

get_total_bytes() -> int

Get total memory_bytes across all entries.

Source code in vllm_mlx/ssd_cache.py

def get_total_bytes(self) -> int:
    """Get total memory_bytes across all entries."""
    with self._db_lock:
        cur = self._conn.execute(
            "SELECT COALESCE(SUM(memory_bytes), 0) FROM entries"
        )
        return cur.fetchone()[0]

vllm_mlx.ssd_cache.SSDIndex.get_entry_count

get_entry_count() -> int

Get number of entries in the index.

Source code in vllm_mlx/ssd_cache.py

def get_entry_count(self) -> int:
    """Get number of entries in the index."""
    with self._db_lock:
        cur = self._conn.execute("SELECT COUNT(*) FROM entries")
        return cur.fetchone()[0]

vllm_mlx.ssd_cache.SSDIndex.touch

touch(tokens_key: tuple[int, ...]) -> None

Update accessed_at timestamp for an entry (marks as recently used).

Source code in vllm_mlx/ssd_cache.py

def touch(self, tokens_key: tuple[int, ...]) -> None:
    """Update accessed_at timestamp for an entry (marks as recently used)."""
    token_hash = _tokens_hash(tokens_key)
    with self._db_lock:
        self._conn.execute(
            "UPDATE entries SET accessed_at = ? WHERE token_hash = ?",
            (time.time(), token_hash),
        )
        self._conn.commit()

vllm_mlx.ssd_cache.SSDIndex.all_entries

all_entries() -> list[dict]

Return all entries (for startup reconciliation).

Source code in vllm_mlx/ssd_cache.py

def all_entries(self) -> list[dict]:
    """Return all entries (for startup reconciliation)."""
    with self._db_lock:
        cur = self._conn.execute(
            "SELECT token_hash, tokens_blob, num_tokens, file_path, memory_bytes "
            "FROM entries ORDER BY accessed_at DESC"
        )
        rows = cur.fetchall()
    results = []
    for row in rows:
        results.append(
            {
                "token_hash": row["token_hash"],
                "tokens_blob": row["tokens_blob"],
                "file_path": row["file_path"],
                "memory_bytes": row["memory_bytes"],
                "num_tokens": row["num_tokens"],
            }
        )
    return results

vllm_mlx.ssd_cache.SSDIndex.close

close() -> None

Close the SQLite connection.

Source code in vllm_mlx/ssd_cache.py

def close(self) -> None:
    """Close the SQLite connection."""
    with self._db_lock:
        self._conn.close()

vllm_mlx.ssd_cache.LayerSerializer

Bases: ABC

Interface for per-layer cache serialization.

Spill is split across two threads: snapshot_layer runs on the producer (request handler) thread so the mx→numpy materialization happens where the per-request Stream(gpu, N) is registered; serialize_layer then runs on the SSD writer thread with numpy only.

vllm_mlx.ssd_cache.LayerSerializer.snapshot_layer abstractmethod

snapshot_layer(layer: Any) -> dict[str, Any]

Producer-thread CPU snapshot of an MLX-backed cache layer.

Source code in vllm_mlx/ssd_cache.py

@abstractmethod
def snapshot_layer(self, layer: Any) -> dict[str, Any]:
    """Producer-thread CPU snapshot of an MLX-backed cache layer."""
    ...

vllm_mlx.ssd_cache.LayerSerializer.serialize_layer abstractmethod

serialize_layer(snapshot: dict[str, Any], layer_idx: int, file_path: str) -> dict[str, Any]

Writer-thread: persist a snapshot to safetensors at file_path.

Returns metadata dict with at least 'layer_type'.

Source code in vllm_mlx/ssd_cache.py

@abstractmethod
def serialize_layer(
    self, snapshot: dict[str, Any], layer_idx: int, file_path: str
) -> dict[str, Any]:
    """Writer-thread: persist a snapshot to safetensors at file_path.

    Returns metadata dict with at least 'layer_type'.
    """
    ...

vllm_mlx.ssd_cache.LayerSerializer.deserialize_layer abstractmethod

deserialize_layer(file_path: str, metadata: dict[str, Any]) -> dict

Read a layer back from disk. Returns layer-state dict.

Source code in vllm_mlx/ssd_cache.py

@abstractmethod
def deserialize_layer(self, file_path: str, metadata: dict[str, Any]) -> dict:
    """Read a layer back from disk. Returns layer-state dict."""
    ...

vllm_mlx.ssd_cache.KVCacheSerializer

Bases: LayerSerializer

Serializer for KVCache and RotatingKVCache layers.

Handles layers with .keys, .values, .offset attributes. RotatingKVCache also has .max_size, .keep, .step, ._idx.

vllm_mlx.ssd_cache.KVCacheSerializer._ROTATING_ATTRS class-attribute instance-attribute

_ROTATING_ATTRS = ('max_size', 'keep', 'step', '_idx')

vllm_mlx.ssd_cache.KVCacheSerializer.snapshot_layer

snapshot_layer(layer: Any) -> dict[str, Any]

Copy a KV cache layer into NumPy-backed writer-thread data.

Source code in vllm_mlx/ssd_cache.py

def snapshot_layer(self, layer: Any) -> dict[str, Any]:
    """Copy a KV cache layer into NumPy-backed writer-thread data."""

    keys_np, keys_orig_dtype = _mx_to_numpy_safe(layer.keys)
    values_np, values_orig_dtype = _mx_to_numpy_safe(layer.values)

    # Quantized-spill cast-back sentinel: the enqueue_spill dequant path
    # may have cast bf16 → fp16 before snapshot to dodge numpy's PEP 3118
    # buffer-protocol mismatch. The cast loses the original-dtype signal
    # _mx_to_numpy_safe would otherwise capture (since fp16 IS numpy-
    # supported, _mx_to_numpy_safe returns dtype=None). Honor an explicit
    # sentinel on the layer when present so the reload path can restore
    # bf16 instead of leaving the model with fp16 KV.
    keys_orig_dtype = (
        getattr(layer, "_ssd_keys_original_dtype", None) or keys_orig_dtype
    )
    values_orig_dtype = (
        getattr(layer, "_ssd_values_original_dtype", None) or values_orig_dtype
    )

    snapshot: dict[str, Any] = {
        "keys_np": keys_np,
        "values_np": values_np,
        "offset": layer.offset,
    }
    if keys_orig_dtype is not None:
        snapshot["keys_original_dtype"] = keys_orig_dtype
    if values_orig_dtype is not None:
        snapshot["values_original_dtype"] = values_orig_dtype

    # RotatingKVCache extras (plain Python scalars, no MLX).
    for attr in self._ROTATING_ATTRS:
        if hasattr(layer, attr):
            snapshot[attr] = getattr(layer, attr)
    return snapshot

vllm_mlx.ssd_cache.KVCacheSerializer.serialize_layer

serialize_layer(snapshot: dict[str, Any], layer_idx: int, file_path: str) -> dict[str, Any]

Write one KV layer to safetensors and return reconstruction metadata.

Source code in vllm_mlx/ssd_cache.py

def serialize_layer(
    self, snapshot: dict[str, Any], layer_idx: int, file_path: str
) -> dict[str, Any]:
    """Write one KV layer to safetensors and return reconstruction metadata."""

    from safetensors.numpy import save_file

    tensors = {
        f"layer_{layer_idx}_keys": snapshot["keys_np"],
        f"layer_{layer_idx}_values": snapshot["values_np"],
    }
    save_file(tensors, file_path)

    metadata = {
        "layer_type": "KVCache",
        "layer_idx": layer_idx,
        "offset": snapshot["offset"],
    }
    for k in ("keys_original_dtype", "values_original_dtype"):
        if k in snapshot:
            metadata[k] = snapshot[k]
    for attr in self._ROTATING_ATTRS:
        if attr in snapshot:
            metadata[attr] = snapshot[attr]

    return metadata

vllm_mlx.ssd_cache.KVCacheSerializer.deserialize_layer

deserialize_layer(file_path: str, metadata: dict[str, Any]) -> dict

Load one KV layer as arrays plus cache reconstruction metadata.

Source code in vllm_mlx/ssd_cache.py

def deserialize_layer(self, file_path: str, metadata: dict[str, Any]) -> dict:
    """Load one KV layer as arrays plus cache reconstruction metadata."""

    from safetensors.numpy import load_file

    layer_idx = metadata["layer_idx"]
    tensors = load_file(file_path)

    result = {
        "keys": tensors[f"layer_{layer_idx}_keys"],
        "values": tensors[f"layer_{layer_idx}_values"],
        "offset": metadata["offset"],
    }
    # Dtype hints surfaced so _reconstruct_ssd_layers can cast back.
    for k in ("keys_original_dtype", "values_original_dtype"):
        if k in metadata:
            result[k] = metadata[k]
    for attr in self._ROTATING_ATTRS:
        if attr in metadata:
            result[attr] = metadata[attr]
    return result

vllm_mlx.ssd_cache.ArraysCacheSerializer

Bases: LayerSerializer

Serializer for ArraysCache (Mamba/linear attention) layers.

Handles layers with .state attribute containing a list of arrays.

vllm_mlx.ssd_cache.ArraysCacheSerializer.snapshot_layer

snapshot_layer(layer: Any) -> dict[str, Any]

Copy an arrays-cache state into NumPy-backed writer-thread data.

Source code in vllm_mlx/ssd_cache.py

def snapshot_layer(self, layer: Any) -> dict[str, Any]:
    """Copy an arrays-cache state into NumPy-backed writer-thread data."""

    state_np: list[np.ndarray] = []
    original_dtypes: list[str | None] = []
    for arr in layer.state:
        np_arr, orig = _mx_to_numpy_safe(arr)
        state_np.append(np_arr)
        original_dtypes.append(orig)

    snapshot: dict[str, Any] = {"state_np": state_np}
    # Skip the dtype list in the common (fp16/fp32) case.
    if any(d is not None for d in original_dtypes):
        snapshot["state_original_dtypes"] = original_dtypes
    return snapshot

vllm_mlx.ssd_cache.ArraysCacheSerializer.serialize_layer

serialize_layer(snapshot: dict[str, Any], layer_idx: int, file_path: str) -> dict[str, Any]

Write arrays-cache state to safetensors and return its metadata.

Source code in vllm_mlx/ssd_cache.py

def serialize_layer(
    self, snapshot: dict[str, Any], layer_idx: int, file_path: str
) -> dict[str, Any]:
    """Write arrays-cache state to safetensors and return its metadata."""

    # Writer-thread side: pure numpy + disk.
    from safetensors.numpy import save_file

    state_np = snapshot["state_np"]
    tensors = {
        f"layer_{layer_idx}_state_{i}": arr for i, arr in enumerate(state_np)
    }
    save_file(tensors, file_path)

    metadata = {
        "layer_type": "ArraysCache",
        "layer_idx": layer_idx,
        "num_arrays": len(state_np),
    }
    if "state_original_dtypes" in snapshot:
        metadata["state_original_dtypes"] = snapshot["state_original_dtypes"]
    return metadata

vllm_mlx.ssd_cache.ArraysCacheSerializer.deserialize_layer

deserialize_layer(file_path: str, metadata: dict[str, Any]) -> dict

Load arrays-cache state and any original dtype hints.

Source code in vllm_mlx/ssd_cache.py

def deserialize_layer(self, file_path: str, metadata: dict[str, Any]) -> dict:
    """Load arrays-cache state and any original dtype hints."""

    from safetensors.numpy import load_file

    layer_idx = metadata["layer_idx"]
    num_arrays = metadata["num_arrays"]
    tensors = load_file(file_path)

    state = []
    for i in range(num_arrays):
        state.append(tensors[f"layer_{layer_idx}_state_{i}"])
    result = {"state": state}
    if "state_original_dtypes" in metadata:
        result["state_original_dtypes"] = metadata["state_original_dtypes"]
    return result

vllm_mlx.ssd_cache.SSDCacheTier

SSDCacheTier(config: SSDCacheConfig)

Cold-tier disk cache for KV cache entries.

Manages a SQLite-indexed on-disk cache directory. Evicted RAM entries are spilled here via an async writer thread. Cold-tier fetches reload from disk asynchronously with RAM budget reservation.

Directory layout::

cache_dir/
  index.db           # SQLite metadata index
  data/              # safetensors files per entry
    {hash}/          # one directory per entry
      layer_0.safetensors
      layer_1.safetensors
      manifest.json  # per-entry layer metadata

Source code in vllm_mlx/ssd_cache.py

def __init__(self, config: SSDCacheConfig) -> None:
    self._config = config
    self._closed = True
    self._writer_thread: threading.Thread | None = None

    if config.cache_dir is None:
        raise ValueError("SSDCacheConfig.cache_dir must be set")

    self._cache_dir = config.cache_dir
    self._data_dir = os.path.join(self._cache_dir, "data")

    # Create directory structure
    os.makedirs(self._cache_dir, mode=config.dir_permissions, exist_ok=True)
    os.makedirs(self._data_dir, mode=config.dir_permissions, exist_ok=True)

    try:
        # Open SQLite index
        self._index = SSDIndex(self._cache_dir)

        # Stats
        self._stats = SSDCacheStats()
        self._lock = threading.Lock()

        # Spill queue and writer thread
        self._spill_queue: queue.Queue = queue.Queue(
            maxsize=config.spill_queue_size
        )
        self._writer_stop = threading.Event()
        self._closed = False
    except Exception:
        index = getattr(self, "_index", None)
        if index is not None:
            try:
                index.close()
            except Exception:
                logger.exception(
                    "ssd_cache: failed to close index during init cleanup"
                )
        raise

vllm_mlx.ssd_cache.SSDCacheTier._config instance-attribute

_config = config

vllm_mlx.ssd_cache.SSDCacheTier._writer_thread instance-attribute

_writer_thread: Thread | None = None

vllm_mlx.ssd_cache.SSDCacheTier._cache_dir instance-attribute

_cache_dir = config.cache_dir

vllm_mlx.ssd_cache.SSDCacheTier._data_dir instance-attribute

_data_dir = os.path.join(self._cache_dir, 'data')

vllm_mlx.ssd_cache.SSDCacheTier._index instance-attribute

_index = SSDIndex(self._cache_dir)

vllm_mlx.ssd_cache.SSDCacheTier._stats instance-attribute

_stats = SSDCacheStats()

vllm_mlx.ssd_cache.SSDCacheTier._lock instance-attribute

_lock = threading.Lock()

vllm_mlx.ssd_cache.SSDCacheTier._spill_queue instance-attribute

_spill_queue: Queue = queue.Queue(maxsize=config.spill_queue_size)

vllm_mlx.ssd_cache.SSDCacheTier._writer_stop instance-attribute

_writer_stop = threading.Event()

vllm_mlx.ssd_cache.SSDCacheTier._closed instance-attribute

_closed = False

vllm_mlx.ssd_cache.SSDCacheTier._entry_hash staticmethod

_entry_hash(tokens: tuple[int, ...]) -> str

Compute deterministic hash for a token sequence.

Source code in vllm_mlx/ssd_cache.py

@staticmethod
def _entry_hash(tokens: tuple[int, ...]) -> str:
    """Compute deterministic hash for a token sequence."""
    return _tokens_hash(tokens)

vllm_mlx.ssd_cache.SSDCacheTier.get_stats

get_stats() -> dict

Return current SSD cache statistics.

Source code in vllm_mlx/ssd_cache.py

def get_stats(self) -> dict:
    """Return current SSD cache statistics."""
    return self._stats.to_dict()

vllm_mlx.ssd_cache.SSDCacheTier.start_writer

start_writer() -> None

Start the background spill writer thread.

Source code in vllm_mlx/ssd_cache.py

def start_writer(self) -> None:
    """Start the background spill writer thread."""
    if self._writer_thread is not None:
        return
    self._writer_stop.clear()
    self._writer_thread = threading.Thread(
        target=self._writer_loop, daemon=True, name="ssd-cache-writer"
    )
    self._writer_thread.start()
    logger.info("[ssd_cache] writer thread started")

vllm_mlx.ssd_cache.SSDCacheTier._writer_loop

_writer_loop() -> None

Drain spill queue and persist entries. Numpy-only — no MLX here.

Source code in vllm_mlx/ssd_cache.py

def _writer_loop(self) -> None:
    """Drain spill queue and persist entries. Numpy-only — no MLX here."""
    while not self._writer_stop.is_set():
        try:
            item = self._spill_queue.get(timeout=0.5)
        except queue.Empty:
            continue

        if item is None:  # Poison pill for shutdown
            break

        tokens_key, layer_snapshots, memory_bytes = item
        try:
            self._write_entry(tokens_key, layer_snapshots, memory_bytes)
        except Exception:
            logger.exception(
                f"[ssd_cache] failed to write entry " f"({len(tokens_key)} tokens)"
            )

vllm_mlx.ssd_cache.SSDCacheTier.enqueue_spill

enqueue_spill(tokens: tuple[int, ...], cache: list[Any], memory_bytes: int) -> bool

Enqueue a cache entry for async spill to SSD.

Must be called on the producer thread (the request handler that owns the layer's Stream(gpu, N)) — the snapshot below materializes MLX → numpy here so the writer thread never has to.

Returns True if enqueued, False if queue is full (entry dropped).

Source code in vllm_mlx/ssd_cache.py

def enqueue_spill(
    self,
    tokens: tuple[int, ...],
    cache: list[Any],
    memory_bytes: int,
) -> bool:
    """Enqueue a cache entry for async spill to SSD.

    Must be called on the producer thread (the request handler that
    owns the layer's Stream(gpu, N)) — the snapshot below materializes
    MLX → numpy here so the writer thread never has to.

    Returns True if enqueued, False if queue is full (entry dropped).
    """
    # Dequantize on the CALLER's thread, which owns the MLX GPU stream.
    # mx.dequantize is a GPU compute op; running it on the writer thread
    # aborts the process ("no Stream(gpu,N) in current thread"). Materialize
    # with mx.eval so the writer thread only does a host-side copy. The layer
    # serializers handle plain KVCache/ArraysCache only; any layer whose
    # .keys/.values are a tuple/list of arrays (packed, scales, biases)
    # — either our `_QuantizedCacheWrapper` or mlx-lm's native
    # `QuantizedKVCache` produced when --kv-cache-quantization is on — must
    # be reduced to a single dense array per attribute before queueing.
    import mlx.core as mx
    from .memory_cache import _QuantizedCacheWrapper, _dequantize_cache

    def _is_quantized_layer(layer):
        if isinstance(layer, _QuantizedCacheWrapper):
            return True
        keys = getattr(layer, "keys", None)
        return isinstance(keys, (tuple, list))

    if any(_is_quantized_layer(layer) for layer in cache):
        try:
            from mlx_lm.models.cache import QuantizedKVCache
        except ImportError:
            QuantizedKVCache = ()  # never matches isinstance below

        converted: list = []
        for layer in cache:
            if isinstance(layer, _QuantizedCacheWrapper):
                # Existing path: wrapper → orig_type with dequantized keys.
                converted.extend(_dequantize_cache([layer]))
            elif isinstance(layer, QuantizedKVCache) or (
                hasattr(layer, "keys") and isinstance(layer.keys, (tuple, list))
            ):
                # Native mlx-lm QuantizedKVCache: keys/values are
                # (packed, scales, biases) tuples. Dequantize in place
                # into a fresh KVCache-shaped duck-type for the serializer.
                from mlx_lm.models.cache import KVCache as _KVCache

                bits = getattr(layer, "bits", 8)
                group_size = getattr(layer, "group_size", 64)
                kv = _KVCache.__new__(_KVCache)
                kv.keys = mx.dequantize(
                    *layer.keys, group_size=group_size, bits=bits
                )
                kv.values = mx.dequantize(
                    *layer.values, group_size=group_size, bits=bits
                )
                kv.offset = getattr(layer, "offset", kv.keys.shape[-2])
                if (
                    kv.keys is not None
                    and hasattr(kv.keys, "shape")
                    and len(kv.keys.shape) >= 3
                    and kv.offset < kv.keys.shape[-2]
                ):
                    kv.keys = kv.keys[..., : kv.offset, :]
                    kv.values = kv.values[..., : kv.offset, :]
                converted.append(kv)
            else:
                converted.append(layer)
        cache = converted
        # Numpy's PEP 3118 buffer protocol doesn't understand bfloat16 —
        # it sees the bf16 array as format "B" (uint8) but the buffer items
        # are 2 bytes, so np.array() raises a RuntimeError mismatch. Cast
        # to float16 here on the CALLER's stream: KV values sit well within
        # the fp16 ±65504 range, fp16 has MORE mantissa bits than bf16, and
        # the byte size is identical. Stash the pre-cast dtype as a sentinel
        # so KVCacheSerializer.snapshot_layer can record it and the reload
        # path (scheduler._reconstruct_ssd_layers) can cast back to bf16 —
        # otherwise the model receives fp16 KV where it computed bf16.
        # Then force eval so the writer thread sees materialized host buffers.
        for layer in cache:
            k = getattr(layer, "keys", None)
            v = getattr(layer, "values", None)
            if k is None or v is None or isinstance(k, (tuple, list)):
                continue
            if str(getattr(k, "dtype", "")).endswith("bfloat16"):
                layer._ssd_keys_original_dtype = "bfloat16"
                layer._ssd_values_original_dtype = "bfloat16"
                layer.keys = k.astype(mx.float16)
                layer.values = v.astype(mx.float16)
            mx.eval(layer.keys, layer.values)
        logger.info(
            "[ssd_cache] dequantized %d layers before spill (%d tokens)",
            len(cache),
            len(tokens),
        )

    try:
        layer_snapshots: list[tuple[LayerSerializer, dict[str, Any]]] = []
        for layer in cache:
            serializer = get_serializer_for_layer(layer)
            snapshot = serializer.snapshot_layer(layer)
            layer_snapshots.append((serializer, snapshot))
    except Exception:
        logger.exception(
            "[ssd_cache] failed to snapshot layers for spill "
            f"({len(tokens)} tokens) — entry dropped"
        )
        return False

    try:
        self._spill_queue.put_nowait((tokens, layer_snapshots, memory_bytes))
        return True
    except queue.Full:
        logger.warning(
            f"[ssd_cache] spill queue full, dropping entry "
            f"({len(tokens)} tokens, {memory_bytes} bytes)"
        )
        return False

vllm_mlx.ssd_cache.SSDCacheTier._write_entry

_write_entry(tokens_key: tuple[int, ...], layer_snapshots: list[tuple[LayerSerializer, dict[str, Any]]], memory_bytes: int) -> None

Atomically persist one entry (writer thread; numpy-only input).

Source code in vllm_mlx/ssd_cache.py

def _write_entry(
    self,
    tokens_key: tuple[int, ...],
    layer_snapshots: list[tuple[LayerSerializer, dict[str, Any]]],
    memory_bytes: int,
) -> None:
    """Atomically persist one entry (writer thread; numpy-only input)."""
    import shutil

    entry_hash = self._entry_hash(tokens_key)
    entry_dir = os.path.join(self._data_dir, entry_hash)
    tmp_dir = entry_dir + ".tmp"

    # Clean up any leftover tmp dir from a previous crash
    if os.path.exists(tmp_dir):
        shutil.rmtree(tmp_dir)

    os.makedirs(tmp_dir, mode=self._config.dir_permissions, exist_ok=True)

    layer_manifests = []
    total_file_bytes = 0

    for i, (serializer, snapshot) in enumerate(layer_snapshots):
        layer_path = os.path.join(tmp_dir, f"layer_{i}.safetensors")
        metadata = serializer.serialize_layer(snapshot, i, layer_path)
        layer_manifests.append(metadata)

        # Set file permissions
        os.chmod(layer_path, self._config.file_permissions)
        total_file_bytes += os.path.getsize(layer_path)

    # Write manifest
    manifest = {
        "num_layers": len(layer_snapshots),
        "layers": layer_manifests,
        "memory_bytes": memory_bytes,
        "num_tokens": len(tokens_key),
    }
    manifest_path = os.path.join(tmp_dir, "manifest.json")
    with open(manifest_path, "w") as f:
        json.dump(manifest, f)
    os.chmod(manifest_path, self._config.file_permissions)

    # Save tokens binary
    tokens_path = os.path.join(tmp_dir, "tokens.bin")
    arr = _array.array("i", tokens_key)
    with open(tokens_path, "wb") as f:
        arr.tofile(f)
    os.chmod(tokens_path, self._config.file_permissions)

    # Atomic rename: tmp_dir -> entry_dir
    if os.path.exists(entry_dir):
        shutil.rmtree(entry_dir)
    os.rename(tmp_dir, entry_dir)

    # Update index
    relative_path = entry_hash
    self._index.insert_entry(
        tokens_key=tokens_key,
        file_path=relative_path,
        memory_bytes=memory_bytes,
        num_tokens=len(tokens_key),
    )

    # Update stats
    with self._lock:
        self._stats.spill_count += 1
        self._stats.spill_bytes += total_file_bytes

    logger.debug(
        f"[ssd_cache] spilled entry: {len(tokens_key)} tokens, "
        f"{total_file_bytes} bytes on disk"
    )

    # Enforce capacity after write
    self._enforce_capacity()

vllm_mlx.ssd_cache.SSDCacheTier.lookup_ssd

lookup_ssd(tokens: tuple[int, ...]) -> dict | None

Synchronous check whether tokens exist in SSD tier.

This is fast (SQLite lookup only, no disk I/O for data). Called from synchronous fetch() to report an SSD candidate.

Returns:

• dict | None –

  Dict with entry metadata if found, None otherwise.

Source code in vllm_mlx/ssd_cache.py

def lookup_ssd(self, tokens: tuple[int, ...]) -> dict | None:
    """Synchronous check whether tokens exist in SSD tier.

    This is fast (SQLite lookup only, no disk I/O for data).
    Called from synchronous fetch() to report an SSD candidate.

    Returns:
        Dict with entry metadata if found, None otherwise.
    """
    result = self._index.lookup_exact(tokens)
    if result is not None:
        return result
    return None

vllm_mlx.ssd_cache.SSDCacheTier.lookup_ssd_prefix

lookup_ssd_prefix(tokens: tuple[int, ...]) -> dict | None

Find the longest prefix match in the SSD tier.

Returns the longest-prefix entry metadata or None.

Source code in vllm_mlx/ssd_cache.py

def lookup_ssd_prefix(self, tokens: tuple[int, ...]) -> dict | None:
    """Find the longest prefix match in the SSD tier.

    Returns the longest-prefix entry metadata or None.
    """
    results = self._index.lookup_prefix(tokens)
    if results:
        return results[0]  # Already sorted by num_tokens DESC
    return None

vllm_mlx.ssd_cache.SSDCacheTier.async_promote async

async_promote(tokens: tuple[int, ...], reserve_budget_fn, release_budget_fn) -> list | None

Promote an entry from SSD to RAM asynchronously.

CRITICAL: Reserves RAM budget BEFORE the disk read, to avoid thrash when multiple promotions race.

Parameters:

• tokens (tuple[int, ...]) –

  Token sequence to promote.

• reserve_budget_fn –

  Callable(nbytes) -> bool. Must return True if budget is available and reserved, False otherwise.

• release_budget_fn –

  Callable(nbytes) -> None. Called to release budget on failure.

Returns:

• list | None –

  List of deserialized cache layers, or None if promotion failed.

Source code in vllm_mlx/ssd_cache.py

async def async_promote(
    self,
    tokens: tuple[int, ...],
    reserve_budget_fn,
    release_budget_fn,
) -> list | None:
    """Promote an entry from SSD to RAM asynchronously.

    CRITICAL: Reserves RAM budget BEFORE the disk read, to avoid
    thrash when multiple promotions race.

    Args:
        tokens: Token sequence to promote.
        reserve_budget_fn: Callable(nbytes) -> bool. Must return True
            if budget is available and reserved, False otherwise.
        release_budget_fn: Callable(nbytes) -> None. Called to release
            budget on failure.

    Returns:
        List of deserialized cache layers, or None if promotion failed.
    """
    import asyncio

    # Step 1: Look up metadata (fast, SQLite)
    meta = self._index.lookup_exact(tokens)
    if meta is None:
        with self._lock:
            self._stats.ssd_misses += 1
        return None

    memory_bytes = meta["memory_bytes"]

    # Step 2: Reserve RAM budget BEFORE disk read
    if not reserve_budget_fn(memory_bytes):
        with self._lock:
            self._stats.promotion_failures += 1
        logger.warning(
            f"[ssd_cache] promotion denied: cannot reserve "
            f"{memory_bytes} bytes RAM budget"
        )
        return None

    # Step 3: Read from disk (in thread pool to avoid blocking event loop)
    # Use shield-and-await-on-cancel per CLAUDE.md Golden Rule #4:
    # budget must be released even if the calling task is cancelled.
    t0 = time.time()
    worker = asyncio.ensure_future(
        asyncio.to_thread(self._read_entry, tokens, meta["file_path"])
    )
    try:
        cache_layers = await asyncio.shield(worker)
    except asyncio.CancelledError:
        # Caller cancelled — still need to wait for the disk read
        # to finish, then release the budget
        try:
            await worker
        except Exception:
            pass
        release_budget_fn(memory_bytes)
        raise
    except Exception:
        # Release budget on read failure
        release_budget_fn(memory_bytes)
        with self._lock:
            self._stats.promotion_failures += 1
        logger.exception(
            f"[ssd_cache] failed to read entry from disk "
            f"({meta['num_tokens']} tokens)"
        )
        return None

    if cache_layers is None:
        # Corrupted entry — release budget, quarantine entry
        release_budget_fn(memory_bytes)
        with self._lock:
            self._stats.promotion_failures += 1
        return None

    dt = time.time() - t0
    total_read_bytes = sum(
        os.path.getsize(
            os.path.join(
                self._data_dir, meta["file_path"], f"layer_{i}.safetensors"
            )
        )
        for i in range(len(cache_layers))
        if os.path.exists(
            os.path.join(
                self._data_dir, meta["file_path"], f"layer_{i}.safetensors"
            )
        )
    )

    with self._lock:
        self._stats.ssd_hits += 1
        self._stats.reload_latency_sum += dt
        self._stats.reload_bytes += total_read_bytes

    # Update access time in index
    self._index.touch(tokens)

    logger.info(
        f"[ssd_cache] promoted entry: {meta['num_tokens']} tokens, "
        f"{total_read_bytes} bytes, {dt*1000:.1f}ms"
    )

    return cache_layers

vllm_mlx.ssd_cache.SSDCacheTier._read_entry

_read_entry(tokens: tuple[int, ...], relative_path: str) -> list | None

Read a cache entry from disk. Called from thread pool.

Returns list of deserialized layer dicts, or None on corruption.

Source code in vllm_mlx/ssd_cache.py

def _read_entry(self, tokens: tuple[int, ...], relative_path: str) -> list | None:
    """Read a cache entry from disk. Called from thread pool.

    Returns list of deserialized layer dicts, or None on corruption.
    """
    entry_dir = os.path.join(self._data_dir, relative_path)
    manifest_path = os.path.join(entry_dir, "manifest.json")

    try:
        with open(manifest_path) as f:
            manifest = json.load(f)
    except (json.JSONDecodeError, OSError) as e:
        logger.warning(f"[ssd_cache] corrupt manifest for {relative_path}: {e}")
        self._quarantine_entry(tokens, relative_path)
        return None

    cache_layers = []
    for layer_meta in manifest["layers"]:
        layer_idx = layer_meta["layer_idx"]
        layer_path = os.path.join(entry_dir, f"layer_{layer_idx}.safetensors")
        layer_type = layer_meta["layer_type"]

        try:
            if layer_type in ("KVCache", "RotatingKVCache"):
                serializer = KVCacheSerializer()
            elif layer_type in ("ArraysCache", "MambaCache"):
                serializer = ArraysCacheSerializer()
            else:
                logger.warning(
                    f"[ssd_cache] unknown layer type {layer_type}, skipping"
                )
                self._quarantine_entry(tokens, relative_path)
                return None

            layer_data = serializer.deserialize_layer(layer_path, layer_meta)
            cache_layers.append(layer_data)
        except Exception as e:
            logger.warning(
                f"[ssd_cache] corrupt layer {layer_idx} in {relative_path}: {e}"
            )
            self._quarantine_entry(tokens, relative_path)
            return None

    return cache_layers

vllm_mlx.ssd_cache.SSDCacheTier._quarantine_entry

_quarantine_entry(tokens: tuple[int, ...], relative_path: str) -> None

Move a corrupt entry to quarantine and remove from index.

Source code in vllm_mlx/ssd_cache.py

def _quarantine_entry(self, tokens: tuple[int, ...], relative_path: str) -> None:
    """Move a corrupt entry to quarantine and remove from index."""
    entry_dir = os.path.join(self._data_dir, relative_path)
    quarantine_dir = os.path.join(self._cache_dir, "quarantine", relative_path)

    try:
        if os.path.exists(entry_dir):
            os.makedirs(
                os.path.dirname(quarantine_dir),
                mode=self._config.dir_permissions,
                exist_ok=True,
            )
            os.rename(entry_dir, quarantine_dir)
            logger.warning(
                f"[ssd_cache] quarantined corrupt entry: {relative_path}"
            )
    except OSError as e:
        logger.warning(f"[ssd_cache] failed to quarantine {relative_path}: {e}")

    self._index.delete_entry(tokens)

vllm_mlx.ssd_cache.SSDCacheTier._enforce_capacity

_enforce_capacity() -> None

Evict oldest SSD entries until within capacity limits.

Called after each spill write. Removes entries by LRU order until both entry count and total bytes are within bounds.

Source code in vllm_mlx/ssd_cache.py

def _enforce_capacity(self) -> None:
    """Evict oldest SSD entries until within capacity limits.

    Called after each spill write. Removes entries by LRU order
    until both entry count and total bytes are within bounds.
    """
    import shutil

    while True:
        entry_count = self._index.get_entry_count()
        total_bytes = self._index.get_total_bytes()

        needs_evict = (
            entry_count > self._config.max_entries
            or total_bytes > self._config.max_size_bytes
        )
        if not needs_evict:
            break

        lru = self._index.get_lru(limit=1)
        if not lru:
            break

        victim = lru[0]
        victim_tokens = _blob_to_tokens(victim["tokens_blob"])
        victim_dir = os.path.join(self._data_dir, victim["file_path"])

        # Delete data files
        if os.path.exists(victim_dir):
            shutil.rmtree(victim_dir)

        # Delete from index
        self._index.delete_entry(victim_tokens)

        logger.debug(
            f"[ssd_cache] disk LRU evicted: {victim['num_tokens']} tokens, "
            f"{victim['memory_bytes']} bytes"
        )

vllm_mlx.ssd_cache.SSDCacheTier.reconcile

reconcile() -> int

Reconcile index with files on disk.

Removes index entries whose data files are missing. Removes data directories not in the index.

Returns number of entries cleaned up.

Source code in vllm_mlx/ssd_cache.py

def reconcile(self) -> int:
    """Reconcile index with files on disk.

    Removes index entries whose data files are missing.
    Removes data directories not in the index.

    Returns number of entries cleaned up.
    """
    import shutil

    cleaned = 0

    # Phase 1: Remove index entries with missing data dirs
    all_entries = self._index.all_entries()
    for entry in all_entries:
        entry_dir = os.path.join(self._data_dir, entry["file_path"])
        manifest_path = os.path.join(entry_dir, "manifest.json")
        if not os.path.isdir(entry_dir) or not os.path.exists(manifest_path):
            tokens = _blob_to_tokens(entry["tokens_blob"])
            self._index.delete_entry(tokens)
            cleaned += 1
            logger.info(
                f"[ssd_cache] reconcile: removed orphaned index entry "
                f"({entry['num_tokens']} tokens, path={entry['file_path']})"
            )

    # Phase 2: Remove data directories not in the index
    if os.path.isdir(self._data_dir):
        indexed_hashes = {e["file_path"] for e in self._index.all_entries()}
        for entry_name in os.listdir(self._data_dir):
            entry_path = os.path.join(self._data_dir, entry_name)
            if (
                os.path.isdir(entry_path)
                and entry_name not in indexed_hashes
                and not entry_name.endswith(".tmp")
            ):
                shutil.rmtree(entry_path)
                cleaned += 1
                logger.info(
                    f"[ssd_cache] reconcile: removed orphaned data dir "
                    f"{entry_name}"
                )

    if cleaned > 0:
        logger.info(f"[ssd_cache] reconciliation cleaned {cleaned} entries")

    return cleaned

vllm_mlx.ssd_cache.SSDCacheTier.close

close() -> None

Close the SSD cache tier and release resources.

Source code in vllm_mlx/ssd_cache.py

def close(self) -> None:
    """Close the SSD cache tier and release resources."""
    if self._closed:
        return
    self._closed = True

    # Stop writer thread
    self._writer_stop.set()
    if self._writer_thread is not None:
        try:
            self._spill_queue.put_nowait(None)  # Poison pill
        except queue.Full:
            pass
        self._writer_thread.join(timeout=5.0)
        self._writer_thread = None

    self._index.close()
    logger.info("[ssd_cache] SSDCacheTier closed")

vllm_mlx.ssd_cache._tokens_to_blob

_tokens_to_blob(tokens: tuple[int, ...]) -> bytes

Serialize token tuple to a compact binary blob for SQLite storage.

Uses the full token sequence as a binary blob for prefix matching.

Source code in vllm_mlx/ssd_cache.py

def _tokens_to_blob(tokens: tuple[int, ...]) -> bytes:
    """Serialize token tuple to a compact binary blob for SQLite storage.

    Uses the full token sequence as a binary blob for prefix matching.
    """
    arr = _array.array("i", tokens)
    return arr.tobytes()

vllm_mlx.ssd_cache._blob_to_tokens

_blob_to_tokens(blob: bytes) -> tuple[int, ...]

Deserialize binary blob back to token tuple.

Source code in vllm_mlx/ssd_cache.py

def _blob_to_tokens(blob: bytes) -> tuple[int, ...]:
    """Deserialize binary blob back to token tuple."""
    arr = _array.array("i")
    arr.frombytes(blob)
    return tuple(arr)

vllm_mlx.ssd_cache._tokens_hash

_tokens_hash(tokens: tuple[int, ...]) -> str

Compute SHA-256 hex digest of a token sequence for use as primary key.

Source code in vllm_mlx/ssd_cache.py

def _tokens_hash(tokens: tuple[int, ...]) -> str:
    """Compute SHA-256 hex digest of a token sequence for use as primary key."""
    return hashlib.sha256(_tokens_to_blob(tokens)).hexdigest()

vllm_mlx.ssd_cache._prefix_hash

_prefix_hash(tokens: tuple[int, ...]) -> str

Hash the bounded token prefix used to prefilter prefix lookups.

Source code in vllm_mlx/ssd_cache.py

def _prefix_hash(tokens: tuple[int, ...]) -> str:
    """Hash the bounded token prefix used to prefilter prefix lookups."""
    return _tokens_hash(tokens[:_PREFIX_FILTER_TOKENS])

vllm_mlx.ssd_cache._mx_to_numpy_safe

_mx_to_numpy_safe(arr: Any) -> tuple[ndarray, str | None]

mx.array → np.ndarray, upcasting numpy-unsupported dtypes (bf16) to fp32.

Returns (numpy_array, original_dtype_name_or_None). The name is only set when an upcast happened, so the SSD-promote path can cast back.

Source code in vllm_mlx/ssd_cache.py

def _mx_to_numpy_safe(arr: Any) -> tuple[np.ndarray, str | None]:
    """mx.array → np.ndarray, upcasting numpy-unsupported dtypes (bf16) to fp32.

    Returns (numpy_array, original_dtype_name_or_None). The name is only set
    when an upcast happened, so the SSD-promote path can cast back.
    """
    try:
        return np.array(arr), None
    except RuntimeError as exc:
        # numpy ↔ mlx bf16 buffer-protocol mismatch on mlx ≥ 0.31. Re-raise
        # anything else — don't swallow unrelated errors.
        if "buffer format string" not in str(exc):
            raise
        import mlx.core as mx

        original_dtype = str(arr.dtype).rsplit(".", 1)[-1]
        upcast = arr.astype(mx.float32)
        mx.eval(upcast)  # astype is lazy; force materialization here
        return np.array(upcast), original_dtype

vllm_mlx.ssd_cache.get_serializer_for_layer

get_serializer_for_layer(layer: Any) -> LayerSerializer

Return the appropriate serializer for a cache layer.

Dispatches based on duck-typing: - If layer has .keys and .values and .offset -> KVCacheSerializer - If layer has .state and it's a list -> ArraysCacheSerializer

Raises ValueError for unsupported layer types.

Source code in vllm_mlx/ssd_cache.py

def get_serializer_for_layer(layer: Any) -> LayerSerializer:
    """Return the appropriate serializer for a cache layer.

    Dispatches based on duck-typing:
    - If layer has .keys and .values and .offset -> KVCacheSerializer
    - If layer has .state and it's a list -> ArraysCacheSerializer

    Raises ValueError for unsupported layer types.
    """
    if hasattr(layer, "keys") and hasattr(layer, "values") and hasattr(layer, "offset"):
        return KVCacheSerializer()
    if hasattr(layer, "state") and isinstance(getattr(layer, "state", None), list):
        return ArraysCacheSerializer()
    raise ValueError(
        f"Unsupported cache layer type: {type(layer).__name__}. "
        f"Supported: {list(SERIALIZER_SUPPORT_MATRIX.keys())}"
    )

Complete contract reference

Expand any definition for its exact inputs, annotations, defaults, return contract, directly raised exceptions, source-grounded behavior, and immutable line link. This section includes private and nested definitions that ordinary API generators omit.

vllm_mlx.ssd_cache.SSDCacheConfig · class

vllm_mlx.ssd_cache.SSDCacheConfig(cache_dir: str | None = None, max_size_gb: float = 10.0, max_entries: int = 10000, file_permissions: int = 384, dir_permissions: int = 448, spill_queue_size: int = 64, retention_seconds: int | None = None)

Configuration for SSD cache tier.

Parameters

Name	Type	Required	Default	Description
cache_dir	str \| None	no	None	Optional constructor field; defaults to None.
max_size_gb	float	no	10.0	Optional constructor field; defaults to 10.0.
max_entries	int	no	10000	Optional constructor field; defaults to 10000.
file_permissions	int	no	384	Optional constructor field; defaults to 384.
dir_permissions	int	no	448	Optional constructor field; defaults to 448.
spill_queue_size	int	no	64	Optional constructor field; defaults to 64.
retention_seconds	int \| None	no	None	Optional constructor field; defaults to None.

Returns

• Constructs: vllm_mlx.ssd_cache.SSDCacheConfig

Exceptions and behavior

Class SSDCacheConfig declares 2 direct member(s). No direct raise statement appears in this definition.

View source #L43-L78.

vllm_mlx.ssd_cache.SSDCacheConfig.__post_init__ · method

vllm_mlx.ssd_cache.SSDCacheConfig.__post_init__() -> None

Method SSDCacheConfig.__post_init__ calls ValueError; can raise ValueError.

Parameters

This callable has no explicit inputs.

Returns

• Type: None

Exceptions and behavior

Method SSDCacheConfig.__post_init__ calls ValueError; can raise ValueError. Directly raised exceptions: ValueError.

View source #L65-L73.

vllm_mlx.ssd_cache.SSDCacheConfig.max_size_bytes · method

vllm_mlx.ssd_cache.SSDCacheConfig.max_size_bytes() -> int

Maximum cache size in bytes.

Parameters

This callable has no explicit inputs.

Returns

• Type: int
• Direct return expressions: int(self.max_size_gb * _BYTES_PER_GB)

Exceptions and behavior

Method SSDCacheConfig.max_size_bytes calls int; returns int(self.max_size_gb * _BYTES_PER_GB). No direct raise statement appears in this definition.

View source #L76-L78.

vllm_mlx.ssd_cache.SSDCacheStats · class

vllm_mlx.ssd_cache.SSDCacheStats(spill_count: int = 0, spill_bytes: int = 0, ssd_hits: int = 0, ssd_misses: int = 0, reload_latency_sum: float = 0.0, reload_bytes: int = 0, promotion_failures: int = 0)

Statistics for SSD cache tier — exposed from day one.

Parameters

Name	Type	Required	Default	Description
spill_count	int	no	0	Optional constructor field; defaults to 0.
spill_bytes	int	no	0	Optional constructor field; defaults to 0.
ssd_hits	int	no	0	Optional constructor field; defaults to 0.
ssd_misses	int	no	0	Optional constructor field; defaults to 0.
reload_latency_sum	float	no	0.0	Optional constructor field; defaults to 0.0.
reload_bytes	int	no	0	Optional constructor field; defaults to 0.
promotion_failures	int	no	0	Optional constructor field; defaults to 0.

Returns

• Constructs: vllm_mlx.ssd_cache.SSDCacheStats

Exceptions and behavior

Class SSDCacheStats declares 1 direct member(s). No direct raise statement appears in this definition.

View source #L82-L123.

vllm_mlx.ssd_cache.SSDCacheStats.to_dict · method

vllm_mlx.ssd_cache.SSDCacheStats.to_dict() -> dict

Return spill, lookup, reload, and promotion statistics.

Parameters

This callable has no explicit inputs.

Returns

• Type: dict
• Direct return expressions: {'spill_count': self.spill_count, 'spill_bytes': self.spill_bytes, 'ssd_hits': self.ssd_hits, 'ssd_misses': self.ssd_mi…

Exceptions and behavior

Method SSDCacheStats.to_dict calls round; returns {'spill_count': self.spill_count, 'spill_bytes': self.spill_bytes, 'ssd_hits': self.ssd_hits, 'ssd_misses': self.ssd_mi…. No direct raise statement appears in this definition.

View source #L103-L123.

vllm_mlx.ssd_cache._tokens_to_blob · function

vllm_mlx.ssd_cache._tokens_to_blob(tokens: tuple[int, ...]) -> bytes

Serialize token tuple to a compact binary blob for SQLite storage.

Parameters

Name	Type	Required	Default	Description
tokens	tuple[int, ...]	yes	none	Required positional or keyword input.

Returns

• Type: bytes
• Direct return expressions: arr.tobytes()

Exceptions and behavior

Function _tokens_to_blob calls _array.array, arr.tobytes; returns arr.tobytes(). No direct raise statement appears in this definition.

View source #L126-L132.

vllm_mlx.ssd_cache._blob_to_tokens · function

vllm_mlx.ssd_cache._blob_to_tokens(blob: bytes) -> tuple[int, ...]

Deserialize binary blob back to token tuple.

Parameters

Name	Type	Required	Default	Description
blob	bytes	yes	none	Required positional or keyword input.

Returns

• Type: tuple[int, ...]
• Direct return expressions: tuple(arr)

Exceptions and behavior

Function _blob_to_tokens calls _array.array, arr.frombytes, tuple; returns tuple(arr). No direct raise statement appears in this definition.

View source #L135-L139.

vllm_mlx.ssd_cache._tokens_hash · function

vllm_mlx.ssd_cache._tokens_hash(tokens: tuple[int, ...]) -> str

Compute SHA-256 hex digest of a token sequence for use as primary key.

Parameters

Name	Type	Required	Default	Description
tokens	tuple[int, ...]	yes	none	Required positional or keyword input.

Returns

• Type: str
• Direct return expressions: hashlib.sha256(_tokens_to_blob(tokens)).hexdigest()

Exceptions and behavior

Function _tokens_hash calls hashlib.sha256(_tokens_to_blob(tokens)).hexdigest, hashlib.sha256, _tokens_to_blob; returns hashlib.sha256(_tokens_to_blob(tokens)).hexdigest(). No direct raise statement appears in this definition.

View source #L142-L144.

vllm_mlx.ssd_cache._prefix_hash · function

vllm_mlx.ssd_cache._prefix_hash(tokens: tuple[int, ...]) -> str

Hash the bounded token prefix used to prefilter prefix lookups.

Parameters

Name	Type	Required	Default	Description
tokens	tuple[int, ...]	yes	none	Required positional or keyword input.

Returns

• Type: str
• Direct return expressions: _tokens_hash(tokens[:_PREFIX_FILTER_TOKENS])

Exceptions and behavior

Function _prefix_hash calls _tokens_hash; returns _tokens_hash(tokens[:_PREFIX_FILTER_TOKENS]). No direct raise statement appears in this definition.

View source #L147-L149.

vllm_mlx.ssd_cache.SSDIndex · class

vllm_mlx.ssd_cache.SSDIndex(cache_dir: str)

SQLite-backed index for SSD cache entries.

Parameters

Name	Type	Required	Default	Description
cache_dir	str	yes	none	Required positional or keyword input.

Returns

• Constructs: vllm_mlx.ssd_cache.SSDIndex

Exceptions and behavior

Class SSDIndex declares 14 direct member(s). No direct raise statement appears in this definition.

View source #L152-L405.

vllm_mlx.ssd_cache.SSDIndex.__init__ · method

vllm_mlx.ssd_cache.SSDIndex.__init__(cache_dir: str) -> None

Method SSDIndex.__init__ updates self._cache_dir, self._db_lock, self._conn, self._conn.row_factory; calls threading.Lock, os.path.join, sqlite3.connect, self._conn.execute.

Parameters

Name	Type	Required	Default	Description
cache_dir	str	yes	none	Required positional or keyword input.

Returns

• Type: None

Exceptions and behavior

Method SSDIndex.__init__ updates self._cache_dir, self._db_lock, self._conn, self._conn.row_factory; calls threading.Lock, os.path.join, sqlite3.connect, self._conn.execute. No direct raise statement appears in this definition.

View source #L165-L173.

vllm_mlx.ssd_cache.SSDIndex._create_tables · method

vllm_mlx.ssd_cache.SSDIndex._create_tables() -> None

Method SSDIndex._create_tables calls self._conn.executescript, self._ensure_column, self._conn.execute, cur.fetchone.

Parameters

This callable has no explicit inputs.

Returns

• Type: None

Exceptions and behavior

Method SSDIndex._create_tables calls self._conn.executescript, self._ensure_column, self._conn.execute, cur.fetchone. No direct raise statement appears in this definition.

View source #L175-L213.

vllm_mlx.ssd_cache.SSDIndex._ensure_column · method

vllm_mlx.ssd_cache.SSDIndex._ensure_column(table: str, column: str, definition: str) -> None

Method SSDIndex._ensure_column calls self._conn.execute, cur.fetchall.

Parameters

Name	Type	Required	Default	Description
table	str	yes	none	Required positional or keyword input.
column	str	yes	none	Required positional or keyword input.
definition	str	yes	none	Required positional or keyword input.

Returns

• Type: None

Exceptions and behavior

Method SSDIndex._ensure_column calls self._conn.execute, cur.fetchall. No direct raise statement appears in this definition.

View source #L215-L218.

vllm_mlx.ssd_cache.SSDIndex._backfill_prefix_hashes · method

vllm_mlx.ssd_cache.SSDIndex._backfill_prefix_hashes() -> None

Method SSDIndex._backfill_prefix_hashes calls self._conn.execute, cur.fetchall, _blob_to_tokens, _prefix_hash.

Parameters

This callable has no explicit inputs.

Returns

• Type: None

Exceptions and behavior

Method SSDIndex._backfill_prefix_hashes calls self._conn.execute, cur.fetchall, _blob_to_tokens, _prefix_hash. No direct raise statement appears in this definition.

View source #L220-L230.

vllm_mlx.ssd_cache.SSDIndex.insert_entry · method

vllm_mlx.ssd_cache.SSDIndex.insert_entry(tokens_key: tuple[int, ...], file_path: str, memory_bytes: int, num_tokens: int) -> None

Insert or replace a cache entry in the index.

Parameters

Name	Type	Required	Default	Description
tokens_key	tuple[int, ...]	yes	none	Required positional or keyword input.
file_path	str	yes	none	Required positional or keyword input.
memory_bytes	int	yes	none	Required positional or keyword input.
num_tokens	int	yes	none	Required positional or keyword input.

Returns

• Type: None

Exceptions and behavior

Method SSDIndex.insert_entry calls time.time, _tokens_hash, _prefix_hash, _tokens_to_blob. No direct raise statement appears in this definition.

View source #L232-L263.

vllm_mlx.ssd_cache.SSDIndex.lookup_exact · method

vllm_mlx.ssd_cache.SSDIndex.lookup_exact(tokens_key: tuple[int, ...]) -> dict | None

Look up an exact token sequence.

Parameters

Name	Type	Required	Default	Description
tokens_key	tuple[int, ...]	yes	none	Required positional or keyword input.

Returns

• Type: dict | None
• Direct return expressions: None; {'file_path': row['file_path'], 'memory_bytes': row['memory_bytes'], 'num_tokens': row['num_tokens']}

Exceptions and behavior

Method SSDIndex.lookup_exact calls _tokens_hash, self._conn.execute, cur.fetchone; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L265-L280.

vllm_mlx.ssd_cache.SSDIndex.lookup_prefix · method

vllm_mlx.ssd_cache.SSDIndex.lookup_prefix(query_tokens: tuple[int, ...]) -> list[dict]

Find entries whose token sequence is a prefix of query_tokens.

Parameters

Name	Type	Required	Default	Description
query_tokens	tuple[int, ...]	yes	none	Required positional or keyword input.

Returns

• Type: list[dict]
• Direct return expressions: []; results

Exceptions and behavior

Method SSDIndex.lookup_prefix calls len, _tokens_to_blob, _tokens_hash, range; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L282-L324.

vllm_mlx.ssd_cache.SSDIndex.delete_entry · method

vllm_mlx.ssd_cache.SSDIndex.delete_entry(tokens_key: tuple[int, ...]) -> None

Delete an entry by token sequence.

Parameters

Name	Type	Required	Default	Description
tokens_key	tuple[int, ...]	yes	none	Required positional or keyword input.

Returns

• Type: None

Exceptions and behavior

Method SSDIndex.delete_entry calls _tokens_hash, self._conn.execute, self._conn.commit. No direct raise statement appears in this definition.

View source #L326-L333.

vllm_mlx.ssd_cache.SSDIndex.get_lru · method

vllm_mlx.ssd_cache.SSDIndex.get_lru(limit: int = 10) -> list[dict]

Get the least recently used entries, ordered oldest first.

Parameters

Name	Type	Required	Default	Description
limit	int	no	10	Optional positional or keyword input; defaults to 10.

Returns

• Type: list[dict]
• Direct return expressions: results

Exceptions and behavior

Method SSDIndex.get_lru calls self._conn.execute, cur.fetchall, results.append; returns results. No direct raise statement appears in this definition.

View source #L335-L355.

vllm_mlx.ssd_cache.SSDIndex.get_total_bytes · method

vllm_mlx.ssd_cache.SSDIndex.get_total_bytes() -> int

Get total memory_bytes across all entries.

Parameters

This callable has no explicit inputs.

Returns

• Type: int
• Direct return expressions: cur.fetchone()[0]

Exceptions and behavior

Method SSDIndex.get_total_bytes calls self._conn.execute, cur.fetchone; returns cur.fetchone()[0]. No direct raise statement appears in this definition.

View source #L357-L363.

vllm_mlx.ssd_cache.SSDIndex.get_entry_count · method

vllm_mlx.ssd_cache.SSDIndex.get_entry_count() -> int

Get number of entries in the index.

Parameters

This callable has no explicit inputs.

Returns

• Type: int
• Direct return expressions: cur.fetchone()[0]

Exceptions and behavior

Method SSDIndex.get_entry_count calls self._conn.execute, cur.fetchone; returns cur.fetchone()[0]. No direct raise statement appears in this definition.

View source #L365-L369.

vllm_mlx.ssd_cache.SSDIndex.touch · method

vllm_mlx.ssd_cache.SSDIndex.touch(tokens_key: tuple[int, ...]) -> None

Update accessed_at timestamp for an entry (marks as recently used).

Parameters

Name	Type	Required	Default	Description
tokens_key	tuple[int, ...]	yes	none	Required positional or keyword input.

Returns

• Type: None

Exceptions and behavior

Method SSDIndex.touch calls _tokens_hash, self._conn.execute, time.time, self._conn.commit. No direct raise statement appears in this definition.

View source #L371-L379.

vllm_mlx.ssd_cache.SSDIndex.all_entries · method

vllm_mlx.ssd_cache.SSDIndex.all_entries() -> list[dict]

Return all entries (for startup reconciliation).

Parameters

This callable has no explicit inputs.

Returns

• Type: list[dict]
• Direct return expressions: results

Exceptions and behavior

Method SSDIndex.all_entries calls self._conn.execute, cur.fetchall, results.append; returns results. No direct raise statement appears in this definition.

View source #L381-L400.

vllm_mlx.ssd_cache.SSDIndex.close · method

vllm_mlx.ssd_cache.SSDIndex.close() -> None

Close the SQLite connection.

Parameters

This callable has no explicit inputs.

Returns

• Type: None

Exceptions and behavior

Method SSDIndex.close calls self._conn.close. No direct raise statement appears in this definition.

View source #L402-L405.

vllm_mlx.ssd_cache.LayerSerializer · class

vllm_mlx.ssd_cache.LayerSerializer()

Interface for per-layer cache serialization.

Parameters

This callable has no explicit inputs.

Returns

• Constructs: vllm_mlx.ssd_cache.LayerSerializer

Exceptions and behavior

Class LayerSerializer derives from ABC and declares 3 direct member(s). No direct raise statement appears in this definition.

View source #L419-L446.

vllm_mlx.ssd_cache.LayerSerializer.snapshot_layer · method

vllm_mlx.ssd_cache.LayerSerializer.snapshot_layer(layer: Any) -> dict[str, Any]

Producer-thread CPU snapshot of an MLX-backed cache layer.

Parameters

Name	Type	Required	Default	Description
layer	Any	yes	none	Required positional or keyword input.

Returns

• Type: dict[str, Any]

Exceptions and behavior

Method LayerSerializer.snapshot_layer contains no state mutation, call, raise, return, await, or yield. No direct raise statement appears in this definition.

View source #L429-L431.

vllm_mlx.ssd_cache.LayerSerializer.serialize_layer · method

vllm_mlx.ssd_cache.LayerSerializer.serialize_layer(snapshot: dict[str, Any], layer_idx: int, file_path: str) -> dict[str, Any]

Writer-thread: persist a snapshot to safetensors at file_path.

Parameters

Name	Type	Required	Default	Description
snapshot	dict[str, Any]	yes	none	Required positional or keyword input.
layer_idx	int	yes	none	Required positional or keyword input.
file_path	str	yes	none	Required positional or keyword input.

Returns

• Type: dict[str, Any]

Exceptions and behavior

Method LayerSerializer.serialize_layer contains no state mutation, call, raise, return, await, or yield. No direct raise statement appears in this definition.

View source #L434-L441.

vllm_mlx.ssd_cache.LayerSerializer.deserialize_layer · method

vllm_mlx.ssd_cache.LayerSerializer.deserialize_layer(file_path: str, metadata: dict[str, Any]) -> dict

Read a layer back from disk.

Parameters

Name	Type	Required	Default	Description
file_path	str	yes	none	Required positional or keyword input.
metadata	dict[str, Any]	yes	none	Required positional or keyword input.

Returns

• Type: dict

Exceptions and behavior

Method LayerSerializer.deserialize_layer contains no state mutation, call, raise, return, await, or yield. No direct raise statement appears in this definition.

View source #L444-L446.

vllm_mlx.ssd_cache._mx_to_numpy_safe · function

vllm_mlx.ssd_cache._mx_to_numpy_safe(arr: Any) -> tuple[np.ndarray, str | None]

mx.array → np.ndarray, upcasting numpy-unsupported dtypes (bf16) to fp32.

Parameters

Name	Type	Required	Default	Description
arr	Any	yes	none	Required positional or keyword input.

Returns

• Type: tuple[np.ndarray, str | None]
• Direct return expressions: (np.array(arr), None); (np.array(upcast), original_dtype)

Exceptions and behavior

Function _mx_to_numpy_safe calls np.array, str, str(arr.dtype).rsplit, arr.astype; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L449-L467.

vllm_mlx.ssd_cache.KVCacheSerializer · class

vllm_mlx.ssd_cache.KVCacheSerializer()

Serializer for KVCache and RotatingKVCache layers.

Parameters

This callable has no explicit inputs.

Returns

• Constructs: vllm_mlx.ssd_cache.KVCacheSerializer

Exceptions and behavior

Class KVCacheSerializer derives from LayerSerializer and declares 3 direct member(s). No direct raise statement appears in this definition.

View source #L470-L564.

vllm_mlx.ssd_cache.KVCacheSerializer.snapshot_layer · method

vllm_mlx.ssd_cache.KVCacheSerializer.snapshot_layer(layer: Any) -> dict[str, Any]

Copy a KV cache layer into NumPy-backed writer-thread data.

Parameters

Name	Type	Required	Default	Description
layer	Any	yes	none	Required positional or keyword input.

Returns

• Type: dict[str, Any]
• Direct return expressions: snapshot

Exceptions and behavior

Method KVCacheSerializer.snapshot_layer calls _mx_to_numpy_safe, getattr, hasattr; returns snapshot. No direct raise statement appears in this definition.

View source #L481-L515.

vllm_mlx.ssd_cache.KVCacheSerializer.serialize_layer · method

vllm_mlx.ssd_cache.KVCacheSerializer.serialize_layer(snapshot: dict[str, Any], layer_idx: int, file_path: str) -> dict[str, Any]

Write one KV layer to safetensors and return reconstruction metadata.

Parameters

Name	Type	Required	Default	Description
snapshot	dict[str, Any]	yes	none	Required positional or keyword input.
layer_idx	int	yes	none	Required positional or keyword input.
file_path	str	yes	none	Required positional or keyword input.

Returns

• Type: dict[str, Any]
• Direct return expressions: metadata

Exceptions and behavior

Method KVCacheSerializer.serialize_layer calls save_file; returns metadata. No direct raise statement appears in this definition.

View source #L517-L542.

vllm_mlx.ssd_cache.KVCacheSerializer.deserialize_layer · method

vllm_mlx.ssd_cache.KVCacheSerializer.deserialize_layer(file_path: str, metadata: dict[str, Any]) -> dict

Load one KV layer as arrays plus cache reconstruction metadata.

Parameters

Name	Type	Required	Default	Description
file_path	str	yes	none	Required positional or keyword input.
metadata	dict[str, Any]	yes	none	Required positional or keyword input.

Returns

• Type: dict
• Direct return expressions: result

Exceptions and behavior

Method KVCacheSerializer.deserialize_layer calls load_file; returns result. No direct raise statement appears in this definition.

View source #L544-L564.

vllm_mlx.ssd_cache.ArraysCacheSerializer · class

vllm_mlx.ssd_cache.ArraysCacheSerializer()

Serializer for ArraysCache (Mamba/linear attention) layers.

Parameters

This callable has no explicit inputs.

Returns

• Constructs: vllm_mlx.ssd_cache.ArraysCacheSerializer

Exceptions and behavior

Class ArraysCacheSerializer derives from LayerSerializer and declares 3 direct member(s). No direct raise statement appears in this definition.

View source #L567-L627.

vllm_mlx.ssd_cache.ArraysCacheSerializer.snapshot_layer · method

vllm_mlx.ssd_cache.ArraysCacheSerializer.snapshot_layer(layer: Any) -> dict[str, Any]

Copy an arrays-cache state into NumPy-backed writer-thread data.

Parameters

Name	Type	Required	Default	Description
layer	Any	yes	none	Required positional or keyword input.

Returns

• Type: dict[str, Any]
• Direct return expressions: snapshot

Exceptions and behavior

Method ArraysCacheSerializer.snapshot_layer calls _mx_to_numpy_safe, state_np.append, original_dtypes.append, any; returns snapshot. No direct raise statement appears in this definition.

View source #L573-L587.

vllm_mlx.ssd_cache.ArraysCacheSerializer.serialize_layer · method

vllm_mlx.ssd_cache.ArraysCacheSerializer.serialize_layer(snapshot: dict[str, Any], layer_idx: int, file_path: str) -> dict[str, Any]

Write arrays-cache state to safetensors and return its metadata.

Parameters

Name	Type	Required	Default	Description
snapshot	dict[str, Any]	yes	none	Required positional or keyword input.
layer_idx	int	yes	none	Required positional or keyword input.
file_path	str	yes	none	Required positional or keyword input.

Returns

• Type: dict[str, Any]
• Direct return expressions: metadata

Exceptions and behavior

Method ArraysCacheSerializer.serialize_layer calls enumerate, save_file, len; returns metadata. No direct raise statement appears in this definition.

View source #L589-L610.

vllm_mlx.ssd_cache.ArraysCacheSerializer.deserialize_layer · method

vllm_mlx.ssd_cache.ArraysCacheSerializer.deserialize_layer(file_path: str, metadata: dict[str, Any]) -> dict

Load arrays-cache state and any original dtype hints.

Parameters

Name	Type	Required	Default	Description
file_path	str	yes	none	Required positional or keyword input.
metadata	dict[str, Any]	yes	none	Required positional or keyword input.

Returns

• Type: dict
• Direct return expressions: result

Exceptions and behavior

Method ArraysCacheSerializer.deserialize_layer calls load_file, range, state.append; returns result. No direct raise statement appears in this definition.

View source #L612-L627.

vllm_mlx.ssd_cache.get_serializer_for_layer · function

vllm_mlx.ssd_cache.get_serializer_for_layer(layer: Any) -> LayerSerializer

Return the appropriate serializer for a cache layer.

Parameters

Name	Type	Required	Default	Description
layer	Any	yes	none	Required positional or keyword input.

Returns

• Type: LayerSerializer
• Direct return expressions: KVCacheSerializer(); ArraysCacheSerializer()

Exceptions and behavior

Function get_serializer_for_layer calls hasattr, KVCacheSerializer, isinstance, getattr; can raise ValueError; has 2 explicit return paths. Directly raised exceptions: ValueError.

View source #L630-L646.

vllm_mlx.ssd_cache.SSDCacheTier · class

vllm_mlx.ssd_cache.SSDCacheTier(config: SSDCacheConfig)

Cold-tier disk cache for KV cache entries.

Parameters

Name	Type	Required	Default	Description
config	SSDCacheConfig	yes	none	Required positional or keyword input.

Returns

• Constructs: vllm_mlx.ssd_cache.SSDCacheTier

Exceptions and behavior

Class SSDCacheTier declares 15 direct member(s). No direct raise statement appears in this definition.

View source #L649-L1248.

vllm_mlx.ssd_cache.SSDCacheTier.__init__ · method

vllm_mlx.ssd_cache.SSDCacheTier.__init__(config: SSDCacheConfig) -> None

Method SSDCacheTier.__init__ updates self._config, self._closed, self._writer_thread, self._cache_dir; calls ValueError, os.path.join, os.makedirs, SSDIndex; can raise ValueError.

Parameters

Name	Type	Required	Default	Description
config	SSDCacheConfig	yes	none	Required positional or keyword input.

Returns

• Type: None

Exceptions and behavior

Method SSDCacheTier.__init__ updates self._config, self._closed, self._writer_thread, self._cache_dir; calls ValueError, os.path.join, os.makedirs, SSDIndex; can raise ValueError. Directly raised exceptions: ValueError.

View source #L667-L705.

vllm_mlx.ssd_cache.SSDCacheTier._entry_hash · method

vllm_mlx.ssd_cache.SSDCacheTier._entry_hash(tokens: tuple[int, ...]) -> str

Compute deterministic hash for a token sequence.

Parameters

Name	Type	Required	Default	Description
tokens	tuple[int, ...]	yes	none	Required positional or keyword input.

Returns

• Type: str
• Direct return expressions: _tokens_hash(tokens)

Exceptions and behavior

Method SSDCacheTier._entry_hash calls _tokens_hash; returns _tokens_hash(tokens). No direct raise statement appears in this definition.

View source #L708-L710.

vllm_mlx.ssd_cache.SSDCacheTier.get_stats · method

vllm_mlx.ssd_cache.SSDCacheTier.get_stats() -> dict

Return current SSD cache statistics.

Parameters

This callable has no explicit inputs.

Returns

• Type: dict
• Direct return expressions: self._stats.to_dict()

Exceptions and behavior

Method SSDCacheTier.get_stats calls self._stats.to_dict; returns self._stats.to_dict(). No direct raise statement appears in this definition.

View source #L712-L714.

vllm_mlx.ssd_cache.SSDCacheTier.start_writer · method

vllm_mlx.ssd_cache.SSDCacheTier.start_writer() -> None

Start the background spill writer thread.

Parameters

This callable has no explicit inputs.

Returns

• Type: None
• Direct return expressions: None

Exceptions and behavior

Method SSDCacheTier.start_writer updates self._writer_thread; calls self._writer_stop.clear, threading.Thread, self._writer_thread.start, logger.info; returns None. No direct raise statement appears in this definition.

View source #L716-L725.

vllm_mlx.ssd_cache.SSDCacheTier._writer_loop · method

vllm_mlx.ssd_cache.SSDCacheTier._writer_loop() -> None

Drain spill queue and persist entries.

Parameters

This callable has no explicit inputs.

Returns

• Type: None

Exceptions and behavior

Method SSDCacheTier._writer_loop calls self._writer_stop.is_set, self._spill_queue.get, self._write_entry, logger.exception. No direct raise statement appears in this definition.

View source #L727-L744.

vllm_mlx.ssd_cache.SSDCacheTier.enqueue_spill · method

vllm_mlx.ssd_cache.SSDCacheTier.enqueue_spill(tokens: tuple[int, ...], cache: list[Any], memory_bytes: int) -> bool

Enqueue a cache entry for async spill to SSD.

Parameters

Name	Type	Required	Default	Description
tokens	tuple[int, ...]	yes	none	Required positional or keyword input.
cache	list[Any]	yes	none	Required positional or keyword input.
memory_bytes	int	yes	none	Required positional or keyword input.

Returns

• Type: bool
• Direct return expressions: False; True

Exceptions and behavior

Method SSDCacheTier.enqueue_spill calls any, _is_quantized_layer, isinstance, converted.extend; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L746-L867.

vllm_mlx.ssd_cache.SSDCacheTier.enqueue_spill._is_quantized_layer · nested function

vllm_mlx.ssd_cache.SSDCacheTier.enqueue_spill._is_quantized_layer(layer) -> not annotated

Nested Function SSDCacheTier.enqueue_spill._is_quantized_layer calls isinstance, getattr; has 2 explicit return paths.

Parameters

Name	Type	Required	Default	Description
layer	not annotated	yes	none	Required positional or keyword input.

Returns

• Type: not annotated
• Direct return expressions: True; isinstance(keys, (tuple, list))

Exceptions and behavior

Nested Function SSDCacheTier.enqueue_spill._is_quantized_layer calls isinstance, getattr; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L772-L776.

vllm_mlx.ssd_cache.SSDCacheTier._write_entry · method

vllm_mlx.ssd_cache.SSDCacheTier._write_entry(tokens_key: tuple[int, ...], layer_snapshots: list[tuple[LayerSerializer, dict[str, Any]]], memory_bytes: int) -> None

Atomically persist one entry (writer thread; numpy-only input).

Parameters

Name	Type	Required	Default	Description
tokens_key	tuple[int, ...]	yes	none	Required positional or keyword input.
layer_snapshots	list[tuple[LayerSerializer, dict[str, Any]]]	yes	none	Required positional or keyword input.
memory_bytes	int	yes	none	Required positional or keyword input.

Returns

• Type: None

Exceptions and behavior

Method SSDCacheTier._write_entry updates self._stats.spill_count, self._stats.spill_bytes; calls self._entry_hash, os.path.join, os.path.exists, shutil.rmtree. No direct raise statement appears in this definition.

View source #L869-L944.

vllm_mlx.ssd_cache.SSDCacheTier.lookup_ssd · method

vllm_mlx.ssd_cache.SSDCacheTier.lookup_ssd(tokens: tuple[int, ...]) -> dict | None

Synchronous check whether tokens exist in SSD tier.

Parameters

Name	Type	Required	Default	Description
tokens	tuple[int, ...]	yes	none	Required positional or keyword input.

Returns

• Type: dict | None
• Direct return expressions: result; None

Exceptions and behavior

Method SSDCacheTier.lookup_ssd calls self._index.lookup_exact; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L946-L958.

vllm_mlx.ssd_cache.SSDCacheTier.lookup_ssd_prefix · method

vllm_mlx.ssd_cache.SSDCacheTier.lookup_ssd_prefix(tokens: tuple[int, ...]) -> dict | None

Find the longest prefix match in the SSD tier.

Parameters

Name	Type	Required	Default	Description
tokens	tuple[int, ...]	yes	none	Required positional or keyword input.

Returns

• Type: dict | None
• Direct return expressions: results[0]; None

Exceptions and behavior

Method SSDCacheTier.lookup_ssd_prefix calls self._index.lookup_prefix; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L960-L968.

vllm_mlx.ssd_cache.SSDCacheTier.async_promote · method

async vllm_mlx.ssd_cache.SSDCacheTier.async_promote(tokens: tuple[int, ...], reserve_budget_fn, release_budget_fn) -> list | None

Promote an entry from SSD to RAM asynchronously.

Parameters

Name	Type	Required	Default	Description
tokens	tuple[int, ...]	yes	none	Token sequence to promote.
reserve_budget_fn	not annotated	yes	none	Callable(nbytes) -> bool. Must return True if budget is available and reserved, False otherwise.
release_budget_fn	not annotated	yes	none	Callable(nbytes) -> None. Called to release budget on failure.

Returns

• Type: list | None
• Direct return expressions: None; cache_layers

Exceptions and behavior

Method SSDCacheTier.async_promote updates self._stats.ssd_misses, self._stats.promotion_failures, self._stats.ssd_hits, self._stats.reload_latency_sum; calls self._index.lookup_exact, reserve_budget_fn, logger.warning, time.time; awaits asynchronous work; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L970-L1076.

vllm_mlx.ssd_cache.SSDCacheTier._read_entry · method

vllm_mlx.ssd_cache.SSDCacheTier._read_entry(tokens: tuple[int, ...], relative_path: str) -> list | None

Read a cache entry from disk.

Parameters

Name	Type	Required	Default	Description
tokens	tuple[int, ...]	yes	none	Required positional or keyword input.
relative_path	str	yes	none	Required positional or keyword input.

Returns

• Type: list | None
• Direct return expressions: None; cache_layers

Exceptions and behavior

Method SSDCacheTier._read_entry calls os.path.join, open, json.load, logger.warning; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L1078-L1121.

vllm_mlx.ssd_cache.SSDCacheTier._quarantine_entry · method

vllm_mlx.ssd_cache.SSDCacheTier._quarantine_entry(tokens: tuple[int, ...], relative_path: str) -> None

Move a corrupt entry to quarantine and remove from index.

Parameters

Name	Type	Required	Default	Description
tokens	tuple[int, ...]	yes	none	Required positional or keyword input.
relative_path	str	yes	none	Required positional or keyword input.

Returns

• Type: None

Exceptions and behavior

Method SSDCacheTier._quarantine_entry calls os.path.join, os.path.exists, os.makedirs, os.path.dirname. No direct raise statement appears in this definition.

View source #L1123-L1142.

vllm_mlx.ssd_cache.SSDCacheTier._enforce_capacity · method

vllm_mlx.ssd_cache.SSDCacheTier._enforce_capacity() -> None

Evict oldest SSD entries until within capacity limits.

Parameters

This callable has no explicit inputs.

Returns

• Type: None

Exceptions and behavior

Method SSDCacheTier._enforce_capacity calls self._index.get_entry_count, self._index.get_total_bytes, self._index.get_lru, _blob_to_tokens. No direct raise statement appears in this definition.

View source #L1144-L1181.

vllm_mlx.ssd_cache.SSDCacheTier.reconcile · method

vllm_mlx.ssd_cache.SSDCacheTier.reconcile() -> int

Reconcile index with files on disk.

Parameters

This callable has no explicit inputs.

Returns

• Type: int
• Direct return expressions: cleaned

Exceptions and behavior

Method SSDCacheTier.reconcile calls self._index.all_entries, os.path.join, os.path.isdir, os.path.exists; returns cleaned. No direct raise statement appears in this definition.

View source #L1183-L1229.

vllm_mlx.ssd_cache.SSDCacheTier.close · method

vllm_mlx.ssd_cache.SSDCacheTier.close() -> None

Close the SSD cache tier and release resources.

Parameters

This callable has no explicit inputs.

Returns

• Type: None
• Direct return expressions: None

Exceptions and behavior

Method SSDCacheTier.close updates self._closed, self._writer_thread; calls self._writer_stop.set, self._spill_queue.put_nowait, self._writer_thread.join, self._index.close; returns None. No direct raise statement appears in this definition.

View source #L1231-L1248.

Complete symbol map

This map also includes private definitions and nested helpers. The signature column exposes every explicit input even when an internal helper has no dedicated parameter prose.

Symbol	Kind	Signature and inputs	What it does	Source
SSDCacheConfig	class	SSDCacheConfig(cache_dir: str \| None = None, max_size_gb: float = 10.0, max_entries: int = 10000, file_permissions: int = 384, dir_permissions: int = 448, spill_queue_size: int = 64, retention_seconds: int \| None = None)	Configuration for SSD cache tier.	#L43-L78
SSDCacheConfig.__post_init__	method	SSDCacheConfig.__post_init__() -> None	Method SSDCacheConfig.__post_init__ calls ValueError; can raise ValueError.	#L65-L73
SSDCacheConfig.max_size_bytes	method	SSDCacheConfig.max_size_bytes() -> int	Maximum cache size in bytes.	#L76-L78
SSDCacheStats	class	SSDCacheStats(spill_count: int = 0, spill_bytes: int = 0, ssd_hits: int = 0, ssd_misses: int = 0, reload_latency_sum: float = 0.0, reload_bytes: int = 0, promotion_failures: int = 0)	Statistics for SSD cache tier — exposed from day one.	#L82-L123
SSDCacheStats.to_dict	method	SSDCacheStats.to_dict() -> dict	Return spill, lookup, reload, and promotion statistics.	#L103-L123
_tokens_to_blob	function	_tokens_to_blob(tokens: tuple[int, ...]) -> bytes	Serialize token tuple to a compact binary blob for SQLite storage.	#L126-L132
_blob_to_tokens	function	_blob_to_tokens(blob: bytes) -> tuple[int, ...]	Deserialize binary blob back to token tuple.	#L135-L139
_tokens_hash	function	_tokens_hash(tokens: tuple[int, ...]) -> str	Compute SHA-256 hex digest of a token sequence for use as primary key.	#L142-L144
_prefix_hash	function	_prefix_hash(tokens: tuple[int, ...]) -> str	Hash the bounded token prefix used to prefilter prefix lookups.	#L147-L149
SSDIndex	class	SSDIndex(cache_dir: str)	SQLite-backed index for SSD cache entries.	#L152-L405
SSDIndex.__init__	method	SSDIndex.__init__(cache_dir: str) -> None	Method SSDIndex.__init__ updates self._cache_dir, self._db_lock, self._conn, self._conn.row_factory; calls threading.Lock, os.path.join, sqlite3.connect, self._conn.execute.	#L165-L173
SSDIndex._create_tables	method	SSDIndex._create_tables() -> None	Method SSDIndex._create_tables calls self._conn.executescript, self._ensure_column, self._conn.execute, cur.fetchone.	#L175-L213
SSDIndex._ensure_column	method	SSDIndex._ensure_column(table: str, column: str, definition: str) -> None	Method SSDIndex._ensure_column calls self._conn.execute, cur.fetchall.	#L215-L218
SSDIndex._backfill_prefix_hashes	method	SSDIndex._backfill_prefix_hashes() -> None	Method SSDIndex._backfill_prefix_hashes calls self._conn.execute, cur.fetchall, _blob_to_tokens, _prefix_hash.	#L220-L230
SSDIndex.insert_entry	method	SSDIndex.insert_entry(tokens_key: tuple[int, ...], file_path: str, memory_bytes: int, num_tokens: int) -> None	Insert or replace a cache entry in the index.	#L232-L263
SSDIndex.lookup_exact	method	SSDIndex.lookup_exact(tokens_key: tuple[int, ...]) -> dict \| None	Look up an exact token sequence.	#L265-L280
SSDIndex.lookup_prefix	method	SSDIndex.lookup_prefix(query_tokens: tuple[int, ...]) -> list[dict]	Find entries whose token sequence is a prefix of query_tokens.	#L282-L324
SSDIndex.delete_entry	method	SSDIndex.delete_entry(tokens_key: tuple[int, ...]) -> None	Delete an entry by token sequence.	#L326-L333
SSDIndex.get_lru	method	SSDIndex.get_lru(limit: int = 10) -> list[dict]	Get the least recently used entries, ordered oldest first.	#L335-L355
SSDIndex.get_total_bytes	method	SSDIndex.get_total_bytes() -> int	Get total memory_bytes across all entries.	#L357-L363
SSDIndex.get_entry_count	method	SSDIndex.get_entry_count() -> int	Get number of entries in the index.	#L365-L369
SSDIndex.touch	method	SSDIndex.touch(tokens_key: tuple[int, ...]) -> None	Update accessed_at timestamp for an entry (marks as recently used).	#L371-L379
SSDIndex.all_entries	method	SSDIndex.all_entries() -> list[dict]	Return all entries (for startup reconciliation).	#L381-L400
SSDIndex.close	method	SSDIndex.close() -> None	Close the SQLite connection.	#L402-L405
LayerSerializer	class	LayerSerializer()	Interface for per-layer cache serialization.	#L419-L446
LayerSerializer.snapshot_layer	method	LayerSerializer.snapshot_layer(layer: Any) -> dict[str, Any]	Producer-thread CPU snapshot of an MLX-backed cache layer.	#L429-L431
LayerSerializer.serialize_layer	method	LayerSerializer.serialize_layer(snapshot: dict[str, Any], layer_idx: int, file_path: str) -> dict[str, Any]	Writer-thread: persist a snapshot to safetensors at file_path.	#L434-L441
LayerSerializer.deserialize_layer	method	LayerSerializer.deserialize_layer(file_path: str, metadata: dict[str, Any]) -> dict	Read a layer back from disk.	#L444-L446
_mx_to_numpy_safe	function	_mx_to_numpy_safe(arr: Any) -> tuple[np.ndarray, str \| None]	mx.array → np.ndarray, upcasting numpy-unsupported dtypes (bf16) to fp32.	#L449-L467
KVCacheSerializer	class	KVCacheSerializer()	Serializer for KVCache and RotatingKVCache layers.	#L470-L564
KVCacheSerializer.snapshot_layer	method	KVCacheSerializer.snapshot_layer(layer: Any) -> dict[str, Any]	Copy a KV cache layer into NumPy-backed writer-thread data.	#L481-L515
KVCacheSerializer.serialize_layer	method	KVCacheSerializer.serialize_layer(snapshot: dict[str, Any], layer_idx: int, file_path: str) -> dict[str, Any]	Write one KV layer to safetensors and return reconstruction metadata.	#L517-L542
KVCacheSerializer.deserialize_layer	method	KVCacheSerializer.deserialize_layer(file_path: str, metadata: dict[str, Any]) -> dict	Load one KV layer as arrays plus cache reconstruction metadata.	#L544-L564
ArraysCacheSerializer	class	ArraysCacheSerializer()	Serializer for ArraysCache (Mamba/linear attention) layers.	#L567-L627
ArraysCacheSerializer.snapshot_layer	method	ArraysCacheSerializer.snapshot_layer(layer: Any) -> dict[str, Any]	Copy an arrays-cache state into NumPy-backed writer-thread data.	#L573-L587
ArraysCacheSerializer.serialize_layer	method	ArraysCacheSerializer.serialize_layer(snapshot: dict[str, Any], layer_idx: int, file_path: str) -> dict[str, Any]	Write arrays-cache state to safetensors and return its metadata.	#L589-L610
ArraysCacheSerializer.deserialize_layer	method	ArraysCacheSerializer.deserialize_layer(file_path: str, metadata: dict[str, Any]) -> dict	Load arrays-cache state and any original dtype hints.	#L612-L627
get_serializer_for_layer	function	get_serializer_for_layer(layer: Any) -> LayerSerializer	Return the appropriate serializer for a cache layer.	#L630-L646
SSDCacheTier	class	SSDCacheTier(config: SSDCacheConfig)	Cold-tier disk cache for KV cache entries.	#L649-L1248
SSDCacheTier.__init__	method	SSDCacheTier.__init__(config: SSDCacheConfig) -> None	Method SSDCacheTier.__init__ updates self._config, self._closed, self._writer_thread, self._cache_dir; calls ValueError, os.path.join, os.makedirs, SSDIndex; can raise ValueError.	#L667-L705
SSDCacheTier._entry_hash	method	SSDCacheTier._entry_hash(tokens: tuple[int, ...]) -> str	Compute deterministic hash for a token sequence.	#L708-L710
SSDCacheTier.get_stats	method	SSDCacheTier.get_stats() -> dict	Return current SSD cache statistics.	#L712-L714
SSDCacheTier.start_writer	method	SSDCacheTier.start_writer() -> None	Start the background spill writer thread.	#L716-L725
SSDCacheTier._writer_loop	method	SSDCacheTier._writer_loop() -> None	Drain spill queue and persist entries.	#L727-L744
SSDCacheTier.enqueue_spill	method	SSDCacheTier.enqueue_spill(tokens: tuple[int, ...], cache: list[Any], memory_bytes: int) -> bool	Enqueue a cache entry for async spill to SSD.	#L746-L867
SSDCacheTier.enqueue_spill._is_quantized_layer	nested function	SSDCacheTier.enqueue_spill._is_quantized_layer(layer) -> not annotated	Nested Function SSDCacheTier.enqueue_spill._is_quantized_layer calls isinstance, getattr; has 2 explicit return paths.	#L772-L776
SSDCacheTier._write_entry	method	SSDCacheTier._write_entry(tokens_key: tuple[int, ...], layer_snapshots: list[tuple[LayerSerializer, dict[str, Any]]], memory_bytes: int) -> None	Atomically persist one entry (writer thread; numpy-only input).	#L869-L944
SSDCacheTier.lookup_ssd	method	SSDCacheTier.lookup_ssd(tokens: tuple[int, ...]) -> dict \| None	Synchronous check whether tokens exist in SSD tier.	#L946-L958
SSDCacheTier.lookup_ssd_prefix	method	SSDCacheTier.lookup_ssd_prefix(tokens: tuple[int, ...]) -> dict \| None	Find the longest prefix match in the SSD tier.	#L960-L968
SSDCacheTier.async_promote	method	async SSDCacheTier.async_promote(tokens: tuple[int, ...], reserve_budget_fn, release_budget_fn) -> list \| None	Promote an entry from SSD to RAM asynchronously.	#L970-L1076
SSDCacheTier._read_entry	method	SSDCacheTier._read_entry(tokens: tuple[int, ...], relative_path: str) -> list \| None	Read a cache entry from disk.	#L1078-L1121
SSDCacheTier._quarantine_entry	method	SSDCacheTier._quarantine_entry(tokens: tuple[int, ...], relative_path: str) -> None	Move a corrupt entry to quarantine and remove from index.	#L1123-L1142
SSDCacheTier._enforce_capacity	method	SSDCacheTier._enforce_capacity() -> None	Evict oldest SSD entries until within capacity limits.	#L1144-L1181
SSDCacheTier.reconcile	method	SSDCacheTier.reconcile() -> int	Reconcile index with files on disk.	#L1183-L1229
SSDCacheTier.close	method	SSDCacheTier.close() -> None	Close the SSD cache tier and release resources.	#L1231-L1248