diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index 178eb42b93f..4945942ba3a 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/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 ────────────────────────────────────────────── diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index e3388a9f85a..3d3d8c1573a 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/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]]: diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index bde6af0e66b..d6f3c3ad0e3 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/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) diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index de47ce0a2c7..24404839094 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/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