[WIP][Spec Decode] DSpark capacity reallocation without sampler padding

[LucasWilkinson]

This PR implements DSpark confidence-scheduled verification: a capacity manager that prunes low-confidence draft tokens from target verification, with two enforcement modes — mask (pruned rows kept as padding) and varlen (verifier batch physically compacted). The varlen path keeps FULL CUDA graphs for both the target verify step and the DSpark draft step.

Based on #47469 (SM100 DSA varlen decode); comparison point is #47513 (padded/full-CG approach). Replaces #47694 (opened from the wrong fork).

How it works

• Capacities come from the paper's Algorithm 1: greedy global admission by prefix-survival score, sum(capacities) == spent budget (the earlier threshold-recount formulation blew through the budget on tied/saturated scores).
• The stopping rule is the paper's theta = tau * SPS(B) argmax over a profiled step-rate curve, supplied via dspark_sps_curve (measure it with benchmarks/profile_dspark_sps_curve.py); dspark_budget_frac still acts as a hard admission cap.
• The released confidence head ships uncalibrated (saturated 0/1 sigmoids whose exact zeros would silently remove tokens from the candidate set), so dspark_online_sts fits the paper's Sequential Temperature Scaling online: per-(position, verdict) empirical conditional acceptance from the rejection sampler's own outcomes. dspark_confidence_temperature is a simpler static alternative.

Performance

DeepSeek-V4-Flash-NVFP4 + DSpark draft, TP4 B200, k=7, speed_bench (all categories), 256 prompts x 1024 output tokens, temperature 1.0, --ignore-eos (equal work per cell). Cells: output tok/s / mean ITL / acceptance length (delta vs no capacity):

c	No capacity	Mask + SPS	Varlen + SPS	Mask + SPS + STS	Varlen + SPS + STS
16	1,397 / 14.69 / 1.324	1,136 / 15.63 / 1.132 (-18.7%)	1,109 / 15.79 / 1.116 (-20.6%)	1,172 / 14.78 / 1.099 (-16.1%)	1,303 / 13.79 / 1.151 (-6.7%)
32	2,202 / 18.35 / 1.321	1,760 / 20.09 / 1.135 (-20.1%)	1,871 / 18.48 / 1.113 (-15.0%)	1,729 / 18.91 / 1.042 (-21.5%)	2,105 / 15.47 / 1.051 (-4.4%)
64	3,388 / 23.59 / 1.327	2,706 / 26.13 / 1.139 (-20.1%)	2,936 / 23.49 / 1.115 (-13.4%)	2,664 / 23.78 / 1.000 (-21.4%)	3,503 / 20.07 / 1.137 (+3.4%)
128	4,832 / 32.39 / 1.334	3,977 / 35.24 / 1.144 (-17.7%)	4,606 / 29.57 / 1.116 (-4.7%)	3,783 / 33.41 / 1.000 (-21.7%)	6,050 / 20.69 / 1.000 (+25.2%)
256	6,211 / 49.01 / 1.321	5,204 / 52.66 / 1.128 (-16.2%)	6,893 / 39.09 / 1.117 (+11.0%)	4,807 / 52.56 / 1.000 (-22.6%)	9,089 / 27.52 / 1.000 (+46.3%)

• varlen turns the same pruning decisions into a physically smaller batch; mask only loses acceptance (identical calibrated capacities, opposite outcomes).
• With online calibration, varlen wins from c=64 up (+46% at c=256, where the scheduler correctly verifies ~nothing at saturation); below the contention knee the verify-only SPS curve still over-prunes slightly (see TODO).

GSM8K 5-shot: varlen 0.933-0.939, mask 0.950, invalid <= 0.002 (eval config included as tests/evals/gsm8k/configs/DeepSeek-V4-Flash-DSpark-varlen-TP4.yaml).

Exact reproduction commands

TARGET=nvidia/DeepSeek-V4-Flash-NVFP4
DRAFT=deepseek-ai/DeepSeek-V4-Flash-DSpark

# 1. profile the SPS curve (one breakpoint per cudagraph capture size)
python benchmarks/profile_dspark_sps_curve.py "$TARGET" \
  --speculative-config '{"method":"dspark","model":"'"$DRAFT"'","attention_backend":"FLASH_ATTN","num_speculative_tokens":7,"draft_sample_method":"probabilistic","dspark_confidence_threshold":0.0,"dspark_capacity_verification_mode":"varlen"}' \
  --engine-args '{"tokenizer":"'"$TARGET"'","tokenizer_mode":"deepseek_v4","trust_remote_code":true,"dtype":"bfloat16","max_model_len":8192,"tensor_parallel_size":4,"enable_expert_parallel":true,"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":2048},"kernel_config":{"moe_backend":"flashinfer_trtllm"}}' \
  --output sps_curve.json
CURVE=$(python -c 'import json; print(json.dumps(json.load(open("sps_curve.json"))["dspark_sps_curve"]))')

# 2. server (same engine flags for every config; SPEC varies per column)
vllm serve "$TARGET" --port 8011 \
  --tokenizer "$TARGET" --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": 2048}' \
  --moe-backend flashinfer_trtllm --speculative-config "$SPEC"

# no capacity:
SPEC='{"method":"dspark","model":"'"$DRAFT"'","attention_backend":"FLASH_ATTN","num_speculative_tokens":7,"draft_sample_method":"probabilistic"}'
# mask/varlen + SPS: append to the object
#   ,"dspark_confidence_threshold":0.0,"dspark_capacity_verification_mode":"varlen","dspark_sps_curve":'"$CURVE"'
# + online STS: additionally append
#   ,"dspark_online_sts":true

# 3. bench (C in 16 32 64 128 256; STS configs get one warmup pass first:
#    same command with --num-prompts 128 --max-concurrency 32 --speed-bench-output-len 256)
vllm bench serve --model "$TARGET" --tokenizer "$TARGET" --tokenizer-mode deepseek_v4 \
  --base-url http://127.0.0.1:8011 \
  --dataset-name speed_bench --dataset-path "$SPEED_BENCH" \
  --num-prompts 256 --save-result --disable-shuffle --max-concurrency $C \
  --temperature 1.0 --speed-bench-output-len 1024 --ignore-eos \
  --backend openai-chat --endpoint /v1/chat/completions --skip-chat-template

# 4. accuracy
python tests/evals/gsm8k/gsm8k_eval.py --num-shots 5 --port 8011

TODO

• Regenerate the SPS curve with a measured per-step overhead (--overhead-ms): the verify-only curve over-prunes below the contention knee (~-7% vs no capacity at c=16)
• Dynamically disable drafting under saturation (at c=256 the scheduler verifies ~nothing but still pays the full draft block — the paper's noted limitation)
• Auto-profile the SPS curve at engine init
• Human line-review before undrafting

AI assistance

Built with AI assistance (Codex + Claude Code); all benchmarks above were executed as part of this work. Draft until human-reviewed.

[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