[SM12x] Default the fused MHC post+pre path on

[zhendonghua]

Motivation

On sm120/sm121 the SM120 block in server_args.py sets SGLANG_OPT_USE_TILELANG_MHC_PRE=False, so DeepSeek-V4's hc_pre falls through to hc_pre_torch_impl — an fp32 F.linear of shape [M, 16384] × [16384, 24]. cuBLAS serves that with cutlass_80_simt_sgemm, i.e. plain CUDA cores, no tensor cores at all.

On 2× DGX Spark (GB10 / sm_121, TP=2) running deepseek-ai/DeepSeek-V4-Flash-0731 + DSPARK that fallback costs 93 calls/step, 9.54 ms/step — about 14% of the decode step, for a projection whose weights are only 1.5 MB.

mhc_fused_post_pre covers the same work, and _is_fused_mhc_post_pre_enabled() already special-cases this arch:

# SM120 disables the standalone TileLang pre path. mhc_fused_post_pre does
# not read that flag and dispatches independently for both small and large
# token batches, so the standalone pre flag must not veto the fused opt-in.

But SGLANG_OPT_FUSE_MHC_POST_PRE defaults to False, so on sm12x the combination the comment anticipates is never reached out of the box — the SIMT fallback is what everyone gets.

Modifications

Default SGLANG_OPT_FUSE_MHC_POST_PRE=True inside the existing is_sm120_supported() block, guarded by is_set() so an explicit setting still wins. No other arch is affected; no change to the flag's global default.

Accuracy Tests

This change is not bit-identical. The fused kernel is fp32 throughout — same precision as the fallback, verified from the TileLang kernel's dtype declarations (pre_fn, mixes_partial_out, sqrsum_partial_out and every accumulator are T.float32) — but it uses a split-K reduction, so the accumulation order differs. MHC feeds a discrete cluster assignment, so a last-bit difference can flip an assignment and diverge the whole generation at temperature 0. Greedy output is therefore not reproducible against the fallback.

Task accuracy is unaffected. GSM8K 5-shot, --parallel 16, 2× DGX Spark TP=2:

build	n	Accuracy	Invalid
baseline	200	0.970	0.000
this change	200	0.980	0.000
baseline (older base, n=1000)	1000	0.943	0.000
this change (older base, n=1000)	1000	0.945	0.000

Both arms are within noise of each other at both sample sizes. Note accuracy is not comparable across different --num-questions: the same build scores 0.970 at n=200 and 0.943 at n=1000, because the first 200 questions are easier — so an apparent "drop" from 0.970 to 0.945 across a resize is an artifact, not a regression.

Anyone who relies on bit-exact reproduction against the previous default should set SGLANG_OPT_FUSE_MHC_POST_PRE=0.

Benchmarking and Profiling

sglang.bench_one_batch_server, bs=1, SGLANG_SIMULATE_ACC_LEN=5 so the accept length is pinned and the comparison is pure step cost (acc_length: 5.00 confirmed in both arms). Same container, unlocked clocks, RoCE/RDMA available to both arms.

in / out	baseline	this change + FP8 wo_a (#34018)
1024 / 512	58.48 tok/s	72.67 tok/s	+24.3%
8192 / 1024	61.88 tok/s	66.76 tok/s	+7.9%

The two changes were measured together, but a 20-step profile separates them cleanly by kernel family:

kernel family	baseline	patched	attributable to
fp32 cuBLAS / SIMT (hc_pre)	9.54 ms / 93 calls	0.86 ms / 8 calls	this PR
MHC TileLang (fused kernel)	0.36 ms / 92 calls	1.75 ms / 180 calls	this PR
bf16 SM80-WMMA	20.18 ms / 136 calls	12.74 ms / 93 calls	FP8 wo_a (#34018)
deep_gemm sm120_fp8_fp4_gemm_1d1d	0	3.29 ms / 43 calls	FP8 wo_a (#34018)
decode step	70.31 ms	63.19 ms	−10.1%

Net for this PR alone: −8.68 ms of SIMT fp32 GEMM for +1.39 ms of fused TileLang ≈ −7.3 ms/step.

The 8 remaining fp32 calls are the first layer of each forward, where prev_residual is None and the fused post+pre has nothing to fuse against.

Test environment caveat

DSPARK on sm120/sm121 cannot boot on stock flashinfer today: the draft's index width is
ceil_align(swa_window + gamma, 64), which is 192 for every gamma in [1, 64], while
flashinfer 0.6.15.post1 only instantiates decode top-k widths {128, 512, 1024}. Both known
fixes are still open — flashinfer-ai/flashinfer#4309 (adds the 192 instantiation) and
#33407 (dispatch fallback in SGLang). The numbers above were collected with
flashinfer#4309 carried as a local patch; SGLANG_SM120_FLASHMLA_BACKEND=triton is the only
stock alternative and it routes all sparse-MLA calls through Triton, which costs throughput
on the hot path (measured 4.8x slower end to end here), so it is not a usable baseline.

On tests

No unit test is included. The change is arch-gated to sm120/sm121, and CI has no such runner,
so a test would not execute where the behaviour differs. Validation was done on hardware
instead: GSM8K for accuracy and a kernel-level profile (reported above) to confirm the intended
kernels actually run rather than inferring it from throughput. Happy to add a test if there is
a preferred pattern for arch-gated paths.

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, see Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, see Benchmark and Profiling and Accuracy Results.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
• Please feel free to join our Slack channel at https://slack.sglang.ai to discuss your PR.

---

CI States

Latest PR Test (Base): ⏳ Run #31660094177
Latest PR Test (Extra): ❌ Run #31660094118

[zhendonghua]

/rerun-failed-ci

[ormandj]

Independent SM120 TP2 profiling of the behavior enabled by this PR:

• Hardware: 2× RTX PRO 6000 Blackwell Max-Q (SM120), TP2 over PCIe Gen4 x16.
• Model: DeepSeek-V4-Flash-0731 with DSPARK speculative decoding.
• Input: 16,384 tokens.
• The two captures used the same image, source trees, model arguments, graph sizes, token IDs, and profiler settings.
• The only configuration change was SGLANG_OPT_FUSE_MHC_POST_PRE unset versus SGLANG_OPT_FUSE_MHC_POST_PRE=1.
• Each arm retained 127 complete target-plus-draft scheduler steps on device 1.

Measured results:

• Target verifier graph span: 14.872 ms → 13.403 ms (-9.88%).
• Target busy-union time: 13.643 ms → 12.458 ms (-8.69%).
• Target internal idle time: 1.220 ms → 0.943 ms (-22.68%).
• Draft graph span: 1.596 ms → 1.605 ms (+0.61%).
• Component-median complete step: 16.694 ms → 15.237 ms (-8.73%).
• Target static graph nodes: 3,094 → 2,587.

The FP32 SIMT MHC-pre GEMM, mean, HC split, and standalone post operations each fell from 86 launches to either one or four launches. The enabled capture contained 85 fused post-pre launches totaling 0.390 ms and 85 large fused-with-norm launches totaling 0.377 ms.

A later configuration with the selector enabled completed five repetitions at every decode concurrency from C1 through C32, five cache-cold prefill repetitions at 8K, 32K, 64K, and 128K, GSM8K with 1,261/1,319 correct and zero request errors, and 8/8 long-output requests. These later results are integration coverage, not isolated attribution to this selector.

[zhendonghua]

/rerun-failed-ci