perf(gemm): optimize SM103 CuTe DSL mm_fp4 kernel

[ashxudev]

📌 Description

In Issue #2621 and its related PR #2888, changes from the upstream CUTLASS CuTe DSL SM103 mm_fp4 kernel were ported to flashinfer. However, the ported SM103 mm_fp4 kernel underperformed the SM100 kernel across all shapes when both were run on a B300. This PR does the following to fix performance:

1. Carry over changes from PR#2888

This carries over the SM103 kernel changes from PR#2888 excluding the alpha scaling sections which are dealt with below.

2. Alpha epilogue bug

The PR#2888 alpha epilogue unnecessarily cast the FP32 accumulator to BF16 before scaling, promoted it back to FP32 for multiplication, then cast it to BF16 again.

This extra type conversion was fixed by PR #4526, which added dedicated TMA and non-TMA helpers in epilogue_utils.py that apply alpha directly to the FP32 accumulator before a single conversion to the output dtype.

This PR makes a small edit to epilogue_utils.py by replacing cute.nvgpu.CopyStgOp() with the supported explicit register-to-global operation cute.nvgpu.CopyR2GOp(). The CopyStgOp does not appear to be exposed by the supported CuTe DSL wheels.

3. Restore 256-bit stores in the N-major non-TMA epilogue

Upstream CUTLASS gives CuTe explicit compile-time guarantees that C is compact N-major, N is divisible by 64 and its pointer is 32-byte aligned. FlashInfer’s TVM-FFI wrapper rebuilt C as a generic tensor and discarded those guarantees, so CuTe could prove only 16-bit BF16 stores were safe.

This PR reconstructs the same layout, divisibility and alignment contract, guarded by can_implement(), allowing the intended 256-bit vectorized stores.

Performance

SM103 now wins five of the six largest shapes, with up to 14.8% higher throughput. Measured on an NVIDIA B300 SXM6 using independently exhaustive tactic selection, CUPTI CUDA-Graph timing and cold-L2 conditions.

M	N	K	SM100 ms	SM100 TFLOPS	SM103 ms	SM103 TFLOPS	Speedup
1	4096	7168	0.0100	5.9	0.0115	5.1	0.87x
4	4096	7168	0.0099	23.8	0.0116	20.3	0.85x
8	4096	7168	0.0096	48.9	0.0115	41.0	0.84x
16	4096	7168	0.0097	97.2	0.0114	82.7	0.85x
32	4096	7168	0.0097	194.4	0.0114	165.4	0.85x
64	4096	7168	0.0094	399.5	0.0114	329.4	0.82x
128	4096	7168	0.0096	786.6	0.0115	653.3	0.83x
128	7168	2048	0.0061	620.3	0.0071	531.4	0.86x
256	4096	7168	0.0091	1645.3	0.0115	1304.9	0.79x
256	14336	4096	0.0129	2322.7	0.0126	2378.4	1.02x
512	14336	4096	0.0174	3450.8	0.0177	3388.7	0.98x
1024	4096	7168	0.0149	4032.3	0.0134	4479.3	1.11x
2048	4096	7168	0.0237	5078.5	0.0215	5592.4	1.10x
4096	4096	7168	0.0421	5707.1	0.0367	6549.9	1.15x
4096	4096	4096	0.0274	5008.7	0.0251	5474.7	1.09x

🔍 Related Issues

• Addresses issue Enable sm103 cute_dsl mm_fp4 kernel #2621
• Ports the SM103 kernel improvements from PR [WIP][Do not review] feat: enable sm103 fp4 gemm #2888

🚀 Pull Request Checklist

✅ Pre-commit Checks

• I have installed pre-commit by running pip install pre-commit (or used your preferred method).
• I have installed the hooks with pre-commit install.
• I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

🧪 Tests

• Tests have been added or updated as needed.
• All tests are passing (unittest, etc.).
• Ran tests/gemm/test_mm_fp4_sm103.py and tests/gemm/test_mm_fp4.py on a B300.

Reviewer Notes

• the SM103 kernel should complement but not replace SM100 kernel when running on B300s/SM103 hardware. SM100 kernel remains substantially faster at small M, likely because it allows for narrower output tiles.
• production enablement of the mm_fp4 tactic selection path is deliberately left to be handled separately

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}