[Bugfix][Spec Decode] Keep EAGLE cache registration on the partial-hash-hit path (#52419)
mispa-ms committed
Aug 15, 2026 at 21:03 UTC
6593754e61a45b27a93e40fa0479568945d38661
2 files changed
+92
-20
tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py
+74
@@ -261,6 +261,80 @@ def test_hybrid_mamba_align_partial_hash_hit():
261
assert manager.get_blocks("1").blocks[1][1].block_hash_num_tokens == 8
262
263
264
+def test_eagle_group_registers_unaligned_tail_under_partial_hash_hits():
265
+ """An EAGLE group must not re-floor what partial hash hits leaves un-floored.
266
+
267
+ ``cache_blocks`` decides once how far a request may be registered, and with
268
+ fine-grained partial hash hits that bound is the raw token count. The EAGLE
269
+ branch then re-derives its own bound for the lookahead block; if it rounds
270
+ down to ``scheduler_block_size`` again, everything between the last aligned
271
+ boundary and the tail stops being registered -- ``(n % scheduler_block_size)
272
+ - manager.block_size`` tokens per call, which is most of a segment whenever
273
+ the group's own block is much smaller than the scheduler block.
274
+ """
275
+ hash_block_size = 2
276
+ mamba_block_size = 4 * hash_block_size
277
+ kv_cache_config = KVCacheConfig(
278
+ num_blocks=40,
279
+ kv_cache_tensors=[],
280
+ kv_cache_groups=[
281
+ KVCacheGroupSpec(
282
+ ["full"],
283
+ FullAttentionSpec(
284
+ block_size=hash_block_size,
285
+ num_kv_heads=1,
286
+ head_size=1,
287
+ dtype=torch.float32,
288
+ ),
289
+ ),
290
+ KVCacheGroupSpec(
291
+ ["mamba"],
292
+ MambaSpec(
293
+ block_size=mamba_block_size,
294
+ shapes=(1, 1),
295
+ dtypes=(torch.float32,),
296
+ mamba_cache_mode="align",
297
+ ),
298
+ ),
299
+ ],
300
+ )
301
+ manager = make_kv_cache_manager(
302
+ kv_cache_config=kv_cache_config,
303
+ max_model_len=8192,
304
+ enable_caching=True,
305
+ hash_block_size=hash_block_size,
306
+ )
307
+ coordinator = manager.coordinator
308
+ assert coordinator.enable_partial_hash_hits
309
+ # The full-attention group is the EAGLE one, and its block is smaller than
310
+ # the scheduler block -- the geometry where re-flooring loses tokens.
311
+ eagle_manager = coordinator.single_type_managers[0]
312
+ eagle_manager.use_eagle = True
313
+ assert eagle_manager.block_size < coordinator.scheduler_block_size
314
+
315
+ # Deliberately not a multiple of the scheduler block, so the two bounds
316
+ # differ: floor(22/8)*8 + 2 = 18 against 22.
317
+ num_tokens = coordinator.scheduler_block_size * 2 + hash_block_size * 3
318
+ req = make_request("0", list(range(num_tokens)), hash_block_size, sha256)
319
+
320
+ recorded: list[int] = []
321
+ for single_type_manager in coordinator.single_type_managers:
322
+ original = single_type_manager.cache_blocks
323
+
324
+ def spy(request, num_tokens_to_cache, *args, _orig=original, **kwargs):
325
+ recorded.append(num_tokens_to_cache)
326
+ return _orig(request, num_tokens_to_cache, *args, **kwargs)
327
+
328
+ single_type_manager.cache_blocks = spy
329
+
330
+ # allocate_slots caches on the way out, so this exercises the real path.
331
+ computed_blocks, num_computed, _ = manager.get_computed_blocks(req)
332
+ assert manager.allocate_slots(req, num_tokens, num_computed, computed_blocks)
333
+
334
+ # Every group, EAGLE or not, may register the whole unaligned tail.
335
+ assert recorded == [num_tokens] * len(coordinator.single_type_managers)
336
+
337
+
338
def test_hybrid_mamba_partial_tail_owner_uses_cow_on_continue():
339
hash_block_size = 2
340
block_size = 2 * hash_block_size
vllm/v1/core/kv_cache_coordinator.py
+18
-20
@@ -6,7 +6,7 @@ from typing import NamedTuple
6
7
from vllm import envs
8
from vllm.logger import init_logger
9
-from vllm.utils.math_utils import cdiv
9
+from vllm.utils.math_utils import cdiv, round_down
10
from vllm.v1.core.block_pool import BlockPool
11
from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector
12
from vllm.v1.core.kv_cache_utils import (
@@ -704,39 +704,37 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
704
for gid in group.group_ids:
705
self.single_type_managers[gid].use_eagle = True
706
707
- def cache_blocks(self, request: Request, num_computed_tokens: int) -> None:
707
+ def _align_cacheable(self, num_tokens: int) -> int:
708
+ """Largest prefix of ``num_tokens`` a future cache hit could match.
709
+
710
+ Hits are ``scheduler_block_size``-aligned (see
711
+ ``find_longest_cache_hit``) unless fine-grained partial hash hits are
712
+ enabled, in which case no rounding applies -- rounding even to
713
+ ``hash_block_size`` would re-register a privatized Mamba tail.
714
+ """
715
if self.enable_partial_hash_hits:
709
- aligned_num_computed_tokens = num_computed_tokens
710
- else:
711
- # Cache hits in this coordinator are always a multiple of
712
- # ``scheduler_block_size`` tokens (see ``find_longest_cache_hit``).
713
- # Within an aligned region, SWA groups may only consult a subset of
714
- # blocks per ``scheduler_block_size``-segment so the unused blocks
715
- # also stay out of the prefix-cache hash map.
716
- aligned_num_computed_tokens = (
717
- num_computed_tokens
718
- // self.scheduler_block_size
719
- * self.scheduler_block_size
720
- )
716
+ return num_tokens
717
+ return round_down(num_tokens, self.scheduler_block_size)
718
+
719
+ def cache_blocks(self, request: Request, num_computed_tokens: int) -> None:
720
+ cached_num_computed_tokens = self._align_cacheable(num_computed_tokens)
721
for manager in self.single_type_managers:
722
- num_tokens_to_cache = aligned_num_computed_tokens
722
+ num_tokens_to_cache = cached_num_computed_tokens
723
# EAGLE groups match one block past each aligned boundary and drop
724
# it, so make that lookahead block eligible to be cached.
725
- if manager.use_eagle and aligned_num_computed_tokens > 0:
725
+ if manager.use_eagle and cached_num_computed_tokens > 0:
726
# Only cache tokens with finalized KV. The last
727
# num_reprefillable_tokens tokens can be re-prefilled during
728
# multi-module MTP.
729
num_finalized_computed_tokens = max(
730
0, num_computed_tokens - self.num_reprefillable_tokens
731
)
732
- aligned_num_finalized_computed_tokens = (
732
+ cached_num_finalized_computed_tokens = self._align_cacheable(
733
num_finalized_computed_tokens
734
- // self.scheduler_block_size
735
- * self.scheduler_block_size
734
)
735
num_tokens_to_cache = min(
736
num_finalized_computed_tokens,
739
- aligned_num_finalized_computed_tokens + manager.block_size,
737
+ cached_num_finalized_computed_tokens + manager.block_size,
738
)
739
# The manager already knows the fine hit granularity
740
# (``scheduler_block_size``); retention is passed separately so it