Skip to content

[SpecDecode] Preserve reject-only residual verifier padding - #54102

Open
GirasoleY wants to merge 4 commits into
vllm-project:mainfrom
GirasoleY:specdecode-residual-verifier-padding
Open

GirasoleY wants to merge 4 commits into
vllm-project:mainfrom
GirasoleY:specdecode-residual-verifier-padding

Conversation

@GirasoleY

@GirasoleY GirasoleY commented Aug 27, 2026

Copy link
Copy Markdown
Contributor

Summary

  • Keep the real one-token prompt tail through Mamba scheduling alignment, then restore the synthetic speculative width before allocation.
  • Track reject-only synthetic drafts separately from GPU-safe placeholder token IDs and restore their -1 sentinel only for rejection sampling.
  • Skip uniform residual padding for adaptive DSpark, which already uses a variable verification budget and variable-length decode graphs.
  • Select target-verifier FULL graphs only when every request carries the complete verifier draft layout, and describe dummy captures with the same topology.

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_tokens so 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

  1. Align residual verifier padding after Mamba splits.
  2. Preserve synthetic draft rejection sentinels and exclude them from acceptance accounting.
  3. Skip uniform padding for adaptive verification.
  4. Select FULL graphs by verifier topology.

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 check on all changed Python files.
  • uvx ruff format --check on all changed Python files.
  • uv run --no-project --python 3.12 python -m py_compile on 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.

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>
@mergify mergify Bot added speculative-decoding mrv2 Model Runner V2 specific scheduler labels Aug 27, 2026
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
GirasoleY marked this pull request as ready for review August 28, 2026 22:16

@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.

@zixi-qi zixi-qi added the ready ONLY add when PR is ready to merge/full CI is needed label Aug 29, 2026
@zixi-qi

zixi-qi commented Aug 29, 2026

Copy link
Copy Markdown
Collaborator

/ci run

@github-actions

Copy link
Copy Markdown

✅ Triggered Buildkite CI #86174 for commit fad62d22ed2a.

@zixi-qi zixi-qi left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

The fixes LGTM! cc @njhill

@mergify

mergify Bot commented Aug 31, 2026

Copy link
Copy Markdown
Contributor

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @GirasoleY.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

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

Labels

mrv2 Model Runner V2 specific needs-rebase ready ONLY add when PR is ready to merge/full CI is needed scheduler speculative-decoding

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants