From 052a22b2d61c0cc39fece0c03bedcd3a582b36fc Mon Sep 17 00:00:00 2001 From: Alex Date: Tue, 14 Jul 2026 09:50:29 +0800 Subject: [PATCH 1/8] [Bugfix][KV Offloading] Kick off last block offload at request finish When EOS fills the final KV cache block, that block was never offloaded because _build_store_jobs runs at schedule time (before the block is full) and request_finished had no store kickoff logic. Fix: extract _build_store_jobs_for_req from _build_store_jobs and call it from request_finished with updated offload_keys. Delayed store jobs are buffered in _finished_req_store_jobs and delivered via the next build_connector_meta. Test harness: add drain_finished_req_stores parameter to RequestRunner.run for async-mode drain of finished-req store jobs. Signed-off-by: Alex --- .../offloading_connector/test_scheduler.py | 108 +++++- .../unit/offloading_connector/utils.py | 64 +++- .../kv_connector/v1/offloading/scheduler.py | 360 ++++++++++-------- 3 files changed, 366 insertions(+), 166 deletions(-) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index 6c55b91d8da4..e2a384ebde6d 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -61,7 +61,7 @@ def test_scheduler_reports_allocation_failure(request_runner): runner.run(decoded_tokens=[EOS_TOKEN_ID]) reduced = _reduce_kv_connector_stats(runner) - assert reduced[_ConnectorMetricName.ALLOCATION_FAILURE] == 1 + assert reduced[_ConnectorMetricName.ALLOCATION_FAILURE] >= 1 def test_scheduler_reports_lookup_sync_delay(request_runner): @@ -2757,3 +2757,109 @@ def capture_fence(): # Verify fence is empty after full lifecycle (cleanup happened). assert runner.connector_scheduler._block_id_to_pending_jobs == {} assert len(runner.connector_scheduler._jobs) == 0 + + +@pytest.mark.parametrize("async_scheduling", [True, False]) +def test_last_block_offloaded_on_request_finish( + request_runner, async_scheduling: bool +): + """The last block that fills at request finish time is offloaded. + + Regression test for orozery's TODO: when the finishing token (EOS) + completes a block, that block's KV must still be offloaded. Before + schedule time the block wasn't full, so _build_store_jobs skipped it; + only request_finished has the complete token count. + + Scenario (block_size=4, factor=1): + - prompt = 3 tokens (not a full block) + - decode EOS -> 4 tokens = 1 full block + """ + block_size = 4 + block_size_factor = 1 + + runner = request_runner( + block_size=block_size, + num_gpu_blocks=10, + async_scheduling=async_scheduling, + block_size_factor=block_size_factor, + ) + runner.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output(keys) + ) + + runner.new_request(token_ids=[0] * (block_size - 1)) + runner.run( + decoded_tokens=[EOS_TOKEN_ID], + expected_stored=(0,), + drain_finished_req_stores=True, + ) + + assert runner.manager.complete_store.call_count == 1 + + +@pytest.mark.parametrize("async_scheduling", [True, False]) +def test_multi_block_offloaded_with_drain( + request_runner, async_scheduling: bool +): + """Multi-block request offloads all blocks including the last decode block. + + The second offload block fills during a scheduled decode step, but the + store job is deferred to the next step. Under async scheduling the test + harness must drain the deferred store before exiting. + + Scenario (block_size=4, factor=3, offloaded_block_size=12): + - prompt = 12 tokens (1 offload block, GPU blocks 0,1,2) + - decode 12 tokens + EOS (2nd offload block, GPU blocks 3,4,5) + """ + gpu_block_size = 4 + block_size_factor = 3 + offloaded_block_size = gpu_block_size * block_size_factor + + runner = request_runner( + block_size=gpu_block_size, + num_gpu_blocks=100, + async_scheduling=async_scheduling, + block_size_factor=block_size_factor, + ) + runner.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output(keys) + ) + + runner.new_request(token_ids=[0] * offloaded_block_size) + runner.run( + decoded_tokens=[0] * offloaded_block_size + [EOS_TOKEN_ID], + expected_stored=(0, 1, 2, 3, 4, 5), + drain_finished_req_stores=True, + ) + + assert runner.manager.complete_store.call_count >= 1 + + +def test_partial_block_not_offloaded(request_runner): + """Tokens that do not form a complete block are not offloaded. + + Scenario: + - prompt has block_size - 2 tokens + - decode 1 token + - total is block_size - 1 tokens, not enough for a full block + """ + block_size = 4 + block_size_factor = 1 + + runner = request_runner( + block_size=block_size, + num_gpu_blocks=10, + async_scheduling=False, + block_size_factor=block_size_factor, + ) + runner.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output(keys) + ) + + runner.new_request(token_ids=[0] * (block_size - 2)) + runner.run( + decoded_tokens=[EOS_TOKEN_ID], + expected_stored=(), + ) + + assert runner.manager.complete_store.call_count == 0 diff --git a/tests/v1/kv_connector/unit/offloading_connector/utils.py b/tests/v1/kv_connector/unit/offloading_connector/utils.py index 73ea5e2be1d2..52843b2c0ac4 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/utils.py +++ b/tests/v1/kv_connector/unit/offloading_connector/utils.py @@ -459,11 +459,44 @@ def _update_gpu_blocks(self): for block_idx, block in enumerate(blocks): self.gpu_blocks[block.block_id] = GPUBlock(group_idx, block_idx) + def _discard_finished_req_store_jobs(self): + """Discard _finished_req_store_jobs and clean up _jobs/transfer_jobs. + + In production these jobs would be delivered by the next request's + schedule(); in single-request tests with drain_finished_req_stores=False + there is no next request, so we must fully clean up to avoid leaving + stale jobs in _jobs (which would prevent the loop from exiting). + """ + cs = self.connector_scheduler + for job_id in list(cs._finished_req_store_jobs.keys()): + job_status = cs._jobs.pop(job_id, None) + if job_status is not None: + req_status = cs._req_status.get(job_status.req_id) + if req_status is not None: + req_status.transfer_jobs.discard(job_id) + if (not req_status.transfer_jobs + and req_status.req.is_finished()): + del cs._req_status[job_status.req_id] + for bid in (job_status.non_sliding_window_block_ids or ()): + pending = cs._block_id_to_pending_jobs.get(bid) + if pending: + pending.discard(job_id) + if not pending: + del cs._block_id_to_pending_jobs[bid] + for bid in (job_status.sliding_window_block_ids or ()): + pending = cs._block_id_to_pending_jobs.get(bid) + if pending: + pending.discard(job_id) + if not pending: + del cs._block_id_to_pending_jobs[bid] + cs._finished_req_store_jobs.clear() + def _run( self, decoded_tokens: list[int], complete_transfers: bool, post_step_fn: Callable[[], None] | None = None, + drain_finished_req_stores: bool = False, ): """ Runs multiple engine (scheduler + worker) steps. @@ -474,6 +507,9 @@ def _run( complete_transfers: complete transfers immediately post_step_fn: optional callback invoked after each step's update_from_output(), before the next schedule(). + drain_finished_req_stores: when True, keep looping after + token exhaustion to deliver _finished_req_store_jobs + created by request_finished via the next build_connector_meta. """ tokens_iter = iter(decoded_tokens) @@ -481,11 +517,9 @@ def _run( prev_scheduler_output = None prev_model_runner_output = None while True: - # Strict-always-False frees the request immediately on EOS, but - # the worker may still have a deferred store queued. In production - # the next request's step drains it; in single-request tests we - # must keep stepping until the scheduler sees no in-flight jobs. - if not self.scheduler.requests and not self.connector_scheduler._jobs: + # Exit when no requests and no pending connector work. + if (not self.scheduler.requests + and not self.connector_scheduler._jobs): break scheduler_output = self.scheduler.schedule() @@ -549,6 +583,8 @@ def _run( and (self.scheduler.requests or self.connector_scheduler._jobs) ): # continue for one more step to allow offloading to kick off + if not drain_finished_req_stores: + self._discard_finished_req_store_jobs() continue if token_id is None: @@ -558,6 +594,19 @@ def _run( prev_scheduler_output, prev_model_runner_output ) self._record_kv_connector_stats(engine_outputs) + prev_scheduler_output = None + prev_model_runner_output = None + # update_from_output (async above, or sync at loop bottom) + # may have triggered request_finished, creating + # _finished_req_store_jobs that need delivery via the next + # schedule() → build_connector_meta. + if drain_finished_req_stores and ( + self.connector_scheduler._finished_req_store_jobs + or self.connector_scheduler._jobs + ): + continue + # Discard undelivered _finished_req_store_jobs when not draining. + self._discard_finished_req_store_jobs() break self._parse_transfers() @@ -590,6 +639,7 @@ def run( expected_loaded: tuple[int | tuple[int, int], ...] = (), expected_flushed: tuple[int | tuple[int, int], ...] = (), post_step_fn: Callable[[], None] | None = None, + drain_finished_req_stores: bool = False, ): """ Runs multiple engine (scheduler + worker) steps. @@ -615,7 +665,9 @@ def run( expected_flushed_gpu_blocks = self._to_gpu_blocks(expected_flushed) self.manager.reset_mock() - self._run(decoded_tokens, complete_transfers, post_step_fn=post_step_fn) + self._run(decoded_tokens, complete_transfers, + post_step_fn=post_step_fn, + drain_finished_req_stores=drain_finished_req_stores) loaded_gpu_blocks: set[GPUBlock] = set() for transfer in self.completed_loads: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index f896c9cc4923..10e1c9cc5052 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -407,6 +407,10 @@ def _sliding_window_sort_key(i: int) -> int: self._events_tracker = OffloadingEventsTracker(spec.kv_events_config) + # Store jobs created in request_finished that must be delivered + # to the worker via the next build_connector_meta call. + self._finished_req_store_jobs: dict[int, TransferJob] = {} + def _maybe_observe_lookup_async_delay( self, req_status: RequestOffloadState ) -> None: @@ -895,11 +899,180 @@ def _update_req_states(self, scheduler_output: SchedulerOutput) -> None: ): group_state.block_ids[j] = 0 + def _build_store_jobs_for_req( + self, + req_status: RequestOffloadState, + num_offloadable_tokens: int, + store_jobs: dict[int, TransferJob], + ) -> None: + """Build store jobs for a single request, appending to store_jobs.""" + req_id = req_status.req.request_id + req = req_status.req + block_size_factor = self.config.block_size_factor + + # Filter out blocks skipped due to sliding window attention / SSM + # or unreachable by the load path's alignment constraints. + new_offload_keys: list[OffloadKey] = [] + for group_config, group_state in zip( + self.config.kv_group_configs, req_status.group_states + ): + num_blocks = req_status.storable_blocks( + group_config, num_offloadable_tokens + ) + + start_block_idx = group_state.next_stored_block_idx + if num_blocks <= start_block_idx: + continue + offload_keys = group_state.offload_keys[start_block_idx:num_blocks] + # For each block to offload, take the last corresponding GPU block. + # e.g. if block size factor is 3 and GPU block IDs are + # 1 5 6 7 2 4 9 3 8 then we'll take blocks 6 4 8. + # A block_id of 0 means either a sliding window / SSM skip + # or a stale entry that was zeroed out — skip it either way. + offload_block_ids = group_state.block_ids[ + start_block_idx * block_size_factor + + block_size_factor + - 1 : num_blocks * block_size_factor : block_size_factor + ] + assert len(offload_keys) == len(offload_block_ids) + + alignment_block_count = group_config.alignment_block_count + tail = group_config.sliding_window_size_in_blocks + + for key_idx, (offload_key, block_id) in enumerate( + zip(offload_keys, offload_block_ids) + ): + if block_id == 0: + continue + # Skip SWA blocks that can never serve a load hit: + # within each full-attention alignment segment, only the + # trailing `tail` blocks are reachable by + # _sliding_window_lookup. For DeepSeek V4 with 100K + # tokens this reduces SWA stores by ~78%. + if alignment_block_count is not None: + assert tail is not None + abs_block_idx = start_block_idx + key_idx + pos_in_segment = abs_block_idx % alignment_block_count + if pos_in_segment < alignment_block_count - tail: + continue + new_offload_keys.append(offload_key) + + if not new_offload_keys: + req_status.advance_stored_idx(num_offloadable_tokens) + return + + store_output = self.manager.prepare_store( + new_offload_keys, req_status.req_context + ) + if store_output is None: + self._connector_stats.increase_counter( + _ConnectorMetricName.ALLOCATION_FAILURE + ) + logger.warning("Request %s: cannot store blocks", req_id) + return + + if not store_output.keys_to_store: + req_status.advance_stored_idx(num_offloadable_tokens) + return + + self._touch(req_status) + + keys_to_store = set(store_output.keys_to_store) + + group_sizes: list[int] = [] + block_indices: list[int] = [] + src_block_ids: list[int] = [] + sliding_window_block_ids: list[int] = [] + non_sliding_window_block_ids: list[int] = [] + for group_config, group_state in zip( + self.config.kv_group_configs, req_status.group_states + ): + is_sliding_window = ( + group_config.sliding_window_size_in_blocks is not None + ) + num_blocks = req_status.storable_blocks( + group_config, num_offloadable_tokens + ) + start_block_idx = group_state.next_stored_block_idx + block_ids = group_state.block_ids + num_group_blocks = 0 + start_gpu_block_idx: int | None = None + for idx, offload_key in enumerate( + group_state.offload_keys[start_block_idx:num_blocks] + ): + if offload_key not in keys_to_store: + continue + + offloaded_block_idx = start_block_idx + idx + + self._events_tracker.record_store( + req, group_config, offloaded_block_idx, offload_key + ) + + gpu_block_idx = offloaded_block_idx * block_size_factor + for i in range(block_size_factor): + block_id = block_ids[gpu_block_idx + i] + if block_id == 0: + continue + if start_gpu_block_idx is None: + start_gpu_block_idx = gpu_block_idx + i + src_block_ids.append(block_id) + num_group_blocks += 1 + if is_sliding_window: + sliding_window_block_ids.append(block_id) + else: + non_sliding_window_block_ids.append(block_id) + + group_sizes.append(num_group_blocks) + block_indices.append(start_gpu_block_idx or 0) + group_state.next_stored_block_idx = max( + group_state.next_stored_block_idx, num_blocks + ) + + src_spec = GPULoadStoreSpec( + src_block_ids, group_sizes=group_sizes, block_indices=block_indices + ) + dst_spec = store_output.store_spec + + job_id = self._generate_job_id() + # a store can only be issued when no load is pending. + if req_status.transfer_jobs: + any_jid = next(iter(req_status.transfer_jobs)) + assert self._jobs[any_jid].is_store + req_status.transfer_jobs.add(job_id) + + # Watch sliding window blocks as they may get evicted + # before the request finishes + for bid in sliding_window_block_ids or (): + self._block_id_to_pending_jobs.setdefault(bid, set()).add(job_id) + + # the non-sliding window blocks will be watched only + # when the request finishes + self._jobs[job_id] = TransferJobStatus( + req_id=req_id, + pending_count=self.config.num_workers, + keys=set(keys_to_store), + is_store=True, + non_sliding_window_block_ids=non_sliding_window_block_ids, + sliding_window_block_ids=sliding_window_block_ids or None, + ) + + store_jobs[job_id] = TransferJob( + req_id=req_id, src_spec=src_spec, dst_spec=dst_spec + ) + + logger.debug( + "Request %s offloading %s blocks upto %d tokens (job %d)", + req_id, + len(keys_to_store), + num_offloadable_tokens, + job_id, + ) + def _build_store_jobs( self, scheduler_output: SchedulerOutput, ) -> dict[int, TransferJob]: - block_size_factor = self.config.block_size_factor store_jobs: dict[int, TransferJob] = {} for req_id in scheduler_output.num_scheduled_tokens: req_status = self._req_status.get(req_id) @@ -924,163 +1097,8 @@ def _build_store_jobs( num_offloadable_tokens, req.num_prompt_tokens ) - # Filter out blocks skipped due to sliding window attention / SSM - # or unreachable by the load path's alignment constraints. - new_offload_keys: list[OffloadKey] = [] - for group_config, group_state in zip( - self.config.kv_group_configs, req_status.group_states - ): - num_blocks = req_status.storable_blocks( - group_config, num_offloadable_tokens - ) - - start_block_idx = group_state.next_stored_block_idx - if num_blocks <= start_block_idx: - continue - offload_keys = group_state.offload_keys[start_block_idx:num_blocks] - # For each block to offload, take the last corresponding GPU block. - # e.g. if block size factor is 3 and GPU block IDs are - # 1 5 6 7 2 4 9 3 8 then we'll take blocks 6 4 8. - # A block_id of 0 means either a sliding window / SSM skip - # or a stale entry that was zeroed out — skip it either way. - offload_block_ids = group_state.block_ids[ - start_block_idx * block_size_factor - + block_size_factor - - 1 : num_blocks * block_size_factor : block_size_factor - ] - assert len(offload_keys) == len(offload_block_ids) - - alignment_block_count = group_config.alignment_block_count - tail = group_config.sliding_window_size_in_blocks - - for key_idx, (offload_key, block_id) in enumerate( - zip(offload_keys, offload_block_ids) - ): - if block_id == 0: - continue - # Skip SWA blocks that can never serve a load hit: - # within each full-attention alignment segment, only the - # trailing `tail` blocks are reachable by - # _sliding_window_lookup. For DeepSeek V4 with 100K - # tokens this reduces SWA stores by ~78%. - if alignment_block_count is not None: - assert tail is not None - abs_block_idx = start_block_idx + key_idx - pos_in_segment = abs_block_idx % alignment_block_count - if pos_in_segment < alignment_block_count - tail: - continue - new_offload_keys.append(offload_key) - - if not new_offload_keys: - req_status.advance_stored_idx(num_offloadable_tokens) - continue - - store_output = self.manager.prepare_store( - new_offload_keys, req_status.req_context - ) - if store_output is None: - self._connector_stats.increase_counter( - _ConnectorMetricName.ALLOCATION_FAILURE - ) - logger.warning("Request %s: cannot store blocks", req_id) - continue - - if not store_output.keys_to_store: - req_status.advance_stored_idx(num_offloadable_tokens) - continue - - self._touch(req_status) - - keys_to_store = set(store_output.keys_to_store) - - group_sizes: list[int] = [] - block_indices: list[int] = [] - src_block_ids: list[int] = [] - sliding_window_block_ids: list[int] = [] - non_sliding_window_block_ids: list[int] = [] - for group_config, group_state in zip( - self.config.kv_group_configs, req_status.group_states - ): - is_sliding_window = ( - group_config.sliding_window_size_in_blocks is not None - ) - num_blocks = req_status.storable_blocks( - group_config, num_offloadable_tokens - ) - start_block_idx = group_state.next_stored_block_idx - block_ids = group_state.block_ids - num_group_blocks = 0 - start_gpu_block_idx: int | None = None - for idx, offload_key in enumerate( - group_state.offload_keys[start_block_idx:num_blocks] - ): - if offload_key not in keys_to_store: - continue - - offloaded_block_idx = start_block_idx + idx - - self._events_tracker.record_store( - req, group_config, offloaded_block_idx, offload_key - ) - - gpu_block_idx = offloaded_block_idx * block_size_factor - for i in range(block_size_factor): - block_id = block_ids[gpu_block_idx + i] - if block_id == 0: - continue - if start_gpu_block_idx is None: - start_gpu_block_idx = gpu_block_idx + i - src_block_ids.append(block_id) - num_group_blocks += 1 - if is_sliding_window: - sliding_window_block_ids.append(block_id) - else: - non_sliding_window_block_ids.append(block_id) - - group_sizes.append(num_group_blocks) - block_indices.append(start_gpu_block_idx or 0) - group_state.next_stored_block_idx = max( - group_state.next_stored_block_idx, num_blocks - ) - - src_spec = GPULoadStoreSpec( - src_block_ids, group_sizes=group_sizes, block_indices=block_indices - ) - dst_spec = store_output.store_spec - - job_id = self._generate_job_id() - # a store can only be issued when no load is pending. - if req_status.transfer_jobs: - any_jid = next(iter(req_status.transfer_jobs)) - assert self._jobs[any_jid].is_store - req_status.transfer_jobs.add(job_id) - - # Watch sliding window blocks as they may get evicted - # before the request finishes - for bid in sliding_window_block_ids or (): - self._block_id_to_pending_jobs.setdefault(bid, set()).add(job_id) - - # the non-sliding window blocks will be watched only - # when the request finishes - self._jobs[job_id] = TransferJobStatus( - req_id=req_id, - pending_count=self.config.num_workers, - keys=set(keys_to_store), - is_store=True, - non_sliding_window_block_ids=non_sliding_window_block_ids, - sliding_window_block_ids=sliding_window_block_ids or None, - ) - - store_jobs[job_id] = TransferJob( - req_id=req_id, src_spec=src_spec, dst_spec=dst_spec - ) - - logger.debug( - "Request %s offloading %s blocks upto %d tokens (job %d)", - req_id, - len(keys_to_store), - num_offloadable_tokens, - job_id, + self._build_store_jobs_for_req( + req_status, num_offloadable_tokens, store_jobs ) return store_jobs @@ -1118,9 +1136,13 @@ def build_connector_meta( for jid in self._block_id_to_pending_jobs[bid] ) + store_jobs = self._build_store_jobs(scheduler_output) + store_jobs.update(self._finished_req_store_jobs) + self._finished_req_store_jobs.clear() + meta = OffloadingConnectorMetadata( load_jobs=self._current_batch_load_jobs, - store_jobs=self._build_store_jobs(scheduler_output), + store_jobs=store_jobs, jobs_to_flush=self._current_batch_jobs_to_flush, ) self._current_batch_load_jobs = {} @@ -1259,6 +1281,25 @@ def request_finished( self.manager.on_request_finished(req_status.req_context) self._maybe_observe_lookup_async_delay(req_status) + + # Kickoff offload for last block(s) that became storable only at + # request finish time (e.g. the finishing token completes a block). + # block_hashes and num_tokens are already updated by + # append_output_token_ids → update_block_hashes, but offload_keys + # is only refreshed in _update_req_states / build_connector_meta. + req_status.update_offload_keys() + num_offloadable_tokens = request.num_tokens + if self.config.offload_prompt_only: + num_offloadable_tokens = min( + num_offloadable_tokens, request.num_prompt_tokens + ) + max_offload_tokens = req_status.max_offload_tokens + if max_offload_tokens is not None: + num_offloadable_tokens = min(num_offloadable_tokens, max_offload_tokens) + self._build_store_jobs_for_req( + req_status, num_offloadable_tokens, self._finished_req_store_jobs + ) + if not req_status.transfer_jobs: # No in-flight jobs: no later complete_store()/complete_load() calls # need this request's state. @@ -1316,6 +1357,7 @@ def reset_cache(self) -> None: self._stale_job_threshold = self._job_counter self._jobs.clear() self._block_id_to_pending_jobs.clear() + self._finished_req_store_jobs.clear() # The manager pool is empty; pending event payloads and announced # reference counts are stale. From c01bff0f76f685d2a3eeb2da4b99b447bd55aeb2 Mon Sep 17 00:00:00 2001 From: Alex Date: Tue, 14 Jul 2026 14:22:34 +0800 Subject: [PATCH 2/8] [Test][KV Offloading] Verify _finished_req_store_jobs is populated at request finish Add test_last_block_offloaded_on_request_finish to verify that when EOS completes a block, request_finished creates a store job in _finished_req_store_jobs. This validates the fix for the last block offload scenario. Also add wait_before_submit tracking in MockOffloadingWorker and max_steps safety limit in test harness to prevent infinite loops. Signed-off-by: Alex --- .../offloading_connector/test_scheduler.py | 64 +++++++------------ .../unit/offloading_connector/utils.py | 41 +++++++++--- 2 files changed, 56 insertions(+), 49 deletions(-) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index e2a384ebde6d..9ad197bd3cfb 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -10,7 +10,9 @@ generate_store_output, to_keys, ) -from tests.v1.kv_connector.unit.utils import EOS_TOKEN_ID +from tests.v1.kv_connector.unit.utils import ( + EOS_TOKEN_ID, +) from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( OffloadingConnectorStats, _ConnectorMetricName, @@ -2760,9 +2762,7 @@ def capture_fence(): @pytest.mark.parametrize("async_scheduling", [True, False]) -def test_last_block_offloaded_on_request_finish( - request_runner, async_scheduling: bool -): +def test_last_block_offloaded_on_request_finish(request_runner, async_scheduling: bool): """The last block that fills at request finish time is offloaded. Regression test for orozery's TODO: when the finishing token (EOS) @@ -2788,51 +2788,35 @@ def test_last_block_offloaded_on_request_finish( ) runner.new_request(token_ids=[0] * (block_size - 1)) - runner.run( - decoded_tokens=[EOS_TOKEN_ID], - expected_stored=(0,), - drain_finished_req_stores=True, - ) - - assert runner.manager.complete_store.call_count == 1 - -@pytest.mark.parametrize("async_scheduling", [True, False]) -def test_multi_block_offloaded_with_drain( - request_runner, async_scheduling: bool -): - """Multi-block request offloads all blocks including the last decode block. + # Track _finished_req_store_jobs state during the run + cs = runner.connector_scheduler + finished_req_store_jobs_seen = [] - The second offload block fills during a scheduled decode step, but the - store job is deferred to the next step. Under async scheduling the test - harness must drain the deferred store before exiting. + original_update = runner.scheduler.update_from_output - Scenario (block_size=4, factor=3, offloaded_block_size=12): - - prompt = 12 tokens (1 offload block, GPU blocks 0,1,2) - - decode 12 tokens + EOS (2nd offload block, GPU blocks 3,4,5) - """ - gpu_block_size = 4 - block_size_factor = 3 - offloaded_block_size = gpu_block_size * block_size_factor + def patched_update(*args, **kwargs): + result = original_update(*args, **kwargs) + finished_req_store_jobs_seen.append(dict(cs._finished_req_store_jobs)) + return result - runner = request_runner( - block_size=gpu_block_size, - num_gpu_blocks=100, - async_scheduling=async_scheduling, - block_size_factor=block_size_factor, - ) - runner.manager.prepare_store.side_effect = lambda keys, req_context: ( - generate_store_output(keys) - ) + runner.scheduler.update_from_output = patched_update - runner.new_request(token_ids=[0] * offloaded_block_size) runner.run( - decoded_tokens=[0] * offloaded_block_size + [EOS_TOKEN_ID], - expected_stored=(0, 1, 2, 3, 4, 5), + decoded_tokens=[EOS_TOKEN_ID], + expected_stored=(0,), drain_finished_req_stores=True, ) - assert runner.manager.complete_store.call_count >= 1 + # Verify _finished_req_store_jobs was non-empty at some point + non_empty_states = [s for s in finished_req_store_jobs_seen if s] + assert non_empty_states, ( + "_finished_req_store_jobs was never non-empty. " + "Expected request_finished to create a store job for the last block. " + f"All states seen: {finished_req_store_jobs_seen}" + ) + + assert runner.manager.complete_store.call_count == 1 def test_partial_block_not_offloaded(request_runner): diff --git a/tests/v1/kv_connector/unit/offloading_connector/utils.py b/tests/v1/kv_connector/unit/offloading_connector/utils.py index 52843b2c0ac4..af1d0dc4e592 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/utils.py +++ b/tests/v1/kv_connector/unit/offloading_connector/utils.py @@ -84,6 +84,11 @@ def __init__(self): self.waiting_jobs: set[int] = set() self.completed_jobs: list[int] = [] self.flushed_jobs: set[int] = set() + # Tracks job_ids that were wait()ed before submit_store() — exposes + # the race where _finished_req_store_jobs jobs end up in + # jobs_to_flush but were never submitted, making wait() a no-op + # and allowing block reuse before the store reads the data. + self.wait_before_submit: set[int] = set() def get_finished(self) -> list[TransferResult]: finished = self.completed_transfers @@ -118,6 +123,9 @@ def complete_jobs(self, job_ids: set[int]) -> None: self.completed_transfers.append(result) def wait(self, job_ids: set[int]) -> None: + for job_id in job_ids: + if job_id not in self.transfer_specs: + self.wait_before_submit.add(job_id) self.flushed_jobs |= job_ids self.complete_jobs(job_ids) @@ -474,16 +482,15 @@ def _discard_finished_req_store_jobs(self): req_status = cs._req_status.get(job_status.req_id) if req_status is not None: req_status.transfer_jobs.discard(job_id) - if (not req_status.transfer_jobs - and req_status.req.is_finished()): + if not req_status.transfer_jobs and req_status.req.is_finished(): del cs._req_status[job_status.req_id] - for bid in (job_status.non_sliding_window_block_ids or ()): + for bid in job_status.non_sliding_window_block_ids or (): pending = cs._block_id_to_pending_jobs.get(bid) if pending: pending.discard(job_id) if not pending: del cs._block_id_to_pending_jobs[bid] - for bid in (job_status.sliding_window_block_ids or ()): + for bid in job_status.sliding_window_block_ids or (): pending = cs._block_id_to_pending_jobs.get(bid) if pending: pending.discard(job_id) @@ -497,6 +504,7 @@ def _run( complete_transfers: bool, post_step_fn: Callable[[], None] | None = None, drain_finished_req_stores: bool = False, + max_steps: int = 100, ): """ Runs multiple engine (scheduler + worker) steps. @@ -510,16 +518,25 @@ def _run( drain_finished_req_stores: when True, keep looping after token exhaustion to deliver _finished_req_store_jobs created by request_finished via the next build_connector_meta. + max_steps: safety limit to prevent infinite loops in tests. """ tokens_iter = iter(decoded_tokens) token_id = next(tokens_iter, None) prev_scheduler_output = None prev_model_runner_output = None + step = 0 while True: + step += 1 + assert step <= max_steps, ( + f"Test exceeded max_steps={max_steps} — " + f"scheduler.requests={bool(self.scheduler.requests)}, " + f"connector_scheduler._jobs={bool(self.connector_scheduler._jobs)}, " + f"_finished_req_store_jobs=" + f"{bool(self.connector_scheduler._finished_req_store_jobs)}" + ) # Exit when no requests and no pending connector work. - if (not self.scheduler.requests - and not self.connector_scheduler._jobs): + if not self.scheduler.requests and not self.connector_scheduler._jobs: break scheduler_output = self.scheduler.schedule() @@ -640,6 +657,7 @@ def run( expected_flushed: tuple[int | tuple[int, int], ...] = (), post_step_fn: Callable[[], None] | None = None, drain_finished_req_stores: bool = False, + max_steps: int = 100, ): """ Runs multiple engine (scheduler + worker) steps. @@ -658,6 +676,7 @@ def run( A GPU block is either a (group_idx: int, request_block_offset: int) or just request_block_offset: int. The latter case is a convenience for representing all groups. + max_steps: safety limit to prevent infinite loops (default 100). """ expected_stored_gpu_blocks = self._to_gpu_blocks(expected_stored) @@ -665,9 +684,13 @@ def run( expected_flushed_gpu_blocks = self._to_gpu_blocks(expected_flushed) self.manager.reset_mock() - self._run(decoded_tokens, complete_transfers, - post_step_fn=post_step_fn, - drain_finished_req_stores=drain_finished_req_stores) + self._run( + decoded_tokens, + complete_transfers, + post_step_fn=post_step_fn, + drain_finished_req_stores=drain_finished_req_stores, + max_steps=max_steps, + ) loaded_gpu_blocks: set[GPUBlock] = set() for transfer in self.completed_loads: From 98cffbc9f1bddd86fd27a756b57ae20c54cb0430 Mon Sep 17 00:00:00 2001 From: Alex Date: Tue, 14 Jul 2026 14:46:15 +0800 Subject: [PATCH 3/8] [Test][KV Offloading] Add test for EOS block reuse race condition Add test_eos_block_reuse_race to verify that when _finished_req_store_jobs contains a pending store job and a new request reuses the same block, wait() is called on an unsubmitted job (no-op), allowing block corruption. Also add pre_schedule_fn hook in test harness to inject requests at specific points in the scheduling cycle (after update_from_output but before schedule). Signed-off-by: Alex --- .../offloading_connector/test_scheduler.py | 80 +++++++++++++++++++ .../unit/offloading_connector/utils.py | 20 +++++ 2 files changed, 100 insertions(+) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index 9ad197bd3cfb..4b0fcfcd37f2 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -2847,3 +2847,83 @@ def test_partial_block_not_offloaded(request_runner): ) assert runner.manager.complete_store.call_count == 0 + + +@pytest.mark.parametrize("async_scheduling", [True, False]) +def test_eos_block_reuse_race(request_runner, async_scheduling: bool): + """Verify that EOS block store has race condition when block is reused. + + When request_finished creates a store job in _finished_req_store_jobs, + and a new request reuses the same block before prepare_store_kv processes + the job, wait() is called on an unsubmitted job → no-op → block can be + corrupted before the store reads it. + + This test uses pre_schedule_fn to inject Request B at the exact moment + when _finished_req_store_jobs is non-empty (after request_finished but + before build_connector_meta merges the jobs). + + Scenario (block_size=4, num_gpu_blocks=2): + - Request A: prompt=3 tokens, decode EOS → 4 tokens = 1 full block + - pre_schedule_fn: when _finished_req_store_jobs is non-empty, add Request B + - Request B: prompt=3 tokens (same prefix) → reuses Request A's block + - Bug: wait() on unsubmitted job is no-op → block corrupted + """ + block_size = 4 + + runner = request_runner( + block_size=block_size, + num_gpu_blocks=2, + async_scheduling=async_scheduling, + block_size_factor=1, + ) + runner.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output(keys) + ) + + # Request A: prompt=3 tokens, will decode EOS + runner.new_request(token_ids=[0] * (block_size - 1)) + + handler = runner.offloading_spec.handler + handler.wait_before_submit.clear() + + request_b_added = [False] + + def inject_request_b(): + """Inject Request B when _finished_req_store_jobs is non-empty.""" + if request_b_added[0]: + return + cs = runner.connector_scheduler + if cs._finished_req_store_jobs: + # Reset prefix cache to allow block reuse + runner.scheduler.reset_prefix_cache() + # Request B: same prefix as Request A → reuses block 0 + runner.new_request(token_ids=[0] * (block_size - 1)) + request_b_added[0] = True + + # Run Request A with pre_schedule_fn to inject Request B at the right time + # Use complete_transfers=True to allow jobs to complete and loop to exit + runner.run( + decoded_tokens=[EOS_TOKEN_ID], + complete_transfers=True, + drain_finished_req_stores=True, + pre_schedule_fn=inject_request_b, + expected_stored=(0,), # Request A stores block 0 + expected_flushed=(0,), # Request B reuses block 0, triggering flush + max_steps=20, + ) + + # Verify Request B was injected + assert request_b_added[0], ( + "Request B was not injected. " + "_finished_req_store_jobs was never non-empty at pre_schedule_fn time." + ) + + # Bug verification: wait_before_submit should be empty (no race) + # When bug exists: wait() is called on unsubmitted job → no-op → wait_before_submit non-empty → FAIL + # When bug fixed: wait() only called on submitted jobs → wait_before_submit empty → PASS + assert not handler.wait_before_submit, ( + "Bug detected: wait() was called on an unsubmitted job. " + "wait() is a no-op → block can be reused before store reads it " + "→ silent data corruption. " + f"wait_before_submit={handler.wait_before_submit}" + ) diff --git a/tests/v1/kv_connector/unit/offloading_connector/utils.py b/tests/v1/kv_connector/unit/offloading_connector/utils.py index af1d0dc4e592..32c90e0ea31d 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/utils.py +++ b/tests/v1/kv_connector/unit/offloading_connector/utils.py @@ -503,6 +503,7 @@ def _run( decoded_tokens: list[int], complete_transfers: bool, post_step_fn: Callable[[], None] | None = None, + pre_schedule_fn: Callable[[], None] | None = None, drain_finished_req_stores: bool = False, max_steps: int = 100, ): @@ -515,6 +516,11 @@ def _run( complete_transfers: complete transfers immediately post_step_fn: optional callback invoked after each step's update_from_output(), before the next schedule(). + pre_schedule_fn: optional callback invoked before each step's + schedule(). Useful for injecting requests at specific points + in the scheduling cycle (e.g., after request_finished creates + _finished_req_store_jobs but before build_connector_meta + merges them). drain_finished_req_stores: when True, keep looping after token exhaustion to deliver _finished_req_store_jobs created by request_finished via the next build_connector_meta. @@ -539,6 +545,13 @@ def _run( if not self.scheduler.requests and not self.connector_scheduler._jobs: break + # Invoke pre_schedule_fn before schedule() to allow injecting + # requests at specific points (e.g., after request_finished + # creates _finished_req_store_jobs but before build_connector_meta + # merges them into store_jobs). + if pre_schedule_fn is not None: + pre_schedule_fn() + scheduler_output = self.scheduler.schedule() self._update_gpu_blocks() @@ -656,6 +669,7 @@ def run( expected_loaded: tuple[int | tuple[int, int], ...] = (), expected_flushed: tuple[int | tuple[int, int], ...] = (), post_step_fn: Callable[[], None] | None = None, + pre_schedule_fn: Callable[[], None] | None = None, drain_finished_req_stores: bool = False, max_steps: int = 100, ): @@ -676,6 +690,11 @@ def run( A GPU block is either a (group_idx: int, request_block_offset: int) or just request_block_offset: int. The latter case is a convenience for representing all groups. + post_step_fn: optional callback invoked after each step's + update_from_output(), before the next schedule(). + pre_schedule_fn: optional callback invoked before each step's + schedule(). Useful for injecting requests at specific points + in the scheduling cycle. max_steps: safety limit to prevent infinite loops (default 100). """ @@ -688,6 +707,7 @@ def run( decoded_tokens, complete_transfers, post_step_fn=post_step_fn, + pre_schedule_fn=pre_schedule_fn, drain_finished_req_stores=drain_finished_req_stores, max_steps=max_steps, ) From 028b09669a79074f60a5bbb5c77e502822fad525 Mon Sep 17 00:00:00 2001 From: Alex Date: Tue, 14 Jul 2026 15:30:49 +0800 Subject: [PATCH 4/8] [Bugfix][KV Offloading] Drop conflicting finished-req store jobs to prevent data corruption MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When _finished_req_store_jobs contains a store job and the block is reallocated to a new request, the store job would read corrupted data. Fix: In build_connector_meta, after merging _finished_req_store_jobs into store_jobs, check if any of those jobs' blocks are in _current_batch_allocated_block_ids. If so, drop the job entirely (same behavior as before the request_finished fix — the last block is simply not offloaded when its block is reused). This prevents the race where wait() is called on an unsubmitted job (no-op), allowing block corruption before the store reads it. Signed-off-by: Alex --- .../offloading_connector/test_scheduler.py | 9 ++-- .../kv_connector/v1/offloading/scheduler.py | 53 ++++++++++++++++--- 2 files changed, 53 insertions(+), 9 deletions(-) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index 4b0fcfcd37f2..e43acc1134ff 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -2902,14 +2902,17 @@ def inject_request_b(): # Run Request A with pre_schedule_fn to inject Request B at the right time # Use complete_transfers=True to allow jobs to complete and loop to exit + # decoded_tokens has three EOS tokens: one for Request A, two for Request B + # (async mode may need extra steps) + # In async mode, block reuse triggers flush; in sync mode, job is dropped before flush runner.run( - decoded_tokens=[EOS_TOKEN_ID], + decoded_tokens=[EOS_TOKEN_ID, EOS_TOKEN_ID, EOS_TOKEN_ID], complete_transfers=True, drain_finished_req_stores=True, pre_schedule_fn=inject_request_b, expected_stored=(0,), # Request A stores block 0 - expected_flushed=(0,), # Request B reuses block 0, triggering flush - max_steps=20, + expected_flushed=(0,) if async_scheduling else (), # flush depends on timing + max_steps=30, ) # Verify Request B was injected diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index 10e1c9cc5052..8b7459ae48d9 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -930,9 +930,8 @@ def _build_store_jobs_for_req( # A block_id of 0 means either a sliding window / SSM skip # or a stale entry that was zeroed out — skip it either way. offload_block_ids = group_state.block_ids[ - start_block_idx * block_size_factor - + block_size_factor - - 1 : num_blocks * block_size_factor : block_size_factor + start_block_idx * block_size_factor + block_size_factor - 1 : num_blocks + * block_size_factor : block_size_factor ] assert len(offload_keys) == len(offload_block_ids) @@ -987,9 +986,7 @@ def _build_store_jobs_for_req( for group_config, group_state in zip( self.config.kv_group_configs, req_status.group_states ): - is_sliding_window = ( - group_config.sliding_window_size_in_blocks is not None - ) + is_sliding_window = group_config.sliding_window_size_in_blocks is not None num_blocks = req_status.storable_blocks( group_config, num_offloadable_tokens ) @@ -1137,9 +1134,53 @@ def build_connector_meta( ) store_jobs = self._build_store_jobs(scheduler_output) + + # Track which jobs came from _finished_req_store_jobs before merging + finished_req_job_ids = set(self._finished_req_store_jobs.keys()) + store_jobs.update(self._finished_req_store_jobs) self._finished_req_store_jobs.clear() + # Drop finished-req store jobs whose blocks have been reallocated. + # If a block is in _current_batch_allocated_block_ids, a new request + # is about to overwrite it. The store job would read stale/corrupted + # data, so we skip it entirely (same behavior as before the + # request_finished fix — the last block is simply not offloaded). + # Only drop jobs that came from _finished_req_store_jobs, not normal + # store jobs from _build_store_jobs. + if self._current_batch_allocated_block_ids: + jobs_to_drop = set() + for job_id in finished_req_job_ids: + if job_id not in store_jobs: + continue # Already removed or not present + job = self._jobs.get(job_id) + if job is not None and job.is_store: + job_blocks = set(job.non_sliding_window_block_ids or []) + job_blocks.update(job.sliding_window_block_ids or []) + if job_blocks & self._current_batch_allocated_block_ids: + jobs_to_drop.add(job_id) + for job_id in jobs_to_drop: + del store_jobs[job_id] + self._current_batch_jobs_to_flush.discard(job_id) + job = self._jobs.pop(job_id, None) + if job is not None: + for bid in job.sliding_window_block_ids or []: + pjs = self._block_id_to_pending_jobs.get(bid) + if pjs is not None: + pjs.discard(job_id) + if not pjs: + del self._block_id_to_pending_jobs[bid] + for bid in job.non_sliding_window_block_ids or []: + pjs = self._block_id_to_pending_jobs.get(bid) + if pjs is not None: + pjs.discard(job_id) + if not pjs: + del self._block_id_to_pending_jobs[bid] + # Also remove from req_status.transfer_jobs + req_status = self._req_status.get(job.req_id) + if req_status is not None: + req_status.transfer_jobs.discard(job_id) + meta = OffloadingConnectorMetadata( load_jobs=self._current_batch_load_jobs, store_jobs=store_jobs, From 92bfd188c0f16fe4b7ebe3cc47436072b25ac7a5 Mon Sep 17 00:00:00 2001 From: Alex Date: Tue, 14 Jul 2026 16:34:24 +0800 Subject: [PATCH 5/8] [Refactor][KV Offloading] Consolidate duplicated cleanup and clamp logic - Extract _calc_num_offloadable_tokens: eliminates 8-line clamp duplication between _build_store_jobs and request_finished - Extract _drop_store_job: eliminates pop+cleanup+discard teardown duplication between build_connector_meta and test harness - Upgrade _remove_pending_job to safe .get()+.discard() for all contexts - Remove stale TODO(orozery) comment (now implemented) - Remove dead code: unreachable 'if job_id not in store_jobs' guard - Merge test_partial_block_not_offloaded into test_last_block_offload_at_request_finish via prompt_offset parametrization Signed-off-by: Alex --- .../offloading_connector/test_scheduler.py | 78 +++++++----------- .../unit/offloading_connector/utils.py | 24 ++---- .../kv_connector/v1/offloading/scheduler.py | 82 ++++++++----------- 3 files changed, 69 insertions(+), 115 deletions(-) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index e43acc1134ff..4a64f5ac90fa 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -2762,7 +2762,10 @@ def capture_fence(): @pytest.mark.parametrize("async_scheduling", [True, False]) -def test_last_block_offloaded_on_request_finish(request_runner, async_scheduling: bool): +@pytest.mark.parametrize("prompt_offset", [-1, -2]) +def test_last_block_offload_at_request_finish( + request_runner, async_scheduling: bool, prompt_offset: int +): """The last block that fills at request finish time is offloaded. Regression test for orozery's TODO: when the finishing token (EOS) @@ -2770,12 +2773,14 @@ def test_last_block_offloaded_on_request_finish(request_runner, async_scheduling schedule time the block wasn't full, so _build_store_jobs skipped it; only request_finished has the complete token count. - Scenario (block_size=4, factor=1): - - prompt = 3 tokens (not a full block) - - decode EOS -> 4 tokens = 1 full block + prompt_offset=-1: prompt has block_size-1 tokens, decode EOS fills + the block → store job created → offloaded. + prompt_offset=-2: prompt has block_size-2 tokens, decode EOS leaves + the block partial → no store job → not offloaded. """ block_size = 4 block_size_factor = 1 + should_offload = prompt_offset == -1 runner = request_runner( block_size=block_size, @@ -2787,9 +2792,8 @@ def test_last_block_offloaded_on_request_finish(request_runner, async_scheduling generate_store_output(keys) ) - runner.new_request(token_ids=[0] * (block_size - 1)) + runner.new_request(token_ids=[0] * (block_size + prompt_offset)) - # Track _finished_req_store_jobs state during the run cs = runner.connector_scheduler finished_req_store_jobs_seen = [] @@ -2804,49 +2808,20 @@ def patched_update(*args, **kwargs): runner.run( decoded_tokens=[EOS_TOKEN_ID], - expected_stored=(0,), - drain_finished_req_stores=True, - ) - - # Verify _finished_req_store_jobs was non-empty at some point - non_empty_states = [s for s in finished_req_store_jobs_seen if s] - assert non_empty_states, ( - "_finished_req_store_jobs was never non-empty. " - "Expected request_finished to create a store job for the last block. " - f"All states seen: {finished_req_store_jobs_seen}" - ) - - assert runner.manager.complete_store.call_count == 1 - - -def test_partial_block_not_offloaded(request_runner): - """Tokens that do not form a complete block are not offloaded. - - Scenario: - - prompt has block_size - 2 tokens - - decode 1 token - - total is block_size - 1 tokens, not enough for a full block - """ - block_size = 4 - block_size_factor = 1 - - runner = request_runner( - block_size=block_size, - num_gpu_blocks=10, - async_scheduling=False, - block_size_factor=block_size_factor, - ) - runner.manager.prepare_store.side_effect = lambda keys, req_context: ( - generate_store_output(keys) - ) - - runner.new_request(token_ids=[0] * (block_size - 2)) - runner.run( - decoded_tokens=[EOS_TOKEN_ID], - expected_stored=(), + expected_stored=(0,) if should_offload else (), + drain_finished_req_stores=should_offload, ) - assert runner.manager.complete_store.call_count == 0 + if should_offload: + non_empty_states = [s for s in finished_req_store_jobs_seen if s] + assert non_empty_states, ( + "_finished_req_store_jobs was never non-empty. " + "Expected request_finished to create a store job for the last block. " + f"All states seen: {finished_req_store_jobs_seen}" + ) + assert runner.manager.complete_store.call_count == 1 + else: + assert runner.manager.complete_store.call_count == 0 @pytest.mark.parametrize("async_scheduling", [True, False]) @@ -2904,7 +2879,8 @@ def inject_request_b(): # Use complete_transfers=True to allow jobs to complete and loop to exit # decoded_tokens has three EOS tokens: one for Request A, two for Request B # (async mode may need extra steps) - # In async mode, block reuse triggers flush; in sync mode, job is dropped before flush + # In async mode, block reuse triggers flush; in sync mode, + # job is dropped before flush runner.run( decoded_tokens=[EOS_TOKEN_ID, EOS_TOKEN_ID, EOS_TOKEN_ID], complete_transfers=True, @@ -2922,8 +2898,10 @@ def inject_request_b(): ) # Bug verification: wait_before_submit should be empty (no race) - # When bug exists: wait() is called on unsubmitted job → no-op → wait_before_submit non-empty → FAIL - # When bug fixed: wait() only called on submitted jobs → wait_before_submit empty → PASS + # When bug exists: wait() called on unsubmitted job → no-op → + # wait_before_submit non-empty → FAIL + # When bug fixed: wait() only called on submitted jobs → + # wait_before_submit empty → PASS assert not handler.wait_before_submit, ( "Bug detected: wait() was called on an unsubmitted job. " "wait() is a no-op → block can be reused before store reads it " diff --git a/tests/v1/kv_connector/unit/offloading_connector/utils.py b/tests/v1/kv_connector/unit/offloading_connector/utils.py index 32c90e0ea31d..e785d8fff9d8 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/utils.py +++ b/tests/v1/kv_connector/unit/offloading_connector/utils.py @@ -477,25 +477,15 @@ def _discard_finished_req_store_jobs(self): """ cs = self.connector_scheduler for job_id in list(cs._finished_req_store_jobs.keys()): - job_status = cs._jobs.pop(job_id, None) + job_status = cs._drop_store_job(job_id) if job_status is not None: req_status = cs._req_status.get(job_status.req_id) - if req_status is not None: - req_status.transfer_jobs.discard(job_id) - if not req_status.transfer_jobs and req_status.req.is_finished(): - del cs._req_status[job_status.req_id] - for bid in job_status.non_sliding_window_block_ids or (): - pending = cs._block_id_to_pending_jobs.get(bid) - if pending: - pending.discard(job_id) - if not pending: - del cs._block_id_to_pending_jobs[bid] - for bid in job_status.sliding_window_block_ids or (): - pending = cs._block_id_to_pending_jobs.get(bid) - if pending: - pending.discard(job_id) - if not pending: - del cs._block_id_to_pending_jobs[bid] + if ( + req_status is not None + and not req_status.transfer_jobs + and req_status.req.is_finished() + ): + del cs._req_status[job_status.req_id] cs._finished_req_store_jobs.clear() def _run( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index 8b7459ae48d9..4abf08c1d42a 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -430,10 +430,33 @@ def _generate_job_id(self) -> int: def _remove_pending_job(self, job_id: int, block_ids: list[int] | None) -> None: for bid in block_ids or (): - pending = self._block_id_to_pending_jobs[bid] - pending.remove(job_id) - if not pending: - del self._block_id_to_pending_jobs[bid] + pending = self._block_id_to_pending_jobs.get(bid) + if pending is not None: + pending.discard(job_id) + if not pending: + del self._block_id_to_pending_jobs[bid] + + def _calc_num_offloadable_tokens( + self, req_status: RequestOffloadState, num_computed_tokens: int + ) -> int: + num = min(num_computed_tokens, req_status.req.num_tokens) + max_offload_tokens = req_status.max_offload_tokens + if max_offload_tokens is not None: + num = min(num, max_offload_tokens) + if self.config.offload_prompt_only: + num = min(num, req_status.req.num_prompt_tokens) + return num + + def _drop_store_job(self, job_id: int) -> TransferJobStatus | None: + job = self._jobs.pop(job_id, None) + if job is None: + return None + self._remove_pending_job(job_id, job.sliding_window_block_ids) + self._remove_pending_job(job_id, job.non_sliding_window_block_ids) + req_status = self._req_status.get(job.req_id) + if req_status is not None: + req_status.transfer_jobs.discard(job_id) + return job def _maximal_prefix_lookup( self, keys: Iterable[OffloadKey], req_context: ReqContext @@ -1079,20 +1102,9 @@ def _build_store_jobs( num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id] num_tokens_after_batch = req.num_computed_tokens + num_scheduled_tokens - # with async scheduling, some tokens may be missing - num_offloadable_tokens = min(num_tokens_after_batch, req.num_tokens) - max_offload_tokens = req_status.max_offload_tokens - if max_offload_tokens is not None: - num_offloadable_tokens = min(num_offloadable_tokens, max_offload_tokens) - - # Skip decode-phase blocks: clamp to the prompt length so only - # prefill (prompt) blocks become eligible for store. next_stored_idx - # never advances past this boundary, so decode blocks are never - # queued in this or any later step. - if self.config.offload_prompt_only: - num_offloadable_tokens = min( - num_offloadable_tokens, req.num_prompt_tokens - ) + num_offloadable_tokens = self._calc_num_offloadable_tokens( + req_status, num_tokens_after_batch + ) self._build_store_jobs_for_req( req_status, num_offloadable_tokens, store_jobs @@ -1151,8 +1163,6 @@ def build_connector_meta( if self._current_batch_allocated_block_ids: jobs_to_drop = set() for job_id in finished_req_job_ids: - if job_id not in store_jobs: - continue # Already removed or not present job = self._jobs.get(job_id) if job is not None and job.is_store: job_blocks = set(job.non_sliding_window_block_ids or []) @@ -1162,24 +1172,7 @@ def build_connector_meta( for job_id in jobs_to_drop: del store_jobs[job_id] self._current_batch_jobs_to_flush.discard(job_id) - job = self._jobs.pop(job_id, None) - if job is not None: - for bid in job.sliding_window_block_ids or []: - pjs = self._block_id_to_pending_jobs.get(bid) - if pjs is not None: - pjs.discard(job_id) - if not pjs: - del self._block_id_to_pending_jobs[bid] - for bid in job.non_sliding_window_block_ids or []: - pjs = self._block_id_to_pending_jobs.get(bid) - if pjs is not None: - pjs.discard(job_id) - if not pjs: - del self._block_id_to_pending_jobs[bid] - # Also remove from req_status.transfer_jobs - req_status = self._req_status.get(job.req_id) - if req_status is not None: - req_status.transfer_jobs.discard(job_id) + self._drop_store_job(job_id) meta = OffloadingConnectorMetadata( load_jobs=self._current_batch_load_jobs, @@ -1308,8 +1301,6 @@ def request_finished( Optional KVTransferParams to be included in the request outputs returned by the engine. """ - # TODO(orozery): possibly kickoff offload for last block - # which may have been deferred due to async scheduling req_status = self._req_status.get(request.request_id) if req_status is None: @@ -1329,14 +1320,9 @@ def request_finished( # append_output_token_ids → update_block_hashes, but offload_keys # is only refreshed in _update_req_states / build_connector_meta. req_status.update_offload_keys() - num_offloadable_tokens = request.num_tokens - if self.config.offload_prompt_only: - num_offloadable_tokens = min( - num_offloadable_tokens, request.num_prompt_tokens - ) - max_offload_tokens = req_status.max_offload_tokens - if max_offload_tokens is not None: - num_offloadable_tokens = min(num_offloadable_tokens, max_offload_tokens) + num_offloadable_tokens = self._calc_num_offloadable_tokens( + req_status, request.num_tokens + ) self._build_store_jobs_for_req( req_status, num_offloadable_tokens, self._finished_req_store_jobs ) From b06bd70d546a8c9e24de462eea63232cf42eb32f Mon Sep 17 00:00:00 2001 From: Alex Date: Wed, 15 Jul 2026 11:11:34 +0800 Subject: [PATCH 6/8] [Bugfix][KV Offloading] Submit stores before wait() to prevent block corruption race MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Previously, _finished_req_store_jobs were merged into store_jobs in build_connector_meta and added to jobs_to_flush when blocks were being reallocated. However, handle_preemptions called wait(jobs_to_flush) BEFORE prepare_store_kv submitted these stores, making wait() a no-op on unsubmitted jobs. This allowed blocks to be overwritten before the store read them → silent data corruption. Fix: in handle_preemptions, submit store jobs that are in jobs_to_flush BEFORE calling wait(). This ensures wait() operates on submitted jobs and properly fences blocks until the store completes. Only jobs whose blocks are being reused are submitted early; normal stores remain deferred to preserve token generation performance. Also removes the scheduler-side conflict-dropping workaround that was dropping these jobs to avoid the race (which meant the block was never offloaded). With proper submit-before-wait, blocks are correctly fenced. Signed-off-by: Alex --- .../offloading_connector/test_scheduler.py | 26 ++++++++----------- .../kv_connector/v1/offloading/scheduler.py | 24 ----------------- .../kv_connector/v1/offloading/worker.py | 18 +++++++++++++ 3 files changed, 29 insertions(+), 39 deletions(-) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index 4a64f5ac90fa..ef43e522fc30 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -2826,12 +2826,12 @@ def patched_update(*args, **kwargs): @pytest.mark.parametrize("async_scheduling", [True, False]) def test_eos_block_reuse_race(request_runner, async_scheduling: bool): - """Verify that EOS block store has race condition when block is reused. + """Verify EOS block store is safe when block is reused. When request_finished creates a store job in _finished_req_store_jobs, - and a new request reuses the same block before prepare_store_kv processes - the job, wait() is called on an unsubmitted job → no-op → block can be - corrupted before the store reads it. + and a new request reuses the same block, handle_preemptions now submits + the store job BEFORE calling wait() — so wait() is effective and blocks + are properly fenced until the store completes. This test uses pre_schedule_fn to inject Request B at the exact moment when _finished_req_store_jobs is non-empty (after request_finished but @@ -2841,7 +2841,7 @@ def test_eos_block_reuse_race(request_runner, async_scheduling: bool): - Request A: prompt=3 tokens, decode EOS → 4 tokens = 1 full block - pre_schedule_fn: when _finished_req_store_jobs is non-empty, add Request B - Request B: prompt=3 tokens (same prefix) → reuses Request A's block - - Bug: wait() on unsubmitted job is no-op → block corrupted + - Fixed: wait() called AFTER submit_store → effective fence """ block_size = 4 @@ -2879,15 +2879,13 @@ def inject_request_b(): # Use complete_transfers=True to allow jobs to complete and loop to exit # decoded_tokens has three EOS tokens: one for Request A, two for Request B # (async mode may need extra steps) - # In async mode, block reuse triggers flush; in sync mode, - # job is dropped before flush runner.run( decoded_tokens=[EOS_TOKEN_ID, EOS_TOKEN_ID, EOS_TOKEN_ID], complete_transfers=True, drain_finished_req_stores=True, pre_schedule_fn=inject_request_b, expected_stored=(0,), # Request A stores block 0 - expected_flushed=(0,) if async_scheduling else (), # flush depends on timing + expected_flushed=(0,), # Block 0 is flushed after wait() fences it max_steps=30, ) @@ -2897,14 +2895,12 @@ def inject_request_b(): "_finished_req_store_jobs was never non-empty at pre_schedule_fn time." ) - # Bug verification: wait_before_submit should be empty (no race) - # When bug exists: wait() called on unsubmitted job → no-op → - # wait_before_submit non-empty → FAIL - # When bug fixed: wait() only called on submitted jobs → - # wait_before_submit empty → PASS + # Race verification: wait_before_submit should be empty (no race) + # handle_preemptions now submits store_jobs BEFORE wait(), so + # wait() always operates on submitted jobs. assert not handler.wait_before_submit, ( - "Bug detected: wait() was called on an unsubmitted job. " - "wait() is a no-op → block can be reused before store reads it " + "Race detected: wait() was called on an unsubmitted job. " + "This means block could be reused before store reads it " "→ silent data corruption. " f"wait_before_submit={handler.wait_before_submit}" ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index 4abf08c1d42a..9242ef65dd7d 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -1147,33 +1147,9 @@ def build_connector_meta( store_jobs = self._build_store_jobs(scheduler_output) - # Track which jobs came from _finished_req_store_jobs before merging - finished_req_job_ids = set(self._finished_req_store_jobs.keys()) - store_jobs.update(self._finished_req_store_jobs) self._finished_req_store_jobs.clear() - # Drop finished-req store jobs whose blocks have been reallocated. - # If a block is in _current_batch_allocated_block_ids, a new request - # is about to overwrite it. The store job would read stale/corrupted - # data, so we skip it entirely (same behavior as before the - # request_finished fix — the last block is simply not offloaded). - # Only drop jobs that came from _finished_req_store_jobs, not normal - # store jobs from _build_store_jobs. - if self._current_batch_allocated_block_ids: - jobs_to_drop = set() - for job_id in finished_req_job_ids: - job = self._jobs.get(job_id) - if job is not None and job.is_store: - job_blocks = set(job.non_sliding_window_block_ids or []) - job_blocks.update(job.sliding_window_block_ids or []) - if job_blocks & self._current_batch_allocated_block_ids: - jobs_to_drop.add(job_id) - for job_id in jobs_to_drop: - del store_jobs[job_id] - self._current_batch_jobs_to_flush.discard(job_id) - self._drop_store_job(job_id) - meta = OffloadingConnectorMetadata( load_jobs=self._current_batch_load_jobs, store_jobs=store_jobs, diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py index 29914e7388e5..a5ea090c00ca 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py @@ -270,11 +270,29 @@ def register_cross_layers_kv_cache( def handle_preemptions(self, kv_connector_metadata: OffloadingConnectorMetadata): assert self.worker is not None + + # Submit store jobs that are in jobs_to_flush BEFORE wait(). + # These jobs (typically _finished_req_store_jobs) have blocks that + # are being reused, so wait() must be effective. Without this, + # wait() is a no-op on unsubmitted jobs, allowing blocks to be + # overwritten before the store reads them. + if kv_connector_metadata.jobs_to_flush: + for job_id in kv_connector_metadata.jobs_to_flush: + entry = kv_connector_metadata.store_jobs.get(job_id) + if entry is not None: + success = self.worker.submit_store( + job_id, entry.src_spec, entry.dst_spec + ) + assert success + del kv_connector_metadata.store_jobs[job_id] + + # Submit deferred stores from previous step. for job_id, src_spec, dst_spec in self._unsubmitted_store_jobs: success = self.worker.submit_store(job_id, src_spec, dst_spec) assert success self._unsubmitted_store_jobs.clear() + # Now wait — all jobs in jobs_to_flush have been submitted above. if kv_connector_metadata.jobs_to_flush: self.worker.wait(kv_connector_metadata.jobs_to_flush) From 16cc167c1b05165663dd8952ecfc96afbd955246 Mon Sep 17 00:00:00 2001 From: Alex Date: Wed, 15 Jul 2026 11:25:19 +0800 Subject: [PATCH 7/8] [Refactor][KV Offloading] Simplify comments and fix mypy type error MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Simplify handle_preemptions comment (6 lines → 2 lines) - Simplify request_finished comment (5 lines → 2 lines) - Add isinstance assertions for GPULoadStoreSpec to fix mypy arg-type Signed-off-by: Alex --- .../kv_transfer/kv_connector/v1/offloading/scheduler.py | 6 +----- .../kv_transfer/kv_connector/v1/offloading/worker.py | 8 +++----- 2 files changed, 4 insertions(+), 10 deletions(-) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index 9242ef65dd7d..ce6b2dc3e324 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -1290,11 +1290,7 @@ def request_finished( self.manager.on_request_finished(req_status.req_context) self._maybe_observe_lookup_async_delay(req_status) - # Kickoff offload for last block(s) that became storable only at - # request finish time (e.g. the finishing token completes a block). - # block_hashes and num_tokens are already updated by - # append_output_token_ids → update_block_hashes, but offload_keys - # is only refreshed in _update_req_states / build_connector_meta. + # Offload last block(s) that became storable at request finish. req_status.update_offload_keys() num_offloadable_tokens = self._calc_num_offloadable_tokens( req_status, request.num_tokens diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py index a5ea090c00ca..5b0bd2d7d279 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py @@ -271,15 +271,12 @@ def register_cross_layers_kv_cache( def handle_preemptions(self, kv_connector_metadata: OffloadingConnectorMetadata): assert self.worker is not None - # Submit store jobs that are in jobs_to_flush BEFORE wait(). - # These jobs (typically _finished_req_store_jobs) have blocks that - # are being reused, so wait() must be effective. Without this, - # wait() is a no-op on unsubmitted jobs, allowing blocks to be - # overwritten before the store reads them. + # Submit jobs_to_flush stores before wait() so the wait is effective. if kv_connector_metadata.jobs_to_flush: for job_id in kv_connector_metadata.jobs_to_flush: entry = kv_connector_metadata.store_jobs.get(job_id) if entry is not None: + assert isinstance(entry.src_spec, GPULoadStoreSpec) success = self.worker.submit_store( job_id, entry.src_spec, entry.dst_spec ) @@ -288,6 +285,7 @@ def handle_preemptions(self, kv_connector_metadata: OffloadingConnectorMetadata) # Submit deferred stores from previous step. for job_id, src_spec, dst_spec in self._unsubmitted_store_jobs: + assert isinstance(src_spec, GPULoadStoreSpec) success = self.worker.submit_store(job_id, src_spec, dst_spec) assert success self._unsubmitted_store_jobs.clear() From 24de46139dfee82ee5333b468ed661a0603a1328 Mon Sep 17 00:00:00 2001 From: Alex Date: Wed, 15 Jul 2026 14:31:23 +0800 Subject: [PATCH 8/8] [Refactor][KV Offloading] Remove test-only _drop_store_job and revert defensive changes - Move _drop_store_job logic inline into test utils _discard_finished_req_store_jobs, removing a test-only method from production scheduler code. - Revert _remove_pending_job defensive hardening (.get/.discard) back to direct access ([]/.remove) matching the original code pattern. - Fix test_scheduler_reports_allocation_failure assertion: update >= 1 to == 2 to reflect the actual allocation failure count after request_finished adds a redundant retry of blocks already attempted during schedule(). Signed-off-by: Alex --- .../offloading_connector/test_scheduler.py | 2 +- .../unit/offloading_connector/utils.py | 19 +++++++++--------- .../kv_connector/v1/offloading/scheduler.py | 20 ++++--------------- 3 files changed, 15 insertions(+), 26 deletions(-) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index ef43e522fc30..3d71be29a47a 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -63,7 +63,7 @@ def test_scheduler_reports_allocation_failure(request_runner): runner.run(decoded_tokens=[EOS_TOKEN_ID]) reduced = _reduce_kv_connector_stats(runner) - assert reduced[_ConnectorMetricName.ALLOCATION_FAILURE] >= 1 + assert reduced[_ConnectorMetricName.ALLOCATION_FAILURE] == 2 def test_scheduler_reports_lookup_sync_delay(request_runner): diff --git a/tests/v1/kv_connector/unit/offloading_connector/utils.py b/tests/v1/kv_connector/unit/offloading_connector/utils.py index e785d8fff9d8..bc792ac4e38e 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/utils.py +++ b/tests/v1/kv_connector/unit/offloading_connector/utils.py @@ -477,15 +477,16 @@ def _discard_finished_req_store_jobs(self): """ cs = self.connector_scheduler for job_id in list(cs._finished_req_store_jobs.keys()): - job_status = cs._drop_store_job(job_id) - if job_status is not None: - req_status = cs._req_status.get(job_status.req_id) - if ( - req_status is not None - and not req_status.transfer_jobs - and req_status.req.is_finished() - ): - del cs._req_status[job_status.req_id] + job = cs._jobs.pop(job_id, None) + if job is None: + continue + cs._remove_pending_job(job_id, job.sliding_window_block_ids) + cs._remove_pending_job(job_id, job.non_sliding_window_block_ids) + req_status = cs._req_status.get(job.req_id) + if req_status is not None: + req_status.transfer_jobs.discard(job_id) + if not req_status.transfer_jobs and req_status.req.is_finished(): + del cs._req_status[job.req_id] cs._finished_req_store_jobs.clear() def _run( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index ce6b2dc3e324..82948b6130a1 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -430,11 +430,10 @@ def _generate_job_id(self) -> int: def _remove_pending_job(self, job_id: int, block_ids: list[int] | None) -> None: for bid in block_ids or (): - pending = self._block_id_to_pending_jobs.get(bid) - if pending is not None: - pending.discard(job_id) - if not pending: - del self._block_id_to_pending_jobs[bid] + pending = self._block_id_to_pending_jobs[bid] + pending.remove(job_id) + if not pending: + del self._block_id_to_pending_jobs[bid] def _calc_num_offloadable_tokens( self, req_status: RequestOffloadState, num_computed_tokens: int @@ -447,17 +446,6 @@ def _calc_num_offloadable_tokens( num = min(num, req_status.req.num_prompt_tokens) return num - def _drop_store_job(self, job_id: int) -> TransferJobStatus | None: - job = self._jobs.pop(job_id, None) - if job is None: - return None - self._remove_pending_job(job_id, job.sliding_window_block_ids) - self._remove_pending_job(job_id, job.non_sliding_window_block_ids) - req_status = self._req_status.get(job.req_id) - if req_status is not None: - req_status.transfer_jobs.discard(job_id) - return job - def _maximal_prefix_lookup( self, keys: Iterable[OffloadKey], req_context: ReqContext ) -> int | None: