diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 258917d07ab..59e8d392c1d 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -2929,14 +2929,27 @@ def check_availability(self, req: DynamicInferenceRequest) -> Tuple[bool, bool, self.total_request_count < self.max_requests and self.paused_request_count == 0 ) - (_, num_blocks_from_pool, _, _, _, effective_prefill_chunk_length) = ( + (matched_block_ids, num_blocks_from_pool, _, _, _, effective_prefill_chunk_length) = ( self._compute_prefix_match(req, req.remaining_prompt_length) ) request_tokens_can_be_added = ( self.active_token_count + effective_prefill_chunk_length <= self.max_tokens ) - kv_cache_available = self.kv_block_allocator.is_memory_available(num_blocks_from_pool) + # add_request pins the matched blocks before allocating. Only matches that + # are currently evictable (ref_count == 0) count against the evictable + # pool; matches already pinned by another in-flight request are not in + # get_evictable_block_count() and pinning them frees nothing. Reserve only + # the ref_count == 0 matches so availability is not under-reported. + potential_matched_count = 0 + if matched_block_ids: + matched_tensor = torch.tensor(matched_block_ids, dtype=torch.int32, device='cpu') + potential_matched_count = int( + (self.kv_block_allocator.block_ref_counts[matched_tensor] == 0).sum() + ) + kv_cache_available = self.kv_block_allocator.is_memory_available( + num_blocks_from_pool, potential_matched_count=potential_matched_count + ) return request_can_be_added, request_tokens_can_be_added, kv_cache_available def _find_kv_match_count( @@ -3037,19 +3050,31 @@ def add_request( # Slice tokens to skip matched prefix this_round_tokens = req.remaining_prompt_tokens[prefix_skip_tokens:prefill_chunk_length] - new_block_ids = None - if num_blocks_from_pool > 0: - new_block_ids = self.kv_block_allocator.allocate_memory_blocks(num_blocks_from_pool) - if new_block_ids is None or len(new_block_ids) != num_blocks_from_pool: - raise BlockOverflowError(req.request_id) - - # Increment ref counts and update timestamps for matched (shared) blocks + # Pin matched (shared) blocks BEFORE allocation. allocate_memory_blocks() + # may trigger LRU eviction, and a matched block still at ref_count == 0 + # would be an eviction candidate — descendant-first LRU could evict a + # matched leaf and immediately reuse its ID for a new block, leaving the + # block table with a duplicate ID and a dangling parent. Incrementing ref + # counts first removes matched blocks from the evictable set (see + # evict_lru_blocks / get_evictable_block_count), so eviction falls back to + # genuinely unused cached blocks. + matched_tensor = None if num_matched_blocks > 0: matched_tensor = torch.tensor(matched_block_ids, dtype=torch.int32, device='cpu') self.kv_block_allocator.block_ref_counts[matched_tensor] += 1 if self.prefix_caching_eviction_policy == PrefixCachingEvictionPolicy.LRU: self.kv_block_allocator.update_timestamps(matched_tensor) + new_block_ids = None + if num_blocks_from_pool > 0: + new_block_ids = self.kv_block_allocator.allocate_memory_blocks(num_blocks_from_pool) + if new_block_ids is None or len(new_block_ids) != num_blocks_from_pool: + # Roll back the pin so a failed add does not leak ref counts on + # the matched blocks (which would make them permanently unevictable). + if matched_tensor is not None: + self.kv_block_allocator.block_ref_counts[matched_tensor] -= 1 + raise BlockOverflowError(req.request_id) + # Note that we decremented the total_request_count for the chunked prefill request # in update_requests, so setting current_id to the total_request_count will again # make the last request the continuing chunked prefill request if one exists. @@ -3151,8 +3176,14 @@ def _register_range(start: int, end: int): return block_ids_to_hash = self.request_to_kv_block_ids[current_id][start:end].tolist() block_hashes_slice = req.precomputed_block_hashes[start:end] + # Parent hash of block k is the hash of block k-1 in the chain; + # block 0 is a root (parent hash 0). Enables LRU eviction to keep + # parents cached until their children are gone. + parent_hashes_slice = [ + req.precomputed_block_hashes[k - 1] if k > 0 else 0 for k in range(start, end) + ] self.kv_block_allocator.register_kv_block_hashes( - block_ids_to_hash, block_hashes_slice + block_ids_to_hash, block_hashes_slice, parent_hashes_slice ) # Range 1: prior-chunk partial block that this chunk just completed diff --git a/megatron/core/inference/contexts/kv_block_allocator.py b/megatron/core/inference/contexts/kv_block_allocator.py index d555c925c93..6711cb2e606 100644 --- a/megatron/core/inference/contexts/kv_block_allocator.py +++ b/megatron/core/inference/contexts/kv_block_allocator.py @@ -1,5 +1,6 @@ # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +import heapq from collections import deque from typing import Callable, Dict, Optional @@ -70,6 +71,24 @@ def __init__( (self.total_count,), dtype=torch.int64, device='cpu' ) + # Persisted prefix-chain bookkeeping for LRU eviction, maintained + # incrementally on register/deregister. Block hashes are + # parent-chained: a cached block that is another cached block's + # parent must not be evicted before its child (see evict_lru_blocks). + # + # block_parent_id[b] = block id of b's parent in the prefix chain, + # or -1 when b is a root block or its parent is not registered. + self.block_parent_id = torch.full( + (self.total_count,), -1, dtype=torch.int64, device='cpu' + ) + # block_child_count[b] = number of currently-registered children of b. + # For a cached block all of its children are cached too, so this + # equals its cached-child count and b is an evictable leaf exactly + # when it reaches 0. + self.block_child_count = torch.zeros( + (self.total_count,), dtype=torch.int64, device='cpu' + ) + # Per-block MoE routing storage (populated when routing replay is enabled) self.block_routing: Dict[int, np.ndarray] = {} @@ -128,13 +147,20 @@ def get_paused_avail(self): """Compute number of paused blocks available.""" return self.paused_count - self.get_paused_used() - def is_memory_available(self, num_blocks: int) -> bool: + def is_memory_available(self, num_blocks: int, potential_matched_count: int = 0) -> bool: """Check if memory blocks are available. Includes both free pool blocks and evictable cached blocks (ref_count == 0). Args: 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``. Return: (bool) Is memory available? @@ -146,8 +172,8 @@ def is_memory_available(self, num_blocks: int) -> bool: return False if self.prefix_caching_eviction_policy == PrefixCachingEvictionPolicy.REF_ZERO: return False # RZ: no cached blocks to evict - # Also count evictable cached blocks - evictable_count = self.get_evictable_block_count() + # Also count evictable cached blocks, excluding those the caller will pin. + evictable_count = int(self.get_evictable_block_count()) - potential_matched_count return (self.total_avail + evictable_count) >= num_blocks def allocate_memory_blocks(self, num_blocks: int) -> Optional[Tensor]: @@ -256,6 +282,8 @@ def reset(self) -> None: self.block_ref_counts.fill_(0) if self.prefix_caching_eviction_policy == PrefixCachingEvictionPolicy.LRU: self.block_timestamps.fill_(0) + self.block_parent_id.fill_(-1) + self.block_child_count.fill_(0) # Clear per-block routing storage self.block_routing.clear() @@ -264,20 +292,54 @@ def reset(self) -> None: # Prefix caching methods # ========================================================================= - def register_kv_block_hashes(self, block_ids: list[int], block_hashes: list[int]) -> None: + def register_kv_block_hashes( + self, + 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). Args: 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. """ if not block_ids: return id_tensor = torch.tensor(block_ids, dtype=torch.int64, device=self.block_hashes.device) hash_tensor = torch.tensor(block_hashes, dtype=torch.int64, device=self.block_hashes.device) self.block_hashes[id_tensor] = hash_tensor + if parent_hashes is not None: + assert len(parent_hashes) == len(block_ids) + # Add the new blocks to the hash map first so that a block whose parent is + # elsewhere in this same batch (block k's parent is block k-1) resolves. self.kv_hash_to_block_id.update(zip(block_hashes, block_ids)) + if self.prefix_caching_eviction_policy == PrefixCachingEvictionPolicy.LRU: + # Persist the resolved parent block id and bump each parent's child count. + # Parents are earlier in the prefix chain and already registered + # (a matched block or a prior chunk / earlier entry in this batch), + # so a valid parent hash resolves; 0 marks a root and an unknown hash + # falls back to -1. + if parent_hashes is None: + parent_hashes = [0] * len(block_ids) + parent_ids = [ + self.kv_hash_to_block_id.get(ph, -1) if ph != 0 else -1 for ph in parent_hashes + ] + parent_id_tensor = torch.tensor(parent_ids, dtype=torch.int64, device=id_tensor.device) + self.block_parent_id[id_tensor] = parent_id_tensor + has_parent = parent_id_tensor >= 0 + if has_parent.any(): + self.block_child_count.scatter_add_( + 0, + parent_id_tensor[has_parent], + torch.ones(int(has_parent.sum()), dtype=torch.int64), + ) + def _deregister_blocks(self, block_ids: Tensor) -> None: """Remove blocks from prefix caching state and return to free pool. @@ -306,10 +368,23 @@ def _deregister_blocks(self, block_ids: Tensor) -> None: self.on_blocks_deregistered(block_ids.tolist(), keys_to_delete) # Reset block state (batched tensor ops) - self.block_hashes[block_ids] = -1 - self.block_ref_counts[block_ids] = 0 if self.prefix_caching_eviction_policy == PrefixCachingEvictionPolicy.LRU: + # Drop these blocks from their parents' child counts before clearing + # their own bookkeeping, keeping block_child_count in sync so a parent + # becomes an evictable leaf once its last child is deregistered. + parent_ids = self.block_parent_id[block_ids_i64] + has_parent = parent_ids >= 0 + if has_parent.any(): + self.block_child_count.scatter_add_( + 0, + parent_ids[has_parent], + torch.full((int(has_parent.sum()),), -1, dtype=torch.int64), + ) + self.block_parent_id[block_ids] = -1 + self.block_child_count[block_ids] = 0 self.block_timestamps[block_ids] = 0 + self.block_hashes[block_ids] = -1 + self.block_ref_counts[block_ids] = 0 # Return blocks to free pool self.block_bag[self.total_avail : self.total_avail + num_blocks] = block_ids @@ -340,7 +415,53 @@ def get_evictable_block_count(self) -> Tensor: def evict_lru_blocks(self, num_blocks_needed: int) -> bool: """Evict LRU cached blocks to free up space in the pool. - Evicts blocks with ref_count == 0, starting with oldest timestamps. + 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. Args: num_blocks_needed: Number of blocks to evict. @@ -352,14 +473,49 @@ def evict_lru_blocks(self, num_blocks_needed: int) -> bool: cached_mask = (self.block_ref_counts == 0) & (self.block_hashes != -1) cached_block_ids = torch.nonzero(cached_mask, as_tuple=True)[0] - if cached_block_ids.numel() < num_blocks_needed: + num_cached = cached_block_ids.numel() + if num_cached < num_blocks_needed: return False # Not enough cached blocks to evict + if num_blocks_needed <= 0: + return True - # Sort by timestamp (ascending = oldest first) - cached_timestamps = self.block_timestamps[cached_block_ids] - sorted_indices = torch.argsort(cached_timestamps) - blocks_to_evict = cached_block_ids[sorted_indices[:num_blocks_needed]] + ts = self.block_timestamps[cached_block_ids].tolist() + bid = cached_block_ids.tolist() + parent_global = self.block_parent_id[cached_block_ids].tolist() + child_count = self.block_child_count[cached_block_ids].tolist() + + # Map a cached block's global id to its local index so the peel can find a + # parent's slot to decrement. Parents that are not cached (root, or a + # parent still in use) are absent and are simply treated as peel roots. + global_to_local = {bid[i]: i for i in range(num_cached)} + parent_local = [global_to_local.get(p, -1) for p in parent_global] + + # Min-heap of currently-evictable leaves keyed by (own timestamp, block + # id). Block ids are unique, so the tie-break is total and deterministic. + heap = [(ts[i], bid[i], i) for i in range(num_cached) if child_count[i] == 0] + heapq.heapify(heap) + + evicted_local = [] + while heap and len(evicted_local) < num_blocks_needed: + _, _, i = heapq.heappop(heap) + evicted_local.append(i) + p = parent_local[i] + if p >= 0: + child_count[p] -= 1 + if child_count[p] == 0: + heapq.heappush(heap, (ts[p], bid[p], p)) + + # A forest is always fully peelable, so the heap always exposes enough + # leaves to collect num_blocks_needed (guaranteed by the num_cached >= + # num_blocks_needed check above). Falling short means the parent graph is + # cyclic — only possible under a hash collision, which we treat as a bug. + assert len(evicted_local) == num_blocks_needed, ( + f"leaf peel evicted {len(evicted_local)} of {num_blocks_needed} " + f"requested from {num_cached} cached blocks; parent graph is not a " + f"forest (likely a block-hash collision)" + ) + blocks_to_evict = cached_block_ids[torch.tensor(evicted_local, dtype=torch.int64)] self._deregister_blocks(blocks_to_evict) return True diff --git a/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py b/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py index 66d3a963786..c1f6bcbfa1d 100644 --- a/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py +++ b/tests/unit_tests/inference/contexts/test_dynamic_prefix_caching.py @@ -429,6 +429,113 @@ def test_ref_count_lru(self): for bid in active_blocks: assert alloc3.block_ref_counts[bid.item()].item() == 1 + @pytest.mark.internal + def test_add_request_full_cache_partial_hit_pins_matched_blocks(self): + """On a partial prefix hit against a FULL cache, the matched + blocks must be pinned before allocation so LRU eviction cannot reclaim + one of them for the new (non-matched) block. + + Scenario (mirrors the descendant-first LRU edge case): a cached chain + H0/S0 -> H1/S1 (older) plus an unrelated cached root HX/SX (newer) fill + the pool. An incoming prompt H0 -> H1 -> H2 matches [S0, S1] and needs one + new block. If S0/S1 are not pinned first, descendant-first LRU evicts the + older leaf S1 and immediately reuses it, yielding block_table [S0, S1, S1] + and a dangling H2 -> missing H1 chain. Correct behavior evicts SX and + yields [S0, S1, SX] with a contiguous H0 -> H1 -> H2 chain. + """ + ctx = self._ctx() + bs = ctx.block_size_tokens + alloc = ctx.kv_block_allocator + + # Cached chain H0/S0 -> H1/S1, seeded with an OLD timestamp. + ctx.prefix_cache_lru_clock = 1 + req_chain = self._req(ctx, self._prompt(bs * 2)) + ctx.add_request(req_chain) + s0, s1 = self._block_ids(ctx, 0, 2) + h0, h1 = req_chain.precomputed_block_hashes[0], req_chain.precomputed_block_hashes[1] + + # Unrelated cached root HX/SX, seeded with a NEWER timestamp, so a naive + # oldest-first / descendant-first eviction would prefer the chain leaf. + ctx.prefix_cache_lru_clock = 10 + ctx.add_request(self._req(ctx, self._prompt(bs, offset=9000), request_id=2)) + (sx,) = self._block_ids(ctx, 1, 1) + + # All three slots are distinct and now cached (ref_count drops to 0). + assert len({s0, s1, sx}) == 3 + ctx.release_memory_blocks_from_request_indexes(torch.tensor([0, 1])) + ctx.total_request_count = 0 + assert alloc.block_ref_counts[s0].item() == 0 + assert alloc.block_ref_counts[s1].item() == 0 + assert alloc.block_ref_counts[sx].item() == 0 + + # Force a full pool: the new block for H2 can only come from eviction. + alloc.total_avail = 0 + + # Incoming prompt H0 -> H1 -> H2: first two blocks match the cached chain, + # the third (H2) is new and must trigger a single eviction. + ctx.prefix_cache_lru_clock = 20 + req_new = self._req(ctx, self._prompt(bs * 3), request_id=3) + ctx.add_request(req_new) + h2 = req_new.precomputed_block_hashes[2] + + block_table = self._block_ids(ctx, 0, 3) + + # Matched blocks are preserved and SX (the unrelated root) is evicted/reused. + assert block_table == [s0, s1, sx] + # All three block IDs are distinct — no duplicate from a reclaimed match. + assert len(set(block_table)) == 3 + # Matched blocks stay pinned for the new request. + assert alloc.block_ref_counts[s0].item() == 1 + assert alloc.block_ref_counts[s1].item() == 1 + assert alloc.block_ref_counts[sx].item() == 1 + # Contiguous H0 -> H1 -> H2 hash chain over [S0, S1, SX]. + assert alloc.block_hashes[s0].item() == h0 + assert alloc.block_hashes[s1].item() == h1 + assert alloc.block_hashes[sx].item() == h2 + # Parent bookkeeping is stored as resolved block ids: S1's parent is S0 + # and SX's parent is S1 along the H0 -> H1 -> H2 chain. + assert alloc.block_parent_id[s1].item() == s0 + assert alloc.block_parent_id[sx].item() == s1 + assert alloc.kv_hash_to_block_id[h1] == s1 + + @pytest.mark.internal + def test_check_availability_excludes_already_pinned_matches(self): + """check_availability reserves only matched blocks that are currently + evictable (ref_count == 0). A matched prefix already pinned by an + in-flight request frees no capacity when re-pinned, so reserving it would + under-report availability and needlessly defer shared-prefix requests.""" + ctx = self._ctx() + bs = ctx.block_size_tokens + alloc = ctx.kv_block_allocator + + # Request A stays active, pinning the shared prefix H0/S0 -> H1/S1. + ctx.add_request(self._req(ctx, self._prompt(bs * 2))) + s0, s1 = self._block_ids(ctx, 0, 2) + assert alloc.block_ref_counts[s0].item() == 1 + assert alloc.block_ref_counts[s1].item() == 1 + + # One unrelated block is cached and evictable (ref_count == 0). + ctx.add_request(self._req(ctx, self._prompt(bs, offset=9000), request_id=2)) + (sx,) = self._block_ids(ctx, 1, 1) + ctx.release_memory_blocks_from_request_indexes(torch.tensor([1])) + assert alloc.block_ref_counts[sx].item() == 0 + assert int(alloc.get_evictable_block_count()) == 1 + + # Free pool exhausted: the one new block B needs (H2) can only come from + # evicting SX. The already-pinned matches S0/S1 must not be reserved. + alloc.total_avail = 0 + + # Request B shares H0/H1 with A and needs one new block for H2. + req_b = self._req(ctx, self._prompt(bs * 3), request_id=3) + matched, num_from_pool, *_ = ctx._compute_prefix_match(req_b, req_b.remaining_prompt_length) + assert matched == [s0, s1] + assert num_from_pool == 1 + + _, _, kv_cache_available = ctx.check_availability(req_b) + # SX (the sole evictable block) can satisfy H2; reserving the pinned + # matches would wrongly report the request as un-addable. + assert kv_cache_available is True + @pytest.mark.internal def test_ref_count_refzero(self): bs = 32 diff --git a/tests/unit_tests/inference/contexts/test_kv_block_allocator.py b/tests/unit_tests/inference/contexts/test_kv_block_allocator.py index 57b58230171..da068087579 100644 --- a/tests/unit_tests/inference/contexts/test_kv_block_allocator.py +++ b/tests/unit_tests/inference/contexts/test_kv_block_allocator.py @@ -19,6 +19,7 @@ def _make_context( total_request_count=0, request_kv_block_counts=None, request_to_kv_block_ids=None, + prefix_cache_lru_clock=0, ): """Build a minimal DynamicInferenceContext-like fake for the allocator.""" if request_kv_block_counts is None: @@ -30,6 +31,7 @@ def _make_context( total_request_count=total_request_count, request_kv_block_counts=request_kv_block_counts, request_to_kv_block_ids=request_to_kv_block_ids, + prefix_cache_lru_clock=prefix_cache_lru_clock, ) @@ -112,7 +114,9 @@ def test_block_usage_counts_no_prefix_caching( ) def test_prefix_caching_state_layout(policy, expect_timestamps): """Prefix-caching mode allocates block_hashes (initially -1) and ref_counts - (initially 0). LRU policy also allocates timestamps; REF_ZERO does not.""" + (initially 0). LRU policy also allocates timestamps and the persisted + prefix-forest bookkeeping (block_parent_id / block_child_count); REF_ZERO + does not.""" a = KVBlockAllocator( _make_context(), total_count=8, @@ -124,6 +128,11 @@ def test_prefix_caching_state_layout(policy, expect_timestamps): assert (a.block_ref_counts == 0).all().item() assert a.kv_hash_to_block_id == {} assert hasattr(a, "block_timestamps") is expect_timestamps + assert hasattr(a, "block_parent_id") is expect_timestamps + assert hasattr(a, "block_child_count") is expect_timestamps + if expect_timestamps: + assert (a.block_parent_id == -1).all().item() + assert (a.block_child_count == 0).all().item() def test_prefix_caching_allocate_and_hash_registration(): @@ -143,15 +152,25 @@ def test_prefix_caching_allocate_and_hash_registration(): ids = a.allocate_memory_blocks(2) assert (a.block_ref_counts[ids] == 1).all().item() - # Hash registration populates both the tensor and the dict. + # Hash registration populates both the tensor and the dict. Parent hashes are + # ignored under REF_ZERO (they only drive LRU eviction ordering), so this mode + # keeps no per-block parent bookkeeping. a.register_kv_block_hashes(block_ids=[1, 3], block_hashes=[111, 333]) assert a.block_hashes[1].item() == 111 assert a.block_hashes[3].item() == 333 + assert not hasattr(a, "block_parent_id") assert a.kv_hash_to_block_id == {111: 1, 333: 3} + # Supplying parent hashes is accepted (and ignored) under REF_ZERO. + a.register_kv_block_hashes(block_ids=[2, 4], block_hashes=[222, 444], parent_hashes=[111, 222]) + + # Mismatched parent-hash length is rejected regardless of policy. + with pytest.raises(AssertionError): + a.register_kv_block_hashes(block_ids=[5], block_hashes=[555], parent_hashes=[1, 2]) + # Empty inputs are a no-op (avoids zero-element tensor construction). a.register_kv_block_hashes(block_ids=[], block_hashes=[]) - assert a.kv_hash_to_block_id == {111: 1, 333: 3} + assert a.kv_hash_to_block_id == {111: 1, 333: 3, 222: 2, 444: 4} # REF_ZERO has no eviction path when the free pool is short. small = KVBlockAllocator( @@ -189,3 +208,316 @@ def test_block_usage_counts_with_prefix_caching( a = KVBlockAllocator(ctx, total_count=TOTAL_COUNT, paused_count=3, enable_prefix_caching=True) assert a.get_active_used() == expected_active assert a.get_paused_used() == expected_paused + + +# --------------------------------------------------------------------------- +# LRU eviction: parent-chain safety +# --------------------------------------------------------------------------- + + +def _lru_allocator(total_count=16, paused_count=1): + """LRU-mode prefix-caching allocator over a fresh fake context.""" + return KVBlockAllocator( + _make_context(), + total_count=total_count, + paused_count=paused_count, + enable_prefix_caching=True, + prefix_caching_eviction_policy=PrefixCachingEvictionPolicy.LRU, + ) + + +def _seed_cached_chain(a, block_ids, hashes, parents, timestamps): + """Register a chain of cached (ref_count == 0) blocks with explicit LRU + timestamps, bypassing the allocation path to control the layout directly.""" + a.register_kv_block_hashes(block_ids=block_ids, block_hashes=hashes, parent_hashes=parents) + ids = torch.tensor(block_ids, dtype=torch.int64) + a.block_ref_counts[ids] = 0 # cached / evictable + a.block_timestamps[ids] = torch.tensor(timestamps, dtype=torch.int64) + # Mark the blocks as out of the free pool so _deregister_blocks (which pushes + # them back) keeps total_avail bookkeeping consistent. + a.total_avail -= len(block_ids) + + +def _assert_prefix_invariant(a): + """Every cached block must have its parent cached too (or be a root). This is + exactly the invariant _find_kv_match_count relies on.""" + cached_ids = set(a.kv_hash_to_block_id.values()) + for block_hash, block_id in a.kv_hash_to_block_id.items(): + parent_id = a.block_parent_id[block_id].item() + if parent_id >= 0: + assert parent_id in cached_ids, ( + f"dangling child: block {block_id} (hash {block_hash}) parent " + f"block {parent_id} not cached" + ) + + +def test_evict_lru_never_orphans_a_child(): + """Regression: with chunked prefill an ancestor block can end up OLDER than + its descendant. A naive oldest-first eviction would evict the parent and leave + a dangling child; leaf-only eviction must evict the child instead.""" + a = _lru_allocator() + # Chain b0 -> b1 -> b2. Parent b1 (ts=1) is older than child b2 (ts=5). + _seed_cached_chain( + a, block_ids=[0, 1, 2], hashes=[10, 20, 30], parents=[0, 10, 20], timestamps=[1, 1, 5] + ) + + assert a.evict_lru_blocks(1) is True + # The leaf (b2, hash 30) is evicted, not the older parent b1 (hash 20). + assert a.kv_hash_to_block_id == {10: 0, 20: 1} + assert a.block_hashes[2].item() == -1 + assert a.block_parent_id[2].item() == -1 + # Evicting the leaf drops it from its parent's child count. + assert a.block_child_count[1].item() == 0 + _assert_prefix_invariant(a) + + +def test_evict_lru_cascades_up_the_chain(): + """Evicting more blocks than there are leaves walks up the chain from the + deepest descendant, always keeping the retained set descendant-closed.""" + a = _lru_allocator() + _seed_cached_chain( + a, block_ids=[0, 1, 2], hashes=[10, 20, 30], parents=[0, 10, 20], timestamps=[1, 1, 5] + ) + + assert a.evict_lru_blocks(2) is True + # b2 then b1 evicted; only the root b0 remains. + assert a.kv_hash_to_block_id == {10: 0} + _assert_prefix_invariant(a) + + +def test_evict_lru_normal_lru_order_when_leaf_is_oldest(): + """When the oldest block is already a leaf (the common partial-match case, + where ancestors are refreshed and descendants are stale), plain LRU order + applies and the oldest leaf is evicted first.""" + a = _lru_allocator() + # Ancestors refreshed (ts=9); descendant stale (ts=3) and is the leaf. + _seed_cached_chain( + a, block_ids=[0, 1, 2], hashes=[10, 20, 30], parents=[0, 10, 20], timestamps=[9, 9, 3] + ) + + assert a.evict_lru_blocks(1) is True + assert a.kv_hash_to_block_id == {10: 0, 20: 1} + _assert_prefix_invariant(a) + + +def test_evict_lru_branching_prefix_tree(): + """A shared parent with two divergent children (branching prefixes) must keep + the parent cached until BOTH children are evicted.""" + a = _lru_allocator() + # b0 is the parent of both b1 and b2 (e.g. prompts "P+X" and "P+Y"). + _seed_cached_chain( + a, block_ids=[0, 1, 2], hashes=[10, 20, 30], parents=[0, 10, 10], timestamps=[1, 2, 8] + ) + + # Evicting one block takes a leaf (b1, the older child), never the parent. + assert a.evict_lru_blocks(1) is True + assert a.kv_hash_to_block_id == {10: 0, 30: 2} + _assert_prefix_invariant(a) + + # Evicting the second child leaves only the parent. + assert a.evict_lru_blocks(1) is True + assert a.kv_hash_to_block_id == {10: 0} + _assert_prefix_invariant(a) + + +def test_evict_lru_cached_child_with_pinned_parent_treated_as_root(): + """Multi-turn / agentic case: a shared prefix block S0 stays pinned by an + active request (ref_count > 0) while a descendant S1 from a finished turn is + cached (ref_count == 0). S0 is not in the candidate set, so S1's parent + resolves to -1 (S1 is treated as a forest root) and is safely evicted by + normal LRU. The pinned parent must never be touched, even when it is the + oldest block of all.""" + a = _lru_allocator() + # Chain S0 -> S1, plus an unrelated cached root SX. + a.register_kv_block_hashes( + block_ids=[0, 1, 2], block_hashes=[10, 20, 30], parent_hashes=[0, 10, 0] + ) + ids = torch.tensor([0, 1, 2], dtype=torch.int64) + # S0 pinned (active request), S1 and SX cached/evictable. S0 is the OLDEST + # (ts=0) — a pin-blind oldest-first eviction would wrongly take it and orphan + # nothing here, but in general orphan its children. + a.block_ref_counts[ids] = torch.tensor([1, 0, 0], dtype=torch.int32) + a.block_timestamps[ids] = torch.tensor([0, 1, 9], dtype=torch.int64) + a.total_avail -= 3 + + # Only S1 and SX are candidates; the pinned S0 is excluded. + assert int(a.get_evictable_block_count()) == 2 + + # Evict one: S1 (ts=1) is the oldest candidate and a leaf; evicted first. + assert a.evict_lru_blocks(1) is True + assert a.kv_hash_to_block_id == {10: 0, 30: 2} # S0 (pinned) + SX survive + assert a.block_ref_counts[0].item() == 1 # parent still pinned + assert a.block_hashes[0].item() == 10 # parent hash intact + assert a.block_hashes[1].item() == -1 # child deregistered + _assert_prefix_invariant(a) + + # Evict again: only SX remains as a candidate; S0 stays pinned throughout. + assert a.evict_lru_blocks(1) is True + assert a.kv_hash_to_block_id == {10: 0} + assert a.block_ref_counts[0].item() == 1 + # The pinned parent can never be evicted, so a third eviction fails. + assert a.evict_lru_blocks(1) is False + + +def test_evict_lru_partial_chain_eviction_peels_from_leaf_keeping_root(): + """Evicting fewer blocks than a chain's length peels from the leaf end, even + when the root is the least-recently-used block. + + Chain A -> B -> C with the root A oldest (ts 1 < 2 < 3); evict 2. Eviction + proceeds leaf-first, so C then B are removed and the root A is retained. The + retained cache stays descendant-closed (no cached block is left with an + evicted parent). + """ + a = _lru_allocator() + _seed_cached_chain( + a, block_ids=[0, 1, 2], hashes=[10, 20, 30], parents=[0, 10, 20], timestamps=[1, 2, 3] + ) + + assert a.evict_lru_blocks(2) is True + # Leaf C and its parent B are evicted; the root A survives despite being oldest. + assert a.kv_hash_to_block_id == {10: 0} + assert a.block_hashes[0].item() == 10 # root A retained + assert a.block_hashes[1].item() == -1 # B deregistered + assert a.block_hashes[2].item() == -1 # C deregistered + _assert_prefix_invariant(a) + + +def test_evict_lru_insufficient_cached_blocks_returns_false(): + """When fewer cached blocks exist than requested, eviction fails without + touching the cache.""" + a = _lru_allocator() + _seed_cached_chain(a, block_ids=[0, 1], hashes=[10, 20], parents=[0, 10], timestamps=[1, 2]) + assert a.evict_lru_blocks(3) is False + assert a.kv_hash_to_block_id == {10: 0, 20: 1} + + +def test_evict_lru_keeps_hottest_leaf_over_cold_interior_parent(): + """Optimality: leaf-peeling must retain the single most-recently-used block + even when reaching it means evicting a colder interior parent elsewhere. A + block is kept only for its own recency, never because a hot descendant props + it up, so the hot leaf E survives while the colder interior block B is evicted. + + A(ts 1) -> B(ts 2) -> C(ts 5) + \-> F(ts 3) + \-> D(ts 3) -> E(ts 5) + """ + a = _lru_allocator(total_count=8) + # hashes: A=10, B=20, C=30, F=40, D=50, E=60 + _seed_cached_chain( + a, + block_ids=[0, 1, 2, 3, 4, 5], + hashes=[10, 20, 30, 40, 50, 60], + parents=[0, 10, 20, 20, 10, 50], + timestamps=[1, 2, 5, 3, 3, 5], + ) + + assert a.evict_lru_blocks(3) is True + # Evicted F(3), C(5), then B(2) once childless. Retains A, D, and the hottest + # block E -- never evicting E in favor of the colder interior B. + assert a.kv_hash_to_block_id == {10: 0, 50: 4, 60: 5} + assert a.block_hashes[5].item() == 60 # hottest leaf E retained + assert a.block_hashes[1].item() == -1 # cold interior B evicted + _assert_prefix_invariant(a) + + +def test_evict_lru_asserts_on_cyclic_parent_graph(): + """The parent graph is assumed acyclic (a forest). A hash collision producing + a cycle exposes no leaf, so the peel cannot collect enough blocks; this is a + bug and must fail loudly rather than silently under-evict.""" + a = _lru_allocator() + # 2-cycle: block 0's parent hash is 20 (block 1) and block 1's parent hash is + # 10 (block 0). register_kv_block_hashes never produces this — we seed it + # directly to model the pathological collision case. + _seed_cached_chain(a, block_ids=[0, 1], hashes=[10, 20], parents=[20, 10], timestamps=[1, 2]) + assert int(a.get_evictable_block_count()) == 2 + + with pytest.raises(AssertionError): + a.evict_lru_blocks(1) + + +def test_is_memory_available_excludes_soon_to_be_pinned_blocks(): + """potential_matched_count removes soon-to-be-pinned cached blocks from the + evictable capacity, so availability matches what allocation can satisfy + once those blocks (e.g. prefix matches) are pinned.""" + a = _lru_allocator(total_count=6, paused_count=1) + # Drain the free pool: every block is allocated (ref_count == 1), none free. + a.allocate_memory_blocks(a.total_avail) + assert a.total_avail == 0 + # Mark two blocks as cached/evictable, mirroring an LRU release: ref_count + # drops to 0 and the hash is retained, but the block stays out of the free + # pool (total_avail unchanged). + a.register_kv_block_hashes(block_ids=[0, 1], block_hashes=[10, 20], parent_hashes=[0, 10]) + a.block_ref_counts[torch.tensor([0, 1])] = 0 + assert a.total_avail == 0 + assert int(a.get_evictable_block_count()) == 2 + + # Both evictable blocks count toward availability by default. + assert a.is_memory_available(2) is True + # Excluding one (it will be pinned) leaves only one usable for the request. + assert a.is_memory_available(2, potential_matched_count=1) is False + assert a.is_memory_available(1, potential_matched_count=1) is True + # Excluding all evictable blocks leaves nothing to satisfy a new block. + assert a.is_memory_available(1, potential_matched_count=2) is False + + +def _reference_leaf_peel(block_ids, hashes, parents, timestamps, k_evict): + """Independent, straightforward greedy reference: repeatedly evict the + currently-evictable leaf with the oldest (timestamp, block_id). Returns the + set of evicted block ids. Used to pin the optimal eviction choice.""" + hash_to_id = dict(zip(hashes, block_ids)) + ts = dict(zip(block_ids, timestamps)) + child_count = {b: 0 for b in block_ids} + parent_of = {} + for b, p in zip(block_ids, parents): + pid = hash_to_id.get(p) + parent_of[b] = pid + if pid is not None: + child_count[pid] += 1 + + import heapq as _heapq + + heap = [(ts[b], b) for b in block_ids if child_count[b] == 0] + _heapq.heapify(heap) + evicted = set() + while heap and len(evicted) < k_evict: + _, b = _heapq.heappop(heap) + evicted.add(b) + pid = parent_of[b] + if pid is not None: + child_count[pid] -= 1 + if child_count[pid] == 0: + _heapq.heappush(heap, (ts[pid], pid)) + return evicted + + +def test_evict_lru_preserves_invariant_under_random_chains(): + """Property test: across many randomized multi-chain layouts and eviction + counts, eviction (a) preserves the parent-chain invariant and (b) evicts + exactly the optimal leaf-peel set (matched against an independent + reference).""" + torch.manual_seed(0) + for _ in range(50): + n = int(torch.randint(2, 10, (1,)).item()) + a = _lru_allocator(total_count=n + 4) + block_ids = list(range(n)) + # Build a forest: block k's parent is a random earlier block or a root. + hashes = [100 + k for k in range(n)] + parents = [] + for k in range(n): + if k == 0 or int(torch.randint(0, 2, (1,)).item()) == 0: + parents.append(0) # root + else: + parents.append(hashes[int(torch.randint(0, k, (1,)).item())]) + # Distinct timestamps so the optimal evicted set is unique and the + # reference comparison is exact (no tie-break ambiguity). + timestamps = torch.randperm(50)[:n].add(1).tolist() + _seed_cached_chain(a, block_ids, hashes, parents, timestamps) + + k_evict = int(torch.randint(1, n + 1, (1,)).item()) + expected_evicted = _reference_leaf_peel(block_ids, hashes, parents, timestamps, k_evict) + + assert a.evict_lru_blocks(k_evict) is True + retained = set(a.kv_hash_to_block_id.values()) + assert retained == set(block_ids) - expected_evicted + assert len(retained) == n - k_evict + _assert_prefix_invariant(a)