core.inference.contexts.prefix_cache_block_state

Module Contents

Classes

PromptLogprobsBlock	Prompt log-probs cached on one KV block, indexed by source token position.
PrefixCacheBlockState	Per-block CPU shadow state for prefix caching.

Data

MAX_CACHED_PROMPT_TOP_N_LOGPROBS

API

core.inference.contexts.prefix_cache_block_state.MAX_CACHED_PROMPT_TOP_N_LOGPROBS

100

class core.inference.contexts.prefix_cache_block_state.PromptLogprobsBlock(block_size: int, top_n: int)

Prompt log-probs cached on one KV block, indexed by source token position.

Slot i holds the scores computed at the block’s token i, i.e. the distribution over the token that follows it. The final slot scores the first token of the next block, which this block’s hash does not cover, so that token is recorded in next_token_id.

Parameters:

• block_size (int) – Tokens per KV block.

• top_n (int) – Number of top alternatives stored per slot.

Initialization

__slots__

(‘top_n’, ‘log_probs’, ‘top_n_log_probs’, ‘top_n_token_ids’, ‘valid’, ‘next_token_id’)

class core.inference.contexts.prefix_cache_block_state.PrefixCacheBlockState(

pool_size: int,

eviction_policy: megatron.core.inference.config.PrefixCachingEvictionPolicy,

)

Per-block CPU shadow state for prefix caching.

Owns the per-block bookkeeping (hashes, ref counts, MTP successor tokens, cached prompt log-probs, LRU timestamps and the LRU parent chain) and exposes mutation primitives that return the hashes a caller must forward to PrefixCacheRegistry. The state class itself is registry-free so that the registry boundary stays the only seam between in-process and remote prefix-cache deployments.

Initialization

property is_lru: bool

True when the eviction policy is LRU (block_timestamps is allocated).

reset() → None

Reset all per-block state. Caller is responsible for clearing the registry.

stamp_block_hashes(

block_ids: List[int],

block_hashes: List[int],

) → List[int]

Write block_hashes into the per-block shadow tensor.

Returns:

Positions into the inputs of the blocks that were newly stamped, in input order.

record_parent_chain(

block_ids: List[int],

parent_ids: List[int],

) → None

Persist the LRU-eviction parent chain for freshly-registered blocks.

parent_ids[k] is the block id of block_ids[k]’s parent (resolved from the parent hash by the caller via the registry), or -1 for a root or not-yet-registered parent. Bumps each resolved parent’s child count so a parent only becomes an evictable leaf once its last child is gone. No-op unless LRU; see find_lru_evictable for how the chain constrains eviction.

update_timestamps(block_ids: torch.Tensor, lru_clock: int) → None

Stamp block_timestamps[block_ids] = lru_clock. No-op in REF_ZERO mode.

on_allocate(block_ids: torch.Tensor, lru_clock: int) → None

Initialize ref counts (and timestamps under LRU) for newly-allocated blocks.

Called by the allocator immediately after popping block_ids from the free pool.

retain(block_ids: torch.Tensor, lru_clock: int) → None

Add one reference per occurrence in block_ids (and refresh LRU timestamps).

get_evictable_block_count() → torch.Tensor

Count of blocks that are cached (ref_count == 0) and have a registered hash.

extra_blocks_available() → int

How many cached blocks could be made available via eviction.

Zero under REF_ZERO: in that policy a block whose ref count hits zero is immediately deregistered and pushed to the pool; the cached reservoir is always empty by construction. Equals the evictable count under LRU.

try_lru_evict_for_pool(

num_blocks_needed: int,

) → Optional[Tuple[torch.Tensor, List[int]]]

Pick + reset the LRU-oldest victims; return them plus their hashes.

Returns:

• None under REF_ZERO (no cached reservoir to evict from).

• None under LRU when fewer than num_blocks_needed evictable blocks exist.

• (victims, hashes) on success: victims is a tensor of exactly num_blocks_needed block IDs, hashes is the list of hashes the caller must drop from the registry. The caller (the allocator) is also responsible for pushing victims onto the free pool.

find_lru_evictable(

num_blocks_needed: int,

) → Optional[torch.Tensor]

Return exactly num_blocks_needed evictable block IDs, or None.

Least-recently-used first, but never a parent before its children. Block hashes are parent-chained and prefix matching relies on the invariant that a cached child always has all of its ancestors cached too; a naive oldest-first eviction breaks it (with chunked prefill an ancestor chunk can be allocated first, hence older, yet outlive a younger descendant). We peel the cached forest from its leaves inward with a min-heap: only a leaf (a cached block with no cached children) is evictable, and among evictable leaves we 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. Keying each block by its own recency (and reconsidering a parent only once its children are gone) is what makes this the optimal LRU generalization.

None if fewer than num_blocks_needed evictable blocks exist. Callers must push the returned blocks back onto the free pool via the allocator after deregister_blocks clears their state.

deregister_blocks(block_ids: torch.Tensor) → List[int]

Reset per-block state for block_ids; return their hashes.

Does NOT touch the host hash registry and does NOT push the blocks back to the free pool – both are the caller’s responsibility.

Used by both the LRU eviction path and the REF_ZERO release path.

on_release_compute_pool_returns(

blocks: torch.Tensor,

) → Tuple[torch.Tensor, List[int]]

Drop one reference per occurrence; return (blocks_for_pool, hashes_to_drop).

Duplicate ids (a shared block released by several finishing owners in one batch) each drop one reference, but a block reaching zero is returned at most once.

Policy-specific:

• REF_ZERO: every block whose ref count reaches zero is reset; both its block ID and its hash are returned for the caller to push to the pool and drop from the registry, respectively.

• LRU: only unregistered (hash == -1) zero-ref blocks are returned for the pool; their hashes are empty (already -1). Blocks with a registered hash stay cached for later reuse and are evicted on demand.

The allocator pushes blocks_for_pool onto the block bag and forwards hashes_to_drop to the registry.