Is this page helpful?

core.inference.contexts.prefix_cache_registry#

Module Contents#

Classes#

PrefixCacheRegistry

Pure-CPU prefix-cache hash registry.

API#

class core.inference.contexts.prefix_cache_registry.PrefixCacheRegistry#

Pure-CPU prefix-cache hash registry.

Owns the host-side hash -> block_id mappings for KV cache blocks and Mamba state blocks, and exposes the match / register / evict primitives that drive prefix caching.

The interface is deliberately torch-free: only Python int / list / set cross the boundary.

Physical resources (the per-block shadow tensors and ref counts in PrefixCacheBlockState, the slot pool and state tensors on the Mamba allocator) stay where they are; their owners call into this registry at register / evict time to keep the host view consistent with their physical state.

Initialization

set_mamba_evict_callback(
cb: Callable[[List[int]], None],
) → None#

Wire a callback fired when Mamba dict entries are evicted.

The callback receives the list of block IDs whose Mamba state should be released. Used by the Mamba allocator to free GPU slots; remote deployments leave this unset.

match_kv_prefix(
hashes: List[int],
) → Tuple[List[int], int]#

Longest cached prefix of hashes.

Parent-chained hashing guarantees that if hashes[N] is cached, all earlier hashes are too. Backwards scan finds the longest cached prefix in one pass.

Returns (matched_block_ids, parent_hash) parent_hash is the hash of the last matched block, 0 if no matches.

register_kv(
block_ids: List[int],
hashes: List[int],
) → None#

Bulk add (hash -> block_id) entries.

evict_kv(hashes: Iterable[int]) → set#

Drop entries whose hashes appear in hashes (filtering -1).

Cascades to the Mamba dict: a deregistered KV block can never retain valid Mamba state. Returns the set of hashes that were actually removed from the KV dict.

match_mamba_farthest(hashes: List[int]) → int#

Index (one past) the farthest cached Mamba block in hashes.

Mamba state is cumulative; only the farthest cached block needs to be restored. Backwards scan finds it in one pass. Returns 0 if nothing matches.

find_mamba_backoff(hashes: List[int], search_len: int) → int#

Farthest cached Mamba block within hashes[:search_len].

Used when a raw Mamba match overruns the chunk and the caller needs to back off to a previous cached state. match_mamba_farthest scans the whole list; this scans a bounded prefix.

register_mamba(
block_ids: List[int],
hashes: List[int],
) → None#

Bulk add (hash -> block_id) entries; skips hash <= 0.

evict_mamba(hashes: Iterable[int]) → set#

Drop entries whose hashes appear in hashes (filtering -1).

Returns the set of hashes that were actually removed. Fires the Mamba-evict callback so the allocator can free GPU slots.

_cascade_mamba_evict(hashes: set) → None#

Internal: KV eviction implies Mamba eviction for the same hashes.

clear_kv() → None#

Drop all KV entries.

clear_mamba() → None#

Drop all Mamba entries.

reset() → None#

Drop all entries (KV and Mamba).