[Test] Regression test for hybrid-Mamba eagle cache-peek in Mooncake connector (#43559) (#48361)
Signed-off-by: Rishi Puri <riship@nvidia.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Rishi Puri committed
Jul 27, 2026 at 16:03 UTC
fd10e8946d7b27d1f3070d4f58febd8eac1c3a4b
2 files changed
+39
tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py
+36
@@ -531,3 +531,39 @@ def test_load_mask_without_eagle_unchanged():
531
assert hit == 64
532
masks = coord.load_mask(hs, token_len=hit)
533
assert masks[0] == [True, True, True, True]
534
+
535
+
536
+def _mamba(block_size=16):
537
+ return MambaSpec(
538
+ block_size=block_size,
539
+ shapes=((1, 1),),
540
+ dtypes=(torch.float32,),
541
+ mamba_cache_mode="align",
542
+ )
543
+
544
+
545
+def test_lookup_with_eagle_hybrid_full_plus_mamba_no_overrun():
546
+ """Full+Mamba with eagle must not overrun the attention-verified hit.
547
+
548
+ ``MambaManager`` ignores ``drop_eagle_block`` (a Mamba block at position
549
+ p IS the recurrent state after (p + 1) * block_size tokens; there is
550
+ nothing to recompute), so granting the Mamba group the one-block eagle
551
+ peek margin lets it match one block PAST the eagle-pruned full-attention
552
+ hit, and adopting that length resumes the recurrent state ahead of the
553
+ verified token prefix (#43559). Gating the margin on ``not
554
+ isinstance(spec, MambaSpec)`` pins the hit to the attention-verified 48.
555
+ """
556
+ groups = [
557
+ KVCacheGroupSpec(["L0"], _full(16)),
558
+ KVCacheGroupSpec(["L1"], _mamba(16)),
559
+ ]
560
+ coord = _make_coord(groups, hash_block_size=16, use_eagle=True)
561
+ hs = _hashes(4)
562
+ exists = {(g, bytes(h)) for g in (0, 1) for h in hs}
563
+ cmap = ExternalCachedBlockPool(16, exists)
564
+ _masks, hit = coord.find_longest_cache_hit(
565
+ hs, max_length=64, cached_block_pool=cmap
566
+ )
567
+ # FullAttn matches 4 blocks, eagle pops 1 -> 48 verified tokens. The
568
+ # Mamba group must serve its state@48 snapshot, not peek to state@64.
569
+ assert hit == 48
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py
+3
@@ -338,6 +338,9 @@ class MooncakeStoreCoordinator:
338
339
drop_eagle_block = idx in eagle_indices and idx not in eagle_verified
340
_max_length = curr_hit_length
341
+ # No eagle peek margin for a recurrent (Mamba) group: its finder
342
+ # never drops a block, so a widened bound would match past the
343
+ # attention-verified hit and resume from speculative state (#43559).
344
if drop_eagle_block and not isinstance(spec, MambaSpec):
345
eagle_margin = (
346
self.hash_block_size