[Bugfix][Spec Decode] Keep EAGLE cache registration on the partial-hash-hit path (#52419)

Author: mispa-ms

tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py

@@ -261,6 +261,80 @@ def test_hybrid_mamba_align_partial_hash_hit():
     assert manager.get_blocks("1").blocks[1][1].block_hash_num_tokens == 8
 
 
+def test_eagle_group_registers_unaligned_tail_under_partial_hash_hits():
+    """An EAGLE group must not re-floor what partial hash hits leaves un-floored.
+
+    ``cache_blocks`` decides once how far a request may be registered, and with
+    fine-grained partial hash hits that bound is the raw token count. The EAGLE
+    branch then re-derives its own bound for the lookahead block; if it rounds
+    down to ``scheduler_block_size`` again, everything between the last aligned
+    boundary and the tail stops being registered -- ``(n % scheduler_block_size)
+    - manager.block_size`` tokens per call, which is most of a segment whenever
+    the group's own block is much smaller than the scheduler block.
+    """
+    hash_block_size = 2
+    mamba_block_size = 4 * hash_block_size
+    kv_cache_config = KVCacheConfig(
+        num_blocks=40,
+        kv_cache_tensors=[],
+        kv_cache_groups=[
+            KVCacheGroupSpec(
+                ["full"],
+                FullAttentionSpec(
+                    block_size=hash_block_size,
+                    num_kv_heads=1,
+                    head_size=1,
+                    dtype=torch.float32,
+                ),
+            ),
+            KVCacheGroupSpec(
+                ["mamba"],
+                MambaSpec(
+                    block_size=mamba_block_size,
+                    shapes=(1, 1),
+                    dtypes=(torch.float32,),
+                    mamba_cache_mode="align",
+                ),
+            ),
+        ],
+    )
+    manager = make_kv_cache_manager(
+        kv_cache_config=kv_cache_config,
+        max_model_len=8192,
+        enable_caching=True,
+        hash_block_size=hash_block_size,
+    )
+    coordinator = manager.coordinator
+    assert coordinator.enable_partial_hash_hits
+    # The full-attention group is the EAGLE one, and its block is smaller than
+    # the scheduler block -- the geometry where re-flooring loses tokens.
+    eagle_manager = coordinator.single_type_managers[0]
+    eagle_manager.use_eagle = True
+    assert eagle_manager.block_size < coordinator.scheduler_block_size
+
+    # Deliberately not a multiple of the scheduler block, so the two bounds
+    # differ: floor(22/8)*8 + 2 = 18 against 22.
+    num_tokens = coordinator.scheduler_block_size * 2 + hash_block_size * 3
+    req = make_request("0", list(range(num_tokens)), hash_block_size, sha256)
+
+    recorded: list[int] = []
+    for single_type_manager in coordinator.single_type_managers:
+        original = single_type_manager.cache_blocks
+
+        def spy(request, num_tokens_to_cache, *args, _orig=original, **kwargs):
+            recorded.append(num_tokens_to_cache)
+            return _orig(request, num_tokens_to_cache, *args, **kwargs)
+
+        single_type_manager.cache_blocks = spy
+
+    # allocate_slots caches on the way out, so this exercises the real path.
+    computed_blocks, num_computed, _ = manager.get_computed_blocks(req)
+    assert manager.allocate_slots(req, num_tokens, num_computed, computed_blocks)
+
+    # Every group, EAGLE or not, may register the whole unaligned tail.
+    assert recorded == [num_tokens] * len(coordinator.single_type_managers)
+
+
 def test_hybrid_mamba_partial_tail_owner_uses_cow_on_continue():
     hash_block_size = 2
     block_size = 2 * hash_block_size

vllm/v1/core/kv_cache_coordinator.py

@@ -6,7 +6,7 @@
 
 from vllm import envs
 from vllm.logger import init_logger
-from vllm.utils.math_utils import cdiv
+from vllm.utils.math_utils import cdiv, round_down
 from vllm.v1.core.block_pool import BlockPool
 from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector
 from vllm.v1.core.kv_cache_utils import (
@@ -704,39 +704,37 @@ def verify_and_split_kv_cache_groups(self) -> None:
                 for gid in group.group_ids:
                     self.single_type_managers[gid].use_eagle = True
 
-    def cache_blocks(self, request: Request, num_computed_tokens: int) -> None:
+    def _align_cacheable(self, num_tokens: int) -> int:
+        """Largest prefix of ``num_tokens`` a future cache hit could match.
+
+        Hits are ``scheduler_block_size``-aligned (see
+        ``find_longest_cache_hit``) unless fine-grained partial hash hits are
+        enabled, in which case no rounding applies -- rounding even to
+        ``hash_block_size`` would re-register a privatized Mamba tail.
+        """
         if self.enable_partial_hash_hits:
-            aligned_num_computed_tokens = num_computed_tokens
-        else:
-            # Cache hits in this coordinator are always a multiple of
-            # ``scheduler_block_size`` tokens (see ``find_longest_cache_hit``).
-            # Within an aligned region, SWA groups may only consult a subset of
-            # blocks per ``scheduler_block_size``-segment so the unused blocks
-            # also stay out of the prefix-cache hash map.
-            aligned_num_computed_tokens = (
-                num_computed_tokens
-                // self.scheduler_block_size
-                * self.scheduler_block_size
-            )
+            return num_tokens
+        return round_down(num_tokens, self.scheduler_block_size)
+
+    def cache_blocks(self, request: Request, num_computed_tokens: int) -> None:
+        cached_num_computed_tokens = self._align_cacheable(num_computed_tokens)
         for manager in self.single_type_managers:
-            num_tokens_to_cache = aligned_num_computed_tokens
+            num_tokens_to_cache = cached_num_computed_tokens
             # EAGLE groups match one block past each aligned boundary and drop
             # it, so make that lookahead block eligible to be cached.
-            if manager.use_eagle and aligned_num_computed_tokens > 0:
+            if manager.use_eagle and cached_num_computed_tokens > 0:
                 # Only cache tokens with finalized KV. The last
                 # num_reprefillable_tokens tokens can be re-prefilled during
                 # multi-module MTP.
                 num_finalized_computed_tokens = max(
                     0, num_computed_tokens - self.num_reprefillable_tokens
                 )
-                aligned_num_finalized_computed_tokens = (
+                cached_num_finalized_computed_tokens = self._align_cacheable(
                     num_finalized_computed_tokens
-                    // self.scheduler_block_size
-                    * self.scheduler_block_size
                 )
                 num_tokens_to_cache = min(
                     num_finalized_computed_tokens,
-                    aligned_num_finalized_computed_tokens + manager.block_size,
+                    cached_num_finalized_computed_tokens + manager.block_size,
                 )
             # The manager already knows the fine hit granularity
             # (``scheduler_block_size``); retention is passed separately so it