[Spec][DSV4] perf: Compile draft-extend hc_head during CUDA graph capture

[mattteochen]

Motivation

The DSV4 MTP draft-extend CUDA graph currently captures hc_head as a sequence of small ATen elementwise and reduction kernels between the cuBLASLt split-K reduction and the shared-head RMSNorm. These kernels are replayed on every draft-extend step, so their launch overhead is significant for the low token counts used by speculative decoding.

This PR takes the selective capture-time torch.compile idea from #32220 and applies it to the current DSV4 MTP NextN draft-extend path. It allows Inductor to fuse the hc_head pointwise and reduction chain before the generated kernels are recorded in the CUDA graph, without enabling whole-model torch compilation.

Modifications

• Replace the duplicated eager math in DeepseekV4ModelNextN.hc_head with the shared hc_head_torch implementation.
• Use compile_in_capture_mode for non-empty CUDA inputs so hc_head_torch is compiled during CUDA-graph warmup and the generated kernels are captured for replay.
• Preserve the eager implementation outside capture and for empty or non-CUDA inputs.
• Keep the existing CUDA-graph runners, capture backend, and global torch-compile configuration unchanged.

Accuracy Tests

sgl-eval run gsm8k \
  --base-url http://localhost:30000/v1 \
  --num-threads 32
Run directory: /root/.sgl_eval/sgl_eval_gsm8k_20260727-003926
gsm8k: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1319/1319 [03:56<00:00,  5.59it/s, acc=97.04%]
== gsm8k ==
1319 examples (single-shot)  |  236.1s  |  743 tok/s  |  176K tokens

* score           =  97.04%
  stop_rate       =  100.00%
  truncated_rate  =  0.00%
  error_rate      =  0.00%

Speed Tests and Profiling

main:

PR:

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 #30799325097
Latest PR Test (Extra): ❌ Run #30799324933

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.