perf(dspark): add load-aware compact verification capacity

[voipmonitor]

Summary

Add opt-in load-aware physical draft-depth control for DSpark/DFlash while keeping the configured logical width and target acceptance semantics unchanged.

At low load the controller retains exact fixed-K5 behavior. At the configured activation knee it selects a smaller physical depth, uses compact variable-length verification, and keeps the path on full CUDA graphs. TP rank 0 computes and broadcasts one canonical capacity vector so all ranks replay the same padded shape without a host-wide synchronization.

Dependency

Merge #108 and #105 first. #108 removes a duplicate sparse-indexer helper that otherwise makes the combined changed-file lint set fail; #105 provides the corrected DFlash prefix-cache, request-slot, graph-lifetime, and padding behavior consumed by this capacity path. The PR still targets dev/fathomless-firmament; after both prerequisites merge, GitHub will reduce this diff to the capacity feature.

Behavior and scope

• disabled by default
• fixed K remains unchanged below the load knee
• compact variable-length verification activates only at the configured load threshold
• one TP-canonical capacity vector is broadcast before asynchronous D2H transfer
• dynamic-depth warmup and full-graph captures cover every enabled physical depth
• explicit validation rejects inconsistent capacity configuration
• profiling utility and focused unit coverage are included

Measured TP2 Lucifer results

Fixed K5 baseline versus load-aware capacity:

Concurrency	Fixed K5 tok/s	Load-aware tok/s
1	237.2	239.5
8	770.9	764.5
16	1186.1	1180.1
32	1764.6	1799.6
64	2614.6	2718.2

Three-run C64 mean was 2725.1 tok/s with SD 11.2. Coding median was 323.0 tok/s, mean 327.5, maximum 359.2. A clean-cache C1/C64/recovery sequence measured 245.2 / 2700.5 / 241.3 tok/s; a 120-second C64 run measured 2751.4 tok/s.

Rejected approaches

• Masking unused slots while retaining fixed physical width does not save CUTLASS target compute.
• Always-on compact verification regresses C1.
• Global synchronization and event-handle reuse were rejected because they either serialize the hot path or are unsafe across asynchronous replays.
• B12X DMA and NCCL tuning are independent transport experiments and are not part of this change.

Validation

• 68 focused DSpark correctness and capacity tests passed
• 15 B12X sparse-indexer tests passed with current DCP and KV-cache fixtures
• clean canonical integration stack: 233 passed, 714 skipped
• Ruff lint and format checks pass for the complete changed Python set
• git diff --check passes

This is extracted from the old mixed #88 branch so the capacity optimization can be reviewed and merged independently.

Summary by CodeRabbit

• New Features

  • Added DSpark capacity-based speculative decoding with configurable varlen or masked verification modes.
  • Added confidence-based draft-token budgeting, online calibration, dynamic draft depth, and SPS curve support.
  • Added automatic SPS profiling and a standalone tool for generating SPS configuration curves.
  • Added new command-line and environment settings for DSpark capacity behavior.
  • Added DeepSeek and Qwen DSpark confidence-head support.
  • Added a GSM8K evaluation configuration for DeepSeek V4 DSpark.

• Bug Fixes

  • Improved CUDA graph dispatch, prefix-cache handling, padding masks, sampling bounds, and deterministic fallback seeds.
  • Improved sparse attention handling for variable-length and non-causal workloads.

• Tests

  • Expanded coverage for DSpark configuration, capacity scheduling, CUDA graphs, attention metadata, sampling, and prefix caching.

[coderabbitai]

📝 Walkthrough

Walkthrough

Adds DSpark capacity-based speculative decoding with confidence scoring, SPS scheduling, dynamic draft depth, varlen and masked verification, cache-aware DFlash execution, CUDA-graph dispatch constraints, sparse-attention varlen support, profiling utilities, and extensive tests.

Changes

DSpark capacity and execution

Layer / File(s)	Summary
Configuration and confidence scoring vllm/config/speculative.py, vllm/engine/arg_utils.py, vllm/model_executor/models/qwen3_dspark.py, vllm/models/deepseek_v4/nvidia/dspark.py	Adds DSpark capacity settings, CLI support, confidence heads, confidence-weight loading, SPS curves, and validation.
Capacity allocation and verification vllm/v1/worker/gpu/spec_decode/capacity.py, vllm/v1/worker/gpu/spec_decode/dspark/*, vllm/v1/worker/gpu/spec_decode/speculator.py	Adds GPU capacity allocation, online temperature scaling, dynamic draft depth, and varlen or masked batch handling.
DFlash and graph dispatch integration vllm/v1/worker/gpu/spec_decode/dflash/*, vllm/v1/worker/gpu/cudagraph_utils.py, vllm/v1/worker/gpu/model_runner.py, vllm/v1/worker/gpu/input_batch.py	Propagates cached-token and request-token limits through DFlash preparation, batch trimming, CUDA-graph capture, and execution dispatch.
Varlen attention plumbing vllm/v1/attention/*, vllm/utils/deep_gemm.py, vllm/model_executor/layers/sparse_attn_indexer.py	Adds ragged decode indices, varlen DeepGEMM paths, non-causal backend selection, dynamic sparse widths, and deterministic sparse metadata initialization.
Profiling and scheduler behavior benchmarks/profile_dspark_sps_curve.py, vllm/v1/worker/gpu/warmup.py, vllm/v1/core/sched/*	Adds SPS curve profiling and preserves explicit zero speculative-token scheduling values.
Validation coverage tests/v1/worker/test_gpu_model_runner_v2_draft_capacity.py, tests/v1/spec_decode/*, tests/test_config.py, tests/engine/test_arg_utils.py	Adds coverage for capacity kernels, graph dispatch, prefix masking, configuration validation, scheduler semantics, and deterministic sampling.

Estimated code review effort: 5 (Critical) | ~120 minutes

Possibly related issues

• local-inference-lab/rtx6kpro#20 — Covers DSpark speculative-decoding scheduling, placeholder sizing, and capacity behavior addressed by this change.

Suggested reviewers: lukealonso

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 19.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title is concise and matches the main change: load-aware DSpark capacity/compact verification behavior.

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch codex/ff-dspark-load-aware-capacity-stack-20260717

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 6

🧹 Nitpick comments (1)

vllm/v1/attention/backends/mla/indexer.py (1)

305-316: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add Google-style sections to the new helper docstrings.

Document parameters and return values using explicit Args: and Returns: sections.

As per coding guidelines, use Google-style docstrings with Args:/Returns:/Raises: sections.

Also applies to: 511-525, 747-763

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@vllm/v1/attention/backends/mla/indexer.py` around lines 305 - 316, Update the
docstrings for _needs_varlen_decode and the helpers at the referenced additional
locations to use Google-style sections: document every parameter under Args:
with types or descriptions, and document the boolean or other return value under
Returns:. Preserve the existing behavior and wording outside the added
documentation.

Source: Coding guidelines

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@benchmarks/profile_dspark_sps_curve.py`:
- Around line 88-97: Validate profiling parameters before any timing or curve
calculation: require iters > 0, warmup >= 0, and overhead_ms >= 0. Add these
checks at the benchmark entry point covering the timing logic and
curve_breakpoints, and reject invalid values with clear errors before division
or SPS computation.

In `@tests/v1/worker/test_mixed_warmup_gate.py`:
- Around line 24-35: Update the curve in
test_dspark_dynamic_budget_uses_upper_load_sps_knee so its largest overall SPS
drop occurs in the lower half, while retaining an upper-half knee that produces
the expected budget of 160. Keep the assertion and
_derive_dspark_draft_token_budget invocation unchanged.

In `@vllm/model_executor/models/qwen3_dspark.py`:
- Around line 92-94: Update the confidence projection forward path around the
visible forward method to conditionally concatenate markov_embed only when
confidence_head_with_markov is enabled; otherwise pass hidden directly. Ensure
the resulting input dimension matches the projection configuration in both modes
while preserving the existing projection and output squeezing behavior.

In `@vllm/v1/attention/backends/mla/indexer.py`:
- Around line 291-302: Gate varlen CUDA-graph support on the DSpark opt-in by
combining _supports_varlen_paged_mqa_logits() with
_uses_varlen_dspark_capacity(vllm_config) at both hardware capability checks and
the varlen enablement path around the referenced sites. Preserve the existing
hardware checks, but ensure they cannot advertise or enable varlen support
unless the configuration explicitly enables varlen DSpark capacity.

In `@vllm/v1/worker/gpu/model_runner.py`:
- Around line 1171-1174: Make the `self.input_buffers.is_padding` reset in the
surrounding model-runner flow unconditional for every batch, removing the
`verification_capacity_manager`/`VLLM_MOE_SKIP_PADDING`/draft-token guard while
preserving the `[:num_tokens].fill_(False)` operation.

In `@vllm/v1/worker/gpu/spec_decode/dspark/speculator.py`:
- Around line 282-305: Track whether proposal confidence logits are valid for
each request in compute_capacities, marking rows produced under
use_capacity=False or is_profile=True as invalid. Update DSparkOnlineSTS.record
to skip invalid/bypassed/profiled rows so stale confidence buffers are never
calibrated against later verification outcomes; preserve normal recording for
freshly computed confidence logits.

---

Nitpick comments:
In `@vllm/v1/attention/backends/mla/indexer.py`:
- Around line 305-316: Update the docstrings for _needs_varlen_decode and the
helpers at the referenced additional locations to use Google-style sections:
document every parameter under Args: with types or descriptions, and document
the boolean or other return value under Returns:. Preserve the existing behavior
and wording outside the added documentation.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: d5feef63-0a16-4ca8-b1e3-9bb1e1f4c204

📥 Commits

Reviewing files that changed from the base of the PR and between e4879e1 and 7c0d639.

📒 Files selected for processing (58)

• benchmarks/profile_dspark_sps_curve.py
• tests/engine/test_arg_utils.py
• tests/evals/gsm8k/configs/DeepSeek-V4-Flash-DSpark-varlen-TP4.yaml
• tests/test_config.py
• tests/v1/attention/test_deepseek_v4_dspark_metadata.py
• tests/v1/spec_decode/test_acceptance_length_controller.py
• tests/v1/spec_decode/test_dflash_cudagraph_lifetime.py
• tests/v1/spec_decode/test_dflash_prefix_cache_masking.py
• tests/v1/spec_decode/test_dynamic_sd_cug.py
• tests/v1/worker/test_gpu_block_table.py
• tests/v1/worker/test_gpu_model_runner_v2_draft_capacity.py
• tests/v1/worker/test_gpu_sampling_states_seed.py
• tests/v1/worker/test_mixed_warmup_gate.py
• tools/pre_commit/generate_attention_backend_docs.py
• vllm/config/speculative.py
• vllm/engine/arg_utils.py
• vllm/envs.py
• vllm/model_executor/layers/sparse_attn_indexer.py
• vllm/model_executor/models/qwen3_dspark.py
• vllm/models/deepseek_v4/common/ops/cache_utils.py
• vllm/models/deepseek_v4/nvidia/dspark.py
• vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py
• vllm/models/deepseek_v4/sparse_mla.py
• vllm/platforms/cuda.py
• vllm/utils/deep_gemm.py
• vllm/v1/attention/backend.py
• vllm/v1/attention/backends/flashinfer.py
• vllm/v1/attention/backends/mla/indexer.py
• vllm/v1/attention/backends/mla/sparse_swa.py
• vllm/v1/attention/backends/utils.py
• vllm/v1/core/sched/async_scheduler.py
• vllm/v1/core/sched/output.py
• vllm/v1/core/sched/scheduler.py
• vllm/v1/worker/gpu/attn_utils.py
• vllm/v1/worker/gpu/block_table.py
• vllm/v1/worker/gpu/cudagraph_utils.py
• vllm/v1/worker/gpu/dp_utils.py
• vllm/v1/worker/gpu/input_batch.py
• vllm/v1/worker/gpu/model_runner.py
• vllm/v1/worker/gpu/model_states/default.py
• vllm/v1/worker/gpu/sample/bad_words.py
• vllm/v1/worker/gpu/sample/gumbel.py
• vllm/v1/worker/gpu/sample/sampler.py
• vllm/v1/worker/gpu/sample/states.py
• vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py
• vllm/v1/worker/gpu/spec_decode/capacity.py
• vllm/v1/worker/gpu/spec_decode/causal_cascade/speculator.py
• vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py
• vllm/v1/worker/gpu/spec_decode/dflash/speculator.py
• vllm/v1/worker/gpu/spec_decode/dspark/capacity.py
• vllm/v1/worker/gpu/spec_decode/dspark/online_sts.py
• vllm/v1/worker/gpu/spec_decode/dspark/speculator.py
• vllm/v1/worker/gpu/spec_decode/rejection_sampler.py
• vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py
• vllm/v1/worker/gpu/spec_decode/speculator.py
• vllm/v1/worker/gpu/states.py
• vllm/v1/worker/gpu/warmup.py
• vllm/v1/worker/ubatch_utils.py

[voipmonitor]

Canonical-stack update: PR #108 is now a git ancestor of this branch, so the sparse-indexer dedup/lint prerequisite is explicit. The full #107 changed-file set now passes ruff check, ruff format --check, and git diff --check against current FF. Merge #108 before #107; #105 remains the DSpark correctness prerequisite.

[voipmonitor]

Canonical FF + B12X merge handoff (2026-07-17)

Pinned bases, rechecked immediately before this handoff:

• vLLM dev/fathomless-firmament: e4879e195dd9ecf23661803f6775c6a847a18a72
• B12X master: 115926eb6452d1e8b849225d7af112e2bb184030

1. Merge B12X first

1. lukealonso/b12x#28
2. lukealonso/b12x#32
3. lukealonso/b12x#34
4. lukealonso/b12x#35

Then close superseded lukealonso/b12x#33. It is still draft/conflicting because this account cannot close another author’s PR; #35 is its clean replacement.

This exact B12X order was synthesized cleanly and validated with 122 passed, 1 skipped. All four PRs are non-draft, mergeable, clean, and have no unresolved review threads.

2. Merge vLLM into FF in this exact order

1. fix(indexer): remove duplicate RoPE quant helper #108
2. perf: avoid fixed-width MTP draft host synchronization #97
3. perf(spec decode): return variable draft ids through async output #103
4. fix(spec decode): isolate probabilistic draft Gumbel streams #98
5. fix(scheduler): skip lookahead allocation for running prefill chunks #99
6. fix(kv cache): preserve MLA quantization metadata when unifying page sizes #100
7. fix(ds4): harden sparse verifier metadata for DCP1 and padded replays #101
8. fix(ds4): isolate auxiliary-stream events across CUDA graphs #102
9. quant: add native mixed NVFP4/MXFP4 plus NF3 MoE serving #104
10. fix(spec decode): harden DSpark and DFlash edge paths #105
11. perf(dspark): add load-aware compact verification capacity #107
12. perf(dspark): add opt-in rowwise FP8 draft head #106

The exact sequence was rebuilt from the pinned FF head with no manual conflict resolution. Final validation: 248 passed, 714 skipped, Ruff lint and format clean across all 86 changed Python files, and git diff --check clean. There are no unresolved review threads on #97-#108, every PR is non-draft and GitHub reports every PR mergeable. The remaining UNSTABLE labels on a few independent PRs are queued repository pre-commit jobs, not failed checks.

Merge strategy

Use normal merge commits for the stacked pairs:

• perf: avoid fixed-width MTP draft host synchronization #97 before perf(spec decode): return variable draft ids through async output #103
• fix(indexer): remove duplicate RoPE quant helper #108 and fix(spec decode): harden DSpark and DFlash edge paths #105 before perf(dspark): add load-aware compact verification capacity #107

If FF policy requires squash or rebase merges, stop after rewriting #97, #108, or #105 and rebase the dependent PR (#103 or #107) onto the new FF head before continuing. Otherwise GitHub cannot recognize the tested ancestry and may show duplicate changes.

The old mixed #88 and superseded #92/#96 are closed; their canonical replacements are the PRs above.

[voipmonitor]

Superseded by #109. The final code was rebuilt directly on current dev/fathomless-firmament as three clean commits: correctness, optional FP8 draft head, and optional load-aware capacity. #109 preserves the validated behavior while removing the stacked merge history and excluding the unrelated #108 indexer cleanup.