[Bugfix][Core] Exempt exactly the blocks an async KV load writes from zeroing (#59504)

Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: Codex <noreply@openai.com>

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 @@ from vllm.v1.core.single_type_kv_cache_manager import (
     ChunkedLocalAttentionManager,
     CircularBufferManager,
     FullAttentionManager,
+    KpoolTailManager,
     MambaManager,
     RSWAManager,
     SlidingWindowManager,
@@ -24,6 +25,7 @@ from vllm.v1.kv_cache_interface import (
     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 get_fake_allocate_slots_fn(original_allocate_slots_fn: Callable):
         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 get_fake_allocate_slots_fn(original_allocate_slots_fn: Callable):
             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 @@ class KVConnectorBase_V1(ABC):
         """
         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 @@ class MooncakeConnector(KVConnectorBase_V1, SupportsHMA):
     # 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 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):
     # ==============================
     # 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 @@ class NixlBaseConnector(KVConnectorBase_V1, SupportsHMA):
     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 @@ class KVCacheCoordinator(ABC):
         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 @@ class KVCacheCoordinator(ABC):
                 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 @@ class KVCacheCoordinator(ABC):
                 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 @@ class KVCacheManager:
         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 @@ class KVCacheManager:
                 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 @@ class KVCacheManager:
                 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 @@ class KVCacheManager:
             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.

vllm/v1/core/sched/scheduler.py

@@ -353,9 +353,6 @@ class Scheduler(SchedulerInterface):
 
         self.has_mamba_layers = kv_cache_config.has_mamba_layers
         self.needs_kv_cache_zeroing = kv_cache_config.needs_kv_cache_zeroing
-        # Blocks that async KV loads will overwrite this step, skipped from
-        # zeroing since the zeroing could race the out-of-band write.
-        self._skip_zero_block_ids: set[int] = set()
         self.need_mamba_block_aligned_split = (
             self.has_mamba_layers and self.cache_config.mamba_cache_mode == "align"
         )
@@ -1230,6 +1227,11 @@ class Scheduler(SchedulerInterface):
                     num_lookahead_tokens=effective_lookahead_tokens,
                     num_external_computed_tokens=num_external_computed_tokens,
                     delay_cache_blocks=load_kv_async,
+                    skip_zeroing_group_ids=(
+                        self.connector.get_loaded_kv_cache_group_ids(request)
+                        if load_kv_async and self.connector is not None
+                        else ()
+                    ),
                     num_encoder_tokens=num_encoder_tokens,
                     full_sequence_must_fit=self.scheduler_reserve_full_isl,
                     reserved_blocks=reserved_blocks,
@@ -1291,16 +1293,6 @@ class Scheduler(SchedulerInterface):
                     request.num_computed_tokens = num_computed_tokens
                     self._inflight_prefills.add(request)
                     self._set_kv_fetch_stage(request, KV_FETCH_IN_PROGRESS)
-                    if self.needs_kv_cache_zeroing:
-                        # Skip zeroing of the blocks the async load will
-                        # overwrite; the zeroing could race the write.
-                        self._skip_zero_block_ids.update(
-                            self.kv_cache_manager.get_zeroing_block_ids_in_range(
-                                request_id,
-                                num_new_local_computed_tokens,
-                                num_computed_tokens,
-                            )
-                        )
                     skip_request(request_queue)
                     continue
 
@@ -1545,12 +1537,6 @@ class Scheduler(SchedulerInterface):
         new_block_ids_to_zero = self.kv_cache_manager.take_new_block_ids()
         if not self.needs_kv_cache_zeroing:
             return None
-
-        if self._skip_zero_block_ids:
-            skip = self._skip_zero_block_ids
-            new_block_ids_to_zero = [b for b in new_block_ids_to_zero if b not in skip]
-            skip.clear()
-
         return new_block_ids_to_zero or None
 
     def _preempt_request(

vllm/v1/core/single_type_kv_cache_manager.py

@@ -331,6 +331,7 @@ class SingleTypeKVCacheManager(ABC):
         request_id: str,
         num_local_computed_tokens: int,
         num_external_computed_tokens: int,
+        record_for_zeroing: bool = True,
     ) -> None:
         """Allocate new blocks for external (KV-connector) computed tokens.
 
@@ -342,6 +343,8 @@ class SingleTypeKVCacheManager(ABC):
             request_id: The request ID.
             num_local_computed_tokens: The number of local computed tokens.
             num_external_computed_tokens: The number of external computed tokens.
+            record_for_zeroing: Whether the new blocks need zeroing. False when
+                the load writes them after this step, which zeroing would race.
 
         """
         num_total_computed_tokens = (
@@ -363,7 +366,7 @@ class SingleTypeKVCacheManager(ABC):
         )
         allocated_blocks = self.block_pool.get_new_blocks(num_new_blocks)
         req_blocks.extend(allocated_blocks)
-        if self._record_new_block_ids:
+        if self._record_new_block_ids and record_for_zeroing:
             self.new_block_ids.extend(b.block_id for b in allocated_blocks)
 
     def allocate_new_blocks(
@@ -1201,13 +1204,15 @@ class CircularBufferManager(FullAttentionManager):
 
     supports_fine_grained_hash_lookup: ClassVar[bool] = False
 
-    def _claim_ring_block(self, request_id: str) -> list[KVCacheBlock]:
+    def _claim_ring_block(
+        self, request_id: str, record_for_zeroing: bool = True
+    ) -> list[KVCacheBlock]:
         req_blocks = self.req_to_blocks[request_id]
         if req_blocks:
             return []
         new_blocks = self.block_pool.get_new_blocks(1)
         req_blocks.extend(new_blocks)
-        if self._record_new_block_ids:
+        if self._record_new_block_ids and record_for_zeroing:
             self.new_block_ids.extend(block.block_id for block in new_blocks)
         return new_blocks
 
@@ -1234,8 +1239,9 @@ class CircularBufferManager(FullAttentionManager):
         request_id: str,
         num_local_computed_tokens: int,
         num_external_computed_tokens: int,
+        record_for_zeroing: bool = True,
     ) -> None:
-        self._claim_ring_block(request_id)
+        self._claim_ring_block(request_id, record_for_zeroing)
 
     @classmethod
     def find_longest_cache_hit(
@@ -2166,6 +2172,7 @@ class CrossAttentionManager(SingleTypeKVCacheManager):
         request_id: str,
         num_local_computed_tokens: int,
         num_external_computed_tokens: int,
+        record_for_zeroing: bool = True,
     ) -> None:
         # Cross-attention does not use prefix caching / external KV loads.
         return
@@ -2311,6 +2318,7 @@ class HiSparseSourceManager(FullAttentionManager):
         request_id: str,
         num_local_computed_tokens: int,
         num_external_computed_tokens: int,
+        record_for_zeroing: bool = True,
     ) -> None:
         if num_external_computed_tokens <= 0:
             return
@@ -2321,7 +2329,10 @@ class HiSparseSourceManager(FullAttentionManager):
             request_id, num_local_computed_tokens + num_external_computed_tokens
         )
         super().allocate_external_computed_blocks(
-            request_id, num_local_computed_tokens, num_external_computed_tokens
+            request_id,
+            num_local_computed_tokens,
+            num_external_computed_tokens,
+            record_for_zeroing,
         )
 
     def cache_blocks(
@@ -2456,6 +2467,7 @@ class HiSparseHotManager(_HiSparseAuxiliaryManager):
         request_id: str,
         num_local_computed_tokens: int,
         num_external_computed_tokens: int,
+        record_for_zeroing: bool = True,
     ) -> None:
         self.require_hot(request_id)
 
@@ -2518,6 +2530,7 @@ class HiSparseResidentManager(_HiSparseAuxiliaryManager):
         request_id: str,
         num_local_computed_tokens: int,
         num_external_computed_tokens: int,
+        record_for_zeroing: bool = True,
     ) -> None:
         """Reserve a writable final page; earlier imported pages stay on host."""
         assert num_external_computed_tokens > 0

vllm/v1/kv_offload/cpu/gpu_worker.py

@@ -689,10 +689,10 @@ class SingleDirectionOffloadingHandler:
 
         # Stores must wait for the model to finish writing the KV they read.
         # Loads must wait for pending writes (including zeroing) to their
-        # destination blocks: the scheduler's _skip_zero_block_ids only edits
-        # the step being scheduled and cannot retract zeroing shipped in an
-        # earlier step for a since-reallocated block; with async scheduling
-        # nothing else orders that zeroing against this copy.
+        # destination blocks: the scheduler leaves a block unzeroed only in
+        # the step that allocates it for the load and cannot retract zeroing
+        # shipped in an earlier step for a since-reallocated block; with async
+        # scheduling nothing else orders that zeroing against this copy.
         stream.wait_stream(current_platform.current_stream())
         if self._transfers:
             last_transfer: Transfer = self._transfers[-1]