[Bugfix][Core] Exempt exactly the blocks an async KV load writes from zeroing

tests/v1/core/test_prefix_replay.py

@@ -4,6 +4,8 @@
 the hit to rebuild the non-cacheable sliding-window group, keeps the hit's
 blocks, and hands the worker the replay start."""
 
+import dataclasses
+
 import pytest
 import torch
 
@@ -277,3 +279,53 @@ def test_hit_no_longer_than_window_is_ignored(via_connector):
     assert new_req.num_computed_tokens == 0
     assert new_req.replay_start == 0
     assert out.num_scheduled_tokens[request.request_id] == NUM_PROMPT_TOKENS
+
+
+@pytest.mark.parametrize("is_async", [True, False])
+@pytest.mark.parametrize("loads_window", [True, False])
+def test_kv_load_zeroing_respects_loaded_groups(is_async, loads_window):
+    """Only groups restored by an async load skip zeroing. P/D loads include
+    the non-cacheable window; stores leave it to be initialized locally."""
+    local = 5 * BLOCK_SIZE
+    scheduler = _replay_scheduler(use_kv_connector=MockKVConfig())
+    if loads_window:
+        scheduler.connector.get_loaded_kv_cache_group_ids = lambda request: (FULL, SWA)
+    assert scheduler.needs_kv_cache_zeroing
+    first = create_requests(
+        num_requests=1, num_tokens=local, same_prompt=True, block_size=BLOCK_SIZE
+    )[0]
+    _prefill(scheduler, first)
+    scheduler.connector.config = dataclasses.replace(
+        scheduler.connector.config,
+        matched_tokens=NUM_PROMPT_TOKENS - 1 - local,
+        is_async=is_async,
+    )
+    request = create_requests(
+        num_requests=1,
+        num_tokens=NUM_PROMPT_TOKENS,
+        same_prompt=True,
+        block_size=BLOCK_SIZE,
+        req_ids=["load"],
+    )[0]
+    scheduler.add_request(request)
+    out = scheduler.schedule()
+
+    manager = scheduler.kv_cache_manager
+    hit_ids = {
+        b.block_id for g in manager.get_blocks(first.request_id).blocks for b in g
+    }
+    swa_blocks = manager.get_blocks(request.request_id).blocks[SWA]
+    assert any(
+        not b.is_null and b.block_id not in hit_ids
+        for b in swa_blocks[: local // BLOCK_SIZE]
+    )
+    zeroed = set(out.new_block_ids_to_zero or ())
+    for group_id, blocks in enumerate(manager.get_blocks(request.request_id).blocks):
+        new_ids = {
+            b.block_id for b in blocks if not b.is_null and b.block_id not in hit_ids
+        }
+        assert new_ids
+        if is_async and (group_id == FULL or loads_window):
+            assert not new_ids & zeroed
+        else:
+            assert new_ids <= zeroed

tests/v1/core/test_single_type_kv_cache_manager.py

@@ -16,6 +16,7 @@
     ChunkedLocalAttentionManager,
     CircularBufferManager,
     FullAttentionManager,
+    KpoolTailManager,
     MambaManager,
     RSWAManager,
     SlidingWindowManager,
@@ -24,6 +25,7 @@
     ChunkedLocalAttentionSpec,
     CircularBufferSpec,
     FullAttentionSpec,
+    KpoolTailSpec,
     MambaSpec,
     RSWASpec,
     SlidingWindowSpec,
@@ -332,6 +334,37 @@ def test_circular_buffer_allocates_one_block_for_the_request_lifetime():
     assert manager.req_to_blocks[request_id] == blocks
 
 
+@pytest.mark.parametrize("record_for_zeroing", [True, False])
+def test_external_kpool_tail_zeroing(record_for_zeroing):
+    """A transferred ring at index zero skips zeroing even with a local hit."""
+    spec = KpoolTailSpec(
+        block_size=4,
+        num_kv_heads=2,
+        head_size=8,
+        dtype=torch.bfloat16,
+        sliding_window=4,
+    )
+    manager = KpoolTailManager(
+        spec,
+        block_pool=BlockPool(10, enable_caching=True, hash_block_size=4),
+        enable_caching=True,
+        kv_cache_group_id=0,
+        scheduler_block_size=4,
+        needs_kv_cache_zeroing=True,
+        max_admission_blocks_per_request=1,
+    )
+    manager.allocate_external_computed_blocks(
+        "request", 16, 4, record_for_zeroing=record_for_zeroing
+    )
+    blocks = manager.req_to_blocks["request"]
+    assert len(blocks) == 1
+    assert manager.take_new_block_ids() == (
+        [blocks[0].block_id] if record_for_zeroing else []
+    )
+    assert manager.allocate_new_blocks("request", 20, 20) == []
+    assert manager.take_new_block_ids() == []
+
+
 def test_sliding_window_records_new_blocks_for_zeroing():
     block_size = 2
     spec = SlidingWindowSpec(

tests/v1/e2e/general/test_mamba_prefix_cache.py

@@ -192,6 +192,7 @@ def fake_allocate_slots_fn(
         full_sequence_must_fit: bool = False,
         reserved_blocks: int = 0,
         has_scheduled_reqs: bool = True,
+        skip_zeroing_group_ids: tuple[int, ...] = (),
     ):
         ret = original_allocate_slots_fn(
             self,
@@ -206,6 +207,7 @@ def fake_allocate_slots_fn(
             full_sequence_must_fit,
             reserved_blocks,
             has_scheduled_reqs,
+            skip_zeroing_group_ids,
         )
         if cur_step_action is not None:
             cur_block_ids = self.coordinator.single_type_managers[0].req_to_blocks[

tests/v1/kv_connector/unit/test_error_propagation.py

@@ -111,6 +111,7 @@ def test_error_propagation_async_load(fail_scheduler: Scheduler):
     }
 
     fail_scheduler.connector = Mock()
+    fail_scheduler.connector.get_loaded_kv_cache_group_ids.return_value = (0,)
     fail_scheduler.connector.get_num_new_matched_tokens.side_effect = (
         _make_get_num_new_matched_tokens(req_num_new_matched_tokens, True)
     )

tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py

@@ -327,6 +327,7 @@ def test_async_recompute_blocks_not_cached_when_invalid(
     recompute_scheduler.connector.get_num_new_matched_tokens.side_effect = (
         _make_get_num_new_matched_tokens(req_num_new_matched_tokens, True)
     )
+    recompute_scheduler.connector.get_loaded_kv_cache_group_ids.return_value = (0,)
     recompute_scheduler.connector.request_finished.return_value = (False, None)
     recompute_scheduler.connector.take_events.return_value = ()
 

tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py

@@ -42,6 +42,7 @@ def test_failed_receive_completion_honors_failure_policy(policy):
     request = create_request(num_tokens=3 * scheduler.block_size)
     scheduler.add_request(request)
     scheduler.connector = Mock()
+    scheduler.connector.get_loaded_kv_cache_group_ids.return_value = (0,)
     scheduler.connector.get_num_new_matched_tokens.side_effect = [
         (2 * scheduler.block_size, True),
         (0, False),
@@ -103,6 +104,7 @@ def test_async_load_failure(
     }
 
     scheduler.connector = Mock()
+    scheduler.connector.get_loaded_kv_cache_group_ids.return_value = (0,)
     scheduler.connector.get_num_new_matched_tokens.side_effect = (
         _make_get_num_new_matched_tokens(req_num_new_matched_tokens, async_load=True)
     )
@@ -332,6 +334,7 @@ def test_async_progressive_load_failure(
     }
 
     scheduler.connector = Mock()
+    scheduler.connector.get_loaded_kv_cache_group_ids.return_value = (0,)
     scheduler.connector.get_num_new_matched_tokens.side_effect = (
         _make_get_num_new_matched_tokens(req_num_new_matched_tokens, async_load=True)
     )

tests/v1/kv_connector/unit/test_multi_connector.py

@@ -202,6 +202,20 @@ def mc() -> MultiConnector:
     return mc
 
 
+@pytest.mark.parametrize("chosen", [0, 1])
+def test_loaded_groups_follow_the_connector_serving_the_request(mc, chosen):
+    request = MagicMock(request_id="request")
+    for i, connector in enumerate(mc._connectors):
+        connector.get_num_new_matched_tokens.return_value = (
+            (16, True) if i == chosen else (0, False)
+        )
+        connector.get_loaded_kv_cache_group_ids.return_value = (i,)
+
+    assert mc.get_num_new_matched_tokens(request, 0) == (16, True)
+    assert mc.get_loaded_kv_cache_group_ids(request) == (chosen,)
+    mc._connectors[1 - chosen].get_loaded_kv_cache_group_ids.assert_not_called()
+
+
 def test_multi_connector_mem_pool_context_none(mc: MultiConnector):
     assert mc.get_mem_pool_context() is None
 

tests/v1/kv_connector/unit/test_nixl_connector_hma.py

@@ -1522,34 +1522,6 @@ def _make_fake_kv_cache_manager():
     return manager
 
 
-@pytest.mark.cpu_test
-def test_zeroing_block_ids_cover_only_loaded_attention_blocks():
-    """Only zero-recorded (attention) groups contribute, sliced to the
-    externally-loaded token range; Mamba state blocks are never zeroed."""
-    manager = _make_fake_kv_cache_manager()
-
-    # Tokens [0, 16) are locally cached; the load covers tokens [16, 56).
-    assert manager.get_zeroing_block_ids_in_range("req-1", 16, 56) == [11, 12, 13]
-
-
-@pytest.mark.cpu_test
-def test_scheduler_filters_connector_loaded_blocks_from_zeroing():
-    """Blocks that will be loaded by the connector must not be zeroed."""
-    from vllm.v1.core.sched.scheduler import Scheduler
-
-    class FakeKVCacheManager:
-        def take_new_block_ids(self):
-            return [9, 10, 11, 12]
-
-    scheduler = object.__new__(Scheduler)
-    scheduler.needs_kv_cache_zeroing = True
-    scheduler.kv_cache_manager = FakeKVCacheManager()
-    scheduler._skip_zero_block_ids = {10, 12}
-
-    assert scheduler._get_new_block_ids_to_zero() == [9, 11]
-    assert not scheduler._skip_zero_block_ids
-
-
 @pytest.mark.cpu_test
 def test_failed_load_rezeroes_unwritten_skipped_blocks():
     """A failed async load leaves zeroing-skipped blocks unwritten beyond
@@ -1575,7 +1547,6 @@ def test_failed_load_rezeroes_unwritten_skipped_blocks():
 
     # Attention blocks covering tokens >= 48 are re-recorded for zeroing
     # and flow into the next step's zero list; Mamba blocks are not.
-    scheduler._skip_zero_block_ids = set()
     assert scheduler._get_new_block_ids_to_zero() == [13, 14, 15]
 
 

vllm/distributed/kv_transfer/kv_connector/v1/base.py

@@ -482,6 +482,16 @@ def bind_gpu_block_pool(self, gpu_block_pool: "BlockPool") -> None:
         """
         return
 
+    def get_loaded_kv_cache_group_ids(self, request: "Request") -> tuple[int, ...]:
+        """KV cache groups restored by this connector's load for ``request``.
+
+        Called after ``get_num_new_matched_tokens`` returned a positive count.
+        Defaults to the prefix-cacheable groups, all a hash-addressed store
+        holds; a connector transferring the request's own blocks restores
+        every transfer group.
+        """
+        return self._kv_cache_config.prefix_cacheable_group_ids
+
     @abstractmethod
     def get_num_new_matched_tokens(
         self,

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py

@@ -748,6 +748,9 @@ def get_required_kvcache_layout(cls, vllm_config: VllmConfig):
     # Scheduler Side Methods
     ############################################################
 
+    def get_loaded_kv_cache_group_ids(self, request: "Request") -> tuple[int, ...]:
+        return self._kv_cache_config.transfer_group_ids
+
     def get_num_new_matched_tokens(
         self, request: "Request", num_computed_tokens: int
     ) -> tuple[int, bool]:

vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py

@@ -426,6 +426,10 @@ def build_connector_worker_meta(self) -> KVConnectorWorkerMetadata | None:
     # ==============================
     # Scheduler-side methods
     # ==============================
+    def get_loaded_kv_cache_group_ids(self, request: "Request") -> tuple[int, ...]:
+        connector = self._connectors[self._requests_to_connector[request.request_id]]
+        return connector.get_loaded_kv_cache_group_ids(request)
+
     def get_num_new_matched_tokens(
         self,
         request: "Request",

vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py

@@ -164,6 +164,9 @@ def get_required_kvcache_layout(cls, vllm_config: VllmConfig):
     def get_finished_count(self) -> int | None:
         return None
 
+    def get_loaded_kv_cache_group_ids(self, request: "Request") -> tuple[int, ...]:
+        return self._kv_cache_config.transfer_group_ids
+
     def get_num_new_matched_tokens(
         self, request: "Request", num_computed_tokens: int
     ) -> tuple[int | None, bool]:

vllm/v1/core/kv_cache_coordinator.py

@@ -241,6 +241,7 @@ def allocate_new_computed_blocks(
         new_computed_blocks: tuple[Sequence[KVCacheBlock], ...],
         num_local_computed_tokens: int,
         num_external_computed_tokens: int,
+        skip_zeroing_group_ids: tuple[int, ...] = (),
     ) -> None:
         """Add the new computed blocks to the request. Optionally allocate new
             blocks for external computed tokens (if any).
@@ -251,6 +252,8 @@ def allocate_new_computed_blocks(
                 prefix cache.
             num_local_computed_tokens: The number of local computed tokens.
             num_external_computed_tokens: The number of external computed tokens.
+            skip_zeroing_group_ids: Groups whose external-token blocks are
+                written by an async load and must not be zeroed.
 
         """
         # A running request is already tracked in num_cached_block and won't
@@ -274,11 +277,12 @@ def allocate_new_computed_blocks(
                 num_external_computed_tokens,
             )
         if num_external_computed_tokens > 0:
-            for manager in self.single_type_managers:
+            for i, manager in enumerate(self.single_type_managers):
                 manager.allocate_external_computed_blocks(
                     request_id,
                     num_local_computed_tokens,
                     num_external_computed_tokens,
+                    record_for_zeroing=i not in skip_zeroing_group_ids,
                 )
 
     def allocate_new_blocks(

vllm/v1/core/kv_cache_manager.py

@@ -381,6 +381,7 @@ def allocate_slots(
         full_sequence_must_fit: bool = False,
         reserved_blocks: int = 0,
         has_scheduled_reqs: bool = True,
+        skip_zeroing_group_ids: tuple[int, ...] = (),
     ) -> KVCacheBlocks | None:
         """Add slots for a request with new tokens to append.
 
@@ -413,6 +414,8 @@ def allocate_slots(
                 blocks an already in-flight (prefilling) sequence is relying on.
             has_scheduled_reqs: Whether any requests are already scheduled to run
                 this step, controls whether watermark is applied.
+            skip_zeroing_group_ids: Groups whose external-token blocks will be
+                written by an async load and must not be zeroed concurrently.
 
         Blocks layout:
         ```
@@ -586,6 +589,7 @@ def allocate_slots(
                 new_computed_blocks=new_computed_block_list,
                 num_local_computed_tokens=num_local_computed_tokens,
                 num_external_computed_tokens=num_external_computed_tokens,
+                skip_zeroing_group_ids=skip_zeroing_group_ids,
             )
 
         new_blocks = self.coordinator.allocate_new_blocks(
@@ -863,20 +867,6 @@ def take_new_block_ids(self) -> list[int]:
             ids.extend(mgr.take_new_block_ids())
         return ids
 
-    def get_zeroing_block_ids_in_range(
-        self, request_id: str, start_token: int, end_token: int
-    ) -> list[int]:
-        """The request's block ids covering [start_token, end_token), from
-        the groups whose new blocks are zeroed by the worker."""
-        ids: list[int] = []
-        for mgr in self.coordinator.single_type_managers:
-            if mgr.records_new_block_ids:
-                start_idx = start_token // mgr.block_size
-                end_idx = cdiv(end_token, mgr.block_size)
-                blocks = mgr.req_to_blocks[request_id]
-                ids.extend(blk.block_id for blk in blocks[start_idx:end_idx])
-        return ids
-
     def record_blocks_for_zeroing(self, request_id: str, start_token: int) -> None:
         """Re-record the request's blocks from start_token onwards for
         zeroing, e.g. blocks a failed async KV load left unwritten.