`core.inference.contexts.kv_block_allocator`#

Module Contents#

Classes#

KVBlockAllocator

Allocator that manages blocks of memory for the KV cache.

Data#

BlocksDeregisteredObserver

API#

core.inference.contexts.kv_block_allocator.BlocksDeregisteredObserver#: None

class core.inference.contexts.kv_block_allocator.KVBlockAllocator( context: DynamicInferenceContext, pool_size: int, paused_limit: int, pc_state: Optional[core.inference.contexts.prefix_cache_block_state.PrefixCacheBlockState] = None, prefix_cache_registry: Optional[core.inference.contexts.prefix_cache_registry.PrefixCacheRegistry] = None, )#

Allocator that manages blocks of memory for the KV cache.

This allocator owns:

The free-pool stack (block_bag, pool_avail).
Allocation, release, retain, and reset orchestration.
The MoE routing-replay per-block storage.

Parameters:

context (DynamicInferenceContext) – Dynamic inference context.
pool_size (int) – Number of blocks in the pool, including the dummy block.
paused_limit (int) – Paused-request block retention limit. Must leave at least one non-dummy block outside the limit.
pc_state (Optional[PrefixCacheBlockState]) – Per-block prefix-caching state. None disables prefix caching entirely on this allocator.
prefix_cache_registry (Optional[PrefixCacheRegistry]) – Host hash registry shared with the Mamba allocator. Given exactly when pc_state is given.

Initialization

property enable_prefix_caching: bool#: True when this allocator carries prefix-caching state.

__str__()#

get_total_used()#: Compute number of physical blocks outside the free pool.

get_active_used()#: Compute number of active blocks used.

get_paused_used()#: Compute number of paused blocks used.

is_memory_available( num_blocks: int, potential_matched_count: int = 0, ) → bool#

Check if memory blocks are available.

Includes both free pool blocks and registered, evictable cached blocks.

Parameters:

num_blocks (int) – Number of blocks to check.
potential_matched_count (int) – Number of currently-evictable cached blocks to subtract from the evictable count because the caller will pin them before allocating (e.g. prefix-matched blocks that get their ref counts bumped in add_request). These blocks are ref_count == 0 now, so they are included in the evictable count, but they will be protected from eviction, so they cannot supply the requested num_blocks.

Returns:

(bool) Is memory available?

allocate_memory_blocks( num_blocks: int, ) → Optional[torch.Tensor]#

Allocate memory blocks if available, else return None.

Under LRU prefix caching, falls back to evicting cached blocks when the free pool is short. Returns None when even eviction cannot satisfy the request.

Parameters:: num_blocks (int) – Number of blocks to allocate.
Returns:: (Optional[Tensor]) Allocated block IDs.

release_memory_blocks(blocks: torch.Tensor) → None#

Release memory blocks.

Without prefix caching: blocks return directly to the free pool. With prefix caching: one reference per occurrence is dropped, and blocks that reach zero are released according to the eviction policy (REF_ZERO deregisters them at once; LRU keeps registered blocks cached and returns only unregistered ones).

Parameters:: blocks (Tensor) – Block IDs to release.

retain_memory_blocks(block_ids: list[int]) → None#

Add one prefix-cache reference to each block.

Parameters:: block_ids – Blocks retained by a new owner.

reset() → None#

Reset the allocator to initial state.

This resets the available block count to the entire memory pool (except for the dummy block).

_push_to_pool(blocks: torch.Tensor) → None#: Push blocks back onto the free-pool stack.

_notify_deregistered( block_ids: torch.Tensor, hashes: list[int], ) → None#: Drop deregistered hashes from the registry, then notify the observers.

register_kv_block_hashes( block_ids: list[int], block_hashes: list[int], parent_hashes: Optional[list[int]] = None, ) → list[int]#

Registration is idempotent: a block that already carries the hash being registered is skipped. Callers may legitimately re-offer an already registered block (a cache-matched block whose block-table slot a later prefill chunk also spans), and the bookkeeping below is one-shot per block – applying it twice adds a second child entry to the block’s parent that no deregistration can ever cancel, leaving that parent permanently short of child_count == 0 and therefore never an evictable leaf (see PrefixCacheBlockState.find_lru_evictable).

Re-registering a live block under a different hash would instead overwrite its recorded parent while leaving the previous parent’s child count raised, so that case is rejected rather than absorbed.

This method never touches reference counts. New blocks are pinned at ref_count == 1 by allocate_memory_blocks, and additional owners of an already registered block are pinned by the caller that matched it.

Parameters:

block_ids – List of block IDs.
block_hashes – List of computed hash values (same length as block_ids).
parent_hashes – Parent hash for each block in the prefix chain (same length as block_ids); 0 marks a root block with no parent. Used by LRU eviction to avoid evicting a parent before its children. If None, parents default to 0.

Returns:

Newly registered block IDs, in input order. Already registered blocks are excluded so callers can preserve their existing metadata.

add_blocks_deregistered_observer( observer: core.inference.contexts.kv_block_allocator.BlocksDeregisteredObserver, ) → None#

Currently used only by DynamoHelper.

get_evictable_block_count() → torch.Tensor#

Get count of cached blocks that can be evicted (ref_count == 0, hash set).

Returns:: Scalar tensor with the number of evictable cached blocks.

get_allocatable_count() → int#

Compute the number of blocks available for allocation.

Includes both blocks in the free pool and, under LRU prefix caching, registered ref-zero blocks that can be evicted.

Returns:: Number of blocks that can currently be allocated.

evict_lru_blocks(num_blocks_needed: int) → bool#

Evict LRU cached blocks to free up space in the pool.

Parameters:: num_blocks_needed – Number of blocks to evict.
Returns:: True if enough blocks were evicted, False otherwise (also False without prefix caching or under REF_ZERO, which keep no evictable reservoir).

store_routing_per_block( flat_routing: numpy.ndarray, block_ids: numpy.ndarray, positions: numpy.ndarray, ) → None#

Scatter flat routing indices into per-block storage.

Matched prefix-cache blocks already contain routing from the original request and are absent from this forward’s token layout.

Parameters:

flat_routing – ndarray of shape [active_token_count, num_layers, topk] aligned with the supplied token layout.
block_ids – KV-block ID for each routing row.
positions – Position within the KV block for each routing row.

reconstruct_routing_from_blocks( block_ids: list[int], total_routing_tokens: int, ) → Optional[numpy.ndarray]#

Reconstruct routing indices from per-block storage.

Concatenates per-block routing ndarrays in block order, trimming the last block to exactly total_routing_tokens entries.

Parameters:

block_ids – Ordered list of block IDs for the request.
total_routing_tokens – Expected number of routing tokens (total_tokens - 1, since the last generated token has no forward-pass routing).

Returns:

ndarray [total_routing_tokens, num_layers, topk] or None if any required block is missing routing data.

store_block_routing( block_id: int, positions: numpy.ndarray, routing: numpy.ndarray, ) → None#

Store routing indices for specific token positions in a block.

Parameters:

block_id – The block ID.
positions – ndarray of token positions within the block (1D, int).
routing – ndarray of routing data [num_positions, num_layers, topk].

get_block_routing(block_id: int) → Optional[numpy.ndarray]#

Get routing indices for a block.

Parameters:: block_id – The block ID.
Returns:: ndarray [block_size_tokens, num_layers, topk] or None if not stored.

store_prompt_logprobs( block_ids: numpy.ndarray, slots: numpy.ndarray, log_probs: numpy.ndarray, top_n_log_probs: numpy.ndarray, top_n_token_ids: numpy.ndarray, next_token_ids: numpy.ndarray, ) → None#

Store prompt log-prob rows on the blocks that hold their source tokens.

A block keeps one top-N width. A wider write replaces it, and a narrower write is dropped, so a block’s rows only ever cover a superset of what an earlier writer requested.

Parameters:

block_ids – Block ID of each row, shape [R].
slots – Source position of each row within its block, shape [R].
log_probs – Selected-token log-prob of each row, shape [R].
top_n_log_probs – Top-N log-probs of each row, shape [R, top_n].
top_n_token_ids – Top-N token IDs of each row, shape [R, top_n].
next_token_ids – Token each row scores, used only for a block’s final slot.

match_prompt_logprob_blocks( block_ids: list[int], first_block: int, prompt_tokens: torch.Tensor, top_n: int, ) → int#

Count leading blocks whose cached prompt log-probs a request can reuse.

Parameters:

block_ids – Matched block IDs, in prefix order.
first_block – Request-relative index of the first matched block.
prompt_tokens – The request’s full prompt.
top_n – Top-N width the request needs.

Returns:

Number of leading blocks with every row the request needs.

Return type:

int

read_prompt_logprobs( block_ids: list[int], first_block: int, start: int, stop: int, top_n: int, ) → Tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]#

Read cached prompt log-prob rows for source positions [start, stop).

Parameters:

block_ids – Block IDs starting at request-relative block first_block.
first_block – Request-relative index of block_ids[0].
start – First source position (inclusive).
stop – Last source position (exclusive).
top_n – Number of top alternatives to return per row.

Returns:

Selected log-probs [N], top-N log-probs [N, top_n], and top-N token IDs.

core.inference.contexts.kv_block_allocator#

Module Contents#

Classes#

Data#

API#

`core.inference.contexts.kv_block_allocator`#