Is this page helpful?

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,
enable_prefix_caching: bool = False,
prefix_caching_eviction_policy: megatron.core.inference.config.PrefixCachingEvictionPolicy = PrefixCachingEvictionPolicy.REF_ZERO,
)#

Allocator that manages blocks of memory for the KV cache.

This allocator is responsible for:

  • Initializing a pool of block IDs

  • Allocating blocks from the pool

  • Releasing blocks back to the pool

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.

Initialization

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

Will attempt LRU eviction of cached blocks if the free pool is insufficient.

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 by decrementing reference counts.

Blocks with ref_count == 0 remain cached (in hash map) for potential reuse. They will be evicted via LRU when space is needed.

Parameters:

blocks (Tensor) – Block IDs to release.

Returns:

None

reset() → None#

Reset the allocator to initial state.

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

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

Register blocks in the hash-to-block mapping for discovery (batch).

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 evict_lru_blocks).

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.

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

Register a callback invoked when cached blocks are deregistered.

Currently used only by DynamoHelper.

_deregister_blocks(block_ids: torch.Tensor) → None#

Remove blocks from prefix caching state and return to free pool.

Shared cleanup logic for both LRU eviction and RZ proactive eviction.

Parameters:

block_ids – Tensor of block IDs to deregister.

update_timestamps(block_ids: torch.Tensor) → None#

Update LRU timestamps for accessed blocks. No-op in RZ mode.

Parameters:

block_ids – Tensor of block IDs that were accessed.

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.

Evicts blocks with ref_count == 0, least-recently-used first, while never evicting a parent before its children. Block hashes are parent-chained, and _find_kv_match_count relies on the invariant that a cached child block always has all of its ancestors cached too. A naive oldest-first eviction breaks this: with chunked prefill, earlier chunks are allocated first (older timestamps) yet are ancestors of later chunks (newer timestamps), so once the request finishes and its blocks are cached, an ancestor can be older than its descendant and get evicted first, leaving a dangling child.

To preserve the invariant while staying optimal we peel the cached forest from its leaves inward with a min-heap: only a leaf (a cached block with no cached children) is ever evictable, and among the currently-evictable leaves we always take the one with the oldest own timestamp. Evicting a leaf can turn its parent into a leaf, which is then pushed onto the heap. Repeating num_blocks_needed times gives, at each step, the globally least-recently-used block that can be removed without orphaning a child — the natural generalization of LRU to the parent-chain constraint. Keying each block by its own recency (and only reconsidering a parent once its children are gone) is what makes this optimal: a block is retained purely because it is recently used, never because a hot descendant props it up, so a colder evictable block is always evicted before a hotter one.

Worked example, evicting 3 from::

A(ts 1) -> B(ts 2) -> C(ts 5)   (C, F are leaves under B)
                  +-> F(ts 3)
        +-> D(ts 3) -> E(ts 5)   (E is a leaf under D)

Leaf-peel evicts F(3), then C(5); B is now childless so it joins the leaves with its own ts=2 and is evicted next -> retains {A, D, E}, keeping the hottest block E(5) rather than the colder interior block B(2).

Note: because a request holds a contiguous block prefix [0..k], any in-use (ref_count > 0) block keeps all of its ancestors in use too. Hence a cached (ref_count == 0) block can only have cached children, and considering the cached set alone is sufficient to avoid dangling children.

The parent block id of each block and its live child count are maintained incrementally on register/deregister (block_parent_id / block_child_count), so this method reads the prefix forest directly rather than rebuilding it from hashes with a per-eviction sort. Only the inherently-sequential leaf peel below is per-element.

The parent graph is assumed acyclic (a forest), which holds for any hashes produced by the prefix-chain builder; an assertion guards against a pathological hash collision wedging the peel.

Parameters:

num_blocks_needed – Number of blocks to evict.

Returns:

True if enough blocks were evicted, False otherwise.

store_routing_per_block(
flat_routing: Optional[numpy.ndarray],
) → None#

Scatter flat routing indices into per-block storage.

Uses the context’s token-to-block mapping to distribute each token’s routing data into the appropriate block. Matched (prefix-cached) blocks already have routing from the original request and are not overwritten here since their tokens are not in the active token layout.

Parameters:

flat_routing – ndarray of shape [active_token_count, num_layers, topk] aligned with the context’s active-token layout, or None.

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