Conversation
Keep the real prompt-tail token count through scheduling alignment, then restore synthetic verification padding before allocation. This prevents Mamba block alignment from clipping uniform verifier batches. Assisted-by: OpenAI Codex Signed-off-by: Summer Yang <girasoleyang@gmail.com>
Track reject-only verifier padding separately from GPU-safe placeholder token IDs. Restore the invalid suffix only in rejection-sampler metadata so synthetic drafts cannot be accepted or counted in acceptance metrics. Carry per-request counts instead of flattened indices so sampler batch sharding preserves the correct request-local layout. Assisted-by: OpenAI Codex Signed-off-by: Summer Yang <girasoleyang@gmail.com>
Adaptive DSpark already selects a variable verification budget and captures variable-length decode graphs. Keep first-decode requests at their real one-token span instead of adding fixed-width synthetic drafts that the adaptive allocator could treat as eligible. Assisted-by: OpenAI Codex Signed-off-by: Summer Yang <girasoleyang@gmail.com>
Require every row in a uniform speculative batch to carry its full draft layout before selecting a target-verifier FULL graph. Describe distributed dummy and idle-rank executions with the same verifier metadata so all ranks select the same graph topology. Assisted-by: OpenAI Codex Signed-off-by: Summer Yang <girasoleyang@gmail.com>
justtestingthingsx
pushed a commit
to meandmyboiclaude/vllm
that referenced
this pull request
Aug 28, 2026
…aph guard + _dummy_run draft layout — scheduler, rejection_sampler, InputBatch and SchedulerOutput halves deliberately skipped Draft PR, partial lift by design. Both pieces land in vllm/v1/worker/gpu/model_runner.py (the executed MRV2 runner); nothing else in the tree is touched. (a) execute_model, inserted after the skip_compiled block and immediately before dispatch_cg_and_sync_dp: when speculative_config is set, compute is_target_verifier_batch and null uniform_tok_count unless every scheduled request carries the full verifier draft layout. Taken verbatim. Our dispatch_cg_and_sync_dp returns a 3-tuple and the existing `batch_desc, dp_sync, _ =` unpack is left exactly as it was — the guard only rewrites the uniform_tok_count argument. (b) _dummy_run, after num_scheduled_tokens is built (:730-732 here): populate scheduled_spec_decode_tokens with [-1] * (n - num_bonus_tokens) under `uniform_decode and self.speculative_config is not None`, so capture and replay describe the same branch topology. The two are one change, not two: without (b) the uniform-decode capture run hits (a) with an empty draft dict, nulls its own uniform_tok_count and the verifier FULL graph is never captured at all. SKIPPED from the PR, as scoped: the scheduler half (scheduler.py num_invalid_spec_tokens plumbing, the enable_adaptive_verification padding opt-out, and the num_new_tokens padding move), the rejection_sampler half (synthetic draft sentinel restore), InputBatch.num_invalid_spec_tokens, the SchedulerOutput comment, the prepare_inputs field pass-through, and all three test files. Neither vendored piece reads num_invalid_spec_tokens, so no half-plumbed field is left behind. Verified against OUR tree rather than the PR's context: - uniform_tok_count comes from gather_batch_req_state (:1547), not from an inline computation as upstream; the guard sits between that and dispatch. - for dummy_run the same helper returns get_uniform_decode_token_count(..., has_prefill=False) (:1120-1124), so the capture path really does reach the guard with a non-None value — hence the coupling above. - decode_query_len = num_speculative_steps + num_new_sampled_tokens_per_step (:425), and _dummy_run's uniform path asserts num_tokens % decode_query_len == 0, so (b) always emits exactly num_speculative_steps drafts; num_speculative_tokens is Field(gt=0) so n - num_bonus_tokens is never 0. - AST scan confirms num_bonus_tokens / draft_tokens / is_target_verifier_batch are each bound once and shadow no existing local in either function; the comprehension's num_tokens is genexp-scoped and does not collide with the outer num_toks. Losslessness: untouched. This changes only which cudagraph descriptor is dispatched; the DFlash K=5 accept test is not on this path. Logic validated out of tree at /tmp/claude-1000/-home-user/9d6988ce-e46e-47b6-b43e-009784106bac/scratchpad/guardtest.py (pure Python, both pieces transcribed, K=5/bonus=1; 8/8 pass): capture keeps uniform_tok_count=6; dropping (b) nulls that same capture; a real verifier batch keeps it; the same-shape short prefill that motivates the PR is now nulled; a mixed draft/no-draft batch is nulled; plain q=1 decode is unaffected because 1 > num_bonus_tokens is false; the non-spec path is inert; and the scheduler's existing [-1]*num_spec reject-only padding still reads as a verifier batch, which is why the skipped scheduler half is not a prerequisite here. NOT VALIDATED HERE: no GPU run. This alters cudagraph dispatch on the serving path and wants a FULL_AND_PIECEWISE boot plus an acceptance leg before the wheel ships. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
GirasoleY
marked this pull request as ready for review
August 28, 2026 22:16
GirasoleY
requested review from
ApostaC,
WoosukKwon,
alexm-redhat,
heheda12345,
ivanium,
njhill,
orozery,
robertgshaw2-redhat,
yewentao256 and
ywang96
as code owners
August 28, 2026 22:16
Collaborator
|
/ci run |
|
✅ Triggered Buildkite CI #86174 for commit |
2 tasks
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
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
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.
Summary
-1sentinel only for rejection sampling.Motivation
#54012 enabled FlashInfer native-CP MLA verification, but the failures exposed here are backend-independent. When a new request with one remaining prompt token joins a running speculative-verification batch, vLLM pads it to
1 + num_speculative_tokensso the target batch remains uniform.Two independent correctness problems occur in that path. Applying Mamba alignment to the padded width can clip the synthetic verifier shape, and the GPU input path replaces negative placeholder IDs with embedding-safe values before rejection sampling, allowing synthetic drafts to be treated as real proposals. Adaptive DSpark should not enter this fixed-width padding path at all because it deliberately reallocates a variable draft budget.
Target-verifier FULL graphs also capture model-wide branch topology, not only tensor dimensions. A same-shaped short prefill must not select a verifier graph when it lacks the complete draft-token layout. Uniform dummy captures must carry equivalent verifier metadata so distributed and idle ranks select the same graph topology.
Commit structure
Relationship to existing work
This is a follow-up to #54012, not an attention-backend change. #53694 avoids a redundant DP synchronization before draft prefill but does not distinguish target-verifier graphs from same-shaped short-prefill execution. I searched open vLLM PRs and issues for
pad_spec_decode, residual verifier padding, invalid draft sentinels, verifier graph topology, and same-shaped short prefills and found no directly overlapping fix. #51508 addresses stale zero-accept recurrent-state rows rather than first-decode residual padding, rejection-sampler sentinel loss, or verifier graph selection.Tests
Passed locally:
uvx ruff checkon all changed Python files.uvx ruff format --checkon all changed Python files.uv run --no-project --python 3.12 python -m py_compileon all changed Python files.git diff --check.Focused pytest and model evaluation are pending because this local worktree does not contain the vLLM runtime/test dependencies. This PR is intentionally opened as a draft. Before marking it ready, I will run the focused scheduler, rejection-sampler, adaptive-verification, and CUDA-graph dispatch tests, then report a mixed-batch Kimi-K3 DCP/EAGLE acceptance and correctness evaluation.
AI assistance disclosure
AI assistance was used to inspect the current upstream data flow, adapt the internal fixes to TP batch sharding and the post-#53694 DP-sync path, write the initial patches and regressions, and run static validation. The submitter will review every changed line and run the required runtime tests and model evaluation before requesting final review.