From ca0df1f1227cdd993835a4d2a62861dd3bd2e893 Mon Sep 17 00:00:00 2001 From: Yifan Qiao Date: Fri, 2 Oct 2026 23:26:12 -0700 Subject: [PATCH] [Bugfix][Core] Exempt exactly the blocks an async KV load writes from zeroing (#59504) Signed-off-by: Yifan Qiao Co-authored-by: Claude Opus 5.5 Co-authored-by: Codex --- tests/v1/core/test_prefix_replay.py | 52 +++++++++++++++++++ .../core/test_single_type_kv_cache_manager.py | 33 ++++++++++++ .../v1/e2e/general/test_mamba_prefix_cache.py | 2 + .../unit/test_error_propagation.py | 1 + .../unit/test_invalid_blocks_correctness.py | 1 + .../unit/test_kv_load_failure_recovery.py | 3 ++ .../kv_connector/unit/test_multi_connector.py | 14 +++++ .../unit/test_nixl_connector_hma.py | 29 ----------- .../kv_transfer/kv_connector/v1/base.py | 10 ++++ .../v1/mooncake/mooncake_connector.py | 3 ++ .../kv_connector/v1/multi_connector.py | 4 ++ .../kv_connector/v1/nixl/connector.py | 3 ++ vllm/v1/core/kv_cache_coordinator.py | 6 ++- vllm/v1/core/kv_cache_manager.py | 18 ++----- vllm/v1/core/sched/scheduler.py | 24 ++------- vllm/v1/core/single_type_kv_cache_manager.py | 23 ++++++-- vllm/v1/kv_offload/cpu/gpu_worker.py | 8 +-- 17 files changed, 162 insertions(+), 72 deletions(-) diff --git a/tests/v1/core/test_prefix_replay.py b/tests/v1/core/test_prefix_replay.py index b2e78e1f92..da57185808 100644 --- a/tests/v1/core/test_prefix_replay.py +++ b/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 diff --git a/tests/v1/core/test_single_type_kv_cache_manager.py b/tests/v1/core/test_single_type_kv_cache_manager.py index f20a96f5cc..dcb7ba923e 100644 --- a/tests/v1/core/test_single_type_kv_cache_manager.py +++ b/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( diff --git a/tests/v1/e2e/general/test_mamba_prefix_cache.py b/tests/v1/e2e/general/test_mamba_prefix_cache.py index 8584d8c6e1..50f3473708 100644 --- a/tests/v1/e2e/general/test_mamba_prefix_cache.py +++ b/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[ diff --git a/tests/v1/kv_connector/unit/test_error_propagation.py b/tests/v1/kv_connector/unit/test_error_propagation.py index 91ec525f5c..35a3749e08 100644 --- a/tests/v1/kv_connector/unit/test_error_propagation.py +++ b/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) ) diff --git a/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py b/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py index 777bada58b..7ea2170584 100644 --- a/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py +++ b/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 = () diff --git a/tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py b/tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py index 108eb049d9..3dbbfdc6b2 100644 --- a/tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py +++ b/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) ) diff --git a/tests/v1/kv_connector/unit/test_multi_connector.py b/tests/v1/kv_connector/unit/test_multi_connector.py index e636aabfd9..67cfdc1ab0 100644 --- a/tests/v1/kv_connector/unit/test_multi_connector.py +++ b/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 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 5a9a10a57d..2a1ca5a877 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/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] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/base.py b/vllm/distributed/kv_transfer/kv_connector/v1/base.py index 6df9d1cfbe..61de1802d5 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/base.py +++ b/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, diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py index 46118f4056..1f1ca517f3 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py +++ b/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]: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py index 73b6b0ccf6..752acb01ec 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py +++ b/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", diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py index ac929f7fcd..27b4d410f0 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py +++ b/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]: diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index daf0cd12c0..9b91bb1920 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/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( diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index 2d39a38651..37ceb4e474 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/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. diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 5585c9c70c..b390f1aadb 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/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( diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 7041600f8b..0ff8059562 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/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 diff --git a/vllm/v1/kv_offload/cpu/gpu_worker.py b/vllm/v1/kv_offload/cpu/gpu_worker.py index 797abed5ea..a56febf702 100644 --- a/vllm/v1/kv_offload/cpu/gpu_worker.py +++ b/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]