[MRV2][Spec Decode] Adaptive Speculative Decoding - Initial Support

[benchislett]

Purpose

This PR implements initial support for Adaptive Speculative Decoding using DSpark's confidence head: we enable fully variable-length per request speculative decoding with FULL CUDA Graph support. This PR does not implement online dynamic scheduling of the verification budget; instead, we rely on user-provided num_speculative_tokens_per_batch_size to determine what size pool of draft tokens to verify at each batch size.

This PR fleshes out a number of fundamental aspects of variable-length speculation, such as:

• Attention metadata flags to declare compatibility with variable length decoding
• FULL CUDA Graph compatibility
• Structured outputs support
• Logprobs support
• Proper integration with statistics capture
• Mixed prefill/decode batch compatibility

Testing

Since only FLASH_ATTENTION is currently functional, I run with Qwen3-8B-FP8 and the corresponding DeepSeek-trained DSpark. Since verification is so cheap on an 8B model, no significant speedup is measurable; that is not the objective of this PR. However, I do observe an increase in acceptance rate under a fixed verification budget:

Acceptance Rate Study: SPEED-Bench, All Categories, Concurrency 128, Shuffled

SPEED-Bench, All Categories, Concurrency 128, Shuffled

Category	Baseline (K=5)	Baseline (K=7)	Adaptive (Max K=7, Effective K=5 Per Request)
coding	4.6267	5.5640	5.0972
humanities	3.2844	3.6392	3.5036
math	4.2535	5.0100	4.5879
multilingual	3.0081	3.4400	3.3648
qa	2.9236	3.2078	3.0896
rag	3.7169	4.2810	4.0475
reasoning	3.9697	4.5843	4.2913
roleplay	2.3373	2.4564	2.4243
stem	3.5963	4.1574	3.9192
summarization	2.9987	3.3062	3.1559
writing	3.1799	3.4470	3.2988
Overall Average	3.4450	3.9176	3.7073

Measured using NVIDIA SPEED-Bench harness, thinking off, vLLM config:

engine_args:
  speculative_config:
    adaptive_verification: true
    num_speculative_tokens_per_batch_size:
      - [1, 128, 5]
    attention_backend: FLASH_ATTN
  gpu_memory_utilization: 0.75
  attention_backend: FLASH_ATTN
  max_num_batched_tokens: 32768
sampling_kwargs:
  temperature: 0.7
  top_p: 0.8
  top_k: 20
dataset_kwargs:
  shuffle_seed: 42

I also implement various unit tests, including E2E coverage of acceptance length, GSM8k correctness, and an assert that we're observing almost-exactly N*K draft tokens at each step.

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[mergify]

Documentation preview: https://vllm--48692.org.readthedocs.build/en/48692/

[benchislett]

Was getting OOM here due to probabilistic drafting taking up lots of memory for logits, not accounted for by profiler. Not sure if this is safe for CI or not

[benchislett]

Bit of a confusing override here; this is so the scheduler treats adaptive verification as a separate mode (and the model runner decides how many decode tokens to run)

[benchislett]

All of this diff is just shuffling things around to fix a bug in the specdec metrics.
Previously we were skipping the metrics observation in cases where request.async_tokens_to_discard == 0, and clearing the entry when num_draft_tokens == 0.

This overcounts the num_drafts since we're only skipping the collection of num_draft_tokens and num_accepted_tokens, throwing off the scores.

[benchislett]

Needed for the test where we assert that num_drafts == K * num_drafted_tokens in the adaptive SD case.

[benchislett]

Much of this here is rearranging, since expanded_idx_mapping depends on cu_num_logits, which we construct explicitly in the adaptive verification pathway

[benchislett]

We update structured outputs so it can apply the bitmask to the varlen batch. mask_stride effectively makes the bitmask tensor into one that includes padding: and then we read the varlen indices in the kernel, using position_is_active to skip any slots for tokens we did not choose to verify

[mergify]

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

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

[cyLi-Tiger]

One thought on the CUDA graph design: SGLang's DSpark integration handles per-request variable verify lengths by packing the ragged batch into a compact varlen buffer and keying the graph purely on the total token count. This seems to allow much finer-grained capture tiers and noticeably less padding waste than the approach here, where decode graphs are rounded up to multiples of the scheduled query length. Might be worth considering whether a similar packed-layout design could reduce compute overhead in vLLM's adaptive verify path.

For example, a verify batch [8, 3, 1, 5] (4 reqs, 17 real tokens) with scheduled qlen 8 needs the 32-token FULL graph here (tiers are multiples of 8, and the 24-tier only has 24/8=3 request slots) — 15 padded rows. SGLang's ragged approach packs tokens compactly and keys the graph on total count only, fitting an 18-token tier — 1 padded row, still FULL graph.

cc @benchislett

[benchislett]

Closing this PR, we will move forward and try to merge #47808.

[benchislett]

Also, @cyLi-Tiger I believe that #47808 does it as you suggest, if I am understanding you properly