feat(moe): add SM120/SM121 dispatch for fp8_blockwise_scaled_grouped_mm

[waynehacking8]

Motivation

fp8_blockwise_scaled_grouped_mm (the CUTLASS blockwise-scaled grouped FP8 GEMM used by the cutlass MoE path) has dispatch cases for SM100/SM103 and SM90 only. On SM120/SM121 (consumer Blackwell — RTX PRO 6000, RTX 50-series, GB10/DGX Spark) the dispatch falls through to:

TORCH_CHECK_NOT_IMPLEMENTED(can_implement,
    "No implemented fp8_blockwise_scaled_grouped_mm for current compute capability: ", sm_version);

so every blockwise-FP8 MoE model crashes at runtime on these GPUs. Note tests/test_fp8_blockwise_moe.py::is_blackwell_supported already returns True for compute capability major 12 (SM120), i.e. the test is expected to run there but currently aborts at this check.

Modifications

Add an SM120 dispatch path for fp8_blockwise_scaled_grouped_mm:

• launch_sm120_fp8_blockwise_scaled_group_mm<OutType, ScheduleConfig, LayoutD> — mirrors the SM100 grouped launcher with ArchTag = cutlass::arch::Sm120. Passing pointer operand layouts makes the SM120 blockwise collective builder select the PtrArray (grouped) kernel automatically (KernelScheduleSm120Blockwise → KernelPtrArrayTmaWarpSpecializedCooperativeBlockwiseScalingSm120). hw_info.sm_count is queried from the device rather than hard-coded.
• sm120_fp8_blockwise_group_mm_dispatch_shape<OutType> — MmaTileShape<128,128,128>, single-SM ClusterShape<1,1,1> (SM120 has no 2-SM TMA cluster), Sm120BlockwiseScaleConfig<1,128,128, K, K> (per-token×128-K activation tiles, 128×128 weight tiles; both scale factors K-major, matching how the grouped caller lays them out and the SM90/SM100 configs).
• Wire sm_version == 120 || sm_version == 121 under CUTLASS_ARCH_MMA_SM120A_SUPPORTED.

No CUTLASS bump: the pinned CUTLASS already provides sm120_mma_array_tma_blockwise_scaling.hpp and the Sm120PtrArrayTmaWarpSpecialized epilogue.

Accuracy Test

Final head d3f801924 was built as a full CUDA 13 AOT wheel on an RTX PRO 6000 Blackwell (SM120) with torch 2.11.0+cu130. The grouped translation unit compiled into both sm90/common_ops and sm100/common_ops; the complete 365 MiB wheel linked successfully.

• Existing public test, deterministic seeds: 10/10 passed across FP16/BF16 and 8/16/32/64/128 experts (calc_diff < 0.001; observed approximately 0.00059-0.00069).
• Fixed boundaries: M=1, N=K=128 and multi-tile M=129, N=K=256, 8 experts, both output dtypes: passed (worst calc_diff = 0.00096).
• Fresh independent verification added mixed zero-token and large-expert rows M=[0,1,257,2049,0,3,128,129], N=K=256: both dtypes matched the quantized reference below 1.1e-9; repeated calls were bit-identical.
• cuobjdump confirmed SM120 cubins in the installed wheel. A separate exact-command compile of this translation unit for compute_121a,code=sm_121a also passed.

SM121 remains compile-only here: I do not have GB10/DGX Spark hardware, so no SM121 runtime result is claimed.

Benchmark

Not a perf change — this enables a path that previously aborted at runtime on SM120/SM121. Functional correctness validated above; perf tuning (shape-specialized configs, pingpong/cooperative selection) can follow.

Checklist

• Format with pre-commit
• Add/extend unit tests (tests/test_fp8_blockwise_moe.py now covers SM120)
• AI assistance was used (Claude) for implementation; changes reviewed and validated on real SM120 hardware by the author.

---

CI States

Latest PR Test (Base): ❌ Run #30548375199
Latest PR Test (Extra): ❌ Run #30548377705

---

Validation scope (SM121): SM121 is compile-only in this PR. Current CUDA 13 aarch64 builds emit sm_121a, and the exact grouped translation unit compiles successfully for that target, but runtime correctness is validated only on physical SM120 hardware.

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[jszzr]

Hardware validation on RTX 6000D (SM120 Blackwell, 84 GB), since this PR appeared stalled awaiting someone with silicon.

Environment: 8× RTX 6000D, driver 580.95.05, CUDA 13.0; built inside lmsysorg/sglang:nightly-dev-cu13-20260416-a4cf2ea1 with make build MAX_JOBS=64 on this branch (8ca1cc80d0).

Build note for CUDA 13: the flashmla extension's host-compiled TU fails with fatal error: cuda/std/utility: No such file because CUDA 13 moved CCCL headers under include/cccl/. Worked around with:

make build CMAKE_ARGS="-DCMAKE_CXX_FLAGS=-isystem /usr/local/cuda/include/cccl"

(unrelated to this PR — happy to file separately if useful.)

Correctness: tests/test_fp8_blockwise_moe.py — 10/10 PASSED on SM120 (0.80 s). These are exactly the 10 cases that currently die with NotImplementedError: No implemented fp8_blockwise_scaled_grouped_mm for current compute capability: 120 on main (cf. #29900). Max diff across all parametrizations ~7e-4 vs the bf16 reference.

Performance (this branch, RTX 6000D, CUDA-event timing over 50 iters after 10 warmups; shapes from the test's random per-expert m_g draws):

experts	M_total	N	K	ms/iter	TFLOP/s
8	808	4608	4608	0.050	686
16	2173	4096	4096	0.322	226
32	3987	3968	3968	1.044	120
64	8406	2048	2048	1.200	59
128	15759	5760	5760	6.571	159

No deep_gemm baseline on SM120 (unsupported upstream), so these are absolute numbers rather than a comparison.

Happy to run additional shapes/configs if the maintainers want more coverage before merging.

[waynehacking8]

Thanks - that's the silicon-side confirmation this was missing, and an RTX 6000D is a different bin from the RTX PRO 6000 I tested on, so it's better coverage than I could give it alone.

Your CUDA 13 build note checks out on my box too: 13.0.3 ships CCCL under include/cccl/ while 12.8 still has include/cuda/std/, so flashmla's host-compiled TU can't find the headers -> worth its own issue, it'll hit anyone building sgl-kernel on cu13.

The perf table is a useful baseline as well; I only claimed correctness in the PR since there's no SM120 deep_gemm to compare against.