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..3d71be29a47a 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, @@ -61,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): @@ -2757,3 +2759,148 @@ 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]) +@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) + 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. + + 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, + 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 + prompt_offset)) + + cs = runner.connector_scheduler + finished_req_store_jobs_seen = [] + + original_update = runner.scheduler.update_from_output + + 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.scheduler.update_from_output = patched_update + + runner.run( + decoded_tokens=[EOS_TOKEN_ID], + expected_stored=(0,) if should_offload else (), + drain_finished_req_stores=should_offload, + ) + + 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]) +def test_eos_block_reuse_race(request_runner, async_scheduling: bool): + """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, 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 + 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 + - Fixed: wait() called AFTER submit_store → effective fence + """ + 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 + # decoded_tokens has three EOS tokens: one for Request A, two for Request B + # (async mode may need extra steps) + 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,), # Block 0 is flushed after wait() fences it + max_steps=30, + ) + + # 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." + ) + + # 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, ( + "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/tests/v1/kv_connector/unit/offloading_connector/utils.py b/tests/v1/kv_connector/unit/offloading_connector/utils.py index 73ea5e2be1d2..bc792ac4e38e 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) @@ -459,11 +467,36 @@ 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 = 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( self, 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, ): """ Runs multiple engine (scheduler + worker) steps. @@ -474,20 +507,42 @@ 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. + 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: - # 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. + 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: 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() @@ -549,6 +604,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 +615,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 +660,9 @@ 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, ): """ Runs multiple engine (scheduler + worker) steps. @@ -608,6 +681,12 @@ 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). """ expected_stored_gpu_blocks = self._to_gpu_blocks(expected_stored) @@ -615,7 +694,14 @@ 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, + pre_schedule_fn=pre_schedule_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: 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..82948b6130a1 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: @@ -431,6 +435,17 @@ def _remove_pending_job(self, job_id: int, block_ids: list[int] | None) -> None: 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 _maximal_prefix_lookup( self, keys: Iterable[OffloadKey], req_context: ReqContext ) -> int | None: @@ -895,192 +910,192 @@ def _update_req_states(self, scheduler_output: SchedulerOutput) -> None: ): group_state.block_ids[j] = 0 - def _build_store_jobs( + def _build_store_jobs_for_req( self, - scheduler_output: SchedulerOutput, - ) -> dict[int, TransferJob]: + 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 - store_jobs: dict[int, TransferJob] = {} - for req_id in scheduler_output.num_scheduled_tokens: - req_status = self._req_status.get(req_id) - if req_status is None: - continue - req = req_status.req - 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 - ) + # 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 + ) - # 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 + 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) ): - 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: + if block_id == 0: 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: + # 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 - # 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 + 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 + 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 ) - if store_output is None: - self._connector_stats.increase_counter( - _ConnectorMetricName.ALLOCATION_FAILURE - ) - logger.warning("Request %s: cannot store blocks", req_id) - continue + 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) - continue + if not store_output.keys_to_store: + req_status.advance_stored_idx(num_offloadable_tokens) + return - self._touch(req_status) + self._touch(req_status) - keys_to_store = set(store_output.keys_to_store) + 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 + 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] ): - is_sliding_window = ( - group_config.sliding_window_size_in_blocks is not None - ) - num_blocks = req_status.storable_blocks( - group_config, num_offloadable_tokens + 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 ) - 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: + + 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 + ) - offloaded_block_idx = start_block_idx + idx + src_spec = GPULoadStoreSpec( + src_block_ids, group_sizes=group_sizes, block_indices=block_indices + ) + dst_spec = store_output.store_spec - self._events_tracker.record_store( - req, group_config, offloaded_block_idx, offload_key - ) + 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) - 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 - ) + # 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) - 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, + ) - # 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 + ) - 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]: + store_jobs: dict[int, TransferJob] = {} + for req_id in scheduler_output.num_scheduled_tokens: + req_status = self._req_status.get(req_id) + if req_status is None: + continue + req = req_status.req + + num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id] + num_tokens_after_batch = req.num_computed_tokens + num_scheduled_tokens + num_offloadable_tokens = self._calc_num_offloadable_tokens( + req_status, num_tokens_after_batch ) - 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 +1133,14 @@ 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 = {} @@ -1245,8 +1265,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: @@ -1259,6 +1277,16 @@ def request_finished( self.manager.on_request_finished(req_status.req_context) self._maybe_observe_lookup_async_delay(req_status) + + # 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 + ) + 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 +1344,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. 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..5b0bd2d7d279 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,27 @@ def register_cross_layers_kv_cache( def handle_preemptions(self, kv_connector_metadata: OffloadingConnectorMetadata): assert self.worker is not None + + # 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 + ) + 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: + assert isinstance(src_spec, GPULoadStoreSpec) 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)