[zixi-qi]

Overview

Continues #48244 by @njhill (co-author). In steps where no scheduled request can use the
drafts proposed after the step, the speculator now proposes with K=0 instead of
drafting. Such requests are non-final prefill chunks and requests sure to reach
max_tokens. This covers every request on a P/D prefill instance, which runs with
max_tokens=1.

Related: #39266 (draft) skips EAGLE drafting on MRV2 for P/D prefill instances and
non-final chunks; this PR decides in the scheduler and also covers requests about to
finish, DFlash and DP. #53426 makes the K=0 draft sync forward optional; this PR keeps
that forward.

Claims

• On Model Runner V2, when no scheduled request on any DP rank can use the drafts, the
  speculator runs with num_speculative_tokens=0:
  • EAGLE/EAGLE3/MTP still run the draft prefill pass, which keeps the drafter's KV cache
    in sync, but skip draft sampling and the draft decode steps.
  • DFlash stores the context K/V and skips the query pass. DSpark, DFlash2 and
    LiLiCorr share this code.
• The drafter's prefill pass still runs, so the gain is modest. Prefill-only throughput
  (Qwen3-8B, 64 x 16,384-token prompts, 1 output token, one GB300 GPU) improves by
  +1.9% with EAGLE3 and +7.3% with DFlash at max_num_batched_tokens=512, and by
  +1.4% and +2.7% at 2048. On a 1P1D prefill instance every step proposes with K=0,
  and prefill throughput improves by +2.1% (EAGLE3) and +3.4% (DFlash) at 512, and by
  +0.6% and +1.9% at 2048.
• Outputs and acceptance do not change. With batch-invariant kernels, outputs and every
  spec-decode counter are identical to main.

Validation

Qwen3-8B with AngelSlim/Qwen3-8B_eagle3 (K=3) and z-lab/Qwen3-8B-DFlash-b16 (K=16),
--attention-backend FLASH_ATTN, GB300. Baseline is main at 532a7d19d2, the base of
this PR; every run below used the rebased commit.

An earlier version of this description reported +9.0% (EAGLE3) and +3.4% (DFlash) at 512.
Those runs were CPU-starved (the benchmark processes were confined to one CPU core),
which inflated host-side costs. The numbers below replace them.

Prefill-only throughput (total tokens/s; mean of 3 reps on one engine after a warm-up
pass). The workload matches vllm bench throughput --dataset-name random --random-input-len 16384 --random-output-len 1 --num-prompts 64 --max-model-len 20480.
main and this PR ran at the same time on two GPUs of one node, then swapped GPUs. The
table shows the GPU whose reps stayed within 1%; on the other GPU, reps of both builds
varied by up to 30%. With the PR, every proposal in these runs used K=0.

Method	max_num_batched_tokens	main	this PR	Delta	Step time saved
EAGLE3	2048	72,100	73,144	+1.4%	0.4 ms
EAGLE3	512	33,834	34,473	+1.9%	0.3 ms
DFlash	2048	72,569	74,508	+2.7%	0.7 ms
DFlash	512	32,367	34,730	+7.3%	1.1 ms

The skipped work is draft sampling and the draft decode steps (EAGLE3), or the query pass
(DFlash). Host time spent in propose() at 512 drops from 1.41 to 0.98 ms per step for
EAGLE3 and from 1.05 to 0.92 ms for DFlash; the rest of the saving is GPU time. The
drafter's prefill pass still runs: it writes the draft KV cache that later steps, and on
a P/D decode instance the transferred KV, depend on.

1P1D prefill throughput. P and D each on one GPU with NixlConnector, the same
drafter on both, P at the given max_num_batched_tokens. Through the toy proxy, 3
phases of 128 random 8,192-token prompts at concurrency 32 (max_tokens=8). P's waiting
queue was non-empty in at least 89% of samples. main and this PR ran side by side on the
node's two GPU pairs, then swapped; tokens/s and deltas are means over both placements.

Method	P max_num_batched_tokens	main	this PR	Delta	P proposals with K=0 (PR)
EAGLE3	2048	56,809	57,213	+0.6%	3,072 / 3,072
EAGLE3	512	28,447	29,043	+2.1%	12,288 / 12,288
DFlash	2048	56,950	57,883	+1.9%	3,072 / 3,072
DFlash	512	28,527	29,524	+3.4%	12,288 / 12,288

Equivalence (offline, greedy). There are 80 requests: 24 long prompts chunked at 512
tokens with max_tokens 1, 2 or 64, run alone and then mixed with 32 short decode
requests.

Method	Mode	Outputs equal	Spec-decode counters	Proposals with K=0 (PR)
EAGLE3	VLLM_BATCH_INVARIANT=1	80/80	identical (3,349 drafts, 1,775 accepted)	19
DFlash	VLLM_BATCH_INVARIANT=1	80/80	identical (2,244 drafts, 2,953 accepted)	18

Without batch invariance, two runs of main already differ on 23/80 EAGLE3 outputs, and
main vs this PR differ on 31/80. DFlash is deterministic across runs of main, but main
vs this PR differ on 31/80: in a K=0 step DFlash stores all context K/V in one eager
pass instead of partly inside the FULL graph, so rounding can differ, and acceptance then
changes batch shapes.

Existing e2e tests give the same results on both:

Test	main	this PR
tests/v1/e2e/spec_decode/eagle/test_eagle_correctness.py -k "medium and qwen3_eagle3 and not transformers"	GSM8K 0.886, pass	GSM8K 0.886, pass
tests/v1/e2e/spec_decode/acceptance_rates/dflash/test_dflash.py (MRV2 reference acceptance)	MT-Bench 4.07, HumanEval 6.42, GSM8K 6.42	same
tests/v1/e2e/spec_decode/acceptance_rates/dflash/test_dflash.py::test_dflash_correctness[qwen3-mrv2]	GSM8K 0.884, AL 4.12	same

Serving equivalence. 1P1D with NixlConnector used EAGLE3 K=3 on both sides and CUDA
graphs on. DP=2 used EAGLE3 K=3. Requests were sent one at a time.

Setup	Outputs equal	Decode-side acceptance length
1P1D	40/40	identical (1.521, 1.598)
DP=2	40/40	identical (1.539, 1.597)

There were no hangs with DP ranks voting to skip.

Unit tests. tests/v1/core/test_scheduler.py, tests/v1/core/test_async_scheduler.py,
tests/v1/worker/test_gpu_dp_utils.py and
tests/v1/worker/test_gpu_target_dependent_ar_speculator.py pass (281 tests), except two
pre-existing PP scheduler tests that fail the same way on main.

Details

• Scheduler. It sets SchedulerOutput.skip_drafts when no scheduled request can use
  the step's drafts, on the V2 runner and for drafters that reserve lookahead slots.
  _drafts_unused counts only the output tokens a request is sure to have, assuming
  every draft is rejected, so it never skips drafts a later step could verify.
• DP. The flag is ANDed across ranks in the batch-descriptor all-reduce that already
  runs, giving DPSyncState.skip_drafts. The runner then proposes with K=0. A rank
  running a dummy batch votes no, so an idle DP rank disables the skip for that step.
• Changes from [Core] Skip redundant draft token alloc + sampling #48244:
  • The per-request lookahead-slot skip is dropped. For a non-final chunk, the lookahead
    slots are the next chunk's first positions, which are allocated one step later
    anyway. For a finishing request it saves at most about one block for one step.
    Keeping it safely means masking every draft KV write, including inside captured
    FULL draft-decode graphs; the fused decode path is the default on FlashAttention,
    FlashInfer and Triton, and its captured graphs compute their own slot mappings.
  • Skipping reuses the existing num_speculative_tokens == 0 path, which already skips
    the AR draft decode steps, instead of a separate flag in each speculator.
    skip_sample and the DFlash early return also apply to K=0 from dynamic speculative
    decoding.
  • The scheduler reports one flag instead of a set of request ids.
• Not covered. Multi-module MTP and the draft-model speculator ignore K, so they
  still draft in these steps and the drafts go unused, as before. N-gram reserves no
  lookahead slots, so the flag is never set for it. MTP, DSpark, DFlash2 and LiLiCorr
  share the changed code but were not run; all validation used Qwen3-8B. MRV1 is
  unchanged.

Throughput script (one engine, warm-up then timed reps)

llm = LLM(model="Qwen/Qwen3-8B", speculative_config=spec, max_model_len=20480,
          max_num_batched_tokens=mbt, enable_prefix_caching=False, seed=42,
          attention_backend="FLASH_ATTN")
rng = np.random.default_rng(42)
prompts = lambda n: [TokensPrompt(prompt_token_ids=r.tolist())
                     for r in rng.integers(1000, 100000, size=(n, 16384))]
params = SamplingParams(temperature=0, max_tokens=1, ignore_eos=True)
llm.generate(prompts(8), params)
for _ in range(3):
    batch = prompts(64)
    t = time.perf_counter(); llm.generate(batch, params)
    print(64 * 16385 / (time.perf_counter() - t))

---

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.

AI assistance (Claude Code) was used for the rebase, the rework and the validation. I
reviewed and can defend every change.

🤖 Generated with Claude Code

[claude]

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.

[TheEpicDolphin]

Very clean, LGTM!

[TheEpicDolphin]

/ci run

[github-actions]

✅ Triggered Buildkite CI #96559 for commit 1366ae26c0c6.

ℹ️ This PR is 85 commits behind upstream main. Merged upstream main into this branch as 1366ae26c0c6 before starting CI. Pull the branch before pushing more changes.

[zixi-qi]

/ci run

[github-actions]

✅ Triggered Buildkite CI #96571 for commit 76387059c411.

[zixi-qi]

/ci retry

[github-actions]

✅ Queued 2 failed job(s) for retry in Buildkite CI #96571.

[TheEpicDolphin]

The failing tests seem related, but should be an easy fix

[zixi-qi]

/ci run

[github-actions]

✅ Triggered Buildkite CI #96715 for commit 7937584f2cd0.

ℹ️ This PR is 51 commits behind upstream main. Merged upstream main into this branch as 7937584f2cd0 before starting CI. Pull the branch before pushing more changes.

[zixi-qi]

/ci retry

[github-actions]

✅ Queued 2 failed job(s) for retry in Buildkite CI #96715.