Is this page helpful?

core.inference.contexts.mamba_slot_allocator#

Module Contents#

Classes#

MambaSlotAllocator

Owns the Mamba prefix cache pool: cached conv/ssm state plus the block-to-slot map.

API#

exception core.inference.contexts.mamba_slot_allocator.MambaSlotCapacityError(required: int, available: int)#

Bases: RuntimeError

Raised when the durable Mamba cache cannot satisfy an allocation.

Initialization

Initialize self. See help(type(self)) for accurate signature.

class core.inference.contexts.mamba_slot_allocator.MambaSlotAllocator(
context: core.inference.contexts.dynamic_context.DynamicInferenceContext,
max_slots: int,
num_mamba_layers: int,
conv_states_shape: tuple,
ssm_states_shape: tuple,
conv_states_dtype: torch.dtype,
ssm_states_dtype: torch.dtype,
prefix_cache_registry: core.inference.contexts.prefix_cache_registry.PrefixCacheRegistry,
)#

Owns the Mamba prefix cache pool: cached conv/ssm state plus the block-to-slot map.

Constructed only when prefix caching is enabled and prefix_caching_mamba_gb > 0. The host hash -> block_id map lives on the shared PrefixCacheRegistry; the intermediate-state extraction buffers and per-request commit bookkeeping live on PrefixCachedMambaMetadata. This allocator owns just the cache pool resources and the commit orchestration.

Parameters:

  • context – The DynamicInferenceContext that owns this allocator.

  • max_slots – Maximum number of cache slots.

  • num_mamba_layers – Number of Mamba layers in the model.

  • conv_states_shape – Shape of per-slot conv state (excluding layer/slot dims).

  • ssm_states_shape – Shape of per-slot SSM state (excluding layer/slot dims).

  • conv_states_dtype – Dtype for conv state tensors.

  • ssm_states_dtype – Dtype for SSM state tensors.

  • prefix_cache_registry – Host hash registry; the Mamba evict callback is wired here.

Initialization

allocate_slots_batch(block_ids: list) → list#

Get free Mamba cache slots for multiple blocks, evicting if necessary.

Handles deduplication: if the same block_id appears multiple times, only one slot is allocated and all occurrences get the same slot.

Parameters:

block_ids – List of KV block IDs to associate with slots.

Returns:

List of allocated slot indices (same length as block_ids).

_evictable_block_ids() → torch.Tensor#

Return blocks whose durable Mamba slots have no live KV owner.

_evict_lru_slots_batch(
num_needed: int,
candidate_ids: torch.Tensor,
) → list#

Evict the least recently used Mamba cache slots.

Does NOT return slots to the free pool — caller takes ownership.

Parameters:

  • num_needed – Number of slots to evict.

  • candidate_ids – Blocks confirmed to be evictable for this allocation.

Returns:

List of freed slot indices.

get_slot(block_id: int) → int#

Return the cache slot for a block, or -1 if none.

Parameters:

block_id – The KV block ID.

Returns:

Slot index or -1.

has_state(block_id: int) → bool#

Check if a block has cached Mamba state.

invalidate_block(block_id: int) → None#

Free cache slot and clear mappings for a block.

Parameters:

block_id – The KV block ID.

_invalidate_blocks_batch(block_ids_list: list) → None#

Free cache slots and clear mappings for multiple blocks at once.

Vectorized version of invalidate_block that avoids per-block .item() GPU syncs. Used by _on_mamba_evicted for bulk slot release.

Parameters:

block_ids_list – List of block IDs to invalidate.

_on_mamba_evicted(block_ids: list) → None#

Registry callback: free GPU slots for the given block IDs.

Fired by PrefixCacheRegistry after the host Mamba dict drops entries: either as a cascade from KV eviction (KV blocks gone means their cached Mamba state is also gone) or from an explicit registry.evict_mamba call.

_invalidate_blocks_batch no-ops on blocks whose block_to_slot is already -1, which lets _evict_lru_slots_batch suppress the free-pool return by clearing mappings before firing the registry call.

store_from_tensors(
block_id: int,
layer_idx: int,
ssm_state: torch.Tensor,
conv_state: torch.Tensor,
) → None#

Write provided state tensors to a cache slot for a specific layer.

Parameters:

  • block_id – The KV block ID.

  • layer_idx – The Mamba layer index.

  • ssm_state – SSM state tensor to store.

  • conv_state – Conv state tensor to store.

store_from_live_batch(slots: list, request_indices: list) → None#

Copy all layers from live per-request buffers to cache slots.

Parameters:

  • slots – List of cache slot indices.

  • request_indices – List of context request indices.

restore_to_live(request_idx: int, block_id: int) → bool#

Copy all layers from cache slot to live request state.

Parameters:

  • request_idx – The context request index.

  • block_id – The KV block ID.

Returns:

True if state was restored, False if block has no cached state.

register_block_hashes_batch(block_ids: list, hashes: list) → None#

Register multiple blocks as having cached Mamba state.

Only registers entries where hash > 0.

Parameters:

  • block_ids – List of block IDs.

  • hashes – List of hash values (same length as block_ids).

commit_intermediate_states() → None#

Move intermediate states recorded this step into the cache pool.

Called after the forward pass (including CUDA graph replay) completes. Pulls per-request commit data from context.mamba_metadata, allocates slots for the new entries, copies the GPU intermediates into the cache pool, copies live EOS state where applicable, and registers Mamba hashes.

_collect_commit_data()#

Pull the metadata’s commit records and the KV-side block hashes.

Returns:

Tuple of (intermediate_bids, src_offsets, eos_bids, eos_ctx_indices, all_hashes) or None if nothing to commit. all_hashes covers intermediate_bids + eos_bids in that order.

_copy_intermediate_to_cache(src_offsets: list, slots: list) → None#

Copy intermediate states from metadata’s GPU buffers to cache slots.

Uses fancy-indexed GPU D2D copy (2 kernel launches instead of 2N).

Parameters:

  • src_offsets – Source indices into metadata’s intermediate buffers.

  • slots – Destination cache slot indices.

reset() → None#

Reset slot pool, cache state, the Mamba host registry, and the metadata’s buffers.