[BugFix] Preserve divergent FA hits with external Mamba state (#51468)
Signed-off-by: Roger Wang <hey@rogerw.io> Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai> Signed-off-by: Jeff Ma <jeffjma@umich.edu> Co-authored-by: Roger Wang <hey@rogerw.io> Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: OpenAI Codex <codex@openai.com>
Jeff (Junze) Ma committed
Aug 8, 2026 at 05:20 UTC
58d3918e3ea0a544ffedadad2ba84559e9c51d8f
3 files changed
+90
-12
tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py
+64
@@ -590,6 +590,70 @@ def test_truncate_computed_blocks_preserves_sparse_prefix_positions():
590
assert [len(group) for group in blocks.blocks] == [3, 2]
591
592
593
+def test_truncate_computed_blocks_allows_short_mamba_group_only():
594
+ """External state may replace a short Mamba hit, but other groups must
595
+ cover the aligned local endpoint."""
596
+ hash_block_size = 2
597
+ kv_cache_config = KVCacheConfig(
598
+ num_blocks=24,
599
+ kv_cache_tensors=[],
600
+ kv_cache_groups=[
601
+ KVCacheGroupSpec(
602
+ ["full"],
603
+ FullAttentionSpec(
604
+ block_size=hash_block_size,
605
+ num_kv_heads=1,
606
+ head_size=1,
607
+ dtype=torch.float32,
608
+ ),
609
+ ),
610
+ KVCacheGroupSpec(
611
+ ["mamba"],
612
+ MambaSpec(
613
+ block_size=2 * hash_block_size,
614
+ shapes=(1, 1),
615
+ dtypes=(torch.float32,),
616
+ mamba_cache_mode="align",
617
+ ),
618
+ ),
619
+ ],
620
+ )
621
+ manager = make_kv_cache_manager(
622
+ kv_cache_config=kv_cache_config,
623
+ max_model_len=8192,
624
+ enable_caching=True,
625
+ hash_block_size=hash_block_size,
626
+ )
627
+ producer = make_request("producer", [0, 0, 1, 1, 2, 2], hash_block_size, sha256)
628
+ blocks, num_computed, _ = manager.get_computed_blocks(producer)
629
+ assert manager.allocate_slots(producer, 6, num_computed, blocks) is not None
630
+ manager.free(producer)
631
+ manager.new_step_starts()
632
+
633
+ consumer = make_request(
634
+ "consumer", [0, 0, 1, 1, 2, 2, 3, 3], hash_block_size, sha256
635
+ )
636
+ blocks, num_computed, _ = manager.get_computed_blocks(consumer)
637
+ assert num_computed == 6
638
+ assert [len(group) for group in blocks.blocks] == [3, 2]
639
+
640
+ short_mamba = manager.create_kv_cache_blocks((list(blocks.blocks[0]), []))
641
+ truncated = manager.truncate_computed_blocks(short_mamba, 4)
642
+ assert [len(group) for group in truncated.blocks] == [2, 0]
643
+
644
+ short_full_attention = manager.create_kv_cache_blocks(
645
+ (list(blocks.blocks[0][:1]), list(blocks.blocks[1]))
646
+ )
647
+ with pytest.raises(AssertionError):
648
+ manager.truncate_computed_blocks(short_full_attention, 4)
649
+
650
+ with pytest.raises(AssertionError):
651
+ manager.truncate_computed_blocks(blocks, 6)
652
+
653
+ # The lookup result itself is never mutated.
654
+ assert [len(group) for group in blocks.blocks] == [3, 2]
655
+
656
+
657
def test_hybrid_mamba_partial_tail_owner_continue_preserves_later_hit():
658
hash_block_size = 2
659
block_size = 2 * hash_block_size
tests/v1/core/test_scheduler.py
+17
-10
@@ -5685,15 +5685,23 @@ def test_hybrid_per_group_hit_divergence_with_connector(
5685
assert replay.num_tokens - num_scheduled == expected_num_computed
5686
5687
5688
-def test_hybrid_per_group_hit_divergence_fa_deeper_no_external():
5689
- """The opposite divergence: the FA prefix survives deeper than the Mamba
5690
- state and the connector supplies nothing (ext == 0). Reporting the deep FA
5691
- hit as locally computed would resume with no valid Mamba state at that
5692
- boundary (silent bad output). The scheduler must fall back to the
5693
- convergent boundary that every group agrees on (block 0's surviving state).
5688
+@pytest.mark.parametrize(
5689
+ ("matched_tokens", "replay_blocks", "expected_num_computed"),
5690
+ [
5691
+ (0, 5, 16),
5692
+ (16, 6, 80),
5693
+ ],
5694
+)
5695
+def test_hybrid_fa_deeper_hit_uses_external_mamba_state(
5696
+ matched_tokens: int,
5697
+ replay_blocks: int,
5698
+ expected_num_computed: int,
5699
+):
5700
+ """A positive external hit supplies the Mamba state at the resume
5701
+ boundary; without one, the scheduler falls back to the common local hit.
5702
"""
5703
block_size = 16
5696
- scheduler = _create_hybrid_mamba_connector_scheduler(matched_tokens=0)
5704
+ scheduler = _create_hybrid_mamba_connector_scheduler(matched_tokens)
5705
manager = scheduler.kv_cache_manager
5706
assert isinstance(manager.coordinator, HybridKVCacheCoordinator)
5707
@@ -5719,7 +5727,7 @@ def test_hybrid_per_group_hit_divergence_fa_deeper_no_external():
5727
5728
[replay] = create_requests(
5729
num_requests=1,
5722
- num_tokens=5 * block_size,
5730
+ num_tokens=replay_blocks * block_size,
5731
max_tokens=1,
5732
same_prompt=True,
5733
block_size=block_size,
@@ -5733,8 +5741,7 @@ def test_hybrid_per_group_hit_divergence_fa_deeper_no_external():
5741
scheduler.add_request(replay)
5742
output = scheduler.schedule()
5743
num_scheduled = output.num_scheduled_tokens[replay.request_id]
5736
- # Must resume at the convergent boundary (block 0), not the deep FA hit.
5737
- assert replay.num_tokens - num_scheduled == block_size
5744
+ assert replay.num_tokens - num_scheduled == expected_num_computed
5745
5746
5747
def _make_encoder_instance_request(scheduler, text_prefix=8, image_tokens=16):
vllm/v1/core/kv_cache_manager.py
+9
-2
@@ -20,6 +20,7 @@ from vllm.v1.kv_cache_interface import (
20
CrossAttentionSpec,
21
EncoderOnlyAttentionSpec,
22
KVCacheConfig,
23
+ MambaSpec,
24
get_kv_cache_spec_kind,
25
get_kv_cache_spec_sliding_window,
26
)
@@ -779,17 +780,23 @@ class KVCacheManager:
780
) -> KVCacheBlocks:
781
"""Return a lookup-result view truncated at an aligned token endpoint.
782
783
+ An external hit can supply the final Mamba state even when the local
784
+ Mamba group ends before this endpoint. Other groups must cover it.
785
Pure slicing: refcounts are untouched and ``blocks`` is not mutated.
786
"""
787
truncated: list[list[KVCacheBlock]] = []
785
- for group_blocks, manager in zip(
788
+ for group_blocks, manager, group in zip(
789
blocks.blocks,
790
self.coordinator.single_type_managers,
791
+ self.kv_cache_config.kv_cache_groups,
792
strict=True,
793
):
794
assert num_computed_tokens % manager.block_size == 0
795
num_blocks = num_computed_tokens // manager.block_size
792
- assert num_blocks <= len(group_blocks)
796
+ if isinstance(group.kv_cache_spec, MambaSpec):
797
+ num_blocks = min(num_blocks, len(group_blocks))
798
+ else:
799
+ assert num_blocks <= len(group_blocks)
800
truncated.append(list(group_blocks[:num_blocks]))
801
return self.create_kv_cache_blocks(tuple(truncated))
802