Repository navigation
[Bugfix][Core] Reserve only the blocks a spec decode step adds when admitting async KV loads - #63136
krishchaudh wants to merge 3 commits into
Conversation
…dmitting async KV loads An async KV load is admitted only if it fits in the free blocks minus reservations, including a margin for the request's first spec decode step. That margin was a flat cdiv(1 + num_spec + lookahead, block_size) per group, so a prompt that fills the pool's last block (or sits at max_model_len) could never be admitted even though its spec step needs no extra block: allocate_slots fails, the FCFS loop breaks, and the request is retried forever with nothing running to free blocks. Compute the margin per request as the blocks between the full sequence and the sequence plus one spec step, capped at max_model_len like allocate_slots, and use the same capped count for in-flight prefill reservations. Signed-off-by: Krish Chaudhary <krishchaudhary1008@gmail.com> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment Once the PR is approved or has the If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. Agent GuidelinesIMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban. 🚀 |
Restart the decode instance of the edge case job with ngram spec decode and a 5-block KV cache, then remote-prefill a prompt that fills the 4 usable blocks with room in its last block for the spec step. On main the decode side never admits it (flat extra-block margin), so the request times out; with the fix it completes. Signed-off-by: Krish Chaudhary <krishchaudhary1008@gmail.com> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Signed-off-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
… blocks Co-authored-by: Codex <noreply@openai.com> Signed-off-by: Chenxi Ma <114001934+algebra-MCX@users.noreply.github.com>
Purpose
With a KV connector doing async loads (P/D decode side, offloading) and spec decode enabled, a request whose remote-hit prompt needs exactly the free KV blocks is never admitted: the scheduler reserves a flat extra block for its first spec step even when that step fits in the block the request already has. The request spins at the head of
waitingforever with no error. Follow-up to #57104 / #59880, found while testing the scheduler (not a production incident).Example
num_speculative_tokens=1, max_model_len 128Step 1 — A arrives
allocate_slots(A, num_new_tokens=126, ..., reserved_blocks=0 inflight + 1 spec margin)A needs
cdiv(126, 32) = 4blocksfree = 4, minus 1 reserved = 3 < 4 →
allocate_slotsreturns None →breakresult:
Step 2 onward — Same cycle
Nothing is running, so no block is ever freed. Peek
waiting→ A → same allocation → fails →break. Spin forever, and every request behind A waits too (FCFS).But A did not need the extra block: after the load, its spec step grows it from 126 to 128 tokens, which is still 4 blocks (2 free slots in the last block). The same happens for any prompt already at
max_model_len(allocate_slotscaps the slots, so the step adds nothing).Root Cause
#57104 fixed the opposite mismatch (reserve too little → promoted request can't get its spec block). It did so with a flat margin precomputed in
Scheduler.__init__:cdiv(1 + num_spec + lookahead, block_size)per non-Mamba group, i.e. 1 block for anynum_spec < block_size, charged against every async load regardless of where the request's sequence ends.Whether the spec step needs a new block depends on the request:
cdiv(min(n + 1 + spec + lookahead, max_model_len), bs) - cdiv(n, bs)is 1 when the step crosses a block boundary (the #57104 case) and 0 when it doesn't. This PR computes that per request (_spec_decode_step_blocks) instead of using the flat 1, so the #57104/#59880 cases still reserve the block and the fitting case is admitted._request_remaining_blocksuses the samemax_model_lencap.Not the same as "blocks exactly match max_model_len": the trigger is prompt length relative to the block boundary (e.g. 126 of 128 slots) or being at
max_model_len; either way the request needs exactly the free blocks.Test Plan
tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py::test_async_load_admits_request_whose_spec_step_fits_its_blocks(num_spec_tokens1/3 × room-in-last-block / capped-by-max_model_len, promptmax_model_len - 1), realScheduler+MockKVConnectortest_async_load_near_max_model_len_does_not_block_another_fitting_load(spec 1/3): a near-max_model_lenin-flight load must not reserve slots past the cap and block a second fitting load (covers the_request_remaining_blockschange)test_async_load_reserves_blocks_for_promotion_margin+ deadlock teststests/v1/kv_connector/nixl_integration/test_spec_decode_admission.py, second phase ofrun_edge_case_test.sh(NixlConnector PD edge caseCI job). Real prefill + decodevllm serve Qwen/Qwen3-0.6BoverNixlConnector; decode restarted with--speculative-config '{"method":"ngram","num_speculative_tokens":1,"prompt_lookup_max":4,"prompt_lookup_min":2}' --max-model-len 128 --block-size 32 --num-gpu-blocks-override 5(1 null + 4 usable). 126-token (fits, 2 output tokens) and 127-token (capped bymax_model_len, 1 output token) prompts remote-prefilled, then sent to decode with the returnedkv_transfer_params(ignore_eos); asserts completion within 60 s. Both phase-1 servers also use--block-size 32so all P/D block sizes match.Test Result
e2e was run locally on CPU (no GPU available;
VLLM_USE_V2_MODEL_RUNNER=0, NIXL over DRAM); the CI job runs it on 2×L4.GPU (2×H200, real NIXL P/D, Qwen3-0.6B, run by @algebra-MCX, 3 runs): base timed out at 60 s on the 126-token request every run (waiting=1, running=0, KV usage 0%); PR completed 126/2 and 127/1 in every run; the two-request cross-block control (#57104 case) completed on both. The test additions in the last commit are hers.
Minimal reproducer (repo root, no fork needed):
repro.py
Tradeoff: for sliding-window groups
get_num_blocks_to_allocatemay cap the real reservation belowcdiv(num_tokens, block_size), so the per-request margin can still over-reserve by one block there; same conservative direction as before, never under-reserves.Not duplicating open work: #61845 fixes a different livelock (
_mamba_block_aligned_splitvs_reserve_prefill_lookahead, Mamba align + MTP L≥2); #62979 is the HiSparse host-page analogue and does not change the GPU-pool margin.AI assistance: developed with Devin (fuzzing, root-cause tracing, drafting); I reviewed every line, ran the tests above, and am responsible for the change.
Pull Request Checklist
I used vLLM's
/pr-checklistskill. (Mandatory for agents, optional for humans).AI assistance was used during the creation of this PR.
Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.
Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.
Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.
Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.