[BugFix] Scope divergent hybrid cache hits to capable connectors (#50344)
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: OpenAI Codex <codex@openai.com>
Yifan Qiao committed
Aug 9, 2026 at 15:38 UTC
d6941300fcb9d4a8bbea19f8b610c2aff9fc5cc3
8 files changed
+130
-27
tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py
+36
@@ -5,6 +5,7 @@
5
on partial hits, and same-step deferral."""
6
7
from types import SimpleNamespace
8
+from unittest.mock import MagicMock
9
10
import pytest
11
import torch
@@ -31,6 +32,41 @@ def _auto_init_hash_fn():
32
init_none_hash(sha256)
33
34
35
+def test_connector_without_divergent_hit_support_uses_common_lookup():
36
+ common_blocks = MagicMock()
37
+ manager = MagicMock()
38
+ manager.get_computed_blocks.return_value = (common_blocks, 0, 0)
39
+ scheduler = SimpleNamespace(
40
+ connector=SimpleNamespace(supports_divergent_local_hybrid_hits=False),
41
+ kv_cache_manager=manager,
42
+ )
43
+
44
+ result = Scheduler._get_local_prefix_cache_hit(scheduler, MagicMock())
45
+
46
+ assert result == (common_blocks, 0, 0, False)
47
+ manager.get_computed_blocks_for_connector.assert_not_called()
48
+
49
+
50
+def test_capable_connector_uses_divergent_partial_hit_lookup():
51
+ per_group_blocks = MagicMock()
52
+ manager = MagicMock()
53
+ manager.get_computed_blocks_for_connector.return_value = (
54
+ per_group_blocks,
55
+ 6,
56
+ 0,
57
+ True,
58
+ )
59
+ scheduler = SimpleNamespace(
60
+ connector=SimpleNamespace(supports_divergent_local_hybrid_hits=True),
61
+ kv_cache_manager=manager,
62
+ )
63
+
64
+ result = Scheduler._get_local_prefix_cache_hit(scheduler, MagicMock())
65
+
66
+ assert result == (per_group_blocks, 6, 0, True)
67
+ manager.get_computed_blocks.assert_not_called()
68
+
69
+
70
def test_mamba_align_split_partial_tail_schedule():
71
"""Chunk ends with partial hits on: block-aligned chunks, one extra stop
72
at the prompt's last hash boundary (registering the partial tail), then
tests/v1/core/test_scheduler.py
+20
-7
@@ -5547,6 +5547,7 @@ def _create_hybrid_mamba_connector_scheduler(
5547
matched_tokens: int,
5548
block_size: int = 16,
5549
num_blocks: int = 100,
5550
+ supports_divergent_hits: bool = True,
5551
) -> Scheduler:
5552
"""FA + Mamba ("all" cache mode) scheduler with a MockKVConnector."""
5553
model_config = ModelConfig(
@@ -5577,6 +5578,7 @@ def _create_hybrid_mamba_connector_scheduler(
5578
kv_connector_extra_config={
5579
"matched_tokens": matched_tokens,
5580
"is_async": False,
5581
+ "supports_divergent_local_hybrid_hits": supports_divergent_hits,
5582
},
5583
),
5584
)
@@ -5686,22 +5688,33 @@ def test_hybrid_per_group_hit_divergence_with_connector(
5688
5689
5690
@pytest.mark.parametrize(
5689
- ("matched_tokens", "replay_blocks", "expected_num_computed"),
5691
+ (
5692
+ "supports_divergent_local_hybrid_hits",
5693
+ "matched_tokens",
5694
+ "replay_blocks",
5695
+ "expected_num_computed",
5696
+ ),
5697
[
5691
- (0, 5, 16),
5692
- (16, 6, 80),
5698
+ (True, 0, 5, 16),
5699
+ (True, 16, 6, 80),
5700
+ (False, 16, 6, 32),
5701
],
5702
)
5695
-def test_hybrid_fa_deeper_hit_uses_external_mamba_state(
5703
+def test_hybrid_fa_deeper_hit_respects_connector_lookup_policy(
5704
+ supports_divergent_local_hybrid_hits: bool,
5705
matched_tokens: int,
5706
replay_blocks: int,
5707
expected_num_computed: int,
5708
):
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.
5709
+ """A capable connector may restore missing Mamba state at the deeper FA
5710
+ boundary. An external miss or an incapable connector uses a locally
5711
+ consistent boundary instead.
5712
"""
5713
block_size = 16
5704
- scheduler = _create_hybrid_mamba_connector_scheduler(matched_tokens)
5714
+ scheduler = _create_hybrid_mamba_connector_scheduler(
5715
+ matched_tokens,
5716
+ supports_divergent_hits=supports_divergent_local_hybrid_hits,
5717
+ )
5718
manager = scheduler.kv_cache_manager
5719
assert isinstance(manager.coordinator, HybridKVCacheCoordinator)
5720
tests/v1/kv_connector/unit/test_multi_connector.py
+29
@@ -53,9 +53,14 @@ class MockConnectorStats(KVConnectorStats):
53
class MockConnector(KVConnectorBase_V1):
54
"""Mock connector for testing."""
55
56
+ _supports_divergent_local_hybrid_hits = False
57
+
58
def __new__(cls, *args, **kwargs):
59
# mock all KVConnectorBase_V1 functions
60
mock = MagicMock(spec_set=KVConnectorBase_V1)
61
+ mock.supports_divergent_local_hybrid_hits = (
62
+ cls._supports_divergent_local_hybrid_hits
63
+ )
64
# Override just build_kv_connector_stats
65
mock.build_kv_connector_stats = cls.build_kv_connector_stats
66
mock.get_kv_connector_stats.return_value = None
@@ -92,8 +97,13 @@ class MockConnector(KVConnectorBase_V1):
97
class MockHMAConnector(KVConnectorBase_V1, SupportsHMA):
98
"""Mock connector that supports HMA for testing."""
99
100
+ _supports_divergent_local_hybrid_hits = False
101
+
102
def __new__(cls, *args, **kwargs):
103
mock = MagicMock(spec_set=cls)
104
+ mock.supports_divergent_local_hybrid_hits = (
105
+ cls._supports_divergent_local_hybrid_hits
106
+ )
107
mock.get_kv_connector_stats.return_value = None
108
return mock
109
@@ -122,11 +132,20 @@ class MockHMAConnector(KVConnectorBase_V1, SupportsHMA):
132
return (False, None)
133
134
135
+class MockDivergentHMAConnector(MockHMAConnector):
136
+ _supports_divergent_local_hybrid_hits = True
137
+
138
+
139
# Register mock connectors
140
KVConnectorFactory.register_connector("MockConnector", __name__, MockConnector.__name__)
141
KVConnectorFactory.register_connector(
142
"MockHMAConnector", __name__, MockHMAConnector.__name__
143
)
144
+KVConnectorFactory.register_connector(
145
+ "MockDivergentHMAConnector",
146
+ __name__,
147
+ MockDivergentHMAConnector.__name__,
148
+)
149
150
151
@pytest.fixture
@@ -1059,6 +1078,16 @@ def test_multi_connector_hma_support_detection():
1078
assert mc_mixed2._all_support_hma is False
1079
1080
1081
+def test_divergent_local_hybrid_hit_capability_is_conservative():
1082
+ all_supported = _make_multi_connector(
1083
+ ["MockDivergentHMAConnector", "MockDivergentHMAConnector"]
1084
+ )
1085
+ assert all_supported.supports_divergent_local_hybrid_hits is True
1086
+
1087
+ mixed = _make_multi_connector(["MockDivergentHMAConnector", "MockHMAConnector"])
1088
+ assert mixed.supports_divergent_local_hybrid_hits is False
1089
+
1090
+
1091
@pytest.mark.skipif(
1092
not torch.cuda.is_available(), reason="Requires GPU to instantiate LLM"
1093
)
tests/v1/kv_connector/unit/utils.py
+8
@@ -363,6 +363,7 @@ class MockKVConfig:
363
matched_tokens: int = 0
364
is_async: bool = False
365
num_defers_before_matching: int = 0
366
+ supports_divergent_local_hybrid_hits: bool = False
367
368
369
class MockKVConnectorMetadata(KVConnectorMetadata):
@@ -388,11 +389,18 @@ class MockKVConnector(KVConnectorBase_V1):
389
num_defers_before_matching=extra_config.get(
390
"num_defers_before_matching", 0
391
),
392
+ supports_divergent_local_hybrid_hits=extra_config.get(
393
+ "supports_divergent_local_hybrid_hits", False
394
+ ),
395
)
396
self._defers_left: defaultdict[str, int] = defaultdict(
397
lambda: self.config.num_defers_before_matching
398
)
399
400
+ @property
401
+ def supports_divergent_local_hybrid_hits(self) -> bool:
402
+ return self.config.supports_divergent_local_hybrid_hits
403
+
404
def get_num_new_matched_tokens(
405
self,
406
request: Request,
vllm/distributed/kv_transfer/kv_connector/v1/base.py
+9
@@ -173,6 +173,15 @@ class KVConnectorBase_V1(ABC):
173
Base class for KV connectors.
174
"""
175
176
+ @property
177
+ def supports_divergent_local_hybrid_hits(self) -> bool:
178
+ """Whether external hits can complete divergent local hybrid hits.
179
+
180
+ A capable connector restores lagging recurrent state when the local
181
+ full-attention group reaches a deeper boundary. Defaults to False.
182
+ """
183
+ return False
184
+
185
@property
186
def prefer_cross_layer_blocks(self) -> bool:
187
"""
vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+6
@@ -203,6 +203,12 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):
203
# Propagated from scheduler to worker side via the connector metadata.
204
self._extra_async_saves: dict[str, int] = {}
205
206
+ @property
207
+ def supports_divergent_local_hybrid_hits(self) -> bool:
208
+ return bool(self._connectors) and all(
209
+ c.supports_divergent_local_hybrid_hits for c in self._connectors
210
+ )
211
+
212
@property
213
def prefer_cross_layer_blocks(self) -> bool:
214
if not self._connectors:
vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py
+4
@@ -79,6 +79,10 @@ logger = init_logger(__name__)
79
class NixlBaseConnector(KVConnectorBase_V1, SupportsHMA):
80
"""Base connector with common logic shared by pull and push modes."""
81
82
+ @property
83
+ def supports_divergent_local_hybrid_hits(self) -> bool:
84
+ return True
85
+
86
@property
87
def prefer_cross_layer_blocks(self) -> bool:
88
if any(
vllm/v1/core/sched/scheduler.py
+18
-20
@@ -423,6 +423,18 @@ class Scheduler(SchedulerInterface):
423
end = min((s for s in stops if start < s < end), default=end)
424
return max(end - start, 0)
425
426
+ def _get_local_prefix_cache_hit(
427
+ self, request: Request
428
+ ) -> tuple[KVCacheBlocks, int, int, bool]:
429
+ connector = self.connector
430
+ if connector is not None and connector.supports_divergent_local_hybrid_hits:
431
+ return self.kv_cache_manager.get_computed_blocks_for_connector(request)
432
+
433
+ blocks, num_local, shared_prefix_boundary = (
434
+ self.kv_cache_manager.get_computed_blocks(request)
435
+ )
436
+ return blocks, num_local, shared_prefix_boundary, False
437
+
438
def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
439
self.current_step += 1
440
# NOTE(woosuk) on the scheduling algorithm:
@@ -739,26 +751,12 @@ class Scheduler(SchedulerInterface):
751
# Get already-cached tokens.
752
if request.num_computed_tokens == 0:
753
did_prefix_cache_lookup = True
742
- hit_diverged = False
743
- # Get locally-cached tokens.
744
- if self.connector is not None:
745
- # A KV connector transfers the missing suffix, which needs a
746
- # hybrid-aware lookup that can diverge across groups.
747
- (
748
- new_computed_blocks,
749
- num_new_local_computed_tokens,
750
- request.shared_prefix_boundary,
751
- hit_diverged,
752
- ) = self.kv_cache_manager.get_computed_blocks_for_connector(
753
- request
754
- )
755
- else:
756
- (
757
- new_computed_blocks,
758
- num_new_local_computed_tokens,
759
- # Marconi shared-prefix junction to pin; 0 if none.
760
- request.shared_prefix_boundary,
761
- ) = self.kv_cache_manager.get_computed_blocks(request)
754
+ (
755
+ new_computed_blocks,
756
+ num_new_local_computed_tokens,
757
+ request.shared_prefix_boundary,
758
+ hit_diverged,
759
+ ) = self._get_local_prefix_cache_hit(request)
760
761
# Get externally-cached tokens if using a KVConnector.
762
if self.connector is not None: