[Spec Decode] DSpark confidence-scheduled verification

[LucasWilkinson]

Adaptively sizes the DSpark draft-verification budget from per-request confidence instead of always verifying every drafted token. Motivation: fixed-k speculation collapses at high concurrency — once the GPU saturates, verifying 7 drafts per request burns more compute than the accepted tokens return, dropping below non-speculative decoding (see table).

Design

• A Triton kernel ranks draft slots by survival probability (cumprod of per-position confidence) and admits the prefix maximizing estimated accepted tokens per millisecond.
• Step costs come from curves profiled with dummy steps at startup: a step function at/below the cudagraph capture limit (padding is physical), linear interpolation above it.
• Batch-level budget sizing runs on CPU from double-buffered stale confidences (no sync); per-request allocation runs on GPU from live values.
• Per-request confidence is smoothed with an EMA (adaptive_verification_ema_alpha, default 0.8 — measured bias-zero crossing across batch sizes).
• Decode cudagraphs become varlen: captured on the token grid, dispatched via max_query_len, every captured slot non-empty. Requires AttentionCGSupport.ALWAYS, which the DSV4 backends report on SM100.

Results

DeepSeek-V4-Flash-DSpark, TP=4, SM100 (B300), speed_bench 256 prompts, 512 output tokens; throughput measured at 32881ec177; subsequent HEAD commits are audit fixes off the captured-graph hot path (decode max_q_len sourcing, CPU upper-bound restoration, non-adaptive request-ordering revert to main), accuracy re-validated at HEAD. Output tok/s:

Adaptive matches fixed-k within ±3% at c≤64 and preserves the spec-decode win at high concurrency where fixed-k goes underwater (at c=256, fixed 7-token is 33% below no-spec). Acceptance length falls from ~3.9 to 3.5 at c=256 — the budget deliberately trims drafts as verification tokens become expensive. Note: the no-spec arm ran with a co-tenant job on the other half of the node; its numbers are directionally correct but modestly depressed. Adaptive/fixed arms ran on an idle node.

Accuracy (re-measured at HEAD after the audit fixes below): GSM8K 0.945 (c=16) / 0.951 (c=64), 0 invalid (bar 0.84); MTBench 80/80 completed, completions manually checked for repetition/degeneration — none found.

Reproduction commands

Server (all measurements; ablations via the noted config deltas):

vllm serve deepseek-ai/DeepSeek-V4-Flash-DSpark \
  --tokenizer-mode deepseek_v4 --trust-remote-code \
  --dtype bfloat16 --max-model-len 8192 \
  --tensor-parallel-size 4 --enable-expert-parallel --block-size 256 \
  --gpu-memory-utilization 0.5 --kv-cache-dtype fp8 \
  --max-num-batched-tokens 16384 --max-num-seqs 256 \
  --compilation-config '{"max_cudagraph_capture_size":1024}' \
  --speculative-config '{"method":"dspark","model":"deepseek-ai/DeepSeek-V4-Flash-DSpark","attention_backend":"FLASH_ATTN","num_speculative_tokens":7,"draft_sample_method":"probabilistic","enable_adaptive_verification":true}'

• fixed 7-token baseline: "enable_adaptive_verification":false
• no-spec baseline: omit --speculative-config

Benchmark harness — vllm bench serve (the Python implementation; the CLI falls back to it automatically when the optional Rust vllm-rs binary is absent). If your install ships a Rust binary that rejects any flag below, invoke the Python implementation directly — still repo-only:

python -c 'from vllm.benchmarks.serve import add_cli_args, main
from vllm.utils.argparse_utils import FlexibleArgumentParser
p = FlexibleArgumentParser(); add_cli_args(p); main(p.parse_args())' <flags...>

GSM8K (in-tree runner; 1319 questions, 5-shot, temperature 0):

python tests/evals/gsm8k/gsm8k_eval.py --port 8000 \
  --num-questions 1319 --num-shots 5 --max-tokens 256 \
  --temperature 0 --max-concurrency 16   # and 64

MTBench coherence (80 prompts, temperature 1; inspect the saved completions for repetition/degeneration, not just the completion count):

vllm bench serve --backend openai-chat --endpoint /v1/chat/completions \
  --base-url http://127.0.0.1:8000 --model dspark-full \
  --tokenizer deepseek-ai/DeepSeek-V4-Flash-DSpark --tokenizer-mode deepseek_v4 \
  --dataset-name hf --dataset-path philschmid/mt-bench \
  --num-prompts 80 --no-oversample --hf-output-len 256 \
  --max-concurrency 16 --request-rate inf --temperature 1 \
  --save-result --save-detailed

Throughput sweep (per concurrency; one warmup pass at --speed-bench-output-len 256 --num-prompts 64 --max-concurrency 32 first):

for c in 1 16 32 64 128 256; do
  vllm bench serve \
    --backend openai-chat --base-url http://127.0.0.1:8000 \
    --endpoint /v1/chat/completions --model dspark-full \
    --tokenizer deepseek-ai/DeepSeek-V4-Flash-DSpark --tokenizer-mode deepseek_v4 \
    --dataset-name speed_bench --dataset-path <speed-bench-dir> \
    --speed-bench-dataset-subset qualitative \
    --skip-chat-template --disable-shuffle --temperature 1.0 \
    --speed-bench-output-len 512 --num-prompts 256 --max-concurrency $c \
    --save-result --result-filename adaptive_on_c${c}.json
done

--disable-shuffle plus the fixed prompt set gives every arm identical prompts in identical order; output_throughput from the result JSON is the tok/s reported above. The prompt set is an internal speed-bench "qualitative" subset; any fixed prompt set of comparable length distribution reproduces the relative behavior, e.g. --dataset-name random --random-input-len 1024 --random-output-len 512.

Limitations

• Output logprobs are rejected when confidence-based verification is enabled (verification compacts logits after the forward pass).
• Non-SM100 falls back to PIECEWISE decode graphs: correct, but without the varlen-FULL-graph throughput win.

Tests

pytest tests/v1/spec_decode/test_adaptive_verification.py \
       tests/v1/spec_decode/test_dynamic_sd_cug.py tests/v1/cudagraph/ -q

All pass; full pre-commit green. GSM8K eval config added at tests/evals/gsm8k/configs/DeepSeek-V4-Flash-DSpark-confidence-TP4.yaml.

Not a duplicate

Canonical PR for DSpark confidence-based verification; no other open PR touches adaptive draft budgeting.

---

AI assistance was used for this change (see commit trailers).

🤖 Generated with Claude Code

[mergify]

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

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

[xhdidi]

@xhdidi what HW are you on?

H20

[weiguihua2]

Hello, I have a question. If cudagraph is changed to AttentionCGSupport.ALWAYS, does the prefill support the full graph mode in common scenarios? This is because the --compilation-config '{"cudagraph_mode": "FULL_AND_PIECEWISE"}' can be configured.

[xhdidi]

VLLM enables chunked prefill by default. In real-world applications with long context lengths and many batches being a mix of prefill and decode batches, the current cost curve, which only applies to pure decode batches, is ineffective, and the performance gain from enable-adaptive-verification is minimal. Have you considered adding cost data to a piecewise CUDA graph?