[Kernel] Route large SM90 row/column-scaled FP8 GEMMs to Torch

[RunFMe]

Summary

Follow-up to #33275, which added MiniMax-H3 and its online FP8 transformer path.

The current sgl_kernel.fp8_scaled_mm is slower than Torch _scaled_mm on the larger H3 dense shapes on SM90. This PR adds both implementations to BaseFusedOp and routes shapes using:

(K >= 5376 and N >= 3584) or (K >= 3584 and M >= 8192)

The AOT kernel remains the fallback for smaller shapes. The Torch path requires FP8 E4M3FN inputs, BF16 output, FP32 row/column scales, compatible layouts, and no bias.

Kernel results

H100 80 GB, median of 10 measurements after 5 warmups:

Operation	M × K × N	SGLang AOT	Torch	Reduction	Torch TFLOP/s
Attention QKV, TP1	8320 × 5376 × 21504	2036.7 µs	1351.1 µs	33.66%	1424
MLP gate/up, TP2	8320 × 5376 × 14336	1321.9 µs	896.4 µs	32.19%	1431
MLP gate/up, TP1	23216 × 5376 × 28672	7823.2 µs	5295.9 µs	32.30%	1351
MLP gate/up, TP1	15760 × 5376 × 28672	4892.1 µs	3570.9 µs	27.01%	1361
MLP down, TP1	8320 × 14336 × 5376	990.0 µs	843.9 µs	14.76%	1520

The sweep covered all 64 H3 dense shapes across TP1/2/4/8 and four production sequence lengths. Torch won 49 shapes, AOT won 15, and the heuristic selected the measured winner for all 64.

Full MiniMax-H3 result

H100 ×8, TP2/Ulysses4, FP8, FlashAttention, 768×1120 FL2VA request with 49 denoising steps. Medians from three matched runs after two warmups:

Metric	AOT only	Routed	Change
Inference	30.213 s	27.955 s	-7.47%
Wall time	31.076 s	29.068 s	-6.46%
Peak memory	35,254 MB	35,254 MB	unchanged

The matched outputs were byte-identical. All eight profiler traces confirmed the routed path while retaining AOT for smaller shapes.

Newer GPUs

I only have access to SM90 hardware, so the capability gate is intentionally limited to SM90. Help benchmarking and tuning these shapes on SM100/SM120 would be appreciated. Torch _scaled_mm may also win there

---

CI States

Latest PR Test (Base): ❌ Run #32853520330
Latest PR Test (Extra): ❌ Run #32853520160
Latest PR Test (AMD ROCm 7.2): ❌ Run #32853520280

[RunFMe]

@mickqian could you take a look pls and help me understand who to ask for review since the change is probably affecting many architectures, not only h3

[RunFMe]

@BBuf would appreciate you taking a look too) thx

[mickqian]

/tag-and-rerun-ci

[mickqian]

unfortunately, the multimodal_gen tests are not triggered

[RunFMe]

@mickqian should I run them myself? I have H200 and H100 nodes

[RunFMe]

multimodal_gen validation

I ran the full H100-compatible/H200-applicable multimodal_gen tess on train_57.

Revisions and environment

• PR head: 83ceab2e13e873be6145fdda5dab74eebd2c2e6d
• Merge base: a1f9508dd41006672710eda29528cccb34d4a52e
• Hardware: 8× NVIDIA H200
• Container image: lmsysorg/sglang@sha256:f7155872f80168da827d36789eb486f85efc9cb42705788921a2e899d6856ac1
• The 5090- and B200-specific lanes were not applicable to this host and were not run.
• Large Hugging Face checkpoints were cached outside the container and mounted into it.

Test scope

The following suites were run from the checked-in sglang/multimodal_gen/test/run_suite.py runner:

• 1-gpu: 5 partitions
• 2-gpu: 3 partitions, including the standalone distributed tests
• 4-gpu-h100
• component-accuracy, including the 1-GPU and 2-GPU cases
• unit
• bcg-diffusion
• kernel-namespace
• Server-case coverage verification

Representative partition commands:

python3 sglang/multimodal_gen/test/run_suite.py \
    --suite 1-gpu \
    --partition-id <0-4> \
    --total-partitions 5 \
    --continue-on-error

python3 sglang/multimodal_gen/test/run_suite.py \
    --suite 2-gpu \
    --partition-id <0-2> \
    --total-partitions 3 \
    --continue-on-error

The server coverage verifier reported 100% coverage of the applicable configured cases: 83/83 entries executed, with 49 passing and 34 failing entries at the PR head.

The standalone distributed files can contain multiple pytest failures/errors under one coverage entry. After expanding those files and including component and unit suites, there were 42 failing/error pytest nodes to compare against the merge base.

Result summary

Category	Failures/errors at PR head	Reproduced at merge base	Contention artifacts
1-GPU server cases	15	13	2
2-GPU server-model cases	17	17	0
Standalone 2-GPU pytest nodes	4	4	0
Component-accuracy cases	5	5	0
Unit tests	1	1	0
Total	42	40	2

Additional clean suites:

• 4-gpu-h100: 1 passed
• bcg-diffusion: 2 passed
• kernel-namespace: 59 passed
• Unit suite overall: 1865 passed, 8 skipped, 173 subtests passed, 1 pre-existing failure

Investigation of the two non-reproducing timing failures

Both cases produced correct outputs and passed their numerical/consistency checks. Only wall-clock performance guards failed while multiple GPU and component partitions were packed onto the same host.

wan2_1_t2v_1.3b_frame_interp_2x_upscaling_4x

The original packed run failed only this timing check:

denoise step 0: 699.5128 ms > allowed 550.05 ms

Observed step-zero times in controlled retries:

• Isolated PR-head retry: 504.37 ms — passed
• Ordered Cosmos → Wan retry: 466.03 ms — passed
• Merge-base warm-cache retry: 112.98 ms — passed

All consistency thresholds passed in the original run. The failure occurred while five 1-GPU partitions plus component-accuracy work were active concurrently. The ordered predecessor retry also passed, ruling out a persistent Cosmos-to-Wan process-state leak.

This is consistent with cold first-step CUDA/Triton initialization being amplified by shared-host contention. Upstream CI normally gives each 1-gpu-h100 partition its own runner.

flux_2_ti2i_multi_image_cache_dit

The original packed run failed only the input-validation timing guard:

InputValidationStage: 333.0807 ms > allowed 204.29 ms

Controlled retries:

• Isolated PR-head retry: 133.19 ms — passed
• Merge-base retry: 130.16 ms — passed
• Ordered Sana → FLUX retry: 103.15 ms — passed

The stage includes two local uploaded-image opens plus PIL/config preprocessing and is measured with wall-clock time.perf_counter, so CPU scheduling and filesystem contention are included. Four other 1-GPU partitions plus component-accuracy work were active at the original failure time.

The ordered predecessor retry passed, ruling out a persistent Sana-to-FLUX state leak.

Pre-existing merge-base failures

13 reproduced 1-GPU server failures

• flux_2_klein_base_image_t2i
• cosmos3_nano_t2v
• zimage_image_t2i_multi_lora
• flux_2_image_t2i_upscaling_4x
• wan2_1_t2v_1.3b
• fast_hunyuan_video
• wan2_1_t2v_1.3b_frame_interp_2x
• hunyuan3d_shape_gen
• turbo_wan2_1_t2v_1.3b
• flux_2_klein_image_t2i
• wan2_1_t2v_1_3b_lora_1gpu
• ltx_2_3_hq_pipeline
• wan2_1_t2v_1.3b_upscaling_4x

17 reproduced 2-GPU server-model failures

• minimax_h3_t2va_2gpu_h100
• minimax_h3_ref2va_video_audio_2gpu_h100
• wan2_2_i2v_a14b_2gpu
• fsdp-inference
• ltx_2.3_one_stage_ti2v
• wan2_2_t2v_a14b_lora_2gpu
• mova_360p_ring1_uly2
• ltx_2_two_stage_t2v
• ltx_2_3_two_stage_ti2v_2gpus
• wan2_1_i2v_14b_480P_2gpu
• qwen_image_t2i_2_gpus
• zimage_image_t2i_2_gpus
• flux_image_t2i_2_gpus
• ltx_2.3_two_stage_t2v_2gpus
• zimage_image_t2i_2_gpus_non_square
• flux_2_image_t2i_2_gpus
• flux2_modelopt_fp8_tp2_t2i

4 reproduced standalone distributed failures/errors

• IPC bitwise all-to-all diverged from the NCCL result and exhausted its retry.
• Three disaggregated-inference tests errored in the denoiser with:
  unsupported operand type(s) for *: 'NoneType' and 'int'

5 reproduced component-accuracy failures

• Wan2.1 1.3B transformer, 1 GPU
• Wan2.2 I2V A14B, 2 GPUs
• Wan2.2 T2V A14B, 2 GPUs
• Wan2.1 T2V 14B, 2 GPUs
• Wan2.1 T2V 1.3B CFG-parallel, 2 GPUs

Unit-test failure

The remaining unit failure was:

test_server_args.py::TestOffloadDefaults::
test_auto_ltx_original_replaces_component_cpu_offload

It fails identically at both the PR head and merge base when the NVIDIA device is exposed:

AssertionError: 'resident' != 'original'

The test mocks CUDA availability and an 80 GiB device but does not mock the device name. On a real H200, the production resolver deliberately selects resident, while the test expects the H100-style original result.

Earlier no-GPU isolated retries passed because the platform reported an X86_64 device name, which hid the H200-specific behavior. The matched GPU reruns confirm this is a pre-existing test-isolation issue, not a change introduced by this PR.

A future test fix should mock an H100-like device name for the original expectation and add an explicit H200 case expecting resident.

PR-relevant checks

The H3 T2VA case produced identical head/base metrics:

similarity: 0.9432
SSIM:       0.8971
PSNR:       32.4219
MAD:        2.7732

The FLUX.2 ModelOpt FP8 case also produced the same consistency result and metrics at the PR head and merge base.

Conclusion

No functional, numerical, or reproducible performance regression unique to this PR was found.

Of the 42 failing/error nodes at the PR head:

• 40 reproduced at the merge base under the same GPU environment.
• 2 were performance-only artifacts caused by packing multiple partitions onto one host; both passed in isolated and ordered PR-head retries.
• 0 were PR-head-only functional or numerical failures.

[RunFMe]

@mickqian wdyt?

[hnyls2002]

cc @RunFMe @BBuf

The routing heuristic here regresses W8A8 FP8 decode on H100: test/registered/quant/test_w8a8_quantization.py::TestW8A8Fp8.test_throughput (Llama-3.1-8B-Instruct-FP8-dynamic, bs=1) dropped from ~215 tok/s to ~190 tok/s and now fails its 200 tok/s threshold on every PR opened since this merged, across several runners (e.g. https://github.com/sgl-project/sglang/actions/runs/33224827665/job/99026386210 ; the last main scheduled run before the merge, https://github.com/sgl-project/sglang/actions/runs/33219371283/job/99021104112 , was at 214.7).

Cause: (k >= 5376 and n >= 3584) has no lower bound on M, so down_proj (K=14336, N=4096) is routed to torch._scaled_mm even at M=1, where the AOT kernel wins. The sweep only covered prefill-size M; suggest requiring a minimum M in that clause too.

[RunFMe]

Thank you for pointing this out, I will make a followup PR to add gating on M too (will take 20 mins to determine which M is the cutoff)

[RunFMe]

here is the fix, update the cutoffs to make sure that we only get the big minimax-related shapes while not including anything smaller
https://github.com/sgl-project/sglang/pull/37018/changes
@mickqian