[SM90] Support MXFP4A8 Group Gemm for the SM90 mixed-input.

[yuyu5333]

Since the SM90 devices gained support for the int4 * fp8 computation pipeline, it has been widely adopted in inference frameworks such as SGLang and vLLM and delivered remarkable performance.

However, with the evolution of large‑language models (LLMs), an increasing number of models employ higher‑precision parameters like MXFP4. Unfortunately, a large number of SM90 devices lack native support for this data type. Similar to the earlier int4 scenario, runtime dequantization is required for adaptation. Most existing schemes dequantize weights to BF16 for inference, which covers the vast majority of use cases. Nevertheless, this is clearly a fallback solution for SM90‑class devices that natively support FP8 GEMM operations.

Accordingly, I have implemented the MXFP4‑to‑FP8 GEMM pipeline. Building upon existing work on WINT4A8, I realized dynamic scale retrieval and the corresponding GEMM computation logic, and removed the traditional fixed‑block‑size‑128 limitation of INT4A8. The implementation has been validated on SGLang. Results show that WMXFP4A8 outperforms existing pipelines in most scenarios.

Summary

This PR adds MXFP4A8 activation block scaling support to the SM90 mixed-input ptr-array mainloop while preserving the existing int4a8 execution path.

The implementation gates the new activation-scale TMA, shared memory, and chunked accumulation logic behind the MXFP4A8-specific EnableActBlockScale condition. Non-MXFP4A8 paths, including int4a8, keep the original mainloop behavior.

Changes

• Added activation block scale TMA descriptor, shared storage, and load path for MXFP4A8.
• Added chunked accumulation with activation scale application for MXFP4A8.
• Preserved the original mma() mainloop for non-act-block-scale paths.
• Added E2M1 to E4M3 numeric conversion coverage.
• Added static tests for int4a8 and MXFP4A8 path separation.

Validation

Tested on machine H20 / SM90, CUDA 12.4.131.

Unit tests

./test/unit/core/cutlass_test_unit_core \
  --gtest_filter="*Mxfp4*:*MXFP4*:*mxfp4*"

Result:

[  PASSED  ] 2 tests.

./test/unit/gemm/device/cutlass_test_unit_gemm_device_tensorop_gmma_rs_warpspecialized_sm90 \
  --gtest_filter="SM90_Device_Gemm_MixedInput_RS.Int4a8AndMxfp4a8StaticProperties"

Result:

[  PASSED  ] 1 test.

int4a8 benchmark

Benchmark target:

examples/69_hopper_mixed_dtype_grouped_gemm/69_hopper_int4_fp8_grouped_gemm

Small shape:

--m=512 --n=2048 --k=2048 --groups=8 --c=512 --mode=1 --warmup=3 --iterations=20 --alpha=1 --beta=0

Version	Avg runtime
baseline	0.317806 ms
patched	0.317676 ms

Delta: -0.041%

Large shape:

--m=2048 --n=5120 --k=8192 --groups=16 --c=512 --mode=1 --warmup=5 --iterations=50 --alpha=1 --beta=0

Version	Avg runtime
baseline	70.65475 ms
patched	70.65260 ms

Delta: -0.003%

Both benchmark runs passed correctness checks. int4a8 performance is unchanged within measurement noise.

int4a8 MNK coverage

The existing CUTLASS int4/fp8 grouped GEMM example was also run on a small shape, a large-group shape, and a large prefill-like shape:

examples/69_hopper_mixed_dtype_grouped_gemm/69_hopper_int4_fp8_grouped_gemm

Shape	Groups	Avg runtime	GFLOPS	Result
m=256, n=128, k=512	6	0.016672 ms	12075.7	Passed
m=128, n=128, k=512	100	0.043360 ms	38692.8	Passed
m=2048, n=5120, k=8192	16	69.4398 ms	39585.0	Passed

This supplements the static int4a8 path-separation test with runtime MNK coverage from the existing CUTLASS example.

MXFP4A8 correctness

MXFP4A8 correctness was validated through a temporary SGLang cutlass_mxfp4a8_moe_mm single-GEMM test on the same H20 machine. The test compares the kernel output against a BF16 dequantized golden reference and exercises:

• activation scale = one identity path;
• full MXFP8 activation path with per-token, per-block activation scales;
• single-expert shapes: m=4/8/16/128, k=256/512/1024, n=512/1024/2048;
• multi-expert grouped GEMM with uneven token counts: [4,4,4,4], [3,5,4,8], [16,8,32,8];
• uniform multi-expert counts from 2 to 32 tokens per expert.

Observed relative mean error was 0.0000 for all tested cases.

MXFP4A8 microbenchmark

The following single-GEMM MXFP4A8 benchmark uses sgl_kernel.cutlass_mxfp4a8_moe_mm with full activation block scaling enabled. For this single-GEMM synthetic benchmark, topk=1 is the clean kernel-level setting; end-to-end MoE top-k behavior is covered separately by full MoE benchmarks.

Shape	Avg runtime	P50 runtime	GFLOPS
m=256, n=128, k=512	0.016179 ms	0.016320 ms	2074.01
m=128, n=128, k=512	0.015639 ms	0.015648 ms	1072.78
m=1024, n=2048, k=4096	0.164708 ms	0.164752 ms	104304.70
m=2048, n=5120, k=8192	1.023284 ms	1.024608 ms	167889.58
m=4096, n=5120, k=8192	1.975428 ms	1.978080 ms	173935.68

For reference, using topk=6/8 in the same single-expert synthetic setup adds non-representative scheduling overhead and should not be used as the single-GEMM microbenchmark number.

[yuyu5333]

@Junkai-Wu @IonThruster @Algy @hwu36 Sorry to bother you. Could you please review this when you have time?