[Diffusion][MiniMax-H3] Add SM120 Sage compute for SubBlock sparse attention

[IPostYellow]

Motivation

Extend MiniMax-H3's SubBlock Sage compute mode from SM90 (#37982) to SM120 using FlashInfer's CuTe-DSL backend. On 8 RTX PRO 5000 GPUs, Sage reduces generation time by 3.86%–14.86% compared with BF16 SubBlock across the evaluated task/duration groups.

Modifications

• Add an SM120 Sage adapter with online INT8-QK / FP8-PV computation and BF16 inputs and outputs.
• Use Q64 × K64 blocks on SM120 with default n_k=4, preserving 16-token key pooling cells.
• Support partial blocks and variable per-row block counts, including mixed sparse/dense query rows.
• Add startup dependency validation, regression tests, and documentation.

The mode is opt-in; BF16 remains the default. SM90 retains its existing Q64 × K128 Sage path, and SM100 remains BF16-only.

--attention-backend subblock_sparse_attn \
--attention-backend-config '{"compute_mode":"sage_fp8"}'

Dependency: requires flashinfer-ai/flashinfer#4691. Benchmarks used FlashInfer source commit 6a84331e. Users enabling SM120 Sage need a source installation containing this change; the default FlashInfer 0.6.18 pin is unchanged.

Accuracy Tests

• 55 tests passed on SM120; three SM90/SM100-specific tests were skipped. Coverage includes dispatch, routing, dense fallback, and kernel correctness against FP32 reference attention with ragged lengths, variable block counts, and empty rows.
• 81/81 end-to-end requests completed: three modes × three tasks (T2VA, FL2VA, Ref2VA_V) × three durations (5/10/15 seconds) × three seeds (0/1/2).

The new SM120 Sage numerical test covers partial blocks, variable block counts, and empty rows against an FP32 reference. It passed on SM120 hardware and remains a manual test pending the required FlashInfer dependency in CI.

All comparisons below use the same inputs and seeds. SB denotes SubBlock BF16, and SB+Sage denotes SubBlock Sage FP8. The dense reference uses Torch SDPA on SM120.

Video SSIM

SSIM is measured against the matching dense video over all frames, with YUV420 weighting (4Y + U + V) / 6. Values are means across three seeds.

Case	Duration	SB	SB+Sage
T2VA	5s	0.830408	0.828118
T2VA	10s	0.858958	0.859222
T2VA	15s	0.898157	0.897310
FL2VA	5s	0.870854	0.871837
FL2VA	10s	0.880160	0.881029
FL2VA	15s	0.872566	0.871813
Ref2VA_V	5s	0.945863	0.945693
Ref2VA_V	10s	0.981771	0.982091
Ref2VA_V	15s	0.968196	0.968183

The difference between SB+Sage and SB in group-mean SSIM against dense ranges from -0.002290 to +0.000983.

Audio similarity

Decoded audio is compared against the matching dense output. Values are means across three seeds, in dB. Higher waveform SNR and lower log-mel MAE indicate closer signals.

The original 32 kHz stereo audio is retained. Log-mel MAE uses 80 Slaney mel bands, a 2048-sample Hann window/FFT, a 512-sample hop, and a shared floor 80 dB below the reference mel peak.

Case	Duration	SB SNR ↑	SB+Sage SNR ↑	SB log-mel MAE ↓	SB+Sage log-mel MAE ↓
T2VA	5s	12.109	11.112	1.951	2.000
T2VA	10s	12.578	12.069	1.812	1.794
T2VA	15s	11.292	11.134	1.857	1.904
FL2VA	5s	12.162	13.718	2.012	1.927
FL2VA	10s	12.876	9.851	1.954	2.593
FL2VA	15s	7.424	7.094	2.332	2.401
Ref2VA_V	5s	15.127	13.786	0.594	0.754
Ref2VA_V	10s	22.004	19.693	0.441	0.567
Ref2VA_V	15s	15.087	14.599	1.197	1.180

Comparing SB+Sage directly against SB, group-mean SNR ranges from 13.93 to 21.44 dB and log-mel MAE from 0.472 to 1.658 dB.

Speed Tests and Profiling

Benchmarked MiniMax-H3 on 8 NVIDIA RTX PRO 5000 GPUs (SM120) with PyTorch 2.13.0+cu130, CUDA 13.0, TP=2, Ulysses SP=4, and text-encoder layerwise offload. All modes use 50 requested inference steps and quality="lossless".

Sparse settings: sparsity=0.75, n_q=n_k=4, skip_first_steps=10, skip_first_layers=0, and min_seq_len=4096.

Generation time

Mean server-side generation time across three seeds, excluding one complete warmup request per mode, task, and duration. The ↓ columns report time reduction relative to dense SDPA: (dense time - method time) / dense time.

Case	Duration	Dense SDPA (s)	SB (s)	SB ↓	SB+Sage (s)	SB+Sage ↓
T2VA	5s	146.14	112.34	23.13%	108.00	26.10%
T2VA	10s	391.57	262.35	33.00%	244.15	37.65%
T2VA	15s	740.95	453.34	38.82%	412.96	44.27%
FL2VA	5s	153.74	120.41	21.68%	114.80	25.33%
FL2VA	10s	405.50	271.74	32.99%	252.60	37.71%
FL2VA	15s	761.29	469.12	38.38%	424.61	44.22%
Ref2VA_V	5s	465.54	306.27	34.21%	283.87	39.02%
Ref2VA_V	10s	1390.90	797.96	42.63%	703.53	49.42%
Ref2VA_V	15s	2724.53	1476.01	45.82%	1256.71	53.87%

Sage is faster than BF16 SubBlock in all 27 matched-seed comparisons, reducing group-mean generation time by 3.86%–14.86%.

Peak reserved memory

Measured by the PyTorch allocator on the output rank. Values are the maximum per-request peak across three seeds, excluding warmup, in GiB.

Case	Duration	Dense SDPA	SB	SB+Sage
T2VA	5s	42.08	42.08	42.04
T2VA	10s	45.35	45.46	45.46
T2VA	15s	49.29	49.57	49.76
FL2VA	5s	42.05	42.05	42.05
FL2VA	10s	46.07	46.07	46.35
FL2VA	15s	49.14	49.14	49.28
Ref2VA_V	5s	44.75	44.75	44.75
Ref2VA_V	10s	50.88	50.88	50.88
Ref2VA_V	15s	56.28	56.28	56.28

Peak reserved memory differs from BF16 SubBlock by -0.04 to +0.27 GiB across the measured groups.

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 #35497377460
Latest PR Test (Extra): ❌ Run #35497377163
Latest PR Test (AMD ROCm 10): ⏳ Run #35497377350

[BBuf]

No blocking issue found in the SM120 adapter. I checked the Q64/K64 geometry, quantized tensor order/layout, and partial-block/per-row-count forwarding against the referenced FlashInfer source. My checks covered source-level dispatch and adapter arguments; I did not independently rerun the SM120 numerical or end-to-end benchmarks. One small documentation suggestion is inline.

[BBuf]

Non-blocking: could we include a source-install link/command or the minimum FlashInfer version providing both SM120 Sage symbols here? The default FlashInfer pin remains 0.6.18; the PR description points to flashinfer-ai/flashinfer#4691 and the tested commit, but users following this README only get the symbol names. Carrying that installation guidance into the README would make the opt-in mode easier to use.

[IPostYellow]

Thanks for the review. I'll add the FlashInfer source-installation guidance and the tested commit to the README so users can enable SM120 Sage without referring to the PR description.

[mickqian]

Exact-head CI follow-up: CPU(0) job106038117734 fails test_disagg_idle_step_counters.py because Scheduler lacks _prev_prefill_end_ts; CPU(6) job106038117740 fails test_pdmux_split_prefill_schedules_auxiliary_output_copy because its batch fixture lacks split_index. These are scheduler-fixture compatibility failures, not evidence of an SM120 Sage regression. Update: main now contains the upstream fixture repair #40411, commit22f02cc3399dcd59380fc546de8d3c89fba6fa3a (including completion timing, burst accounting, and additional regression coverage). Please merge updated main instead of rerunning the stale fixtures or applying the earlier #39206 variants. Maintainer edits are disabled for this head and I have no push permission, so this is an author-action follow-up. No unchanged-fixture rerun requested.