Make LRU prefix caching eviction policy only evict child blocks

megatron/core/inference/contexts/dynamic_context.py

@@ -3078,8 +3078,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

megatron/core/inference/contexts/kv_block_allocator.py

@@ -58,6 +58,15 @@ def __init__(
             # Hash-to-block mapping for O(1) prefix lookup
             self.kv_hash_to_block_id: Dict[int, int] = {}
 
+            # Parent hash per block for prefix-chain bookkeeping: 0 = no parent
+            # (root block or unregistered). Block hashes are parent-chained, so a
+            # cached block whose hash is some other cached block's parent must not
+            # be evicted before its child (see evict_lru_blocks). Valid hashes are
+            # in [1, 2^63-1], so 0 is a safe "no parent" sentinel.
+            self.block_parent_hashes = torch.zeros(
+                (self.total_count,), dtype=torch.int64, device='cpu'
+            )
+
             # Reference count per block: 0 = cached (evictable), >0 = actively used
             self.block_ref_counts = torch.zeros(
                 (self.total_count,), dtype=torch.int32, device='cpu'
@@ -253,6 +262,7 @@ def reset(self) -> None:
 
             # Reset prefix caching state
             self.kv_hash_to_block_id.clear()
+            self.block_parent_hashes.fill_(0)
             self.block_ref_counts.fill_(0)
             if self.prefix_caching_eviction_policy == PrefixCachingEvictionPolicy.LRU:
                 self.block_timestamps.fill_(0)
@@ -264,18 +274,32 @@ 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)
+            self.block_parent_hashes[id_tensor] = torch.tensor(
+                parent_hashes, dtype=torch.int64, device=self.block_hashes.device
+            )
         self.kv_hash_to_block_id.update(zip(block_hashes, block_ids))
 
     def _deregister_blocks(self, block_ids: Tensor) -> None:
@@ -307,6 +331,7 @@ def _deregister_blocks(self, block_ids: Tensor) -> None:
 
         # Reset block state (batched tensor ops)
         self.block_hashes[block_ids] = -1
+        self.block_parent_hashes[block_ids] = 0
         self.block_ref_counts[block_ids] = 0
         if self.prefix_caching_eviction_policy == PrefixCachingEvictionPolicy.LRU:
             self.block_timestamps[block_ids] = 0
@@ -340,7 +365,28 @@ 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, we key each cached block by the *maximum*
+        timestamp over its subtree (itself plus all cached descendants): a cached
+        prefix is as recently used as the most-recently-used block extending it.
+        This makes every parent's key >= all of its descendants' keys, so evicting
+        an ascending-key prefix is always "descendant closed" — a parent is only
+        evicted once all of its children have been. Ties are broken by depth
+        (deeper first) so a parent never precedes a child at equal key.
+
+        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.
 
         Args:
             num_blocks_needed: Number of blocks to evict.
@@ -352,13 +398,45 @@ 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
 
-        # 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]]
+        own_ts = self.block_timestamps[cached_block_ids]
+        hashes = self.block_hashes[cached_block_ids]
+        parents = self.block_parent_hashes[cached_block_ids]
+
+        # Resolve each block's parent hash to the local index of the parent block,
+        # or -1 when the parent is not cached (root block with parent hash 0, or a
+        # parent still in use). Hashes are unique per cached block, so a single
+        # sorted lookup suffices.
+        sorted_hashes, sort_order = torch.sort(hashes)
+        pos = torch.searchsorted(sorted_hashes, parents).clamp(max=num_cached - 1)
+        found = sorted_hashes[pos] == parents
+        parent_idx = torch.where(found, sort_order[pos], torch.full_like(pos, -1))
+        has_parent = parent_idx >= 0
+        parent_safe = parent_idx.clamp(min=0)  # -1 -> 0 for masked-out gathers
+
+        # Propagate subtree-max timestamp and depth up/down the parent chain via
+        # repeated vectorized scatter/gather. Converges in <= max-chain-depth
+        # iterations (bounded by prompt_len / block_size); no per-element Python.
+        subtree_max = own_ts.clone()
+        depth = torch.zeros(num_cached, dtype=torch.int64)
+        while True:
+            new_subtree_max = subtree_max.clone()
+            new_subtree_max.scatter_reduce_(
+                0, parent_idx[has_parent], subtree_max[has_parent], reduce="amax"
+            )
+            new_depth = torch.where(has_parent, depth[parent_safe] + 1, depth)
+            if torch.equal(new_subtree_max, subtree_max) and torch.equal(new_depth, depth):
+                break
+            subtree_max, depth = new_subtree_max, new_depth
+
+        # Order by (subtree_max asc, depth desc) so children precede their parents.
+        # Two stable sorts compose into the desired lexicographic order.
+        by_depth = torch.argsort(depth, descending=True, stable=True)
+        order = by_depth[torch.argsort(subtree_max[by_depth], stable=True)]
+        blocks_to_evict = cached_block_ids[order[:num_blocks_needed]]
 
         self._deregister_blocks(blocks_to_evict)
 

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,
     )
 
 
@@ -122,6 +124,7 @@ def test_prefix_caching_state_layout(policy, expect_timestamps):
     )
     assert (a.block_hashes == -1).all().item()
     assert (a.block_ref_counts == 0).all().item()
+    assert (a.block_parent_hashes == 0).all().item()
     assert a.kv_hash_to_block_id == {}
     assert hasattr(a, "block_timestamps") is expect_timestamps
 
@@ -143,15 +146,27 @@ 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
+    # optional and default to 0 (root) when omitted.
     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 a.block_parent_hashes[1].item() == 0
+    assert a.block_parent_hashes[3].item() == 0
     assert a.kv_hash_to_block_id == {111: 1, 333: 3}
 
+    # When supplied, parent hashes are recorded per block.
+    a.register_kv_block_hashes(block_ids=[2, 4], block_hashes=[222, 444], parent_hashes=[111, 222])
+    assert a.block_parent_hashes[2].item() == 111
+    assert a.block_parent_hashes[4].item() == 222
+
+    # Mismatched parent-hash length is rejected.
+    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 +204,161 @@ 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."""
+    present = set(a.kv_hash_to_block_id.keys())
+    for block_hash, block_id in a.kv_hash_to_block_id.items():
+        parent = a.block_parent_hashes[block_id].item()
+        if parent != 0:
+            assert parent in present, (
+                f"dangling child: block {block_id} (hash {block_hash}) parent "
+                f"{parent} 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_hashes[2].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_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_preserves_invariant_under_random_chains():
+    """Property test: across many randomized multi-chain layouts and eviction
+    counts, the retained cache always satisfies the parent-chain invariant."""
+    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())])
+        timestamps = torch.randint(1, 20, (n,)).tolist()
+        _seed_cached_chain(a, block_ids, hashes, parents, timestamps)
+
+        k_evict = int(torch.randint(1, n + 1, (1,)).item())
+        assert a.evict_lru_blocks(k_evict) is True
+        assert len(a.kv_hash_to_block_id) == n - k_evict
+        _assert_prefix_invariant(a)