[WIP][Spec Decode] DSpark confidence-scheduled verification

[LucasWilkinson]

This PR implements confidence-scheduled speculative verification for
DSpark. The confidence head predicts each
draft token's acceptance probability; the host scheduler prunes verification
work that is not expected to pay for its target-model cost.

confidence_based_verification controls the feature: none and off disable
it, auto uses compact verification when the attention backend supports the
required full CUDA graphs and otherwise falls back to masking, and mask
forces the masked path.

• Masking leaves pruned rows as padding. Supported MoE kernels skip those rows
  through VLLM_MOE_SKIP_PADDING.
• Compact verification rebuilds the variable-length verifier batch and replays
  full CUDA graphs for the target verification and draft query steps.

This builds on #47469 (SM100 DSA varlen decode). #47513 is the padded/full-CUDA-
graph comparison point. It replaces #47694, which was opened from the wrong
fork.

Scheduling and execution

The paper's prefix scheduler runs on the host. Each draft token receives its
cumulative prefix-survival probability. Tokens are considered globally in
score order, and the scheduler chooses the admission count that maximizes
expected accepted tokens times the profiled step rate. dspark_budget_frac
caps admission, while dspark_confidence_threshold provides direct threshold
mode.

Automatic profiling measures overlapping end-to-end step times and fits an
additive request-sized cost and verification-token cost. Their sum is converted
to the step rate used by the scheduler. The request term captures the parallel
DSpark proposal and other work whose shape does not shrink with verifier
pruning; the verification term captures the prunable token work. In a TP4
varlen profile, request-sized cost was 1.408 ms at 32 requests, 4.021 ms at 128,
and 7.398 ms at 256. The fit RMSE was 0.180 ms. Explicit
(verification_tokens, steps_per_second) breakpoints remain supported.

ConfidenceManager owns scheduling and the mode-specific batch transform. The
shared model runner is limited to request initialization, graph dispatch, batch
trimming, and confidence staging hooks.

Confidence logits and sampler outcomes use a fixed two-slot pinned-host ring.
The previous proposal's copy is consumed at a deterministic later dispatch, so
all TP ranks make the same shape decision without per-step payload allocation.
Compact verification uploads one packed metadata buffer and rebuilds positions,
sequence lengths, logit mappings, and verifier token IDs in one Triton kernel.

Online calibration

dspark_online_sts applies per-position Sequential Temperature Scaling from
the rejection sampler's observed accept/reject outcomes. Temperatures are fit
with an order-preserving ECE grid search over binned logits. Data-starved
positions blend toward the pooled fit, and calibration remains the identity
before observations arrive.

Performance

deepseek-ai/DeepSeek-V4-Flash-DSpark, TP4 B300, k=7, probabilistic drafts,
speed_bench, 256 prompts, up to 512 generated tokens, temperature 1.0,
natural EOS termination, and auto-profiled SPS. Cells are output tok/s /
acceptance length:

c	No capacity	Mask	Varlen	Mask + STS	Varlen + STS
16	3,141 / 3.87	3,112 / 3.90	3,042 / 3.73	3,114 / 3.92	2,983 / 3.72
32	4,691 / 3.91	4,647 / 3.94	4,509 / 3.58	4,572 / 3.92	4,304 / 3.58
64	6,993 / 3.90	7,009 / 3.87	6,726 / 3.56	6,968 / 3.95	6,571 / 3.54
128	9,591 / 3.95	9,921 / 3.93	9,826 / 3.48	10,044 / 3.94	9,722 / 3.49
256	12,695 / 3.92	12,517 / 3.90	13,019 / 3.34	12,497 / 3.84	13,226 / 3.33

All 25 cells completed 256/256 requests. The result is workload-dependent. At
c=128, mask + STS improves throughput by 4.7% over baseline. At c=256, varlen
improves throughput by 2.6% and varlen + STS by 4.2%. The table is reported
without smoothing or a fitted trend line.

Correctness and validation

• .venv/bin/python -m pytest tests/test_config.py tests/v1/spec_decode/test_dspark_confidence.py tests/v1/worker/test_gpu_block_table.py -q:
  183 passed.
• .venv/bin/python -m pytest tests/v1/attention/test_dspark_noncausal_sparse_mla.py -q:
  20 passed, 8 skipped.
• .venv/bin/python -m pytest tests/v1/spec_decode/test_dflash_lookahead.py -q:
  3 passed.
• GSM8K 5-shot, all 1,319 questions, TP4 compact verification with budget 0.5:
  0.9454 accuracy (required 0.9200), invalid rate 0.000.
• Pre-commit (Ruff, format, mypy 3.12, and config schema) passed on all changed
  files; the PR pre-commit and DCO checks are green.
• The performance matrix above completed without retries, CUDA errors, or
  tracebacks.

Duplicate-work check

Searched open PRs by the linked PR number and DSpark confidence/capacity
keywords. The only functional overlap is #47513, which implements padded
verification rather than this host confidence scheduler and compact batch path.

AI assistance

Built with AI assistance (OpenAI Codex and Claude). The human submitter must
review and be able to defend every changed line before undrafting.

[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

[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

[sheep94lion]

Hi, I am learning about DSpark and glad to see a working implementation! I wanted to check one design point against the DSpark paper before I rely on this. As I read the paper, the confidence-based prefix scheduler splits the decision in two: the total verify budget for a step is set from the confidence data two steps earlier, but the per-request verify length is then allocated from the current step's confidence.

Instead, we approximate the upcoming verification capacity using the confidence head outputs from two steps prior. Mechanically, the candidate tokens in the current step are still strictly sorted by their actual, up-to-date cumulative confidence scores;

In this PR it looks to me like both quantities come from the same two-steps-earlier snapshot rather than being split. Is that an intentional simplification for now, or is matching the paper's fresher per-request allocation something you're planning?

Possibly I'm misreading the flow — happy to be corrected.