[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>

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

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)
 

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"
 )

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,

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:
         """

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:

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(

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: