Skip to content

[Bugfix][Core] Reserve only the blocks a spec decode step adds when admitting async KV loads - #63136

Open
krishchaudh wants to merge 3 commits into
vllm-project:mainfrom
krishchaudh:fix/async-load-spec-step-reserve
Open

krishchaudh wants to merge 3 commits into
vllm-project:mainfrom
krishchaudh:fix/async-load-spec-step-reserve

Conversation

@krishchaudh

@krishchaudh krishchaudh commented Oct 22, 2026 •

Copy link
Copy Markdown

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 waiting forever with no error. Follow-up to #57104 / #59880, found while testing the scheduler (not a production incident).

Example

  • block_size 32, 4 usable blocks, num_speculative_tokens=1, max_model_len 128
  • 1 request A: 126-token prompt, full remote hit (needs 4 blocks); nothing else running

Step 1 — A arrives

allocate_slots(A, num_new_tokens=126, ..., reserved_blocks=0 inflight + 1 spec margin)

  • A needs cdiv(126, 32) = 4 blocks

  • free = 4, minus 1 reserved = 3 < 4 → allocate_slots returns None → break

  • result:

waiting = [A]
A.status = WAITING   (never reaches WAITING_FOR_REMOTE_KVS)

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_slots caps 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 any num_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_blocks uses the same max_model_len cap.

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

  • unit: tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py::test_async_load_admits_request_whose_spec_step_fits_its_blocks (num_spec_tokens 1/3 × room-in-last-block / capped-by-max_model_len, prompt max_model_len - 1), real Scheduler + MockKVConnector
  • unit: test_async_load_near_max_model_len_does_not_block_another_fitting_load (spec 1/3): a near-max_model_len in-flight load must not reserve slots past the cap and block a second fitting load (covers the _request_remaining_blocks change)
  • existing [BugFix][KV Connector] Fix Deadlock with KVConnector + MTP under KV Pressure #57104/[Bugfix][KV Connector] Fix speculative block reservations across KV cache groups #59880 tests unchanged: test_async_load_reserves_blocks_for_promotion_margin + deadlock tests
  • e2e, no mocks: tests/v1/kv_connector/nixl_integration/test_spec_decode_admission.py, second phase of run_edge_case_test.sh (NixlConnector PD edge case CI job). Real prefill + decode vllm serve Qwen/Qwen3-0.6B over NixlConnector; 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 by max_model_len, 1 output token) prompts remote-prefilled, then sent to decode with the returned kv_transfer_params (ignore_eos); asserts completion within 60 s. Both phase-1 servers also use --block-size 32 so all P/D block sizes match.

Test Result

# unit, main (876d5da): 6 failed
E   assert <RequestStatus.WAITING: 1> == <RequestStatus.WAITING_FOR_REMOTE_KVS: 3>
# unit, this PR:
tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py  22 passed
tests/v1/core/test_scheduler.py tests/v1/core/test_async_scheduler.py  252 passed

# e2e, main (4e59165): decode metrics num_requests_waiting=1 waiting_capacity=1
#   num_requests_running=0 kv_cache_usage_perc=0.0, unchanged after 150 s
FAILED test_spec_decode_admission.py::test_async_load_admits_request_filling_free_blocks
  AssertionError: decode did not admit the 126-token remote-prefilled request within 60s
1 failed in 61.49s
# e2e, this PR:
1 passed in 2.22s   (completion_tokens=2, total_tokens=128)

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
from unittest.mock import patch
from tests.v1.kv_connector.unit.utils import create_vllm_config, create_scheduler, create_request

cfg = create_vllm_config(block_size=16, num_speculative_tokens=1)
sched = create_scheduler(cfg, num_blocks=4)  # 3 usable blocks
req = create_request(request_id=1, num_tokens=40, do_remote_prefill=True, max_tokens=1)
sched.add_request(req)
with patch.object(sched.connector, "get_num_new_matched_tokens", return_value=(40, True)):
    for _ in range(3):
        sched.schedule()
print(req.status)  # main: RequestStatus.WAITING forever; fixed: WAITING_FOR_REMOTE_KVS

Tradeoff: for sliding-window groups get_num_blocks_to_allocate may cap the real reservation below cdiv(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_split vs _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-checklist skill. (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.

…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>

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

@github-actions

Copy link
Copy Markdown

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment /ci run for upstream CI or /amd-ci run for AMD CI only whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use the corresponding /ci run, /ci retry, and /ci cancel commands, or their /amd-ci variants. New commits do not start upstream CI automatically.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: 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.

🚀

krishchaudh and others added 2 commits October 25, 2026 03:56
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>

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working kv-connector scheduler

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants