[RL, Spec] Introduce top-p mask capture for spec and add DFlash/DSpark impl

[jvmncs]

Relationship to existing work

#27408 introduced sampling-mask output support. #35205 restored the finite-top_k requirement and bounded support packing for sampling masks. #33459 added the DFlash SpecV2 logprob path.

This PR extends that machinery to the DFlash family, with worker-specific adapters for DFlash and DSpark, while preserving the finite-top_k contract from #35205.

Motivation

Sampling masks let RL trainers reconstruct the inference policy more accurately and reduce rollout/training policy mismatch. This change makes that metadata available when using DFlash or DSpark speculative decoding, and introduces a contract that other spec implementations can implement for their own support.

Modifications

• Add speculative/spec_sampling_mask.py, the speculative sampling-mask contract shared by DFlash and DSpark:
  • validate_spec_sampling_mask_request rejects configurations whose committed tokens are not exact samples from the captured target distribution: unsupported speculative methods, min_p, simulated acceptance, and for DFlash, acceptance thresholds other than 1.0 or missing sampling-verification support.
  • SpeculativeSamplingMaskCapture.from_logits rebuilds the target policy after acceptance from the final adjusted logits (grammar, penalties, logit bias, custom logit processors), only for batches with mask requests; greedy rows emit singleton masks with logprob 0.
• Keep speculative acceptance untouched; each worker calls the capture once after accept, including DSpark compact/ragged verification after logits are restored to strided layout.
• Align metadata with committed out_tokens / commit_lens via SamplingMaskOutput.num_accept_tokens; the result processor appends one support row per accepted token and trims at finished_len.
• Track per-row mask top_k on SamplingBatchInfo (sampling_mask_top_ks, maintained on filter/merge) to bound the packed support width; unbounded top-k stays rejected per [Sampling] Restore finite top-k requirement for sampling masks #35205.
• Reject runtime acceptance-threshold updates while DFlash mask requests are active, including PD decode queues.
• Rename build_dflash_verify_target_probs to build_speculative_verify_target_probs, as it is shared by DFlash, DSpark, UNO and the capture.
• EAGLE, NGRAM and other speculative methods keep rejecting sampling-mask requests.

Tests

• Add SpecSamplingMaskKit (mixed greedy/sampling batch, selected and support modes) and run it in the DFlash and DSpark basic sanity suites.
• Existing sampling-mask materialization unit tests now build real SamplingMaskOutput objects.

Performance considerations

Capture remains opt-in: batches without mask requests do no extra work. Support is bounded by the largest finite mask top_k in the batch and packed before device-to-host transfer. For batches with mask requests, both DFlash and DSpark recompute the target policy once after verification.

Accuracy Tests

Targeted final implementation:

• 4/4 DFlash sampling-mask unit tests passed in Modal:

  • tensor-backed deferred materialization;
  • filtered target probabilities and selected-token logprobs;
  • greedy singleton masks;
  • whole-request rejection on support overflow.

• Added SpecSamplingMaskKit coverage asserting:

  • seven generated tokens produce seven masks and seven logprobs;
  • every sampled token belongs to its corresponding support;
  • greedy rows return singleton masks with logprob 0.

• Full Qwen3.6 NVFP4+DFlash B200 run:

  • 98,304 measured output tokens;
  • 49,152 requested masks returned, with no missing masks;
  • maximum observed top-p support was 5,306; a repeat reached 7,476;
  • completed without generation or mask-alignment errors.

This feature only returns metadata and is not intended to change sampling decisions, so no downstream task-accuracy benchmark was run.

Speed Tests and Profiling

One B200 with:

• nvidia/Qwen3.6-35B-A3B-NVFP4;
• modal-labs/Qwen3.6-35B-A3B-DFlash;
• DFlash block size 8, FA4 draft attention;
• 32 concurrent requests, 512 generated tokens each;
• warmup followed by six alternating mask/no-mask rounds.

Results:

• Without masks: 4,163.78 tok/s median
• With masks: 4,160.45 tok/s median
• Ratio: 0.9992, or approximately −0.08%

The difference was within run-to-run noise. A complete repeat also showed no slowdown.

These numbers were measured on an earlier revision in which DFlash reused the verification probabilities instead of recomputing them after accept; requests without masks are unaffected either way.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): 🚫 Run #36681727080
Latest PR Test (Extra): ❌ Run #36681726706
Latest PR Test (AMD ROCm 10): ❌ Run #36681727472

[gongy]

/rerun-test test/registered/spec/dflash/test_dflash.py

[github-actions]

⚠️ Rebase Required Before Re-run

A major update has landed on main. Your PR is diverged relative to required base commit aa3f76679930.

Re-run was not dispatched. What to do:

• Rebase your branch onto the latest main and push again
• Follow issue #21065 for context
• CI-fix PRs may request the bypass-maintenance label to skip this check

[dcw02]

/rerun-test test/registered/sampling/test_sampling_mask.py test/registered/spec/dflash/test_dflash.py test/registered/core/test_basic_sanity_dspark.py

[github-actions]

⚠️ Rebase Required Before Re-run

A major update has landed on main. Your PR is diverged relative to required base commit aa3f76679930.

Re-run was not dispatched. What to do:

• Rebase your branch onto the latest main and push again
• Follow issue #21065 for context
• CI-fix PRs may request the bypass-maintenance label to skip this check

[dcw02]

/rerun-test test/registered/sampling/test_sampling_mask.py test/registered/spec/dflash/test_dflash.py test/registered/core/test_basic_sanity_dspark.py

[github-actions]

Results for /rerun-test test/registered/sampling/test_sampling_mask.py test/registered/spec/dflash/test_dflash.py test/registered/core/test_basic_sanity_dspark.py:

🚀 1-gpu-5090 (2 tests): ✅ View workflow run

cd test/ && python3 registered/sampling/test_sampling_mask.py
cd test/ && python3 registered/spec/dflash/test_dflash.py

🚀 1-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/core/test_basic_sanity_dspark.py

[jvmncs]

/rerun-test test/registered/sampling/test_sampling_mask.py test/registered/spec/dflash/test_dflash.py test/registered/core/test_basic_sanity_dspark.py

[github-actions]

Results for /rerun-test test/registered/sampling/test_sampling_mask.py test/registered/spec/dflash/test_dflash.py test/registered/core/test_basic_sanity_dspark.py:

🚀 2-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/sampling/test_sampling_mask.py

🚀 1-gpu-5090 (1 test): ✅ View workflow run

cd test/ && python3 registered/spec/dflash/test_dflash.py

🚀 1-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/core/test_basic_sanity_dspark.py

[hnyls2002]

/rerun-test test_basic_sanity_dflash.py test_basic_sanity_dspark.py test_sampling_mask.py test_scheduler_sampling_mask_validation.py test_batch_result_processor_hidden_states.py test_output_streamer_customized_info.py test_sampling_batch_info.py test_graph_pool_borrow.py

[github-actions]

Results for /rerun-test test_basic_sanity_dflash.py test_basic_sanity_dspark.py test_sampling_mask.py test_scheduler_sampling_mask_validation.py test_batch_result_processor_hidden_states.py test_output_streamer_customized_info.py test_sampling_batch_info.py test_graph_pool_borrow.py:

🚀 1-gpu-5090 (2 tests): ✅ View workflow run

cd test/ && python3 registered/core/test_basic_sanity_dflash.py
cd test/ && python3 registered/unit/model_executor/runner_utils/test_graph_pool_borrow.py

🚀 1-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/core/test_basic_sanity_dspark.py

🚀 2-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/sampling/test_sampling_mask.py

🚀 ubuntu-latest (4 tests): ✅ View workflow run

cd test/ && python3 registered/unit/managers/test_scheduler_sampling_mask_validation.py
cd test/ && python3 registered/unit/managers/test_batch_result_processor_hidden_states.py
cd test/ && python3 registered/unit/managers/test_output_streamer_customized_info.py
cd test/ && python3 registered/unit/sampling/test_sampling_batch_info.py