[Bugfix][MRV2] Require all requests to be decoding for uniform-decode dispatch - #51865
Merged
Merged
Conversation
… dispatch Model runner v2 classified a batch as a uniform decode batch from its shape alone -- num_tokens == num_reqs * max_query_len, `get_uniform_token_count` in vllm/v1/worker/gpu/cudagraph_utils.py -- and nothing about whether the requests were decoding. A prompt chunk of exactly `1 + num_speculative_tokens` tokens has that shape, so a prefill was dispatched with `cg_mode=FULL` and replayed the captured spec-verify decode graph over prompt tokens. Without spec decoding the same collision exists between a 1-token prompt chunk and the q=1 decode graph. Observed on this base with Qwen/Qwen3.5-0.8B (GDN hybrid) + MTP, K=7, FULL_AND_PIECEWISE, VLLM_USE_V2_MODEL_RUNNER=1: an 8-token prompt dispatches `reqs=1 toks=8 uniform_token_count=8 -> FULL` on its prefill step and generates deterministic garbage, while the same request at 4, 5, 6, 7 or 9 prompt tokens dispatches PIECEWISE and is coherent. After this change the 8-token prefill dispatches PIECEWISE and answers the prompt; the control outputs are unchanged byte for byte, and genuine decode steps still replay their FULL graphs (uniform_token_count=8 for the verify step, 1 for the drafter steps). Mechanism, traced at the metadata level on this base: the mis-dispatched batch still builds prefill metadata, because `MambaHybridModelState.prepare_attn` passes is_prefilling and a -1 draft count for the chunk, so `GDNAttentionMetadataBuilder.build` reports num_prefills > 0. Both persistent-buffer refresh blocks in that builder are gated on num_prefills == 0, so neither runs, and the replayed graph reads the state indices its buffers held when the last uniform decode batch was built -- NULL(0) from the capture-time dummy block tables if none has run yet. `state_idx <= 0` makes the fused GDN kernels skip the recurrent state read and the store (third_party/flash_linear_attention/ops/fused_sigmoid_gating.py), so the prompt's conv/SSM state is never created and every later step computes from a zero state. The gate is per batch, not per request, so genuine decodes sharing a batch with such a chunk lose their refresh too: building metadata for 6 spec decodes plus 2 colliding chunks leaves all 8 rows of spec_state_indices_tensor at their capture-time values while the real block table sits elsewhere. Dense models survive the mis-dispatch because spec-verify attention over a fresh sequence is the same computation as prefill attention; only recurrent-state models lose state. Issue vllm-project#49918 REPORTS this as 0/6 correct control requests, reproduced independently of proposer (MTP and ngram), quantization and batch size, with PIECEWISE or enforce-eager as workarounds. The classification rule existed twice and drifted, which is why fixing it in one runner leaves the other broken. `get_uniform_decode_token_count` in vllm/v1/worker/utils.py is now the single source of truth: every request has the same query length AND no request has prompt tokens left to compute. utils.py is already the module that both gpu_model_runner.py and the gpu/ package import from, so neither runner has to depend on the other. Both v2 call sites use it: the target model in `execute_model`, and the draft prefill graph in `AutoRegressiveSpeculator.propose`, whose comment already claimed "when all requests are decoding" while testing only the shape -- that dispatch also replayed FULL for the 8-token prompt above. Capture paths keep the shape-only test. Dummy batches (DP padding, memory profiling, warmup) are uniform by construction and their synthetic request ids have no state to consult, matching `InputBatch.make_dummy`, which marks every dummy request as not prefilling; the v1 runner expresses the same exemption as `force_uniform_decode`. Startup on the run above exercised the profile run, the warmup steps and prefill/decode graph capture with the stricter predicate in place, and the captured decode graphs are still selected at runtime. Deliberately not done: the v1 site, `GPUModelRunner._is_uniform_decode`, is left untouched. PR vllm-project#47123 has been open since 2026-06-30 to fix exactly that call site and this change must not supersede it. The helper is shaped so that `_is_uniform_decode` can collapse onto it by passing `input_batch.num_computed_tokens_cpu` and `input_batch.num_prompt_tokens` and comparing the result with uniform_decode_query_len. vllm/v1/worker/gpu/warmup.py also still pads its warmup prompts to `decode_query_len + 1` to dodge the misclassification; that workaround is now redundant, but it is harmless and it keeps the warmup shapes diverse. The added cost is a per-step `np.fromiter` over the scheduled request ids, one dict lookup each, then a gather of two int32 numpy views. No CPU-GPU sync. It duplicates an index vector `prepare_inputs` rebuilds shortly after; folding the two together is a worthwhile follow-up. Tests: with the decoding clause removed from the predicate (i.e. the pre-fix rule), test_prompt_chunks_shaped_like_spec_decode_miss_the_full_graph and test_prompt_chunk_of_decode_query_len_is_not_uniform_decode fail while the uniform-decode and dummy-run cases still pass. AI assistance was used for this change. Co-authored-by: Janelle Cai <janelle.cai@modal.com> Signed-off-by: Rahul Chalamala <22563365+rchalamala@users.noreply.github.com>
_get_uniform_token_count built a per-request index mapping with a Python dict-lookup loop before calling get_uniform_decode_token_count, which runs the O(1) shape test first and returns None for any mixed prefill/decode batch. Most batches are mixed, so the gather was built and discarded one line later on the common path, where the code it replaced was a single arithmetic comparison. Runs the same shape test in the caller and returns early, so the gather is built only for batches that can still turn out to be uniform decode. The predicate is the one the callee already applies, so classification is unchanged. Co-authored-by: Janelle Cai <janelle.cai@modal.com> Signed-off-by: Rahul Chalamala <22563365+rchalamala@users.noreply.github.com>
The multi-module MTP speculator picks its cudagraph with `get_uniform_token_count`, which tests batch shape alone. A prompt chunk of `1 + num_speculative_tokens` tokens has a spec-decode step's shape, so that test classifies it as a decode batch and replays a decode-captured graph over prompt tokens, whose kernels read per-request state a prefill metadata build does not refresh. Route it through `get_uniform_decode_token_count`, which additionally requires that no request is still prefilling. This speculator is a copy of the autoregressive one, comment included, and arrived after the other call sites were already fixed, so the guard added here is written over every speculator in the package rather than over this one. It fails on `multi_module_mtp/speculator.py:201` without the change and passes with it, and will fail again on the next speculator added by copy. Co-authored-by: Janelle Cai <janelle.cai@modal.com> Signed-off-by: Rahul Chalamala <22563365+rchalamala@users.noreply.github.com>
Keep the main-branch imports required by the autoregressive speculator while removing the replaced shape-only helper. Co-authored-by: Janelle Cai <janelle.cai@modal.com> Signed-off-by: Rahul Chalamala <22563365+rchalamala@users.noreply.github.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
njhill
changed the title
Member
Author
|
/ci run |
|
✅ Triggered Buildkite CI #83430 for commit |
| return num_reqs > 0 and num_tokens == max_query_len * num_reqs | ||
|
|
||
|
|
||
| def get_uniform_decode_token_count( |
Collaborator
There was a problem hiding this comment.
nit: i find having get_uniform_token_count and get_uniform_decode_token_count confusing, why can't we deprecate get_uniform_token_count and use get_uniform_decode_token_count exclusively?
4 tasks
This was referenced Aug 12, 2026
zyp2014
pushed a commit
to zyp2014/vllm
that referenced
this pull request
Aug 21, 2026
… dispatch (vllm-project#51865) Signed-off-by: Rahul Chalamala <22563365+rchalamala@users.noreply.github.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Rahul Chalamala <22563365+rchalamala@users.noreply.github.com> Co-authored-by: Janelle Cai <janelle.cai@modal.com>
This was referenced Aug 21, 2026
1 task
garrygale
added a commit
to garrygale/vllm
that referenced
this pull request
Sep 5, 2026
A chunked-prefill chunk of exactly decode_query_len tokens has the same shape as a uniform spec-decode batch. Shape-only classification then selects a FULL uniform-decode graph for a batch that still contains a prefill, corrupting GDN/Mamba state and causing intermittent acceptance drops whenever a request finishes and a new prefill tail is admitted. Require that every scheduled request is past its prefill before uniform decode dispatch (vLLM vllm-project#51865).
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
4 participants
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.
This is a copy of #50532 from @rchalamala with some additional rework to avoid redundant computation.
Part of the
prepare_inputslogic is split into a separategather_batch_req_statemethod which runs before the dp token count / cuda-graph mode synchronization.