From d6941300fcb9d4a8bbea19f8b610c2aff9fc5cc3 Mon Sep 17 00:00:00 2001 From: Yifan Qiao Date: Sun, 9 Aug 2026 15:38:13 -0700 Subject: [PATCH] [BugFix] Scope divergent hybrid cache hits to capable connectors (#50344) Signed-off-by: Yifan Qiao Co-authored-by: OpenAI Codex --- .../test_partial_prefix_cache_hits.py | 36 ++++++++++++++++++ tests/v1/core/test_scheduler.py | 27 +++++++++---- .../kv_connector/unit/test_multi_connector.py | 29 ++++++++++++++ tests/v1/kv_connector/unit/utils.py | 8 ++++ .../kv_transfer/kv_connector/v1/base.py | 9 +++++ .../kv_connector/v1/multi_connector.py | 6 +++ .../kv_connector/v1/nixl/connector.py | 4 ++ vllm/v1/core/sched/scheduler.py | 38 +++++++++---------- 8 files changed, 130 insertions(+), 27 deletions(-) diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py index 543991a874..5f549ec8df 100644 --- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +++ b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py @@ -5,6 +5,7 @@ on partial hits, and same-step deferral.""" from types import SimpleNamespace +from unittest.mock import MagicMock import pytest import torch @@ -31,6 +32,41 @@ def _auto_init_hash_fn(): init_none_hash(sha256) +def test_connector_without_divergent_hit_support_uses_common_lookup(): + common_blocks = MagicMock() + manager = MagicMock() + manager.get_computed_blocks.return_value = (common_blocks, 0, 0) + scheduler = SimpleNamespace( + connector=SimpleNamespace(supports_divergent_local_hybrid_hits=False), + kv_cache_manager=manager, + ) + + result = Scheduler._get_local_prefix_cache_hit(scheduler, MagicMock()) + + assert result == (common_blocks, 0, 0, False) + manager.get_computed_blocks_for_connector.assert_not_called() + + +def test_capable_connector_uses_divergent_partial_hit_lookup(): + per_group_blocks = MagicMock() + manager = MagicMock() + manager.get_computed_blocks_for_connector.return_value = ( + per_group_blocks, + 6, + 0, + True, + ) + scheduler = SimpleNamespace( + connector=SimpleNamespace(supports_divergent_local_hybrid_hits=True), + kv_cache_manager=manager, + ) + + result = Scheduler._get_local_prefix_cache_hit(scheduler, MagicMock()) + + assert result == (per_group_blocks, 6, 0, True) + manager.get_computed_blocks.assert_not_called() + + def test_mamba_align_split_partial_tail_schedule(): """Chunk ends with partial hits on: block-aligned chunks, one extra stop at the prompt's last hash boundary (registering the partial tail), then diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 30ab5bc09d..75f9402394 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -5547,6 +5547,7 @@ def _create_hybrid_mamba_connector_scheduler( matched_tokens: int, block_size: int = 16, num_blocks: int = 100, + supports_divergent_hits: bool = True, ) -> Scheduler: """FA + Mamba ("all" cache mode) scheduler with a MockKVConnector.""" model_config = ModelConfig( @@ -5577,6 +5578,7 @@ def _create_hybrid_mamba_connector_scheduler( kv_connector_extra_config={ "matched_tokens": matched_tokens, "is_async": False, + "supports_divergent_local_hybrid_hits": supports_divergent_hits, }, ), ) @@ -5686,22 +5688,33 @@ def test_hybrid_per_group_hit_divergence_with_connector( @pytest.mark.parametrize( - ("matched_tokens", "replay_blocks", "expected_num_computed"), + ( + "supports_divergent_local_hybrid_hits", + "matched_tokens", + "replay_blocks", + "expected_num_computed", + ), [ - (0, 5, 16), - (16, 6, 80), + (True, 0, 5, 16), + (True, 16, 6, 80), + (False, 16, 6, 32), ], ) -def test_hybrid_fa_deeper_hit_uses_external_mamba_state( +def test_hybrid_fa_deeper_hit_respects_connector_lookup_policy( + supports_divergent_local_hybrid_hits: bool, matched_tokens: int, replay_blocks: int, expected_num_computed: int, ): - """A positive external hit supplies the Mamba state at the resume - boundary; without one, the scheduler falls back to the common local hit. + """A capable connector may restore missing Mamba state at the deeper FA + boundary. An external miss or an incapable connector uses a locally + consistent boundary instead. """ block_size = 16 - scheduler = _create_hybrid_mamba_connector_scheduler(matched_tokens) + scheduler = _create_hybrid_mamba_connector_scheduler( + matched_tokens, + supports_divergent_hits=supports_divergent_local_hybrid_hits, + ) manager = scheduler.kv_cache_manager assert isinstance(manager.coordinator, HybridKVCacheCoordinator) diff --git a/tests/v1/kv_connector/unit/test_multi_connector.py b/tests/v1/kv_connector/unit/test_multi_connector.py index 14ea6759f4..ae2cd4b8e0 100644 --- a/tests/v1/kv_connector/unit/test_multi_connector.py +++ b/tests/v1/kv_connector/unit/test_multi_connector.py @@ -53,9 +53,14 @@ class MockConnectorStats(KVConnectorStats): class MockConnector(KVConnectorBase_V1): """Mock connector for testing.""" + _supports_divergent_local_hybrid_hits = False + def __new__(cls, *args, **kwargs): # mock all KVConnectorBase_V1 functions mock = MagicMock(spec_set=KVConnectorBase_V1) + mock.supports_divergent_local_hybrid_hits = ( + cls._supports_divergent_local_hybrid_hits + ) # Override just build_kv_connector_stats mock.build_kv_connector_stats = cls.build_kv_connector_stats mock.get_kv_connector_stats.return_value = None @@ -92,8 +97,13 @@ class MockConnector(KVConnectorBase_V1): class MockHMAConnector(KVConnectorBase_V1, SupportsHMA): """Mock connector that supports HMA for testing.""" + _supports_divergent_local_hybrid_hits = False + def __new__(cls, *args, **kwargs): mock = MagicMock(spec_set=cls) + mock.supports_divergent_local_hybrid_hits = ( + cls._supports_divergent_local_hybrid_hits + ) mock.get_kv_connector_stats.return_value = None return mock @@ -122,11 +132,20 @@ class MockHMAConnector(KVConnectorBase_V1, SupportsHMA): return (False, None) +class MockDivergentHMAConnector(MockHMAConnector): + _supports_divergent_local_hybrid_hits = True + + # Register mock connectors KVConnectorFactory.register_connector("MockConnector", __name__, MockConnector.__name__) KVConnectorFactory.register_connector( "MockHMAConnector", __name__, MockHMAConnector.__name__ ) +KVConnectorFactory.register_connector( + "MockDivergentHMAConnector", + __name__, + MockDivergentHMAConnector.__name__, +) @pytest.fixture @@ -1059,6 +1078,16 @@ def test_multi_connector_hma_support_detection(): assert mc_mixed2._all_support_hma is False +def test_divergent_local_hybrid_hit_capability_is_conservative(): + all_supported = _make_multi_connector( + ["MockDivergentHMAConnector", "MockDivergentHMAConnector"] + ) + assert all_supported.supports_divergent_local_hybrid_hits is True + + mixed = _make_multi_connector(["MockDivergentHMAConnector", "MockHMAConnector"]) + assert mixed.supports_divergent_local_hybrid_hits is False + + @pytest.mark.skipif( not torch.cuda.is_available(), reason="Requires GPU to instantiate LLM" ) diff --git a/tests/v1/kv_connector/unit/utils.py b/tests/v1/kv_connector/unit/utils.py index 7d4b17ab39..b1fb433533 100644 --- a/tests/v1/kv_connector/unit/utils.py +++ b/tests/v1/kv_connector/unit/utils.py @@ -363,6 +363,7 @@ class MockKVConfig: matched_tokens: int = 0 is_async: bool = False num_defers_before_matching: int = 0 + supports_divergent_local_hybrid_hits: bool = False class MockKVConnectorMetadata(KVConnectorMetadata): @@ -388,11 +389,18 @@ class MockKVConnector(KVConnectorBase_V1): num_defers_before_matching=extra_config.get( "num_defers_before_matching", 0 ), + supports_divergent_local_hybrid_hits=extra_config.get( + "supports_divergent_local_hybrid_hits", False + ), ) self._defers_left: defaultdict[str, int] = defaultdict( lambda: self.config.num_defers_before_matching ) + @property + def supports_divergent_local_hybrid_hits(self) -> bool: + return self.config.supports_divergent_local_hybrid_hits + def get_num_new_matched_tokens( self, request: Request, diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/base.py b/vllm/distributed/kv_transfer/kv_connector/v1/base.py index 076d40a187..e49dbebb12 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/base.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/base.py @@ -173,6 +173,15 @@ class KVConnectorBase_V1(ABC): Base class for KV connectors. """ + @property + def supports_divergent_local_hybrid_hits(self) -> bool: + """Whether external hits can complete divergent local hybrid hits. + + A capable connector restores lagging recurrent state when the local + full-attention group reaches a deeper boundary. Defaults to False. + """ + return False + @property def prefer_cross_layer_blocks(self) -> bool: """ diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py index 59dad186c6..957b15ea40 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py @@ -203,6 +203,12 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA): # Propagated from scheduler to worker side via the connector metadata. self._extra_async_saves: dict[str, int] = {} + @property + def supports_divergent_local_hybrid_hits(self) -> bool: + return bool(self._connectors) and all( + c.supports_divergent_local_hybrid_hits for c in self._connectors + ) + @property def prefer_cross_layer_blocks(self) -> bool: if not self._connectors: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py index 7d9a34406b..0b7a5d5728 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py @@ -79,6 +79,10 @@ logger = init_logger(__name__) class NixlBaseConnector(KVConnectorBase_V1, SupportsHMA): """Base connector with common logic shared by pull and push modes.""" + @property + def supports_divergent_local_hybrid_hits(self) -> bool: + return True + @property def prefer_cross_layer_blocks(self) -> bool: if any( diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index f00c1f7b24..006da8c7d1 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -423,6 +423,18 @@ class Scheduler(SchedulerInterface): end = min((s for s in stops if start < s < end), default=end) return max(end - start, 0) + def _get_local_prefix_cache_hit( + self, request: Request + ) -> tuple[KVCacheBlocks, int, int, bool]: + connector = self.connector + if connector is not None and connector.supports_divergent_local_hybrid_hits: + return self.kv_cache_manager.get_computed_blocks_for_connector(request) + + blocks, num_local, shared_prefix_boundary = ( + self.kv_cache_manager.get_computed_blocks(request) + ) + return blocks, num_local, shared_prefix_boundary, False + def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: self.current_step += 1 # NOTE(woosuk) on the scheduling algorithm: @@ -739,26 +751,12 @@ class Scheduler(SchedulerInterface): # Get already-cached tokens. if request.num_computed_tokens == 0: did_prefix_cache_lookup = True - hit_diverged = False - # Get locally-cached tokens. - if self.connector is not None: - # A KV connector transfers the missing suffix, which needs a - # hybrid-aware lookup that can diverge across groups. - ( - new_computed_blocks, - num_new_local_computed_tokens, - request.shared_prefix_boundary, - hit_diverged, - ) = self.kv_cache_manager.get_computed_blocks_for_connector( - request - ) - else: - ( - new_computed_blocks, - num_new_local_computed_tokens, - # Marconi shared-prefix junction to pin; 0 if none. - request.shared_prefix_boundary, - ) = self.kv_cache_manager.get_computed_blocks(request) + ( + new_computed_blocks, + num_new_local_computed_tokens, + request.shared_prefix_boundary, + hit_diverged, + ) = self._get_local_prefix_cache_hit(request) # Get externally-cached tokens if using a KVConnector. if self.connector is not None: