[BugFix] Handle per-group prefix-hit divergence for hybrid models with KV connector (#48425)
Nick Hill committed
Jul 23, 2026 at 01:19 UTC
229e01e9e10fa2ed25ad02d689c907a62b73ccf1
4 files changed
+279
-40
tests/v1/core/test_scheduler.py
+197
@@ -25,14 +25,17 @@ from vllm.multimodal.inputs import (
25
from vllm.sampling_params import SamplingParams, StructuredOutputsParams
26
from vllm.utils.hashing import sha256
27
from vllm.v1.core.encoder_cache_manager import EncoderCacheManager
28
+from vllm.v1.core.kv_cache_coordinator import HybridKVCacheCoordinator
29
from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash
30
from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput
31
from vllm.v1.core.sched.scheduler import Scheduler
32
+from vllm.v1.core.single_type_kv_cache_manager import register_all_kvcache_specs
33
from vllm.v1.engine import FinishReason
34
from vllm.v1.kv_cache_interface import (
35
FullAttentionSpec,
36
KVCacheConfig,
37
KVCacheGroupSpec,
38
+ MambaSpec,
39
)
40
from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput
41
from vllm.v1.request import Request, RequestStatus
@@ -5464,3 +5467,197 @@ def test_async_load_reservation_prevents_wedge_e2e():
5467
assert b.status == RequestStatus.WAITING
5468
assert b.num_preemptions == 0
5469
assert b.request_id not in req_to_blocks
5470
+
5471
+
5472
+def _create_hybrid_mamba_connector_scheduler(
5473
+ matched_tokens: int,
5474
+ block_size: int = 16,
5475
+ num_blocks: int = 100,
5476
+) -> Scheduler:
5477
+ """FA + Mamba ("all" cache mode) scheduler with a MockKVConnector."""
5478
+ model_config = ModelConfig(
5479
+ model="facebook/opt-125m",
5480
+ trust_remote_code=True,
5481
+ dtype="float16",
5482
+ seed=42,
5483
+ skip_tokenizer_init=True,
5484
+ )
5485
+ vllm_config = VllmConfig(
5486
+ scheduler_config=SchedulerConfig(
5487
+ max_num_seqs=4,
5488
+ max_num_batched_tokens=8192,
5489
+ max_model_len=8192,
5490
+ enable_chunked_prefill=True,
5491
+ is_encoder_decoder=False,
5492
+ watermark=0.0,
5493
+ ),
5494
+ model_config=model_config,
5495
+ cache_config=CacheConfig(
5496
+ block_size=block_size,
5497
+ enable_prefix_caching=True,
5498
+ mamba_cache_mode="all",
5499
+ ),
5500
+ kv_transfer_config=KVTransferConfig(
5501
+ kv_connector="MockKVConnector",
5502
+ kv_role="kv_both",
5503
+ kv_connector_extra_config={
5504
+ "matched_tokens": matched_tokens,
5505
+ "is_async": False,
5506
+ },
5507
+ ),
5508
+ )
5509
+ vllm_config.cache_config.num_gpu_blocks = num_blocks
5510
+ kv_cache_config = KVCacheConfig(
5511
+ num_blocks=num_blocks,
5512
+ kv_cache_tensors=[],
5513
+ kv_cache_groups=[
5514
+ KVCacheGroupSpec(
5515
+ ["fa"],
5516
+ FullAttentionSpec(
5517
+ block_size=block_size,
5518
+ num_kv_heads=1,
5519
+ head_size=1,
5520
+ dtype=torch.float32,
5521
+ ),
5522
+ ),
5523
+ KVCacheGroupSpec(
5524
+ ["mamba"],
5525
+ MambaSpec(
5526
+ block_size=block_size,
5527
+ shapes=((1, 1),),
5528
+ dtypes=(torch.float32,),
5529
+ mamba_cache_mode="all",
5530
+ ),
5531
+ ),
5532
+ ],
5533
+ )
5534
+ register_all_kvcache_specs(vllm_config)
5535
+ return Scheduler(
5536
+ vllm_config=vllm_config,
5537
+ kv_cache_config=kv_cache_config,
5538
+ structured_output_manager=StructuredOutputManager(vllm_config),
5539
+ block_size=block_size,
5540
+ hash_block_size=block_size,
5541
+ log_stats=True,
5542
+ )
5543
+
5544
+
5545
+@pytest.mark.parametrize(
5546
+ "matched_tokens,expected_num_computed",
5547
+ [
5548
+ # No external hit: resume on the deepest locally-consistent boundary
5549
+ # (block 0's state survives for both groups).
5550
+ (0, 16),
5551
+ # One external block on top of the reconciled local boundary.
5552
+ (16, 32),
5553
+ ],
5554
+)
5555
+def test_hybrid_per_group_hit_divergence_with_connector(
5556
+ matched_tokens: int, expected_num_computed: int
5557
+):
5558
+ """Per-group prefix hits can diverge for hybrid models with a connector
5559
+ (#46453): under block pressure the FA prefix tail is evicted while a
5560
+ deeper Mamba state block survives. The scheduler must not report the
5561
+ deeper hit as locally computed (evicted FA blocks are not resident ->
5562
+ engine crash / dirty KV); it falls back to the reconciled boundary that
5563
+ every group is consistent at.
5564
+ """
5565
+ block_size = 16
5566
+ scheduler = _create_hybrid_mamba_connector_scheduler(matched_tokens)
5567
+ manager = scheduler.kv_cache_manager
5568
+ assert isinstance(manager.coordinator, HybridKVCacheCoordinator)
5569
+
5570
+ # Seed a 4-block prefix so both groups cache all four boundaries
5571
+ # (mamba cache mode "all" caches every block's state densely).
5572
+ [fill] = create_requests(
5573
+ num_requests=1,
5574
+ num_tokens=4 * block_size,
5575
+ max_tokens=1,
5576
+ same_prompt=True,
5577
+ block_size=block_size,
5578
+ req_ids=["fill"],
5579
+ )
5580
+ computed_blocks, num_computed, _ = manager.get_computed_blocks(fill)
5581
+ blocks = manager.allocate_slots(
5582
+ fill, fill.num_tokens, num_computed, computed_blocks
5583
+ )
5584
+ fa_ids = [b.block_id for b in blocks.blocks[0]]
5585
+ mamba_ids = [b.block_id for b in blocks.blocks[1]]
5586
+ manager.free(fill)
5587
+
5588
+ # Evict the FA tail and the middle mamba states; block 0 (both groups)
5589
+ # and the deep mamba state at block 3 survive.
5590
+ manager.block_pool.evict_blocks({fa_ids[2], fa_ids[3], mamba_ids[1], mamba_ids[2]})
5591
+
5592
+ # A replay of the prefix plus one extra block now sees diverged
5593
+ # per-group hits: FA stops at the evicted tail, while the mamba lookup
5594
+ # finds the deeper surviving state.
5595
+ [replay] = create_requests(
5596
+ num_requests=1,
5597
+ num_tokens=5 * block_size,
5598
+ max_tokens=1,
5599
+ same_prompt=True,
5600
+ block_size=block_size,
5601
+ req_ids=["replay"],
5602
+ )
5603
+ _, per_group_hits = manager.coordinator.find_longest_cache_hit_per_group(
5604
+ replay.block_hashes, replay.num_tokens - 1
5605
+ )
5606
+ assert per_group_hits == (2 * block_size, 4 * block_size) # diverged
5607
+
5608
+ scheduler.add_request(replay)
5609
+ output = scheduler.schedule()
5610
+ num_scheduled = output.num_scheduled_tokens[replay.request_id]
5611
+ assert replay.num_tokens - num_scheduled == expected_num_computed
5612
+
5613
+
5614
+def test_hybrid_per_group_hit_divergence_fa_deeper_no_external():
5615
+ """The opposite divergence: the FA prefix survives deeper than the Mamba
5616
+ state and the connector supplies nothing (ext == 0). Reporting the deep FA
5617
+ hit as locally computed would resume with no valid Mamba state at that
5618
+ boundary (silent bad output). The scheduler must fall back to the
5619
+ convergent boundary that every group agrees on (block 0's surviving state).
5620
+ """
5621
+ block_size = 16
5622
+ scheduler = _create_hybrid_mamba_connector_scheduler(matched_tokens=0)
5623
+ manager = scheduler.kv_cache_manager
5624
+ assert isinstance(manager.coordinator, HybridKVCacheCoordinator)
5625
+
5626
+ # Seed a 4-block prefix in both groups.
5627
+ [fill] = create_requests(
5628
+ num_requests=1,
5629
+ num_tokens=4 * block_size,
5630
+ max_tokens=1,
5631
+ same_prompt=True,
5632
+ block_size=block_size,
5633
+ req_ids=["fill"],
5634
+ )
5635
+ computed_blocks, num_computed, _ = manager.get_computed_blocks(fill)
5636
+ blocks = manager.allocate_slots(
5637
+ fill, fill.num_tokens, num_computed, computed_blocks
5638
+ )
5639
+ mamba_ids = [b.block_id for b in blocks.blocks[1]]
5640
+ manager.free(fill)
5641
+
5642
+ # Keep all FA blocks; evict every mamba state but block 0. FA reaches 4
5643
+ # blocks, the mamba hit only reaches 1 -> diverged (FA > Mamba).
5644
+ manager.block_pool.evict_blocks({mamba_ids[1], mamba_ids[2], mamba_ids[3]})
5645
+
5646
+ [replay] = create_requests(
5647
+ num_requests=1,
5648
+ num_tokens=5 * block_size,
5649
+ max_tokens=1,
5650
+ same_prompt=True,
5651
+ block_size=block_size,
5652
+ req_ids=["replay"],
5653
+ )
5654
+ _, per_group_hits = manager.coordinator.find_longest_cache_hit_per_group(
5655
+ replay.block_hashes, replay.num_tokens - 1
5656
+ )
5657
+ assert per_group_hits == (4 * block_size, 1 * block_size) # FA deeper
5658
+
5659
+ scheduler.add_request(replay)
5660
+ output = scheduler.schedule()
5661
+ num_scheduled = output.num_scheduled_tokens[replay.request_id]
5662
+ # Must resume at the convergent boundary (block 0), not the deep FA hit.
5663
+ assert replay.num_tokens - num_scheduled == block_size
vllm/v1/core/kv_cache_coordinator.py
+9
@@ -634,6 +634,15 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
634
key=lambda g: not isinstance(g.spec, FullAttentionSpec)
635
)
636
637
+ # Dense reference group for per-group lookups (None when the model
638
+ # has no full-attention layers): full attention is downward-closed,
639
+ # so any group reporting a longer per-group hit implies the union of
640
+ # per-group hits is not consistent at a single boundary (#46453).
641
+ first = self.attention_groups[0]
642
+ self.full_attention_group_id: int | None = (
643
+ first.group_ids[0] if isinstance(first.spec, FullAttentionSpec) else None
644
+ )
645
+
646
# Propagate the eagle bit to each manager (default to ``use_eagle=False``).
647
for group in self.attention_groups:
648
if group.use_eagle:
vllm/v1/core/kv_cache_manager.py
+51
-1
@@ -9,7 +9,10 @@ from typing import Literal, overload
9
from vllm.distributed.kv_events import BlockStored, KVCacheEvent
10
from vllm.logger import init_logger
11
from vllm.utils.math_utils import cdiv
12
-from vllm.v1.core.kv_cache_coordinator import get_kv_cache_coordinator
12
+from vllm.v1.core.kv_cache_coordinator import (
13
+ HybridKVCacheCoordinator,
14
+ get_kv_cache_coordinator,
15
+)
16
from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector
17
from vllm.v1.core.kv_cache_utils import KVCacheBlock, KVCacheBlockCopy
18
from vllm.v1.kv_cache_interface import (
@@ -287,6 +290,53 @@ class KVCacheManager:
290
blocks = self.create_kv_cache_blocks(computed_blocks)
291
return blocks, num_new_computed_tokens, shared_prefix_boundary
292
293
+ def get_computed_blocks_for_connector(
294
+ self, request: Request
295
+ ) -> tuple[KVCacheBlocks, int, int, bool]:
296
+ """Local prefix-cache lookup for a request scheduled with a KV connector.
297
+
298
+ Hybrid (Mamba + full-attention) models can have per-group prefix hits
299
+ diverge under block pressure: the full-attention tail may be evicted
300
+ while a deeper Mamba state survives, or vice versa. Report the
301
+ full-attention hit as the local prefix - the connector transfers the
302
+ remaining suffix and the Mamba state is transferred unconditionally by
303
+ nixl's ``_apply_prefix_caching`` - and flag when that hit ran deeper
304
+ than a lagging group. Such a hit only has a valid Mamba state at its
305
+ boundary if the connector supplies it, so the caller must fall back to
306
+ ``get_computed_blocks`` to reconcile when no external tokens are found.
307
+
308
+ Non-hybrid models and already-convergent hits use ``get_computed_blocks``.
309
+
310
+ Returns:
311
+ The ``get_computed_blocks`` triple (blocks, number of local computed
312
+ tokens, shared-prefix boundary) plus ``hit_diverged``.
313
+ """
314
+ coordinator = self.coordinator
315
+ if not (
316
+ self.kv_cache_config.has_mamba_layers
317
+ and isinstance(coordinator, HybridKVCacheCoordinator)
318
+ and coordinator.full_attention_group_id is not None
319
+ ):
320
+ return *self.get_computed_blocks(request), False
321
+
322
+ if not self.prefix_cache_lookup_enabled(request):
323
+ return self.empty_kv_cache_blocks, 0, 0, False
324
+
325
+ fa_group_id = coordinator.full_attention_group_id
326
+ computed, per_group_hits = coordinator.find_longest_cache_hit_per_group(
327
+ request.block_hashes, request.num_tokens - 1
328
+ )
329
+ if any(hit > per_group_hits[fa_group_id] for hit in per_group_hits):
330
+ # A lagging group hit deeper than full attention means its
331
+ # full-attention blocks were evicted; use the reconciled boundary
332
+ # that every group agrees on.
333
+ return *self.get_computed_blocks(request), False
334
+
335
+ num_local = per_group_hits[fa_group_id]
336
+ blocks = self.create_kv_cache_blocks(computed)
337
+ # Per-group lookups do not detect an uncached shared prefix (boundary 0).
338
+ return blocks, num_local, 0, min(per_group_hits) < num_local
339
+
340
def allocate_slots(
341
self,
342
request: Request,
vllm/v1/core/sched/scheduler.py
+22
-39
@@ -35,7 +35,6 @@ from vllm.v1.core.encoder_cache_manager import (
35
EncoderCacheManager,
36
EncoderDecoderCacheManager,
37
)
38
-from vllm.v1.core.kv_cache_coordinator import HybridKVCacheCoordinator
38
from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager
39
from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector
40
from vllm.v1.core.kv_cache_utils import KVCacheBlock
@@ -711,50 +710,24 @@ class Scheduler(SchedulerInterface):
710
# Get already-cached tokens.
711
if request.num_computed_tokens == 0:
712
did_prefix_cache_lookup = True
713
+ hit_diverged = False
714
# Get locally-cached tokens.
715
- if (
716
- self.connector is not None
717
- and self.has_mamba_layers
718
- and isinstance(
719
- self.kv_cache_manager.coordinator, HybridKVCacheCoordinator
715
+ if self.connector is not None:
716
+ # A KV connector transfers the missing suffix, which needs a
717
+ # hybrid-aware lookup that can diverge across groups.
718
+ (
719
+ new_computed_blocks,
720
+ num_new_local_computed_tokens,
721
+ request.shared_prefix_boundary,
722
+ hit_diverged,
723
+ ) = self.kv_cache_manager.get_computed_blocks_for_connector(
724
+ request
725
)
721
- ):
722
- # The per-group lookup does not detect an uncached shared
723
- # prefix, so there is no junction to pin in this path.
724
- request.shared_prefix_boundary = 0
725
- kv_cache_manager = self.kv_cache_manager
726
- if not kv_cache_manager.prefix_cache_lookup_enabled(request):
727
- # Mirror the get_computed_blocks() early-out: the
728
- # request must recompute its prompt.
729
- new_computed_blocks = kv_cache_manager.empty_kv_cache_blocks
730
- num_new_local_computed_tokens = 0
731
- else:
732
- computed, per_group_hits = (
733
- self.kv_cache_manager.coordinator.find_longest_cache_hit_per_group(
734
- request.block_hashes, request.num_tokens - 1
735
- )
736
- )
737
- new_computed_blocks = (
738
- self.kv_cache_manager.create_kv_cache_blocks(computed)
739
- )
740
- # NOTE(ZhanqiuHu): For Mamba hybrid models,
741
- # num_new_local_computed_tokens should be the FA hit
742
- # length. This value is passed to the connector's
743
- # get_num_new_matched_tokens which computes:
744
- # external = total - local_computed.
745
- # Using the FA hit skips re-transferring FA blocks
746
- # already cached on D-side. The Mamba state (always
747
- # the last block) is transferred unconditionally by
748
- # _apply_prefix_caching in nixl/worker.py.
749
- num_new_local_computed_tokens = max(per_group_hits)
726
else:
727
(
728
new_computed_blocks,
729
num_new_local_computed_tokens,
754
- # Junction to pin (Marconi-style APC) so its
755
- # sparse-retention state (Mamba block / sliding-window
756
- # tail) survives retention and serves a later hit; 0
757
- # if no uncached shared prefix was detected.
730
+ # Marconi shared-prefix junction to pin; 0 if none.
731
request.shared_prefix_boundary,
732
) = self.kv_cache_manager.get_computed_blocks(request)
733
@@ -776,6 +749,16 @@ class Scheduler(SchedulerInterface):
749
750
num_external_computed_tokens = ext_tokens
751
752
+ if hit_diverged and num_external_computed_tokens == 0:
753
+ # No external tokens back the deeper local hit, so its
754
+ # resume boundary would have no valid Mamba state.
755
+ # Reconcile to the boundary every group agrees on.
756
+ (
757
+ new_computed_blocks,
758
+ num_new_local_computed_tokens,
759
+ request.shared_prefix_boundary,
760
+ ) = self.kv_cache_manager.get_computed_blocks(request)
761
+
762
connector_prefix_cache_queries = (
763
request.num_tokens - num_new_local_computed_tokens
764
)