[KV Connector] Fix PD async scheduling race condition for hybrid attn models (#48481)

Signed-off-by: Artem Perevedentsev <aperevedents@nvidia.com>
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: llx-08 <2596671364@qq.com>
Co-authored-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

tests/v1/kv_connector/unit/test_nixl_connector_hma.py

@@ -752,6 +752,94 @@ def test_nixl_metadata_hybrid_ssm_block_ids():
     assert len(req_meta.remote.block_ids[0]) != len(req_meta.remote.block_ids[1])
 
 
+class _FakeBlock:
+    def __init__(self, block_id):
+        self.block_id = block_id
+
+
+class _FakeSingleTypeManager:
+    def __init__(self, records, block_size, block_ids):
+        self.records_new_block_ids = records
+        self.block_size = block_size
+        self.req_to_blocks = {"req-1": [_FakeBlock(b) for b in block_ids]}
+        self.new_block_ids: list[int] = []
+
+    def take_new_block_ids(self):
+        ids = self.new_block_ids
+        self.new_block_ids = []
+        return ids
+
+
+def _make_fake_kv_cache_manager():
+    from unittest.mock import MagicMock
+
+    from vllm.v1.core.kv_cache_manager import KVCacheManager
+
+    manager = object.__new__(KVCacheManager)
+    manager.coordinator = MagicMock()
+    manager.coordinator.single_type_managers = (
+        _FakeSingleTypeManager(True, 16, [10, 11, 12, 13, 14, 15]),  # attention
+        _FakeSingleTypeManager(False, 16, [20, 21, 22, 23, 24, 25]),  # mamba
+    )
+    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
+    the valid prefix; they must be zeroed before local recompute."""
+    from unittest.mock import MagicMock
+
+    from vllm.v1.core.sched.scheduler import Scheduler
+
+    scheduler = object.__new__(Scheduler)
+    scheduler.connector = MagicMock()
+    scheduler.needs_kv_cache_zeroing = True
+    scheduler.kv_cache_manager = _make_fake_kv_cache_manager()
+    scheduler.kv_cache_manager.cache_blocks = MagicMock()
+    scheduler.failed_recving_kv_req_ids = {"req-1"}
+    scheduler.finished_recving_kv_req_ids = {"req-1"}
+
+    request = MagicMock()
+    request.request_id = "req-1"
+    request.num_computed_tokens = 48  # Truncated at the first invalid block.
+
+    scheduler._update_waiting_for_remote_kv(request)
+
+    # 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]
+
+
 # ── Mamba N-1 prefill tests ──────────────────────────────────────────────
 
 

vllm/v1/core/kv_cache_manager.py

@@ -681,6 +681,34 @@ 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.
+
+        start_token must be block-aligned: zeroing a partially-valid block
+        would wipe its valid prefix.
+        """
+        for mgr in self.coordinator.single_type_managers:
+            if mgr.records_new_block_ids:
+                assert start_token % mgr.block_size == 0
+                start_idx = start_token // mgr.block_size
+                blocks = mgr.req_to_blocks[request_id]
+                mgr.new_block_ids.extend(blk.block_id for blk in blocks[start_idx:])
+
     def take_kv_cache_block_copies(
         self,
     ) -> tuple[list[KVCacheBlockCopy], list[KVCacheBlock]]:

vllm/v1/core/sched/scheduler.py

@@ -297,6 +297,9 @@ 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"
         )
@@ -986,6 +989,16 @@ class Scheduler(SchedulerInterface):
                     # only the successfully loaded tokens.
                     request.num_computed_tokens = num_computed_tokens
                     self._inflight_prefills.add(request)
+                    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.request_id,
+                                num_new_local_computed_tokens,
+                                num_computed_tokens,
+                            )
+                        )
                     continue
 
                 self.running.append(request)
@@ -1098,12 +1111,6 @@ class Scheduler(SchedulerInterface):
             self.prev_step_scheduled_req_ids.clear()
             self.prev_step_scheduled_req_ids.update(num_scheduled_tokens.keys())
 
-        # Drain new attention block ids every step so the manager-side list
-        # does not grow unbounded; only kv-cache zeroing consumes them.
-        new_attn_block_ids = self.kv_cache_manager.take_new_block_ids()
-        new_block_ids_to_zero = (
-            (new_attn_block_ids or None) if self.needs_kv_cache_zeroing else None
-        )
         kv_cache_block_copies, cow_retained_blocks = (
             self.kv_cache_manager.take_kv_cache_block_copies()
         )
@@ -1147,7 +1154,7 @@ class Scheduler(SchedulerInterface):
             # the previous and the current steps.
             finished_req_ids=self.finished_req_ids,
             free_encoder_mm_hashes=self.encoder_cache_manager.get_freed_mm_hashes(),
-            new_block_ids_to_zero=new_block_ids_to_zero,
+            new_block_ids_to_zero=self._get_new_block_ids_to_zero(),
             kv_cache_block_copies=pending_kv_cache_block_copies,
             num_spec_tokens_to_schedule=num_spec_tokens_to_schedule,
         )
@@ -1181,6 +1188,20 @@ class Scheduler(SchedulerInterface):
     ) -> KVConnectorMetadata:
         return connector.build_connector_meta(scheduler_output)
 
+    def _get_new_block_ids_to_zero(self) -> list[int] | None:
+        # Drain new attention block ids every step so the manager-side list
+        # does not grow unbounded; only kv-cache zeroing consumes them.
+        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(self, request: Request, timestamp: float) -> None:
         """Preempt a request and put it back to the waiting queue.
 
@@ -2518,9 +2539,18 @@ class Scheduler(SchedulerInterface):
             if request.num_computed_tokens:
                 # Cache any valid computed tokens.
                 self.kv_cache_manager.cache_blocks(request, request.num_computed_tokens)
+                if self.needs_kv_cache_zeroing:
+                    # The failed load left the blocks beyond the valid
+                    # prefix unwritten and their zeroing was skipped; zero
+                    # them before they are recomputed locally.
+                    self.kv_cache_manager.record_blocks_for_zeroing(
+                        request.request_id, request.num_computed_tokens
+                    )
             else:
                 # No valid computed tokens, release allocated blocks.
                 # There may be a local cache hit on retry.
+                # (Freed blocks are re-recorded for zeroing when
+                # reallocated, so the skipped blocks need no handling.)
                 self.kv_cache_manager.free(request)
 
             self.failed_recving_kv_req_ids.remove(request.request_id)

vllm/v1/core/single_type_kv_cache_manager.py

@@ -359,6 +359,11 @@ class SingleTypeKVCacheManager(ABC):
                 self.new_block_ids.extend(b.block_id for b in new_blocks)
             return cow_blocks + new_blocks
 
+    @property
+    def records_new_block_ids(self) -> bool:
+        """Whether this manager's new blocks are zeroed by the worker."""
+        return self._record_new_block_ids
+
     def take_new_block_ids(self) -> list[int]:
         """Drain and return block IDs allocated since the last call."""
         ids = self.new_block_ids