[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>
Artem Perevedentsev committed
Jul 16, 2026 at 13:42 UTC
530852f9591a822ff4065908778a58fa015f0e69
4 files changed
+158
-7
tests/v1/kv_connector/unit/test_nixl_connector_hma.py
+88
@@ -752,6 +752,94 @@ def test_nixl_metadata_hybrid_ssm_block_ids():
752
assert len(req_meta.remote.block_ids[0]) != len(req_meta.remote.block_ids[1])
753
754
755
+class _FakeBlock:
756
+ def __init__(self, block_id):
757
+ self.block_id = block_id
758
+
759
+
760
+class _FakeSingleTypeManager:
761
+ def __init__(self, records, block_size, block_ids):
762
+ self.records_new_block_ids = records
763
+ self.block_size = block_size
764
+ self.req_to_blocks = {"req-1": [_FakeBlock(b) for b in block_ids]}
765
+ self.new_block_ids: list[int] = []
766
+
767
+ def take_new_block_ids(self):
768
+ ids = self.new_block_ids
769
+ self.new_block_ids = []
770
+ return ids
771
+
772
+
773
+def _make_fake_kv_cache_manager():
774
+ from unittest.mock import MagicMock
775
+
776
+ from vllm.v1.core.kv_cache_manager import KVCacheManager
777
+
778
+ manager = object.__new__(KVCacheManager)
779
+ manager.coordinator = MagicMock()
780
+ manager.coordinator.single_type_managers = (
781
+ _FakeSingleTypeManager(True, 16, [10, 11, 12, 13, 14, 15]), # attention
782
+ _FakeSingleTypeManager(False, 16, [20, 21, 22, 23, 24, 25]), # mamba
783
+ )
784
+ return manager
785
+
786
+
787
+@pytest.mark.cpu_test
788
+def test_zeroing_block_ids_cover_only_loaded_attention_blocks():
789
+ """Only zero-recorded (attention) groups contribute, sliced to the
790
+ externally-loaded token range; Mamba state blocks are never zeroed."""
791
+ manager = _make_fake_kv_cache_manager()
792
+
793
+ # Tokens [0, 16) are locally cached; the load covers tokens [16, 56).
794
+ assert manager.get_zeroing_block_ids_in_range("req-1", 16, 56) == [11, 12, 13]
795
+
796
+
797
+@pytest.mark.cpu_test
798
+def test_scheduler_filters_connector_loaded_blocks_from_zeroing():
799
+ """Blocks that will be loaded by the connector must not be zeroed."""
800
+ from vllm.v1.core.sched.scheduler import Scheduler
801
+
802
+ class FakeKVCacheManager:
803
+ def take_new_block_ids(self):
804
+ return [9, 10, 11, 12]
805
+
806
+ scheduler = object.__new__(Scheduler)
807
+ scheduler.needs_kv_cache_zeroing = True
808
+ scheduler.kv_cache_manager = FakeKVCacheManager()
809
+ scheduler._skip_zero_block_ids = {10, 12}
810
+
811
+ assert scheduler._get_new_block_ids_to_zero() == [9, 11]
812
+ assert not scheduler._skip_zero_block_ids
813
+
814
+
815
+@pytest.mark.cpu_test
816
+def test_failed_load_rezeroes_unwritten_skipped_blocks():
817
+ """A failed async load leaves zeroing-skipped blocks unwritten beyond
818
+ the valid prefix; they must be zeroed before local recompute."""
819
+ from unittest.mock import MagicMock
820
+
821
+ from vllm.v1.core.sched.scheduler import Scheduler
822
+
823
+ scheduler = object.__new__(Scheduler)
824
+ scheduler.connector = MagicMock()
825
+ scheduler.needs_kv_cache_zeroing = True
826
+ scheduler.kv_cache_manager = _make_fake_kv_cache_manager()
827
+ scheduler.kv_cache_manager.cache_blocks = MagicMock()
828
+ scheduler.failed_recving_kv_req_ids = {"req-1"}
829
+ scheduler.finished_recving_kv_req_ids = {"req-1"}
830
+
831
+ request = MagicMock()
832
+ request.request_id = "req-1"
833
+ request.num_computed_tokens = 48 # Truncated at the first invalid block.
834
+
835
+ scheduler._update_waiting_for_remote_kv(request)
836
+
837
+ # Attention blocks covering tokens >= 48 are re-recorded for zeroing
838
+ # and flow into the next step's zero list; Mamba blocks are not.
839
+ scheduler._skip_zero_block_ids = set()
840
+ assert scheduler._get_new_block_ids_to_zero() == [13, 14, 15]
841
+
842
+
843
# ── Mamba N-1 prefill tests ──────────────────────────────────────────────
844
845
vllm/v1/core/kv_cache_manager.py
+28
@@ -681,6 +681,34 @@ class KVCacheManager:
681
ids.extend(mgr.take_new_block_ids())
682
return ids
683
684
+ def get_zeroing_block_ids_in_range(
685
+ self, request_id: str, start_token: int, end_token: int
686
+ ) -> list[int]:
687
+ """The request's block ids covering [start_token, end_token), from
688
+ the groups whose new blocks are zeroed by the worker."""
689
+ ids: list[int] = []
690
+ for mgr in self.coordinator.single_type_managers:
691
+ if mgr.records_new_block_ids:
692
+ start_idx = start_token // mgr.block_size
693
+ end_idx = cdiv(end_token, mgr.block_size)
694
+ blocks = mgr.req_to_blocks[request_id]
695
+ ids.extend(blk.block_id for blk in blocks[start_idx:end_idx])
696
+ return ids
697
+
698
+ def record_blocks_for_zeroing(self, request_id: str, start_token: int) -> None:
699
+ """Re-record the request's blocks from start_token onwards for
700
+ zeroing, e.g. blocks a failed async KV load left unwritten.
701
+
702
+ start_token must be block-aligned: zeroing a partially-valid block
703
+ would wipe its valid prefix.
704
+ """
705
+ for mgr in self.coordinator.single_type_managers:
706
+ if mgr.records_new_block_ids:
707
+ assert start_token % mgr.block_size == 0
708
+ start_idx = start_token // mgr.block_size
709
+ blocks = mgr.req_to_blocks[request_id]
710
+ mgr.new_block_ids.extend(blk.block_id for blk in blocks[start_idx:])
711
+
712
def take_kv_cache_block_copies(
713
self,
714
) -> tuple[list[KVCacheBlockCopy], list[KVCacheBlock]]:
vllm/v1/core/sched/scheduler.py
+37
-7
@@ -297,6 +297,9 @@ class Scheduler(SchedulerInterface):
297
298
self.has_mamba_layers = kv_cache_config.has_mamba_layers
299
self.needs_kv_cache_zeroing = kv_cache_config.needs_kv_cache_zeroing
300
+ # Blocks that async KV loads will overwrite this step, skipped from
301
+ # zeroing since the zeroing could race the out-of-band write.
302
+ self._skip_zero_block_ids: set[int] = set()
303
self.need_mamba_block_aligned_split = (
304
self.has_mamba_layers and self.cache_config.mamba_cache_mode == "align"
305
)
@@ -986,6 +989,16 @@ class Scheduler(SchedulerInterface):
989
# only the successfully loaded tokens.
990
request.num_computed_tokens = num_computed_tokens
991
self._inflight_prefills.add(request)
992
+ if self.needs_kv_cache_zeroing:
993
+ # Skip zeroing of the blocks the async load will
994
+ # overwrite; the zeroing could race the write.
995
+ self._skip_zero_block_ids.update(
996
+ self.kv_cache_manager.get_zeroing_block_ids_in_range(
997
+ request.request_id,
998
+ num_new_local_computed_tokens,
999
+ num_computed_tokens,
1000
+ )
1001
+ )
1002
continue
1003
1004
self.running.append(request)
@@ -1098,12 +1111,6 @@ class Scheduler(SchedulerInterface):
1111
self.prev_step_scheduled_req_ids.clear()
1112
self.prev_step_scheduled_req_ids.update(num_scheduled_tokens.keys())
1113
1101
- # Drain new attention block ids every step so the manager-side list
1102
- # does not grow unbounded; only kv-cache zeroing consumes them.
1103
- new_attn_block_ids = self.kv_cache_manager.take_new_block_ids()
1104
- new_block_ids_to_zero = (
1105
- (new_attn_block_ids or None) if self.needs_kv_cache_zeroing else None
1106
- )
1114
kv_cache_block_copies, cow_retained_blocks = (
1115
self.kv_cache_manager.take_kv_cache_block_copies()
1116
)
@@ -1147,7 +1154,7 @@ class Scheduler(SchedulerInterface):
1154
# the previous and the current steps.
1155
finished_req_ids=self.finished_req_ids,
1156
free_encoder_mm_hashes=self.encoder_cache_manager.get_freed_mm_hashes(),
1150
- new_block_ids_to_zero=new_block_ids_to_zero,
1157
+ new_block_ids_to_zero=self._get_new_block_ids_to_zero(),
1158
kv_cache_block_copies=pending_kv_cache_block_copies,
1159
num_spec_tokens_to_schedule=num_spec_tokens_to_schedule,
1160
)
@@ -1181,6 +1188,20 @@ class Scheduler(SchedulerInterface):
1188
) -> KVConnectorMetadata:
1189
return connector.build_connector_meta(scheduler_output)
1190
1191
+ def _get_new_block_ids_to_zero(self) -> list[int] | None:
1192
+ # Drain new attention block ids every step so the manager-side list
1193
+ # does not grow unbounded; only kv-cache zeroing consumes them.
1194
+ new_block_ids_to_zero = self.kv_cache_manager.take_new_block_ids()
1195
+ if not self.needs_kv_cache_zeroing:
1196
+ return None
1197
+
1198
+ if self._skip_zero_block_ids:
1199
+ skip = self._skip_zero_block_ids
1200
+ new_block_ids_to_zero = [b for b in new_block_ids_to_zero if b not in skip]
1201
+ skip.clear()
1202
+
1203
+ return new_block_ids_to_zero or None
1204
+
1205
def _preempt_request(self, request: Request, timestamp: float) -> None:
1206
"""Preempt a request and put it back to the waiting queue.
1207
@@ -2518,9 +2539,18 @@ class Scheduler(SchedulerInterface):
2539
if request.num_computed_tokens:
2540
# Cache any valid computed tokens.
2541
self.kv_cache_manager.cache_blocks(request, request.num_computed_tokens)
2542
+ if self.needs_kv_cache_zeroing:
2543
+ # The failed load left the blocks beyond the valid
2544
+ # prefix unwritten and their zeroing was skipped; zero
2545
+ # them before they are recomputed locally.
2546
+ self.kv_cache_manager.record_blocks_for_zeroing(
2547
+ request.request_id, request.num_computed_tokens
2548
+ )
2549
else:
2550
# No valid computed tokens, release allocated blocks.
2551
# There may be a local cache hit on retry.
2552
+ # (Freed blocks are re-recorded for zeroing when
2553
+ # reallocated, so the skipped blocks need no handling.)
2554
self.kv_cache_manager.free(request)
2555
2556
self.failed_recving_kv_req_ids.remove(request.request_id)
vllm/v1/core/single_type_kv_cache_manager.py
+5
@@ -359,6 +359,11 @@ class SingleTypeKVCacheManager(ABC):
359
self.new_block_ids.extend(b.block_id for b in new_blocks)
360
return cow_blocks + new_blocks
361
362
+ @property
363
+ def records_new_block_ids(self) -> bool:
364
+ """Whether this manager's new blocks are zeroed by the worker."""
365
+ return self._record_new_block_ids
366
+
367
def take_new_block_ids(self) -> list[int]:
368
"""Drain and return block IDs allocated since the last call."""
369
ids = self.new_block_ids