[Bugfix][KV Offloading] Offload last block at request finish and prevent reuse race#48596
Open
Alex-ai-future wants to merge 8 commits into
Open
[Bugfix][KV Offloading] Offload last block at request finish and prevent reuse race#48596Alex-ai-future wants to merge 8 commits into
Alex-ai-future wants to merge 8 commits into
Conversation
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 <alex.tech.lab@outlook.com>
… 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 <alex.tech.lab@outlook.com>
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 <alex.tech.lab@outlook.com>
…revent data corruption 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 <alex.tech.lab@outlook.com>
Alex-ai-future
requested review from
ApostaC,
NickLucche,
ivanium,
orozery and
xuechendi
as code owners
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
- 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 <alex.tech.lab@outlook.com>
Alex-ai-future
force-pushed
the
fix/last_block
branch
from
65e82de to
92bfd18
Compare
…corruption race 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 <alex.tech.lab@outlook.com>
- 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 <alex.tech.lab@outlook.com>
… 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 <alex.tech.lab@outlook.com>
Contributor
Author
|
@orozery PTAL — this implements your request_finished TODO plus the handle_preemptions wait-before-submit fix. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Purpose
Fix: last KV block that fills at request finish time is never offloaded, causing prefix cache misses for that block. Also fixes a data corruption race when the block is reused.
Bug principle:
_build_store_jobsruns duringschedule()— at that point the finishing token (EOS) hasn't been processed yet, so the last block is still partial and skipped.request_finishedis called after EOS is appended (viaupdate_block_hashes), but it had no store kickoff logic — the last block was silently dropped.Consequence chain: last block not offloaded → next request with the same prefix misses that block → extra prefill compute + no KV reuse for the final block.
Fix: In
request_finished, callupdate_offload_keys+_build_store_jobs_for_reqto create store jobs for blocks that became storable only at finish time. Jobs are placed in_finished_req_store_jobs(a deferred queue), not submitted immediately — they are merged intostore_jobsat the nextbuild_connector_meta, wherejobs_to_flushcan fence the blocks if needed. This matches the approach described in #47107: build store jobs atrequest_finished, fence blocks, flush at nextbuild_connector_meta.Race condition fix:
request_finishedcreates a store job but doesn't submit it. Before the nextbuild_connector_metaruns, the scheduler may reallocate the same block to a new request. On the worker side,handle_preemptionscallswait()onjobs_to_flush— but the new store job hasn't beensubmit_store'd yet, sowait()is a no-op → block gets overwritten → silent data corruption. Fix: inhandle_preemptions, submit store jobs that are injobs_to_flushBEFORE callingwait(). This ensureswait()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.Related issue: #41704 (same class of bug in
SimpleCPUOffloadScheduler, fixed by #41777 — different component, not a duplicate).Test Plan
test_last_block_offload_at_request_finish[-1]request_finishedcreates store job → block is offloadedtest_last_block_offload_at_request_finish[-2]test_eos_block_reuse_raceTests use
wait_before_submittracking to expose the race: ifwait()fires beforesubmit_store(), the test detects the no-op.Test harness improvements (in
utils.py):wait_before_submit: tracks job_ids thatwait()was called on beforesubmit_store()— exposes the race where_finished_req_store_jobsjobs end up injobs_to_flushbut were never submitted, makingwait()a no-oppre_schedule_fn: callback invoked before each step'sschedule()— allows injecting requests at precise timing points (e.g., afterrequest_finishedcreates_finished_req_store_jobsbut beforebuild_connector_metamerges them)drain_finished_req_stores: keeps the test loop running after token exhaustion so_finished_req_store_jobscan be delivered via the nextbuild_connector_meta— without this, store jobs created at request finish would never be submitted_discard_finished_req_store_jobs: cleans up deferred store jobs whendrain_finished_req_stores=False— without this, stale jobs in_jobswould prevent the test loop from exitingmax_steps: safety limit on test loop iterations — prevents infinite loops from hanging CI, and the assertion message includes relevant scheduler state (_jobs,_finished_req_store_jobs) for fast diagnosisExisting test fix:
test_scheduler_reports_allocation_failureassertion updated from== 1to== 2.request_finishedcreates a store job for the same block already attempted duringschedule(), triggering a second allocation failure (the block'snext_stored_block_idxis not advanced on failure, sorequest_finishedretries it)..venv/bin/python -m pytest tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py -v -k "last_block or eos_block_reuse"Test Result