vllm_mlx.memory_cache

Memory-aware prefix cache for vllm-mlx.

View the complete module source at #L1-L1463.

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.memory_cache

Memory-aware prefix cache for vllm-mlx.

This module provides a prefix cache implementation that tracks memory usage and evicts entries based on memory pressure rather than entry count.

Key features: - Automatic memory limit detection based on available system RAM - Accurate memory tracking for MLX array caches - LRU eviction triggered by memory thresholds - No unnecessary deep copies (MLX arrays are immutable)

Example

config = MemoryCacheConfig(max_memory_percent=0.25) cache = MemoryAwarePrefixCache(model, config)

Fetch returns reference (no copy) - safe because MLX arrays are immutable

kv_cache, remaining = cache.fetch(tokens)

Store tracks memory automatically

cache.store(tokens, kv_cache)

vllm_mlx.memory_cache.logger module-attribute

logger = logging.getLogger(__name__)

vllm_mlx.memory_cache._BYTES_PER_MB module-attribute

_BYTES_PER_MB = 1024 * 1024

vllm_mlx.memory_cache._DEFAULT_MEMORY_PERCENT module-attribute

_DEFAULT_MEMORY_PERCENT = 0.2

vllm_mlx.memory_cache._MIN_MEMORY_BYTES module-attribute

_MIN_MEMORY_BYTES = 100 * _BYTES_PER_MB

vllm_mlx.memory_cache._MAX_ENTRIES_FALLBACK module-attribute

_MAX_ENTRIES_FALLBACK = 50

vllm_mlx.memory_cache._CACHE_PERSIST_VERSION module-attribute

_CACHE_PERSIST_VERSION = 3

vllm_mlx.memory_cache.MemoryCacheConfig dataclass

MemoryCacheConfig(max_memory_mb: int | None = None, max_memory_percent: float = _DEFAULT_MEMORY_PERCENT, max_entries: int = 1000, enable_memory_tracking: bool = True, kv_quantize: bool = False, kv_bits: int = 8, kv_group_size: int = 64, kv_min_quantize_tokens: int = 256, min_prefix_tokens: int = 128)

Configuration for memory-aware prefix cache.

Attributes:

• max_memory_mb (int | None) –

  Maximum memory in MB. If None, auto-detects.

• max_memory_percent (float) –

  Fraction of available RAM to use (0.0-1.0).

• max_entries (int) –

  Hard limit on number of entries (safety net).

• enable_memory_tracking (bool) –

  Whether to track per-entry memory.

• kv_quantize (bool) –

  Whether to quantize KV cache layers for reduced memory.

• kv_bits (int) –

  Number of bits for KV cache quantization.

• kv_group_size (int) –

  Group size for KV cache quantization.

• kv_min_quantize_tokens (int) –

  Minimum sequence length for quantization to apply.

• min_prefix_tokens (int) –

  Minimum cached prefix length eligible for reuse.

vllm_mlx.memory_cache.MemoryCacheConfig.max_memory_mb class-attribute instance-attribute

max_memory_mb: int | None = None

vllm_mlx.memory_cache.MemoryCacheConfig.max_memory_percent class-attribute instance-attribute

max_memory_percent: float = _DEFAULT_MEMORY_PERCENT

vllm_mlx.memory_cache.MemoryCacheConfig.max_entries class-attribute instance-attribute

max_entries: int = 1000

vllm_mlx.memory_cache.MemoryCacheConfig.enable_memory_tracking class-attribute instance-attribute

enable_memory_tracking: bool = True

vllm_mlx.memory_cache.MemoryCacheConfig.kv_quantize class-attribute instance-attribute

kv_quantize: bool = False

vllm_mlx.memory_cache.MemoryCacheConfig.kv_bits class-attribute instance-attribute

kv_bits: int = 8

vllm_mlx.memory_cache.MemoryCacheConfig.kv_group_size class-attribute instance-attribute

kv_group_size: int = 64

vllm_mlx.memory_cache.MemoryCacheConfig.kv_min_quantize_tokens class-attribute instance-attribute

kv_min_quantize_tokens: int = 256

vllm_mlx.memory_cache.MemoryCacheConfig.min_prefix_tokens class-attribute instance-attribute

min_prefix_tokens: int = 128

vllm_mlx.memory_cache.MemoryCacheConfig.__post_init__

__post_init__() -> None

Source code in vllm_mlx/memory_cache.py

def __post_init__(self) -> None:
    if not 0.0 < self.max_memory_percent <= 1.0:
        raise ValueError(
            f"max_memory_percent must be in (0, 1], got {self.max_memory_percent}"
        )
    if self.max_entries < 1:
        raise ValueError(f"max_entries must be >= 1, got {self.max_entries}")
    if self.kv_min_quantize_tokens < 0:
        raise ValueError(
            f"kv_min_quantize_tokens must be >= 0, got {self.kv_min_quantize_tokens}"
        )
    if self.min_prefix_tokens < 1:
        raise ValueError(
            f"min_prefix_tokens must be >= 1, got {self.min_prefix_tokens}"
        )

vllm_mlx.memory_cache.MemoryCacheConfig.compute_memory_limit

compute_memory_limit() -> int

Compute the memory limit in bytes.

Returns:

• int –

  Memory limit in bytes.

Source code in vllm_mlx/memory_cache.py

def compute_memory_limit(self) -> int:
    """
    Compute the memory limit in bytes.

    Returns:
        Memory limit in bytes.
    """
    if self.max_memory_mb is not None:
        return self.max_memory_mb * _BYTES_PER_MB

    available = _get_available_memory()
    if available > 0:
        limit = int(available * self.max_memory_percent)
        return max(limit, _MIN_MEMORY_BYTES)

    # Fallback: assume 8GB system, use configured percent
    fallback_total = 8 * 1024 * _BYTES_PER_MB
    return int(fallback_total * self.max_memory_percent)

vllm_mlx.memory_cache.CacheStats dataclass

CacheStats(hits: int = 0, misses: int = 0, evictions: int = 0, tokens_saved: int = 0, current_memory_bytes: int = 0, max_memory_bytes: int = 0, entry_count: int = 0)

Statistics for cache performance monitoring.

vllm_mlx.memory_cache.CacheStats.hits class-attribute instance-attribute

hits: int = 0

vllm_mlx.memory_cache.CacheStats.misses class-attribute instance-attribute

misses: int = 0

vllm_mlx.memory_cache.CacheStats.evictions class-attribute instance-attribute

evictions: int = 0

vllm_mlx.memory_cache.CacheStats.tokens_saved class-attribute instance-attribute

tokens_saved: int = 0

vllm_mlx.memory_cache.CacheStats.current_memory_bytes class-attribute instance-attribute

current_memory_bytes: int = 0

vllm_mlx.memory_cache.CacheStats.max_memory_bytes class-attribute instance-attribute

max_memory_bytes: int = 0

vllm_mlx.memory_cache.CacheStats.entry_count class-attribute instance-attribute

entry_count: int = 0

vllm_mlx.memory_cache.CacheStats.hit_rate property

hit_rate: float

Return successful lookups divided by all completed lookups.

vllm_mlx.memory_cache.CacheStats.memory_utilization property

memory_utilization: float

Return the fraction of the configured memory budget in use.

vllm_mlx.memory_cache.CacheStats.to_dict

to_dict() -> dict[str, Any]

Return rounded cache counters and memory values for APIs and logs.

Source code in vllm_mlx/memory_cache.py

def to_dict(self) -> dict[str, Any]:
    """Return rounded cache counters and memory values for APIs and logs."""

    return {
        "hits": self.hits,
        "misses": self.misses,
        "hit_rate": round(self.hit_rate, 4),
        "evictions": self.evictions,
        "tokens_saved": self.tokens_saved,
        "current_memory_mb": round(self.current_memory_bytes / _BYTES_PER_MB, 2),
        "max_memory_mb": round(self.max_memory_bytes / _BYTES_PER_MB, 2),
        "memory_utilization": round(self.memory_utilization, 4),
        "entry_count": self.entry_count,
    }

vllm_mlx.memory_cache._CacheEntry dataclass

_CacheEntry(tokens: tuple[int, ...], cache: list[Any], memory_bytes: int)

Internal cache entry with memory tracking.

vllm_mlx.memory_cache._CacheEntry.tokens instance-attribute

tokens: tuple[int, ...]

vllm_mlx.memory_cache._CacheEntry.cache instance-attribute

cache: list[Any]

vllm_mlx.memory_cache._CacheEntry.memory_bytes instance-attribute

memory_bytes: int

vllm_mlx.memory_cache._CacheEntry.create classmethod

create(tokens: list[int], cache: list[Any]) -> _CacheEntry

Create a cache entry with memory estimation.

Source code in vllm_mlx/memory_cache.py

@classmethod
def create(cls, tokens: list[int], cache: list[Any]) -> _CacheEntry:
    """Create a cache entry with memory estimation."""
    memory = estimate_kv_cache_memory(cache)
    return cls(
        tokens=tuple(tokens),
        cache=cache,
        memory_bytes=memory,
    )

vllm_mlx.memory_cache._QuantizedCacheWrapper

_QuantizedCacheWrapper(layer: Any, bits: int, group_size: int)

Lightweight wrapper storing quantized KV arrays + original cache metadata.

Unlike QuantizedKVCache, this preserves enough info to reconstruct the original cache type (KVCache, RotatingKVCache, etc.) on dequantize.

Source code in vllm_mlx/memory_cache.py

def __init__(self, layer: Any, bits: int, group_size: int):
    import mlx.core as mx

    self.keys = mx.quantize(layer.keys, group_size=group_size, bits=bits)
    self.values = mx.quantize(layer.values, group_size=group_size, bits=bits)
    self.offset = layer.offset
    self.bits = bits
    self.group_size = group_size
    self.orig_type = type(layer)
    # Preserve RotatingKVCache-specific attrs
    self.orig_attrs = {}
    for attr in ("max_size", "keep", "step", "_idx"):
        if hasattr(layer, attr):
            self.orig_attrs[attr] = getattr(layer, attr)

vllm_mlx.memory_cache._QuantizedCacheWrapper.__slots__ class-attribute instance-attribute

__slots__ = ('keys', 'values', 'offset', 'bits', 'group_size', 'orig_type', 'orig_attrs')

vllm_mlx.memory_cache._QuantizedCacheWrapper.keys instance-attribute

keys = mx.quantize(layer.keys, group_size=group_size, bits=bits)

vllm_mlx.memory_cache._QuantizedCacheWrapper.values instance-attribute

values = mx.quantize(layer.values, group_size=group_size, bits=bits)

vllm_mlx.memory_cache._QuantizedCacheWrapper.offset instance-attribute

offset = layer.offset

vllm_mlx.memory_cache._QuantizedCacheWrapper.bits instance-attribute

bits = bits

vllm_mlx.memory_cache._QuantizedCacheWrapper.group_size instance-attribute

group_size = group_size

vllm_mlx.memory_cache._QuantizedCacheWrapper.orig_type instance-attribute

orig_type = type(layer)

vllm_mlx.memory_cache._QuantizedCacheWrapper.orig_attrs instance-attribute

orig_attrs = {}

vllm_mlx.memory_cache.MemoryAwarePrefixCache

MemoryAwarePrefixCache(model: Any, config: MemoryCacheConfig | None = None)

Prefix cache with memory-based eviction.

This cache tracks memory usage per entry and evicts based on memory pressure rather than entry count. It uses LRU (Least Recently Used) ordering for eviction decisions.

Key design decisions: - No deep copies on fetch: MLX arrays are immutable, so sharing is safe - Memory tracking per entry: Accurate accounting for eviction - Auto-detection of available RAM: Adapts to different systems - OrderedDict for O(1) LRU operations

Thread Safety

This class is NOT thread-safe. Use external locking if needed.

Initialize the memory-aware prefix cache.

Parameters:

• model (Any) –

  The MLX model (used for identification).

• config (MemoryCacheConfig | None, default: None ) –

  Cache configuration. Uses defaults if None.

Source code in vllm_mlx/memory_cache.py

def __init__(
    self,
    model: Any,
    config: MemoryCacheConfig | None = None,
) -> None:
    """
    Initialize the memory-aware prefix cache.

    Args:
        model: The MLX model (used for identification).
        config: Cache configuration. Uses defaults if None.
    """
    self._model_id = id(model)
    self._config = config or MemoryCacheConfig()
    self._model_fingerprint = _compute_model_fingerprint(model)

    # OrderedDict maintains insertion order for LRU
    # Key: tuple(tokens), Value: _CacheEntry
    self._entries: OrderedDict[tuple[int, ...], _CacheEntry] = OrderedDict()

    # Sorted index of token keys for efficient prefix/supersequence lookup.
    # Tuple lexicographic ordering means a prefix key P is always < any
    # extension of P, so bisect gives O(log N) range scans instead of O(N).
    self._sorted_keys: list[tuple[int, ...]] = []

    # Memory tracking
    self._max_memory = self._config.compute_memory_limit()
    self._current_memory = 0
    self._memory_lock = threading.RLock()

    # Statistics
    self._stats = CacheStats(max_memory_bytes=self._max_memory)

    # Track the match type from the last fetch() call
    self._last_match_type: str | None = None

    # Optional SSD cold tier (set via set_ssd_tier())
    self._ssd_tier = None

    logger.info(
        f"MemoryAwarePrefixCache initialized: "
        f"max_memory={self._max_memory / _BYTES_PER_MB:.1f}MB, "
        f"max_entries={self._config.max_entries}"
    )

vllm_mlx.memory_cache.MemoryAwarePrefixCache._model_id instance-attribute

_model_id = id(model)

vllm_mlx.memory_cache.MemoryAwarePrefixCache._config instance-attribute

_config = config or MemoryCacheConfig()

vllm_mlx.memory_cache.MemoryAwarePrefixCache._model_fingerprint instance-attribute

_model_fingerprint = _compute_model_fingerprint(model)

vllm_mlx.memory_cache.MemoryAwarePrefixCache._entries instance-attribute

_entries: OrderedDict[tuple[int, ...], _CacheEntry] = OrderedDict()

vllm_mlx.memory_cache.MemoryAwarePrefixCache._sorted_keys instance-attribute

_sorted_keys: list[tuple[int, ...]] = []

vllm_mlx.memory_cache.MemoryAwarePrefixCache._max_memory instance-attribute

_max_memory = self._config.compute_memory_limit()

vllm_mlx.memory_cache.MemoryAwarePrefixCache._current_memory instance-attribute

_current_memory = 0

vllm_mlx.memory_cache.MemoryAwarePrefixCache._memory_lock instance-attribute

_memory_lock = threading.RLock()

vllm_mlx.memory_cache.MemoryAwarePrefixCache._stats instance-attribute

_stats = CacheStats(max_memory_bytes=self._max_memory)

vllm_mlx.memory_cache.MemoryAwarePrefixCache._last_match_type instance-attribute

_last_match_type: str | None = None

vllm_mlx.memory_cache.MemoryAwarePrefixCache._ssd_tier instance-attribute

_ssd_tier = None

vllm_mlx.memory_cache.MemoryAwarePrefixCache.memory_usage_mb property

memory_usage_mb: float

Current memory usage in MB.

vllm_mlx.memory_cache.MemoryAwarePrefixCache.memory_limit_mb property

memory_limit_mb: float

Memory limit in MB.

vllm_mlx.memory_cache.MemoryAwarePrefixCache.fetch

fetch(tokens: list[int]) -> tuple[list[Any] | None, list[int]]

Find cached KV state for the given tokens.

This method searches for exact matches, prefix matches, supersequence matches, and longest-common-prefix (LCP) matches. Uses a sorted key index for O(log N) lookup instead of scanning all entries.

Returns the cached KV state directly (no copy) since MLX arrays are immutable and safe to share.

Parameters:

• tokens (list[int]) –

  Input token sequence.

Returns:

• list[Any] | None –

  Tuple of (cache, remaining_tokens):

• list[int] –
  • cache: Cached KV state if found, None otherwise
• tuple[list[Any] | None, list[int]] –
  • remaining_tokens: Tokens that still need processing

Source code in vllm_mlx/memory_cache.py

def fetch(self, tokens: list[int]) -> tuple[list[Any] | None, list[int]]:
    """
    Find cached KV state for the given tokens.

    This method searches for exact matches, prefix matches, supersequence
    matches, and longest-common-prefix (LCP) matches.  Uses a sorted key
    index for O(log N) lookup instead of scanning all entries.

    Returns the cached KV state directly (no copy) since MLX arrays
    are immutable and safe to share.

    Args:
        tokens: Input token sequence.

    Returns:
        Tuple of (cache, remaining_tokens):
        - cache: Cached KV state if found, None otherwise
        - remaining_tokens: Tokens that still need processing
    """
    if not tokens:
        self._stats.misses += 1
        self._last_match_type = "miss"
        return None, tokens
    if len(tokens) < self._config.min_prefix_tokens:
        self._stats.misses += 1
        self._last_match_type = "miss_short_prefix"
        return None, tokens

    tokens_key = tuple(tokens)

    # --- O(1) exact match ---
    if tokens_key in self._entries:
        entry = self._entries[tokens_key]
        self._entries.move_to_end(tokens_key)
        self._stats.hits += 1
        self._stats.tokens_saved += len(tokens)
        self._last_match_type = "exact"
        cache_out = (
            _dequantize_cache(entry.cache)
            if self._config.kv_quantize
            else entry.cache
        )
        return cache_out, []

    # --- O(log N) prefix & supersequence match via sorted index ---
    best_match: _CacheEntry | None = None
    best_length = 0
    best_super: _CacheEntry | None = None

    sorted_keys = self._sorted_keys
    if sorted_keys:
        # Find insertion point for tokens_key in the sorted list.
        # Keys that are prefixes of tokens_key or supersequences will be
        # clustered around this position due to lexicographic ordering.
        idx = bisect.bisect_left(sorted_keys, tokens_key)

        # Scan backwards from idx to find cached keys that are PREFIXES
        # of tokens_key (shorter cached sequences).  A prefix P of T
        # satisfies P <= T lexicographically, so P is at idx-1 or earlier.
        for i in range(idx - 1, -1, -1):
            cached_key = sorted_keys[i]
            cached_len = len(cached_key)
            if cached_len >= len(tokens_key):
                continue  # Not a prefix (same length or longer)
            # Check if cached_key is a prefix of tokens_key
            if tokens_key[:cached_len] == cached_key:
                if cached_len > best_length:
                    best_match = self._entries[cached_key]
                    best_length = cached_len
                # Found best prefix — shorter entries can't be longer
                break
            # Once we go past the prefix range, stop
            if cached_key[0] != tokens_key[0]:
                break

        # Scan forward from idx to find cached keys that are SUPERSEQUENCES
        # of tokens_key (longer cached sequences starting with tokens_key).
        for i in range(idx, len(sorted_keys)):
            cached_key = sorted_keys[i]
            cached_len = len(cached_key)
            if cached_len < len(tokens_key):
                continue
            # Check if tokens_key is a prefix of cached_key
            if cached_key[: len(tokens_key)] == tokens_key:
                if best_super is None or cached_len > len(best_super.tokens):
                    best_super = self._entries[cached_key]
            else:
                # Past the supersequence range
                break

    # --- Supersequence match handling ---
    if best_super is not None:
        n_cached = len(best_super.tokens)
        n_requested = len(tokens)
        excess = n_cached - n_requested

        has_non_trimmable = any(
            not _is_cache_layer_trimmable(lc) for lc in best_super.cache
        )

        if excess > 0 and has_non_trimmable:
            logger.debug(
                "[cache_fetch] supersequence match skipped: "
                "non-trimmable cache layers (hybrid model)"
            )
        elif excess > 0:
            trimmed_cache = _trim_cache_offset(best_super.cache, excess)
            self._entries.move_to_end(best_super.tokens)
            self._stats.hits += 1
            self._stats.tokens_saved += n_requested
            self._last_match_type = "supersequence"
            trimmed_cache = (
                _dequantize_cache(trimmed_cache)
                if self._config.kv_quantize
                else trimmed_cache
            )
            return trimmed_cache, []
        else:
            self._entries.move_to_end(best_super.tokens)
            self._stats.hits += 1
            self._stats.tokens_saved += n_requested
            self._last_match_type = "supersequence"
            cache_out = (
                _dequantize_cache(best_super.cache)
                if self._config.kv_quantize
                else best_super.cache
            )
            return cache_out, []

    # --- Prefix match ---
    if best_match is not None:
        self._entries.move_to_end(best_match.tokens)
        self._stats.hits += 1
        self._stats.tokens_saved += best_length
        remaining = tokens[best_length:]
        self._last_match_type = "prefix"
        cache_out = (
            _dequantize_cache(best_match.cache)
            if self._config.kv_quantize
            else best_match.cache
        )
        return cache_out, remaining

    # --- LCP (Longest Common Prefix) for divergent sequences ---
    # This handles the agentic pattern: same system+context prefix
    # but different final user message.  Use the sorted index to find
    # the nearest neighbor which likely shares the longest prefix.
    best_lcp_entry: _CacheEntry | None = None
    best_lcp_length = 0

    if sorted_keys:
        idx = bisect.bisect_left(sorted_keys, tokens_key)
        # Check neighbors around insertion point (they share the most
        # common prefix due to lexicographic ordering).
        for i in (idx - 1, idx):
            if i < 0 or i >= len(sorted_keys):
                continue
            cached_key = sorted_keys[i]
            if cached_key == tokens_key:
                continue  # Skip exact (already handled)
            min_len = min(len(cached_key), len(tokens_key))
            if min_len <= best_lcp_length:
                continue
            # Compute LCP length
            lcp = 0
            for j in range(min_len):
                if cached_key[j] != tokens_key[j]:
                    break
                lcp = j + 1
            if lcp > best_lcp_length:
                best_lcp_entry = self._entries[cached_key]
                best_lcp_length = lcp
                logger.debug(
                    f"[cache_fetch] LCP scan: cached_len={len(cached_key)} "
                    f"req_len={len(tokens_key)} lcp={lcp}"
                )

    if best_lcp_entry is not None and best_lcp_length > 0:
        if best_lcp_length < self._config.min_prefix_tokens:
            logger.debug(
                "[cache_fetch] LCP skipped: shared=%s below min_prefix_tokens=%s",
                best_lcp_length,
                self._config.min_prefix_tokens,
            )
            self._stats.misses += 1
            self._last_match_type = "miss_short_lcp"
            return None, tokens
        excess = len(best_lcp_entry.tokens) - best_lcp_length

        has_non_trimmable = any(
            not _is_cache_layer_trimmable(lc) for lc in best_lcp_entry.cache
        )
        logger.debug(
            f"[cache_fetch] LCP candidate: lcp={best_lcp_length} "
            f"entry_len={len(best_lcp_entry.tokens)} excess={excess} "
            f"non_trimmable={has_non_trimmable} "
            f"cache_layers={len(best_lcp_entry.cache)} "
            f"layer_types={[type(lc).__name__ for lc in best_lcp_entry.cache[:3]]}"
        )

        if has_non_trimmable:
            # Hybrid model (SSM+Attention): SSM state can't be rewound.
            # Block LCP for hybrid models — use think-suffix stripping
            # in the engine layer to get clean PREFIX matches instead.
            logger.debug(
                "[cache_fetch] LCP skipped: non-trimmable cache layers "
                "(hybrid model, SSM state can't be rewound)"
            )
        else:
            trimmed_cache = _trim_cache_offset(best_lcp_entry.cache, excess)
            self._entries.move_to_end(best_lcp_entry.tokens)
            self._stats.hits += 1
            self._stats.tokens_saved += best_lcp_length
            remaining = tokens[best_lcp_length:]
            logger.debug(
                f"[cache_fetch] LCP hit: shared={best_lcp_length} "
                f"trimmed={excess} remaining={len(remaining)}"
            )
            self._last_match_type = "lcp"
            trimmed_cache = (
                _dequantize_cache(trimmed_cache)
                if self._config.kv_quantize
                else trimmed_cache
            )
            return trimmed_cache, remaining

    self._stats.misses += 1
    self._last_match_type = "miss"

    return None, tokens

vllm_mlx.memory_cache.MemoryAwarePrefixCache.store

store(tokens: list[int], cache: list[Any], evict_prefixes: bool = True) -> bool

Store KV cache for future reuse.

This method stores the cache reference directly (no copy) and tracks memory usage. If memory limit is exceeded, LRU entries are evicted until there's room.

Parameters:

• tokens (list[int]) –

  Token sequence that was processed.

• cache (list[Any]) –

  The computed KV cache to store.

• evict_prefixes (bool, default: True ) –

  If True, evict existing entries whose token sequence is a strict prefix of tokens. Set to False when storing prompt+output entries to preserve prompt-only entries created by prompt_cache_save (those are the entries that future requests will actually match).

Returns:

• bool –

  True if stored successfully, False if rejected.

Source code in vllm_mlx/memory_cache.py

def store(
    self, tokens: list[int], cache: list[Any], evict_prefixes: bool = True
) -> bool:
    """
    Store KV cache for future reuse.

    This method stores the cache reference directly (no copy) and
    tracks memory usage. If memory limit is exceeded, LRU entries
    are evicted until there's room.

    Args:
        tokens: Token sequence that was processed.
        cache: The computed KV cache to store.
        evict_prefixes: If True, evict existing entries whose token
            sequence is a strict prefix of ``tokens``.  Set to False
            when storing prompt+output entries to preserve prompt-only
            entries created by prompt_cache_save (those are the entries
            that future requests will actually match).

    Returns:
        True if stored successfully, False if rejected.
    """
    if not tokens or not cache:
        return False
    if len(tokens) < self._config.min_prefix_tokens:
        logger.debug(
            "[cache_store] skipped short prefix: tokens=%s min_prefix_tokens=%s",
            len(tokens),
            self._config.min_prefix_tokens,
        )
        return False

    with self._memory_lock:
        tokens_key = tuple(tokens)

        # If already cached, just update LRU order (skip expensive trim/quantize)
        if tokens_key in self._entries:
            self._entries.move_to_end(tokens_key)
            return True

        # Trim oversized KV arrays to actual used size
        cache = _trim_to_offset(cache)

        # Quantize if enabled and sequence is long enough
        if (
            self._config.kv_quantize
            and len(tokens) >= self._config.kv_min_quantize_tokens
        ):
            cache = _quantize_cache(
                cache, self._config.kv_bits, self._config.kv_group_size
            )

        # Create entry and estimate memory
        entry = _CacheEntry.create(tokens, cache)

        # Check if single entry exceeds limit
        if entry.memory_bytes > self._max_memory:
            logger.warning(
                f"Cache entry too large: {entry.memory_bytes / _BYTES_PER_MB:.1f}MB "
                f"exceeds limit {self._max_memory / _BYTES_PER_MB:.1f}MB"
            )
            return False

        # Prefix-subset eviction: remove entries whose token sequence
        # is a strict prefix of the new entry.  Uses sorted index for
        # O(log N + K) lookup instead of O(N) scan.
        if evict_prefixes and self._sorted_keys:
            to_remove = []
            idx = bisect.bisect_left(self._sorted_keys, tokens_key)
            # Scan backwards — prefixes of tokens_key are immediately before idx
            for i in range(idx - 1, -1, -1):
                key = self._sorted_keys[i]
                klen = len(key)
                if klen >= len(tokens_key):
                    continue
                if tokens_key[:klen] == key:
                    to_remove.append(key)
                elif key[0] != tokens_key[0]:
                    break
            for key in to_remove:
                old = self._entries.pop(key)
                self._current_memory -= old.memory_bytes
                self._stats.evictions += 1
                self._remove_from_sorted(key)
                logger.debug(
                    f"[prefix_evict] removed {len(key)} tokens, "
                    f"freed {old.memory_bytes / _BYTES_PER_MB:.2f}MB, "
                    f"new_entry={len(tokens_key)} tokens"
                )
            if to_remove:
                self._stats.entry_count = len(self._entries)
                self._stats.current_memory_bytes = self._current_memory

        # Evict until we have room
        while (
            self._current_memory + entry.memory_bytes > self._max_memory
            or len(self._entries) >= self._config.max_entries
        ) and self._entries:
            self._evict_lru()

        # Store entry
        self._entries[tokens_key] = entry
        self._current_memory += entry.memory_bytes
        bisect.insort(self._sorted_keys, tokens_key)
        self._stats.entry_count = len(self._entries)
        self._stats.current_memory_bytes = self._current_memory

    logger.debug(
        f"Stored cache: {len(tokens)} tokens, "
        f"{entry.memory_bytes / _BYTES_PER_MB:.2f}MB, "
        f"total={self._current_memory / _BYTES_PER_MB:.1f}MB"
    )

    return True

vllm_mlx.memory_cache.MemoryAwarePrefixCache._remove_from_sorted

_remove_from_sorted(key: tuple[int, ...]) -> None

Remove a key from the sorted index using bisect for O(log N).

Source code in vllm_mlx/memory_cache.py

def _remove_from_sorted(self, key: tuple[int, ...]) -> None:
    """Remove a key from the sorted index using bisect for O(log N)."""
    idx = bisect.bisect_left(self._sorted_keys, key)
    if idx < len(self._sorted_keys) and self._sorted_keys[idx] == key:
        self._sorted_keys.pop(idx)

vllm_mlx.memory_cache.MemoryAwarePrefixCache._evict_lru

_evict_lru() -> None

Evict the least recently used entry.

If an SSD tier is attached, the entry is spilled to disk instead of being discarded.

Source code in vllm_mlx/memory_cache.py

def _evict_lru(self) -> None:
    """Evict the least recently used entry.

    If an SSD tier is attached, the entry is spilled to disk instead
    of being discarded.
    """
    with self._memory_lock:
        if not self._entries:
            return

        # popitem(last=False) removes oldest entry (FIFO order = LRU)
        tokens_key, entry = self._entries.popitem(last=False)
        self._current_memory -= entry.memory_bytes
        self._remove_from_sorted(tokens_key)
        self._stats.evictions += 1
        self._stats.entry_count = len(self._entries)
        self._stats.current_memory_bytes = self._current_memory

    # Spill to SSD tier if available
    if self._ssd_tier is not None:
        self._ssd_tier.enqueue_spill(tokens_key, entry.cache, entry.memory_bytes)

    logger.debug(
        f"[lru_evict] removed {len(tokens_key)} tokens, "
        f"freed {entry.memory_bytes / _BYTES_PER_MB:.2f}MB"
        f"{'  (spilled to SSD)' if self._ssd_tier is not None else ''}"
    )

vllm_mlx.memory_cache.MemoryAwarePrefixCache.remove

remove(tokens: list[int]) -> bool

Remove a specific cache entry.

Parameters:

• tokens (list[int]) –

  Token sequence to remove.

Returns:

• bool –

  True if entry was found and removed.

Source code in vllm_mlx/memory_cache.py

def remove(self, tokens: list[int]) -> bool:
    """
    Remove a specific cache entry.

    Args:
        tokens: Token sequence to remove.

    Returns:
        True if entry was found and removed.
    """
    with self._memory_lock:
        tokens_key = tuple(tokens)
        entry = self._entries.pop(tokens_key, None)
        if entry is not None:
            self._current_memory -= entry.memory_bytes
            self._remove_from_sorted(tokens_key)
            self._stats.entry_count = len(self._entries)
            self._stats.current_memory_bytes = self._current_memory
            return True
        return False

vllm_mlx.memory_cache.MemoryAwarePrefixCache.clear

clear() -> None

Clear all cached entries.

Source code in vllm_mlx/memory_cache.py

def clear(self) -> None:
    """Clear all cached entries."""
    with self._memory_lock:
        self._entries.clear()
        self._sorted_keys.clear()
        self._current_memory = 0
        self._stats = CacheStats(max_memory_bytes=self._max_memory)
    logger.debug("Cache cleared")

vllm_mlx.memory_cache.MemoryAwarePrefixCache.get_stats

get_stats() -> dict[str, Any]

Get cache statistics.

Source code in vllm_mlx/memory_cache.py

def get_stats(self) -> dict[str, Any]:
    """Get cache statistics."""
    return self._stats.to_dict()

vllm_mlx.memory_cache.MemoryAwarePrefixCache.reset_stats

reset_stats() -> None

Reset statistics while preserving cache contents.

Source code in vllm_mlx/memory_cache.py

def reset_stats(self) -> None:
    """Reset statistics while preserving cache contents."""
    with self._memory_lock:
        self._stats = CacheStats(
            max_memory_bytes=self._max_memory,
            current_memory_bytes=self._current_memory,
            entry_count=len(self._entries),
        )

vllm_mlx.memory_cache.MemoryAwarePrefixCache.try_reserve_memory

try_reserve_memory(nbytes: int) -> bool

Tentatively reserve cache memory for an upcoming promotion.

Source code in vllm_mlx/memory_cache.py

def try_reserve_memory(self, nbytes: int) -> bool:
    """Tentatively reserve cache memory for an upcoming promotion."""
    with self._memory_lock:
        if self._current_memory + nbytes > self._max_memory:
            return False
        self._current_memory += nbytes
        self._stats.current_memory_bytes = self._current_memory
        return True

vllm_mlx.memory_cache.MemoryAwarePrefixCache.release_reserved_memory

release_reserved_memory(nbytes: int) -> None

Release memory previously reserved by try_reserve_memory().

Source code in vllm_mlx/memory_cache.py

def release_reserved_memory(self, nbytes: int) -> None:
    """Release memory previously reserved by try_reserve_memory()."""
    with self._memory_lock:
        self._current_memory = max(0, self._current_memory - nbytes)
        self._stats.current_memory_bytes = self._current_memory

vllm_mlx.memory_cache.MemoryAwarePrefixCache.__len__

__len__() -> int

Return number of cached entries.

Source code in vllm_mlx/memory_cache.py

def __len__(self) -> int:
    """Return number of cached entries."""
    return len(self._entries)

vllm_mlx.memory_cache.MemoryAwarePrefixCache.__contains__

__contains__(tokens: list[int]) -> bool

Check if tokens are cached.

Source code in vllm_mlx/memory_cache.py

def __contains__(self, tokens: list[int]) -> bool:
    """Check if tokens are cached."""
    return tuple(tokens) in self._entries

vllm_mlx.memory_cache.MemoryAwarePrefixCache.set_ssd_tier

set_ssd_tier(ssd_tier) -> None

Attach an SSD cache tier for eviction spilling.

When set, evicted entries are spilled to SSD instead of discarded.

Parameters:

• ssd_tier –

  An SSDCacheTier instance (or None to disable).

Source code in vllm_mlx/memory_cache.py

def set_ssd_tier(self, ssd_tier) -> None:
    """Attach an SSD cache tier for eviction spilling.

    When set, evicted entries are spilled to SSD instead of discarded.

    Args:
        ssd_tier: An SSDCacheTier instance (or None to disable).
    """
    self._ssd_tier = ssd_tier
    if ssd_tier is not None:
        logger.info("[memory_cache] SSD tier attached for eviction spilling")

vllm_mlx.memory_cache.MemoryAwarePrefixCache.check_ssd

check_ssd(tokens: list[int]) -> dict | None

Check if tokens have an SSD cache hit (without reading data).

Returns metadata dict with 'match_type' ('exact' or 'prefix') if found in SSD tier, None if not found. For prefix matches, the dict also includes 'matched_tokens' (the count of tokens the SSD entry covers).

This is a fast synchronous call (SQLite lookup only). The actual data read happens via the scheduler handoff.

Source code in vllm_mlx/memory_cache.py

def check_ssd(self, tokens: list[int]) -> dict | None:
    """Check if tokens have an SSD cache hit (without reading data).

    Returns metadata dict with 'match_type' ('exact' or 'prefix') if
    found in SSD tier, None if not found. For prefix matches, the dict
    also includes 'matched_tokens' (the count of tokens the SSD entry
    covers).

    This is a fast synchronous call (SQLite lookup only).
    The actual data read happens via the scheduler handoff.
    """
    if self._ssd_tier is None:
        return None

    tokens_key = tuple(tokens)

    # If already in RAM, no SSD needed
    if tokens_key in self._entries:
        return None

    # Check SSD tier — exact match first, then prefix
    candidate = self._ssd_tier.lookup_ssd(tokens_key)
    if candidate is not None:
        candidate["match_type"] = "exact"
        candidate["matched_tokens"] = len(tokens)
        return candidate

    prefix = self._ssd_tier.lookup_ssd_prefix(tokens_key)
    if prefix is not None:
        prefix["match_type"] = "prefix"
        prefix["matched_tokens"] = prefix["num_tokens"]
        return prefix

    return None

vllm_mlx.memory_cache.MemoryAwarePrefixCache.save_to_disk

save_to_disk(cache_dir: str) -> bool

Save all cache entries to disk using mlx_lm's safetensors format.

Directory layout::

cache_dir/
  index.json          # token keys + metadata per entry
  entry_0.safetensors # KV arrays for entry 0
  entry_1.safetensors
  ...

Returns True if at least one entry was saved.

Source code in vllm_mlx/memory_cache.py

def save_to_disk(self, cache_dir: str) -> bool:
    """Save all cache entries to disk using mlx_lm's safetensors format.

    Directory layout::

        cache_dir/
          index.json          # token keys + metadata per entry
          entry_0.safetensors # KV arrays for entry 0
          entry_1.safetensors
          ...

    Returns True if at least one entry was saved.
    """
    import json
    import os
    import time as _time

    if not self._entries:
        logger.info("[cache_persist] nothing to save (0 entries)")
        return False

    t0 = _time.monotonic()
    os.makedirs(cache_dir, exist_ok=True)

    try:
        from mlx_lm.models.cache import save_prompt_cache
    except ImportError:
        logger.warning("[cache_persist] mlx_lm not available, cannot save")
        return False

    index = {
        "version": _CACHE_PERSIST_VERSION,
        "model_fingerprint": self._model_fingerprint,
        "num_entries": len(self._entries),
        "total_memory_bytes": self._current_memory,
        "entries": [],
    }

    saved = 0
    for i, (tokens_key, entry) in enumerate(self._entries.items()):
        entry_path = os.path.join(cache_dir, f"entry_{i}.safetensors")
        try:
            # Dequantize _QuantizedCacheWrapper layers before saving.
            # save_prompt_cache requires .state and .meta_state which
            # the wrapper does not provide; dequantizing restores the
            # original cache types that do.
            persist_cache = (
                _dequantize_cache(entry.cache)
                if any(isinstance(c, _QuantizedCacheWrapper) for c in entry.cache)
                else entry.cache
            )
            save_prompt_cache(
                entry_path,
                persist_cache,
                metadata={"num_tokens": str(len(tokens_key))},
            )
            # Save tokens separately (can be 100K+ ints → binary is smaller)
            tokens_path = os.path.join(cache_dir, f"entry_{i}_tokens.bin")
            import array as _array

            arr = _array.array("i", tokens_key)  # 32-bit signed ints
            with open(tokens_path, "wb") as f:
                arr.tofile(f)

            index["entries"].append(
                {
                    "index": i,
                    "num_tokens": len(tokens_key),
                    "memory_bytes": entry.memory_bytes,
                }
            )
            saved += 1
            logger.info(
                f"[cache_persist] saved entry {i}: "
                f"{len(tokens_key)} tokens, "
                f"{entry.memory_bytes / _BYTES_PER_MB:.1f}MB KV, "
                f"file={entry_path}"
            )
        except Exception as e:
            logger.warning(f"[cache_persist] failed to save entry {i}: {e}")

    index_path = os.path.join(cache_dir, "index.json")
    with open(index_path, "w") as f:
        json.dump(index, f, indent=2)

    dt = _time.monotonic() - t0
    logger.info(
        f"[cache_persist] SAVED {saved}/{len(self._entries)} entries "
        f"to {cache_dir} in {dt:.1f}s "
        f"({self._current_memory / _BYTES_PER_MB:.0f}MB total)"
    )
    return saved > 0

vllm_mlx.memory_cache.MemoryAwarePrefixCache.load_from_disk

load_from_disk(cache_dir: str) -> int

Load cache entries from disk.

Returns the number of entries successfully loaded.

Source code in vllm_mlx/memory_cache.py

def load_from_disk(self, cache_dir: str) -> int:
    """Load cache entries from disk.

    Returns the number of entries successfully loaded.
    """
    import json
    import os
    import time as _time

    index_path = os.path.join(cache_dir, "index.json")
    if not os.path.exists(index_path):
        logger.info(f"[cache_persist] no index at {index_path}, nothing to load")
        return 0

    t0 = _time.monotonic()

    try:
        from mlx_lm.models.cache import load_prompt_cache
    except ImportError:
        logger.warning("[cache_persist] mlx_lm not available, cannot load")
        return 0

    with open(index_path) as f:
        index = json.load(f)

    version = index.get("version", 1)
    if version != _CACHE_PERSIST_VERSION:
        logger.warning(
            f"[cache_persist] version mismatch: disk={version} "
            f"current={_CACHE_PERSIST_VERSION}, discarding stale cache"
        )
        return 0

    disk_fp = index.get("model_fingerprint", "")
    if disk_fp and disk_fp != self._model_fingerprint:
        logger.warning(
            f"[cache_persist] model fingerprint mismatch: "
            f"disk={disk_fp} current={self._model_fingerprint}, "
            f"discarding incompatible cache"
        )
        return 0

    loaded = 0
    for entry_meta in index.get("entries", []):
        i = entry_meta["index"]
        entry_path = os.path.join(cache_dir, f"entry_{i}.safetensors")
        tokens_path = os.path.join(cache_dir, f"entry_{i}_tokens.bin")

        if not os.path.exists(entry_path) or not os.path.exists(tokens_path):
            logger.warning(f"[cache_persist] missing files for entry {i}, skipping")
            continue

        try:
            # Load tokens from binary
            import array as _array

            arr = _array.array("i")
            with open(tokens_path, "rb") as f:
                arr.fromfile(f, entry_meta["num_tokens"])
            tokens = list(arr)
            if len(tokens) < self._config.min_prefix_tokens:
                logger.info(
                    "[cache_persist] skipping short entry %s: %s tokens < "
                    "min_prefix_tokens=%s",
                    i,
                    len(tokens),
                    self._config.min_prefix_tokens,
                )
                continue

            # Load KV cache
            cache = load_prompt_cache(entry_path)

            # Estimate memory
            memory = estimate_kv_cache_memory(cache)

            with self._memory_lock:
                # Check if it fits
                if self._current_memory + memory > self._max_memory:
                    logger.info(
                        f"[cache_persist] entry {i} would exceed memory limit "
                        f"({(self._current_memory + memory) / _BYTES_PER_MB:.0f}MB > "
                        f"{self._max_memory / _BYTES_PER_MB:.0f}MB), stopping load"
                    )
                    break

                tokens_key = tuple(tokens)
                entry = _CacheEntry(
                    tokens=tokens_key,
                    cache=cache,
                    memory_bytes=memory,
                )
                self._entries[tokens_key] = entry
                self._current_memory += memory
                bisect.insort(self._sorted_keys, tokens_key)
                loaded += 1

            logger.info(
                f"[cache_persist] loaded entry {i}: "
                f"{len(tokens)} tokens, "
                f"{memory / _BYTES_PER_MB:.1f}MB KV"
            )

        except Exception as e:
            logger.warning(f"[cache_persist] failed to load entry {i}: {e}")

    with self._memory_lock:
        self._stats.entry_count = len(self._entries)
        self._stats.current_memory_bytes = self._current_memory

    dt = _time.monotonic() - t0
    logger.info(
        f"[cache_persist] LOADED {loaded} entries from {cache_dir} "
        f"in {dt:.1f}s ({self._current_memory / _BYTES_PER_MB:.0f}MB total)"
    )
    return loaded

vllm_mlx.memory_cache._get_available_memory

_get_available_memory() -> int

Get available system memory in bytes.

Returns:

• int –

  Available memory in bytes, or 0 if detection fails.

Source code in vllm_mlx/memory_cache.py

def _get_available_memory() -> int:
    """
    Get available system memory in bytes.

    Returns:
        Available memory in bytes, or 0 if detection fails.
    """
    try:
        import psutil

        return psutil.virtual_memory().available
    except ImportError:
        logger.warning("psutil not installed, using fallback memory limit")
        return 0
    except Exception as e:
        logger.warning(f"Failed to detect available memory: {e}")
        return 0

vllm_mlx.memory_cache._array_memory

_array_memory(arr) -> int

Estimate array memory from shape+dtype without triggering lazy eval.

Accessing .nbytes on a lazy MLX array forces evaluation of the entire computation graph, causing a VRAM spike. This function uses shape and dtype metadata (which are always available without eval) to compute the same value.

Parameters:

• arr –

  An MLX array or similar object.

Returns:

• int –

  Estimated memory in bytes.

Source code in vllm_mlx/memory_cache.py

def _array_memory(arr) -> int:
    """
    Estimate array memory from shape+dtype without triggering lazy eval.

    Accessing .nbytes on a lazy MLX array forces evaluation of the entire
    computation graph, causing a VRAM spike. This function uses shape and
    dtype metadata (which are always available without eval) to compute
    the same value.

    Args:
        arr: An MLX array or similar object.

    Returns:
        Estimated memory in bytes.
    """
    if hasattr(arr, "shape") and hasattr(arr, "dtype"):
        dtype = arr.dtype
        if hasattr(dtype, "size"):
            return math.prod(arr.shape) * dtype.size
    # Fallback for non-MLX arrays or objects without shape/dtype
    if hasattr(arr, "nbytes"):
        return arr.nbytes
    return 0

vllm_mlx.memory_cache._nested_array_memory

_nested_array_memory(value: Any) -> int

Sum _array_memory over an arbitrarily nested state structure.

Cache state payloads are not always a flat (keys, values) pair: CacheList yields a list of sub-cache states and PoolingCache yields (buf_kv, buf_gate, pooled) with possible None members. Unpacking those as two values raised, was swallowed, and the entry was accounted as zero bytes — so the dashboard showed 0% cache memory and, far worse, the byte-based LRU eviction never fired for such models.

Source code in vllm_mlx/memory_cache.py

def _nested_array_memory(value: Any) -> int:
    """Sum ``_array_memory`` over an arbitrarily nested state structure.

    Cache ``state`` payloads are not always a flat ``(keys, values)`` pair:
    CacheList yields a list of sub-cache states and PoolingCache yields
    ``(buf_kv, buf_gate, pooled)`` with possible ``None`` members. Unpacking
    those as two values raised, was swallowed, and the entry was accounted as
    zero bytes — so the dashboard showed 0% cache memory and, far worse, the
    byte-based LRU eviction never fired for such models.
    """
    if value is None:
        return 0
    if isinstance(value, (list, tuple)):
        return sum(_nested_array_memory(v) for v in value)
    return _array_memory(value)

vllm_mlx.memory_cache.estimate_kv_cache_memory

estimate_kv_cache_memory(cache: list[Any]) -> int

Estimate memory usage of a KV cache in bytes.

This function inspects MLX arrays in the cache and calculates their total memory footprint using shape+dtype metadata to avoid triggering lazy evaluation (which would cause a VRAM spike).

Parameters:

• cache (list[Any]) –

  List of layer cache objects, each containing keys/values tensors.

Returns:

• int –

  Estimated memory usage in bytes.

Source code in vllm_mlx/memory_cache.py

def estimate_kv_cache_memory(cache: list[Any]) -> int:
    """
    Estimate memory usage of a KV cache in bytes.

    This function inspects MLX arrays in the cache and calculates their
    total memory footprint using shape+dtype metadata to avoid triggering
    lazy evaluation (which would cause a VRAM spike).

    Args:
        cache: List of layer cache objects, each containing keys/values tensors.

    Returns:
        Estimated memory usage in bytes.
    """
    if not cache:
        return 0

    total_bytes = 0

    for layer_cache in cache:
        # Handle different cache object types
        # Check dict first since dicts have .keys() method that would match below
        if isinstance(layer_cache, dict) and "state" in layer_cache:
            # Extracted state dict
            keys, values = layer_cache["state"]
            total_bytes += _array_memory(keys)
            total_bytes += _array_memory(values)
        # Handle QuantizedKVCache: keys/values are tuples of (data, scales, biases)
        elif hasattr(layer_cache, "keys") and isinstance(
            getattr(layer_cache, "keys", None), (list, tuple)
        ):
            for arr in layer_cache.keys:
                total_bytes += _array_memory(arr)
            for arr in layer_cache.values:
                total_bytes += _array_memory(arr)
            continue
        elif hasattr(layer_cache, "state") and not isinstance(layer_cache, dict):
            # Cache with a state property. Walk it recursively: the payload may
            # be a plain (keys, values) pair, but CacheList/PoolingCache nest
            # further, and the old two-way unpack silently measured those as 0.
            try:
                total_bytes += _nested_array_memory(layer_cache.state)
            except (TypeError, ValueError):
                pass
        elif hasattr(layer_cache, "keys") and hasattr(layer_cache, "values"):
            # Standard KVCache with keys/values attributes (not dict)
            keys_attr = layer_cache.keys
            values_attr = layer_cache.values
            # Ensure these are arrays, not methods
            if not callable(keys_attr):
                total_bytes += _array_memory(keys_attr)
            if not callable(values_attr):
                total_bytes += _array_memory(values_attr)

    return total_bytes

vllm_mlx.memory_cache._is_cache_layer_trimmable

_is_cache_layer_trimmable(layer_cache: Any) -> bool

Return whether a cache layer can safely be rewound for partial reuse.

Source code in vllm_mlx/memory_cache.py

def _is_cache_layer_trimmable(layer_cache: Any) -> bool:
    """Return whether a cache layer can safely be rewound for partial reuse."""
    if isinstance(layer_cache, _QuantizedCacheWrapper):
        if "max_size" in layer_cache.orig_attrs:
            return False
        return hasattr(layer_cache, "offset") and hasattr(layer_cache, "keys")

    # _trim_cache_offset does not currently rewind container children.
    if hasattr(layer_cache, "caches"):
        return False

    is_trimmable = getattr(layer_cache, "is_trimmable", None)
    if callable(is_trimmable):
        try:
            return bool(is_trimmable())
        except Exception:
            logger.debug(
                "Failed to check cache layer trimmability for %s",
                type(layer_cache).__name__,
                exc_info=True,
            )
            return False

    # Compatibility fallback for simple KV-like cache implementations.
    return hasattr(layer_cache, "offset") and hasattr(layer_cache, "keys")

vllm_mlx.memory_cache._trim_cache_offset

_trim_cache_offset(cache: list[Any], trim_by: int) -> list[Any]

Create copies of cache layers with the last trim_by positions removed.

This is used when returning a cached KV state to the scheduler so that the last N positions are "freed" and the model will recompute them on the next forward pass (preventing duplicate KV entries).

For plain KVCache: reduces offset (surplus data beyond offset is harmless since merge slices to keys[:, :, :offset, :]).

For RotatingKVCache: actually trims the circular buffer — reducing offset alone breaks size() / _temporal_order invariants.

Supports KVCache, RotatingKVCache, and _QuantizedCacheWrapper.

Source code in vllm_mlx/memory_cache.py

def _trim_cache_offset(cache: list[Any], trim_by: int) -> list[Any]:
    """Create copies of cache layers with the last ``trim_by`` positions removed.

    This is used when returning a cached KV state to the scheduler so that
    the last N positions are "freed" and the model will recompute them on the
    next forward pass (preventing duplicate KV entries).

    For plain KVCache: reduces offset (surplus data beyond offset is harmless
    since merge slices to ``keys[:, :, :offset, :]``).

    For RotatingKVCache: actually trims the circular buffer — reducing offset
    alone breaks ``size()`` / ``_temporal_order`` invariants.

    Supports KVCache, RotatingKVCache, and _QuantizedCacheWrapper.
    """
    import mlx.core as mx
    from mlx_lm.models.cache import RotatingKVCache

    trimmed: list[Any] = []
    eval_targets: list[Any] = []
    for layer_cache in cache:
        if isinstance(layer_cache, _QuantizedCacheWrapper):
            # Shallow copy with reduced offset
            tc = _QuantizedCacheWrapper.__new__(_QuantizedCacheWrapper)
            tc.keys = layer_cache.keys
            tc.values = layer_cache.values
            tc.offset = max(layer_cache.offset - trim_by, 0)
            tc.bits = layer_cache.bits
            tc.group_size = layer_cache.group_size
            tc.orig_type = layer_cache.orig_type
            tc.orig_attrs = layer_cache.orig_attrs
            trimmed.append(tc)
        elif isinstance(layer_cache, RotatingKVCache):
            if layer_cache.keys is None or trim_by <= 0:
                trimmed.append(layer_cache)
                continue
            # RotatingKVCache: must trim buffer, not just offset.
            # The buffer stores the last min(offset, max_size) tokens in a
            # circular arrangement.  Trimming excess positions from the END
            # means removing the newest entries (chronologically last).
            old_offset = layer_cache.offset
            new_offset = max(old_offset - trim_by, 0)
            old_size = min(old_offset, layer_cache.max_size)
            entries_to_keep = max(0, old_size - trim_by)

            orig_cls = type(layer_cache)
            tc = orig_cls.__new__(orig_cls)
            tc.offset = new_offset
            tc.max_size = layer_cache.max_size
            tc.keep = getattr(layer_cache, "keep", 0)
            tc.step = getattr(layer_cache, "step", layer_cache.max_size)

            if entries_to_keep <= 0:
                # All buffer content is beyond the trim point — clear
                tc.keys = None
                tc.values = None
                tc._idx = 0
                tc.offset = 0
            elif entries_to_keep < old_size:
                # Reorder to temporal order, keep the oldest entries
                ordered_k = layer_cache._temporal_order(layer_cache.keys)
                ordered_v = layer_cache._temporal_order(layer_cache.values)
                kept_k = ordered_k[:, :, :entries_to_keep, :]
                kept_v = ordered_v[:, :, :entries_to_keep, :]

                if new_offset >= tc.max_size:
                    # Invariant: when offset >= max_size, buffer must be
                    # full (keys.shape[2] == max_size).  Left-pad with
                    # zeros to restore the full buffer.  Zeros represent
                    # positions evicted long ago; _idx = max_size so
                    # _temporal_order returns as-is and _update_in_place
                    # rotates to overwrite zeros first.
                    pad_n = tc.max_size - entries_to_keep
                    pad_k = mx.zeros(
                        (kept_k.shape[0], kept_k.shape[1], pad_n, kept_k.shape[3]),
                        dtype=kept_k.dtype,
                    )
                    pad_v = mx.zeros(
                        (kept_v.shape[0], kept_v.shape[1], pad_n, kept_v.shape[3]),
                        dtype=kept_v.dtype,
                    )
                    tc.keys = mx.concatenate([pad_k, kept_k], axis=2)
                    tc.values = mx.concatenate([pad_v, kept_v], axis=2)
                    tc._idx = tc.max_size
                else:
                    if entries_to_keep < new_offset:
                        # Buffer has fewer entries than offset requires.
                        # This happens when old_offset > max_size (rotating)
                        # and the trim brought new_offset below max_size.
                        # Pad with zeros on the left to maintain the invariant
                        # size() == keys.shape[2], preventing merge crashes.
                        pad_n = new_offset - entries_to_keep
                        pad_k = mx.zeros(
                            (
                                kept_k.shape[0],
                                kept_k.shape[1],
                                pad_n,
                                kept_k.shape[3],
                            ),
                            dtype=kept_k.dtype,
                        )
                        pad_v = mx.zeros(
                            (
                                kept_v.shape[0],
                                kept_v.shape[1],
                                pad_n,
                                kept_v.shape[3],
                            ),
                            dtype=kept_v.dtype,
                        )
                        tc.keys = mx.concatenate([pad_k, kept_k], axis=2)
                        tc.values = mx.concatenate([pad_v, kept_v], axis=2)
                        tc._idx = new_offset
                    else:
                        tc.keys = kept_k
                        tc.values = kept_v
                        tc._idx = entries_to_keep
                eval_targets.extend([tc.keys, tc.values])
            else:
                # No entries removed (trim_by == 0 already handled above,
                # this covers entries_to_keep == old_size edge case)
                tc.keys = layer_cache.keys
                tc.values = layer_cache.values
                tc._idx = layer_cache._idx
            trimmed.append(tc)
        elif (
            hasattr(layer_cache, "offset")
            and hasattr(layer_cache, "keys")
            and not isinstance(layer_cache.keys, (list, tuple))
        ):
            orig_cls = type(layer_cache)
            tc = orig_cls.__new__(orig_cls)
            new_offset = max(layer_cache.offset - trim_by, 0)
            keys = layer_cache.keys
            values = layer_cache.values
            # Slice the arrays down to new_offset rather than just shrinking the
            # offset pointer.  Sharing the original (over-sized) array across
            # requests lets attention paths that read the full underlying
            # buffer (e.g. Gemma 4's KV-shared layers, which read cache.state
            # directly instead of going through update_and_fetch) see stale
            # tokens from the previous owner — issue #384.
            if (
                keys is not None
                and hasattr(keys, "shape")
                and len(keys.shape) >= 3
                and new_offset < keys.shape[-2]
            ):
                tc.keys = keys[..., :new_offset, :]
                tc.values = values[..., :new_offset, :]
            else:
                tc.keys = keys
                tc.values = values
            tc.offset = new_offset
            # Preserve type-specific attrs (max_size, keep, step, _idx)
            for attr in ("max_size", "keep", "step", "_idx"):
                if hasattr(layer_cache, attr):
                    setattr(tc, attr, getattr(layer_cache, attr))
            trimmed.append(tc)
        else:
            trimmed.append(layer_cache)

    if eval_targets:
        mx.eval(*eval_targets)

    return trimmed

vllm_mlx.memory_cache._needs_kv_trim

_needs_kv_trim(layer: Any) -> bool

Check if a cache layer has oversized KV arrays (duck-typed, no MLX import).

Source code in vllm_mlx/memory_cache.py

def _needs_kv_trim(layer: Any) -> bool:
    """Check if a cache layer has oversized KV arrays (duck-typed, no MLX import)."""
    keys = getattr(layer, "keys", None)
    offset = getattr(layer, "offset", None)
    if keys is None or offset is None:
        return False
    if isinstance(keys, (list, tuple)):
        return False  # QuantizedKVCache — skip
    shape = getattr(keys, "shape", None)
    if shape is None or len(shape) < 3:
        return False
    return 0 < offset < shape[2]

vllm_mlx.memory_cache._trim_to_offset

_trim_to_offset(cache: list[Any]) -> list[Any]

Trim KV arrays to their actual used size (offset) before storage.

KV arrays are often pre-allocated larger than needed (e.g. 4096 slots when only 100 are used). This slices them down to offset and evaluates the result so the original large buffer can be freed.

Parameters:

• cache (list[Any]) –

  List of cache layer objects (KVCache or other types).

Returns:

• list[Any] –

  New list with KVCache layers trimmed to their offset.

• list[Any] –

  Non-KVCache layers are passed through unchanged.

Source code in vllm_mlx/memory_cache.py

def _trim_to_offset(cache: list[Any]) -> list[Any]:
    """Trim KV arrays to their actual used size (offset) before storage.

    KV arrays are often pre-allocated larger than needed (e.g. 4096 slots
    when only 100 are used).  This slices them down to ``offset`` and
    evaluates the result so the original large buffer can be freed.

    Args:
        cache: List of cache layer objects (KVCache or other types).

    Returns:
        New list with KVCache layers trimmed to their offset.
        Non-KVCache layers are passed through unchanged.
    """
    if not any(_needs_kv_trim(layer) for layer in cache):
        return cache

    import mlx.core as mx
    from mlx_lm.models.cache import KVCache

    trimmed = []
    eval_targets = []
    for layer in cache:
        if isinstance(layer, KVCache) and layer.keys is not None:
            offset = layer.offset
            if offset <= 0 or offset >= layer.keys.shape[2]:
                trimmed.append(layer)
                continue
            tc = KVCache()
            tc.keys = layer.keys[:, :, :offset, :]
            tc.values = layer.values[:, :, :offset, :]
            tc.offset = offset
            eval_targets.extend([tc.keys, tc.values])
            trimmed.append(tc)
        else:
            trimmed.append(layer)

    if eval_targets:
        mx.eval(*eval_targets)

    return trimmed

vllm_mlx.memory_cache._quantize_cache

_quantize_cache(cache: list[Any], bits: int = 8, group_size: int = 64) -> list[Any]

Quantize KV cache layers to reduce memory.

Only plain KVCache layers are quantized. RotatingKVCache (sliding window) is left as-is because its internal _idx/rotation state is tightly coupled with update_and_fetch logic and cannot survive quantize/dequantize roundtrip. RotatingKVCache is typically small (max_size=1024) so skipping it is fine.

Source code in vllm_mlx/memory_cache.py

def _quantize_cache(cache: list[Any], bits: int = 8, group_size: int = 64) -> list[Any]:
    """Quantize KV cache layers to reduce memory.

    Only plain KVCache layers are quantized. RotatingKVCache (sliding window)
    is left as-is because its internal _idx/rotation state is tightly coupled
    with update_and_fetch logic and cannot survive quantize/dequantize roundtrip.
    RotatingKVCache is typically small (max_size=1024) so skipping it is fine.
    """
    from mlx_lm.models.cache import KVCache

    quantized = []
    for layer in cache:
        if type(layer) is KVCache and getattr(layer, "keys", None) is not None:
            quantized.append(_QuantizedCacheWrapper(layer, bits, group_size))
        else:
            quantized.append(layer)
    return quantized

vllm_mlx.memory_cache._dequantize_cache

_dequantize_cache(cache: list[Any]) -> list[Any]

Dequantize _QuantizedCacheWrapper layers and copy non-quantized layers.

All layers are copied (never returned by reference) so that the model's update_and_fetch mutations don't corrupt the stored cache entry.

Source code in vllm_mlx/memory_cache.py

def _dequantize_cache(cache: list[Any]) -> list[Any]:
    """Dequantize _QuantizedCacheWrapper layers and copy non-quantized layers.

    All layers are copied (never returned by reference) so that the model's
    ``update_and_fetch`` mutations don't corrupt the stored cache entry.
    """
    import mlx.core as mx

    result = []
    for layer in cache:
        if isinstance(layer, _QuantizedCacheWrapper):
            # Reconstruct original cache type from quantized data
            orig_cls = layer.orig_type
            kv = orig_cls.__new__(orig_cls)
            kv.keys = mx.dequantize(
                *layer.keys, group_size=layer.group_size, bits=layer.bits
            )
            kv.values = mx.dequantize(
                *layer.values, group_size=layer.group_size, bits=layer.bits
            )
            kv.offset = layer.offset
            # Slice the dequantized arrays down to offset so that readers
            # which bypass offset (e.g. Gemma 4 KV-shared layers reading
            # cache.state directly) cannot see stale tokens from a previous
            # request.  Mirrors the plain-KVCache slice in
            # _trim_cache_offset — see issue #384.
            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, :]
            # Restore type-specific attrs (max_size, keep, step, _idx)
            for attr, val in layer.orig_attrs.items():
                setattr(kv, attr, val)
            result.append(kv)
        elif hasattr(layer, "keys") and hasattr(layer, "offset"):
            # Deep-copy non-quantized cache layers (e.g. RotatingKVCache)
            # so model's in-place mutations don't corrupt stored entries
            orig_cls = type(layer)
            kv = orig_cls.__new__(orig_cls)
            kv.keys = mx.array(layer.keys) if layer.keys is not None else None
            kv.values = mx.array(layer.values) if layer.values is not None else None
            kv.offset = layer.offset
            for attr in ("max_size", "keep", "step", "_idx"):
                if hasattr(layer, attr):
                    setattr(kv, attr, getattr(layer, attr))
            result.append(kv)
        else:
            result.append(layer)
    return result

vllm_mlx.memory_cache._compute_model_fingerprint

_compute_model_fingerprint(model: Any) -> str

Compute a fingerprint from model architecture for cache compatibility.

Used to reject disk-persisted caches created by a different model or a different quantisation of the same model. The fingerprint is a short hex digest of (num_layers, hidden_size, vocab_size, num_kv_heads, head_dim) — lightweight and deterministic.

Source code in vllm_mlx/memory_cache.py

def _compute_model_fingerprint(model: Any) -> str:
    """Compute a fingerprint from model architecture for cache compatibility.

    Used to reject disk-persisted caches created by a different model or
    a different quantisation of the same model.  The fingerprint is a
    short hex digest of (num_layers, hidden_size, vocab_size, num_kv_heads,
    head_dim) — lightweight and deterministic.
    """
    import hashlib

    parts: list[str] = []
    # Walk model.config / model.args / direct attributes
    for cfg_attr in ("config", "args", "model_config"):
        cfg = getattr(model, cfg_attr, None)
        if cfg is not None:
            break
    if cfg is None:
        cfg = model  # fallback: attributes on the model itself

    for key in (
        "num_hidden_layers",
        "hidden_size",
        "vocab_size",
        "num_key_value_heads",
        "head_dim",
        "intermediate_size",
        "model_type",
    ):
        val = getattr(cfg, key, None)
        if val is not None:
            parts.append(f"{key}={val}")

    fingerprint = hashlib.sha256("|".join(parts).encode()).hexdigest()[:16]
    logger.debug(f"[model_fingerprint] {fingerprint} ({', '.join(parts)})")
    return fingerprint

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.memory_cache._get_available_memory · function

vllm_mlx.memory_cache._get_available_memory() -> int

Get available system memory in bytes.

Parameters

This callable has no explicit inputs.

Returns

• Type: int
• Direct return expressions: psutil.virtual_memory().available; 0

Exceptions and behavior

Function _get_available_memory calls psutil.virtual_memory, logger.warning; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L47-L63.

vllm_mlx.memory_cache._array_memory · function

vllm_mlx.memory_cache._array_memory(arr) -> int

Estimate array memory from shape+dtype without triggering lazy eval.

Parameters

Name	Type	Required	Default	Description
arr	not annotated	yes	none	An MLX array or similar object.

Returns

• Type: int
• Direct return expressions: math.prod(arr.shape) * dtype.size; arr.nbytes; 0

Exceptions and behavior

Function _array_memory calls hasattr, math.prod; has 3 explicit return paths. No direct raise statement appears in this definition.

View source #L66-L88.

vllm_mlx.memory_cache._nested_array_memory · function

vllm_mlx.memory_cache._nested_array_memory(value: Any) -> int

Sum _array_memory over an arbitrarily nested state structure.

Parameters

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

Returns

• Type: int
• Direct return expressions: 0; sum((_nested_array_memory(v) for v in value)); _array_memory(value)

Exceptions and behavior

Function _nested_array_memory calls isinstance, sum, _nested_array_memory, _array_memory; has 3 explicit return paths. No direct raise statement appears in this definition.

View source #L91-L105.

vllm_mlx.memory_cache.estimate_kv_cache_memory · function

vllm_mlx.memory_cache.estimate_kv_cache_memory(cache: list[Any]) -> int

Estimate memory usage of a KV cache in bytes.

Parameters

Name	Type	Required	Default	Description
cache	list[Any]	yes	none	List of layer cache objects, each containing keys/values tensors.

Returns

• Type: int
• Direct return expressions: 0; total_bytes

Exceptions and behavior

Function estimate_kv_cache_memory calls isinstance, _array_memory, hasattr, getattr; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L108-L162.

vllm_mlx.memory_cache.MemoryCacheConfig · class

vllm_mlx.memory_cache.MemoryCacheConfig(max_memory_mb: int | None = None, max_memory_percent: float = _DEFAULT_MEMORY_PERCENT, max_entries: int = 1000, enable_memory_tracking: bool = True, kv_quantize: bool = False, kv_bits: int = 8, kv_group_size: int = 64, kv_min_quantize_tokens: int = 256, min_prefix_tokens: int = 128)

Configuration for memory-aware prefix cache.

Parameters

Name	Type	Required	Default	Description
max_memory_mb	int \| None	no	None	Optional constructor field; defaults to None.
max_memory_percent	float	no	_DEFAULT_MEMORY_PERCENT	Optional constructor field; defaults to _DEFAULT_MEMORY_PERCENT.
max_entries	int	no	1000	Optional constructor field; defaults to 1000.
enable_memory_tracking	bool	no	True	Optional constructor field; defaults to True.
kv_quantize	bool	no	False	Optional constructor field; defaults to False.
kv_bits	int	no	8	Optional constructor field; defaults to 8.
kv_group_size	int	no	64	Optional constructor field; defaults to 64.
kv_min_quantize_tokens	int	no	256	Optional constructor field; defaults to 256.
min_prefix_tokens	int	no	128	Optional constructor field; defaults to 128.

Returns

• Constructs: vllm_mlx.memory_cache.MemoryCacheConfig

Exceptions and behavior

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

View source #L166-L225.

vllm_mlx.memory_cache.MemoryCacheConfig.__post_init__ · method

vllm_mlx.memory_cache.MemoryCacheConfig.__post_init__() -> None

Method MemoryCacheConfig.__post_init__ calls ValueError; can raise ValueError.

Parameters

This callable has no explicit inputs.

Returns

• Type: None

Exceptions and behavior

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

View source #L192-L206.

vllm_mlx.memory_cache.MemoryCacheConfig.compute_memory_limit · method

vllm_mlx.memory_cache.MemoryCacheConfig.compute_memory_limit() -> int

Compute the memory limit in bytes.

Parameters

This callable has no explicit inputs.

Returns

• Type: int
• Direct return expressions: self.max_memory_mb * _BYTES_PER_MB; max(limit, _MIN_MEMORY_BYTES); int(fallback_total * self.max_memory_percent)

Exceptions and behavior

Method MemoryCacheConfig.compute_memory_limit calls _get_available_memory, int, max; has 3 explicit return paths. No direct raise statement appears in this definition.

View source #L208-L225.

vllm_mlx.memory_cache.CacheStats · class

vllm_mlx.memory_cache.CacheStats(hits: int = 0, misses: int = 0, evictions: int = 0, tokens_saved: int = 0, current_memory_bytes: int = 0, max_memory_bytes: int = 0, entry_count: int = 0)

Statistics for cache performance monitoring.

Parameters

Name	Type	Required	Default	Description
hits	int	no	0	Optional constructor field; defaults to 0.
misses	int	no	0	Optional constructor field; defaults to 0.
evictions	int	no	0	Optional constructor field; defaults to 0.
tokens_saved	int	no	0	Optional constructor field; defaults to 0.
current_memory_bytes	int	no	0	Optional constructor field; defaults to 0.
max_memory_bytes	int	no	0	Optional constructor field; defaults to 0.
entry_count	int	no	0	Optional constructor field; defaults to 0.

Returns

• Constructs: vllm_mlx.memory_cache.CacheStats

Exceptions and behavior

Class CacheStats declares 3 direct member(s). No direct raise statement appears in this definition.

View source #L229-L268.

vllm_mlx.memory_cache.CacheStats.hit_rate · method

vllm_mlx.memory_cache.CacheStats.hit_rate() -> float

Return successful lookups divided by all completed lookups.

Parameters

This callable has no explicit inputs.

Returns

• Type: float
• Direct return expressions: self.hits / total if total > 0 else 0.0

Exceptions and behavior

Method CacheStats.hit_rate returns self.hits / total if total > 0 else 0.0. No direct raise statement appears in this definition.

View source #L241-L245.

vllm_mlx.memory_cache.CacheStats.memory_utilization · method

vllm_mlx.memory_cache.CacheStats.memory_utilization() -> float

Return the fraction of the configured memory budget in use.

Parameters

This callable has no explicit inputs.

Returns

• Type: float
• Direct return expressions: 0.0; self.current_memory_bytes / self.max_memory_bytes

Exceptions and behavior

Method CacheStats.memory_utilization has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L248-L253.

vllm_mlx.memory_cache.CacheStats.to_dict · method

vllm_mlx.memory_cache.CacheStats.to_dict() -> dict[str, Any]

Return rounded cache counters and memory values for APIs and logs.

Parameters

This callable has no explicit inputs.

Returns

• Type: dict[str, Any]
• Direct return expressions: {'hits': self.hits, 'misses': self.misses, 'hit_rate': round(self.hit_rate, 4), 'evictions': self.evictions, 'tokens_sa…

Exceptions and behavior

Method CacheStats.to_dict calls round; returns {'hits': self.hits, 'misses': self.misses, 'hit_rate': round(self.hit_rate, 4), 'evictions': self.evictions, 'tokens_sa…. No direct raise statement appears in this definition.

View source #L255-L268.

vllm_mlx.memory_cache._CacheEntry · class

vllm_mlx.memory_cache._CacheEntry(tokens: tuple[int, ...], cache: list[Any], memory_bytes: int)

Internal cache entry with memory tracking.

Parameters

Name	Type	Required	Default	Description
tokens	tuple[int, ...]	yes	none	Required constructor field.
cache	list[Any]	yes	none	Required constructor field.
memory_bytes	int	yes	none	Required constructor field.

Returns

• Constructs: vllm_mlx.memory_cache._CacheEntry

Exceptions and behavior

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

View source #L272-L287.

vllm_mlx.memory_cache._CacheEntry.create · method

vllm_mlx.memory_cache._CacheEntry.create(tokens: list[int], cache: list[Any]) -> _CacheEntry

Create a cache entry with memory estimation.

Parameters

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

Returns

• Type: _CacheEntry
• Direct return expressions: cls(tokens=tuple(tokens), cache=cache, memory_bytes=memory)

Exceptions and behavior

Method _CacheEntry.create calls estimate_kv_cache_memory, cls, tuple; returns cls(tokens=tuple(tokens), cache=cache, memory_bytes=memory). No direct raise statement appears in this definition.

View source #L280-L287.

vllm_mlx.memory_cache._is_cache_layer_trimmable · function

vllm_mlx.memory_cache._is_cache_layer_trimmable(layer_cache: Any) -> bool

Return whether a cache layer can safely be rewound for partial reuse.

Parameters

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

Returns

• Type: bool
• Direct return expressions: False; hasattr(layer_cache, 'offset') and hasattr(layer_cache, 'keys'); bool(is_trimmable())

Exceptions and behavior

Function _is_cache_layer_trimmable calls isinstance, hasattr, getattr, callable; has 3 explicit return paths. No direct raise statement appears in this definition.

View source #L290-L314.

vllm_mlx.memory_cache._trim_cache_offset · function

vllm_mlx.memory_cache._trim_cache_offset(cache: list[Any], trim_by: int) -> list[Any]

Create copies of cache layers with the last trim_by positions removed.

Parameters

Name	Type	Required	Default	Description
cache	list[Any]	yes	none	Required positional or keyword input.
trim_by	int	yes	none	Required positional or keyword input.

Returns

• Type: list[Any]
• Direct return expressions: trimmed

Exceptions and behavior

Function _trim_cache_offset calls isinstance, _QuantizedCacheWrapper.__new__, max, trimmed.append; returns trimmed. No direct raise statement appears in this definition.

View source #L317-L481.

vllm_mlx.memory_cache._needs_kv_trim · function

vllm_mlx.memory_cache._needs_kv_trim(layer: Any) -> bool

Check if a cache layer has oversized KV arrays (duck-typed, no MLX import).

Parameters

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

Returns

• Type: bool
• Direct return expressions: False; 0 < offset < shape[2]

Exceptions and behavior

Function _needs_kv_trim calls getattr, isinstance, len; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L484-L495.

vllm_mlx.memory_cache._trim_to_offset · function

vllm_mlx.memory_cache._trim_to_offset(cache: list[Any]) -> list[Any]

Trim KV arrays to their actual used size (offset) before storage.

Parameters

Name	Type	Required	Default	Description
cache	list[Any]	yes	none	List of cache layer objects (KVCache or other types).

Returns

• Type: list[Any]
• Direct return expressions: cache; trimmed

Exceptions and behavior

Function _trim_to_offset calls any, _needs_kv_trim, isinstance, trimmed.append; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L498-L538.

vllm_mlx.memory_cache._QuantizedCacheWrapper · class

vllm_mlx.memory_cache._QuantizedCacheWrapper(layer: Any, bits: int, group_size: int)

Lightweight wrapper storing quantized KV arrays + original cache metadata.

Parameters

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

Returns

• Constructs: vllm_mlx.memory_cache._QuantizedCacheWrapper

Exceptions and behavior

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

View source #L541-L571.

vllm_mlx.memory_cache._QuantizedCacheWrapper.__init__ · method

vllm_mlx.memory_cache._QuantizedCacheWrapper.__init__(layer: Any, bits: int, group_size: int) -> not annotated

Method _QuantizedCacheWrapper.__init__ updates self.keys, self.values, self.offset, self.bits; calls mx.quantize, type, hasattr, getattr.

Parameters

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

Returns

• Type: not annotated

Exceptions and behavior

Method _QuantizedCacheWrapper.__init__ updates self.keys, self.values, self.offset, self.bits; calls mx.quantize, type, hasattr, getattr. No direct raise statement appears in this definition.

View source #L558-L571.

vllm_mlx.memory_cache._quantize_cache · function

vllm_mlx.memory_cache._quantize_cache(cache: list[Any], bits: int = 8, group_size: int = 64) -> list[Any]

Quantize KV cache layers to reduce memory.

Parameters

Name	Type	Required	Default	Description
cache	list[Any]	yes	none	Required positional or keyword input.
bits	int	no	8	Optional positional or keyword input; defaults to 8.
group_size	int	no	64	Optional positional or keyword input; defaults to 64.

Returns

• Type: list[Any]
• Direct return expressions: quantized

Exceptions and behavior

Function _quantize_cache calls type, getattr, quantized.append, _QuantizedCacheWrapper; returns quantized. No direct raise statement appears in this definition.

View source #L574-L590.

vllm_mlx.memory_cache._dequantize_cache · function

vllm_mlx.memory_cache._dequantize_cache(cache: list[Any]) -> list[Any]

Dequantize _QuantizedCacheWrapper layers and copy non-quantized layers.

Parameters

Name	Type	Required	Default	Description
cache	list[Any]	yes	none	Required positional or keyword input.

Returns

• Type: list[Any]
• Direct return expressions: result

Exceptions and behavior

Function _dequantize_cache calls isinstance, orig_cls.__new__, mx.dequantize, hasattr; returns result. No direct raise statement appears in this definition.

View source #L593-L645.

vllm_mlx.memory_cache._compute_model_fingerprint · function

vllm_mlx.memory_cache._compute_model_fingerprint(model: Any) -> str

Compute a fingerprint from model architecture for cache compatibility.

Parameters

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

Returns

• Type: str
• Direct return expressions: fingerprint

Exceptions and behavior

Function _compute_model_fingerprint calls getattr, parts.append, hashlib.sha256('|'.join(parts).encode()).hexdigest, hashlib.sha256; returns fingerprint. No direct raise statement appears in this definition.

View source #L648-L682.

vllm_mlx.memory_cache.MemoryAwarePrefixCache · class

vllm_mlx.memory_cache.MemoryAwarePrefixCache(model: Any, config: MemoryCacheConfig | None = None)

Prefix cache with memory-based eviction.

Parameters

Name	Type	Required	Default	Description
model	Any	yes	none	The MLX model (used for identification).
config	MemoryCacheConfig \| None	no	None	Cache configuration. Uses defaults if None.

Returns

• Constructs: vllm_mlx.memory_cache.MemoryAwarePrefixCache

Exceptions and behavior

Class MemoryAwarePrefixCache declares 19 direct member(s). No direct raise statement appears in this definition.

View source #L685-L1463.

vllm_mlx.memory_cache.MemoryAwarePrefixCache.__init__ · method

vllm_mlx.memory_cache.MemoryAwarePrefixCache.__init__(model: Any, config: MemoryCacheConfig | None = None) -> None

Initialize the memory-aware prefix cache.

Parameters

Name	Type	Required	Default	Description
model	Any	yes	none	The MLX model (used for identification).
config	MemoryCacheConfig \| None	no	None	Cache configuration. Uses defaults if None.

Returns

• Type: None

Exceptions and behavior

Method MemoryAwarePrefixCache.__init__ updates self._model_id, self._config, self._model_fingerprint, self._entries; calls id, MemoryCacheConfig, _compute_model_fingerprint, OrderedDict. No direct raise statement appears in this definition.

View source #L703-L746.

vllm_mlx.memory_cache.MemoryAwarePrefixCache.fetch · method

vllm_mlx.memory_cache.MemoryAwarePrefixCache.fetch(tokens: list[int]) -> tuple[list[Any] | None, list[int]]

Find cached KV state for the given tokens.

Parameters

Name	Type	Required	Default	Description
tokens	list[int]	yes	none	Input token sequence.

Returns

• Type: tuple[list[Any] | None, list[int]]
• Direct return expressions: (None, tokens); (cache_out, []); (trimmed_cache, []); (cache_out, remaining); (trimmed_cache, remaining)

Exceptions and behavior

Method MemoryAwarePrefixCache.fetch updates self._stats.misses, self._last_match_type, self._stats.hits, self._stats.tokens_saved; calls len, tuple, self._entries.move_to_end, _dequantize_cache; has 5 explicit return paths. No direct raise statement appears in this definition.

View source #L748-L977.

vllm_mlx.memory_cache.MemoryAwarePrefixCache.store · method

vllm_mlx.memory_cache.MemoryAwarePrefixCache.store(tokens: list[int], cache: list[Any], evict_prefixes: bool = True) -> bool

Store KV cache for future reuse.

Parameters

Name	Type	Required	Default	Description
tokens	list[int]	yes	none	Token sequence that was processed.
cache	list[Any]	yes	none	The computed KV cache to store.
evict_prefixes	bool	no	True	If True, evict existing entries whose token sequence is a strict prefix of tokens. Set to False when storing prompt+output entries to preserve prompt-only entries created by prompt_cache_save (those are the entries that future requests will actually match).

Returns

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

Exceptions and behavior

Method MemoryAwarePrefixCache.store updates self._current_memory, self._stats.evictions, self._stats.entry_count, self._stats.current_memory_bytes; calls len, logger.debug, tuple, self._entries.move_to_end; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L979-L1092.

vllm_mlx.memory_cache.MemoryAwarePrefixCache._remove_from_sorted · method

vllm_mlx.memory_cache.MemoryAwarePrefixCache._remove_from_sorted(key: tuple[int, ...]) -> None

Remove a key from the sorted index using bisect for O(log N).

Parameters

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

Returns

• Type: None

Exceptions and behavior

Method MemoryAwarePrefixCache._remove_from_sorted calls bisect.bisect_left, len, self._sorted_keys.pop. No direct raise statement appears in this definition.

View source #L1094-L1098.

vllm_mlx.memory_cache.MemoryAwarePrefixCache._evict_lru · method

vllm_mlx.memory_cache.MemoryAwarePrefixCache._evict_lru() -> None

Evict the least recently used entry.

Parameters

This callable has no explicit inputs.

Returns

• Type: None
• Direct return expressions: None

Exceptions and behavior

Method MemoryAwarePrefixCache._evict_lru updates self._current_memory, self._stats.evictions, self._stats.entry_count, self._stats.current_memory_bytes; calls self._entries.popitem, self._remove_from_sorted, len, self._ssd_tier.enqueue_spill; returns None. No direct raise statement appears in this definition.

View source #L1100-L1126.

vllm_mlx.memory_cache.MemoryAwarePrefixCache.remove · method

vllm_mlx.memory_cache.MemoryAwarePrefixCache.remove(tokens: list[int]) -> bool

Remove a specific cache entry.

Parameters

Name	Type	Required	Default	Description
tokens	list[int]	yes	none	Token sequence to remove.

Returns

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

Exceptions and behavior

Method MemoryAwarePrefixCache.remove updates self._current_memory, self._stats.entry_count, self._stats.current_memory_bytes; calls tuple, self._entries.pop, self._remove_from_sorted, len; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L1128-L1147.

vllm_mlx.memory_cache.MemoryAwarePrefixCache.clear · method

vllm_mlx.memory_cache.MemoryAwarePrefixCache.clear() -> None

Clear all cached entries.

Parameters

This callable has no explicit inputs.

Returns

• Type: None

Exceptions and behavior

Method MemoryAwarePrefixCache.clear updates self._current_memory, self._stats; calls self._entries.clear, self._sorted_keys.clear, CacheStats, logger.debug. No direct raise statement appears in this definition.

View source #L1149-L1156.

vllm_mlx.memory_cache.MemoryAwarePrefixCache.get_stats · method

vllm_mlx.memory_cache.MemoryAwarePrefixCache.get_stats() -> dict[str, Any]

Get cache statistics.

Parameters

This callable has no explicit inputs.

Returns

• Type: dict[str, Any]
• Direct return expressions: self._stats.to_dict()

Exceptions and behavior

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

View source #L1158-L1160.

vllm_mlx.memory_cache.MemoryAwarePrefixCache.reset_stats · method

vllm_mlx.memory_cache.MemoryAwarePrefixCache.reset_stats() -> None

Reset statistics while preserving cache contents.

Parameters

This callable has no explicit inputs.

Returns

• Type: None

Exceptions and behavior

Method MemoryAwarePrefixCache.reset_stats updates self._stats; calls CacheStats, len. No direct raise statement appears in this definition.

View source #L1162-L1169.

vllm_mlx.memory_cache.MemoryAwarePrefixCache.memory_usage_mb · method

vllm_mlx.memory_cache.MemoryAwarePrefixCache.memory_usage_mb() -> float

Current memory usage in MB.

Parameters

This callable has no explicit inputs.

Returns

• Type: float
• Direct return expressions: self._current_memory / _BYTES_PER_MB

Exceptions and behavior

Method MemoryAwarePrefixCache.memory_usage_mb returns self._current_memory / _BYTES_PER_MB. No direct raise statement appears in this definition.

View source #L1172-L1174.

vllm_mlx.memory_cache.MemoryAwarePrefixCache.memory_limit_mb · method

vllm_mlx.memory_cache.MemoryAwarePrefixCache.memory_limit_mb() -> float

Memory limit in MB.

Parameters

This callable has no explicit inputs.

Returns

• Type: float
• Direct return expressions: self._max_memory / _BYTES_PER_MB

Exceptions and behavior

Method MemoryAwarePrefixCache.memory_limit_mb returns self._max_memory / _BYTES_PER_MB. No direct raise statement appears in this definition.

View source #L1177-L1179.

vllm_mlx.memory_cache.MemoryAwarePrefixCache.try_reserve_memory · method

vllm_mlx.memory_cache.MemoryAwarePrefixCache.try_reserve_memory(nbytes: int) -> bool

Tentatively reserve cache memory for an upcoming promotion.

Parameters

Name	Type	Required	Default	Description
nbytes	int	yes	none	Required positional or keyword input.

Returns

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

Exceptions and behavior

Method MemoryAwarePrefixCache.try_reserve_memory updates self._current_memory, self._stats.current_memory_bytes; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L1181-L1188.

vllm_mlx.memory_cache.MemoryAwarePrefixCache.release_reserved_memory · method

vllm_mlx.memory_cache.MemoryAwarePrefixCache.release_reserved_memory(nbytes: int) -> None

Release memory previously reserved by try_reserve_memory().

Parameters

Name	Type	Required	Default	Description
nbytes	int	yes	none	Required positional or keyword input.

Returns

• Type: None

Exceptions and behavior

Method MemoryAwarePrefixCache.release_reserved_memory updates self._current_memory, self._stats.current_memory_bytes; calls max. No direct raise statement appears in this definition.

View source #L1190-L1194.

vllm_mlx.memory_cache.MemoryAwarePrefixCache.__len__ · method

vllm_mlx.memory_cache.MemoryAwarePrefixCache.__len__() -> int

Return number of cached entries.

Parameters

This callable has no explicit inputs.

Returns

• Type: int
• Direct return expressions: len(self._entries)

Exceptions and behavior

Method MemoryAwarePrefixCache.__len__ calls len; returns len(self._entries). No direct raise statement appears in this definition.

View source #L1196-L1198.

vllm_mlx.memory_cache.MemoryAwarePrefixCache.__contains__ · method

vllm_mlx.memory_cache.MemoryAwarePrefixCache.__contains__(tokens: list[int]) -> bool

Check if tokens are cached.

Parameters

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

Returns

• Type: bool
• Direct return expressions: tuple(tokens) in self._entries

Exceptions and behavior

Method MemoryAwarePrefixCache.__contains__ calls tuple; returns tuple(tokens) in self._entries. No direct raise statement appears in this definition.

View source #L1200-L1202.

vllm_mlx.memory_cache.MemoryAwarePrefixCache.set_ssd_tier · method

vllm_mlx.memory_cache.MemoryAwarePrefixCache.set_ssd_tier(ssd_tier) -> None

Attach an SSD cache tier for eviction spilling.

Parameters

Name	Type	Required	Default	Description
ssd_tier	not annotated	yes	none	An SSDCacheTier instance (or None to disable).

Returns

• Type: None

Exceptions and behavior

Method MemoryAwarePrefixCache.set_ssd_tier updates self._ssd_tier; calls logger.info. No direct raise statement appears in this definition.

View source #L1204-L1214.

vllm_mlx.memory_cache.MemoryAwarePrefixCache.check_ssd · method

vllm_mlx.memory_cache.MemoryAwarePrefixCache.check_ssd(tokens: list[int]) -> dict | None

Check if tokens have an SSD cache hit (without reading data).

Parameters

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

Returns

• Type: dict | None
• Direct return expressions: None; candidate; prefix

Exceptions and behavior

Method MemoryAwarePrefixCache.check_ssd calls tuple, self._ssd_tier.lookup_ssd, len, self._ssd_tier.lookup_ssd_prefix; has 3 explicit return paths. No direct raise statement appears in this definition.

View source #L1216-L1249.

vllm_mlx.memory_cache.MemoryAwarePrefixCache.save_to_disk · method

vllm_mlx.memory_cache.MemoryAwarePrefixCache.save_to_disk(cache_dir: str) -> bool

Save all cache entries to disk using mlx_lm's safetensors format.

Parameters

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

Returns

• Type: bool
• Direct return expressions: False; saved > 0

Exceptions and behavior

Method MemoryAwarePrefixCache.save_to_disk calls logger.info, _time.monotonic, os.makedirs, logger.warning; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L1255-L1346.

vllm_mlx.memory_cache.MemoryAwarePrefixCache.load_from_disk · method

vllm_mlx.memory_cache.MemoryAwarePrefixCache.load_from_disk(cache_dir: str) -> int

Load cache entries from disk.

Parameters

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

Returns

• Type: int
• Direct return expressions: 0; loaded

Exceptions and behavior

Method MemoryAwarePrefixCache.load_from_disk updates self._current_memory, self._stats.entry_count, self._stats.current_memory_bytes; calls os.path.join, os.path.exists, logger.info, _time.monotonic; has 2 explicit return paths. No direct raise statement appears in this definition.

View source #L1348-L1463.

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
_get_available_memory	function	_get_available_memory() -> int	Get available system memory in bytes.	#L47-L63
_array_memory	function	_array_memory(arr) -> int	Estimate array memory from shape+dtype without triggering lazy eval.	#L66-L88
_nested_array_memory	function	_nested_array_memory(value: Any) -> int	Sum _array_memory over an arbitrarily nested state structure.	#L91-L105
estimate_kv_cache_memory	function	estimate_kv_cache_memory(cache: list[Any]) -> int	Estimate memory usage of a KV cache in bytes.	#L108-L162
MemoryCacheConfig	class	MemoryCacheConfig(max_memory_mb: int \| None = None, max_memory_percent: float = _DEFAULT_MEMORY_PERCENT, max_entries: int = 1000, enable_memory_tracking: bool = True, kv_quantize: bool = False, kv_bits: int = 8, kv_group_size: int = 64, kv_min_quantize_tokens: int = 256, min_prefix_tokens: int = 128)	Configuration for memory-aware prefix cache.	#L166-L225
MemoryCacheConfig.__post_init__	method	MemoryCacheConfig.__post_init__() -> None	Method MemoryCacheConfig.__post_init__ calls ValueError; can raise ValueError.	#L192-L206
MemoryCacheConfig.compute_memory_limit	method	MemoryCacheConfig.compute_memory_limit() -> int	Compute the memory limit in bytes.	#L208-L225
CacheStats	class	CacheStats(hits: int = 0, misses: int = 0, evictions: int = 0, tokens_saved: int = 0, current_memory_bytes: int = 0, max_memory_bytes: int = 0, entry_count: int = 0)	Statistics for cache performance monitoring.	#L229-L268
CacheStats.hit_rate	method	CacheStats.hit_rate() -> float	Return successful lookups divided by all completed lookups.	#L241-L245
CacheStats.memory_utilization	method	CacheStats.memory_utilization() -> float	Return the fraction of the configured memory budget in use.	#L248-L253
CacheStats.to_dict	method	CacheStats.to_dict() -> dict[str, Any]	Return rounded cache counters and memory values for APIs and logs.	#L255-L268
_CacheEntry	class	_CacheEntry(tokens: tuple[int, ...], cache: list[Any], memory_bytes: int)	Internal cache entry with memory tracking.	#L272-L287
_CacheEntry.create	method	_CacheEntry.create(tokens: list[int], cache: list[Any]) -> _CacheEntry	Create a cache entry with memory estimation.	#L280-L287
_is_cache_layer_trimmable	function	_is_cache_layer_trimmable(layer_cache: Any) -> bool	Return whether a cache layer can safely be rewound for partial reuse.	#L290-L314
_trim_cache_offset	function	_trim_cache_offset(cache: list[Any], trim_by: int) -> list[Any]	Create copies of cache layers with the last trim_by positions removed.	#L317-L481
_needs_kv_trim	function	_needs_kv_trim(layer: Any) -> bool	Check if a cache layer has oversized KV arrays (duck-typed, no MLX import).	#L484-L495
_trim_to_offset	function	_trim_to_offset(cache: list[Any]) -> list[Any]	Trim KV arrays to their actual used size (offset) before storage.	#L498-L538
_QuantizedCacheWrapper	class	_QuantizedCacheWrapper(layer: Any, bits: int, group_size: int)	Lightweight wrapper storing quantized KV arrays + original cache metadata.	#L541-L571
_QuantizedCacheWrapper.__init__	method	_QuantizedCacheWrapper.__init__(layer: Any, bits: int, group_size: int) -> not annotated	Method _QuantizedCacheWrapper.__init__ updates self.keys, self.values, self.offset, self.bits; calls mx.quantize, type, hasattr, getattr.	#L558-L571
_quantize_cache	function	_quantize_cache(cache: list[Any], bits: int = 8, group_size: int = 64) -> list[Any]	Quantize KV cache layers to reduce memory.	#L574-L590
_dequantize_cache	function	_dequantize_cache(cache: list[Any]) -> list[Any]	Dequantize _QuantizedCacheWrapper layers and copy non-quantized layers.	#L593-L645
_compute_model_fingerprint	function	_compute_model_fingerprint(model: Any) -> str	Compute a fingerprint from model architecture for cache compatibility.	#L648-L682
MemoryAwarePrefixCache	class	MemoryAwarePrefixCache(model: Any, config: MemoryCacheConfig \| None = None)	Prefix cache with memory-based eviction.	#L685-L1463
MemoryAwarePrefixCache.__init__	method	MemoryAwarePrefixCache.__init__(model: Any, config: MemoryCacheConfig \| None = None) -> None	Initialize the memory-aware prefix cache.	#L703-L746
MemoryAwarePrefixCache.fetch	method	MemoryAwarePrefixCache.fetch(tokens: list[int]) -> tuple[list[Any] \| None, list[int]]	Find cached KV state for the given tokens.	#L748-L977
MemoryAwarePrefixCache.store	method	MemoryAwarePrefixCache.store(tokens: list[int], cache: list[Any], evict_prefixes: bool = True) -> bool	Store KV cache for future reuse.	#L979-L1092
MemoryAwarePrefixCache._remove_from_sorted	method	MemoryAwarePrefixCache._remove_from_sorted(key: tuple[int, ...]) -> None	Remove a key from the sorted index using bisect for O(log N).	#L1094-L1098
MemoryAwarePrefixCache._evict_lru	method	MemoryAwarePrefixCache._evict_lru() -> None	Evict the least recently used entry.	#L1100-L1126
MemoryAwarePrefixCache.remove	method	MemoryAwarePrefixCache.remove(tokens: list[int]) -> bool	Remove a specific cache entry.	#L1128-L1147
MemoryAwarePrefixCache.clear	method	MemoryAwarePrefixCache.clear() -> None	Clear all cached entries.	#L1149-L1156
MemoryAwarePrefixCache.get_stats	method	MemoryAwarePrefixCache.get_stats() -> dict[str, Any]	Get cache statistics.	#L1158-L1160
MemoryAwarePrefixCache.reset_stats	method	MemoryAwarePrefixCache.reset_stats() -> None	Reset statistics while preserving cache contents.	#L1162-L1169
MemoryAwarePrefixCache.memory_usage_mb	method	MemoryAwarePrefixCache.memory_usage_mb() -> float	Current memory usage in MB.	#L1172-L1174
MemoryAwarePrefixCache.memory_limit_mb	method	MemoryAwarePrefixCache.memory_limit_mb() -> float	Memory limit in MB.	#L1177-L1179
MemoryAwarePrefixCache.try_reserve_memory	method	MemoryAwarePrefixCache.try_reserve_memory(nbytes: int) -> bool	Tentatively reserve cache memory for an upcoming promotion.	#L1181-L1188
MemoryAwarePrefixCache.release_reserved_memory	method	MemoryAwarePrefixCache.release_reserved_memory(nbytes: int) -> None	Release memory previously reserved by try_reserve_memory().	#L1190-L1194
MemoryAwarePrefixCache.__len__	method	MemoryAwarePrefixCache.__len__() -> int	Return number of cached entries.	#L1196-L1198
MemoryAwarePrefixCache.__contains__	method	MemoryAwarePrefixCache.__contains__(tokens: list[int]) -> bool	Check if tokens are cached.	#L1200-L1202
MemoryAwarePrefixCache.set_ssd_tier	method	MemoryAwarePrefixCache.set_ssd_tier(ssd_tier) -> None	Attach an SSD cache tier for eviction spilling.	#L1204-L1214
MemoryAwarePrefixCache.check_ssd	method	MemoryAwarePrefixCache.check_ssd(tokens: list[int]) -> dict \| None	Check if tokens have an SSD cache hit (without reading data).	#L1216-L1249
MemoryAwarePrefixCache.save_to_disk	method	MemoryAwarePrefixCache.save_to_disk(cache_dir: str) -> bool	Save all cache entries to disk using mlx_lm's safetensors format.	#L1255-L1346
MemoryAwarePrefixCache.load_from_disk	method	MemoryAwarePrefixCache.load_from_disk(cache_dir: str) -> int	Load cache entries from disk.	#L1348-L1463