CUTLASS NVFP4 GEMM improvement of SM120

[b8zhong]

Motivation

Profile the NVFP4 x NVFP4 with CUTLASS profiler for SM120 family, I tried the exhaustive combination of various tile size, cooperative vs pingpong, StreamK (which didn't really make a difference for the best configuration). Thus, update some heuristics. For example, the speedup of M=16, N=6144, K=5120 → 1.197x (~20% speedup).

Modifications

• Restruct the code to seperate SM100 and SM120 files, since they have different features (SM120 doesn't support various feature like multicast, certain tile sizes, etc.) We can tune it seperately and seperate it since future architectures will leverage FP4 more.

Benchmarking and Profiling

Thus, we could almost beat cuDNN.

As a followup, we should tune larger M. I only tuned M <= 128, because CUTLASS profiler took a really long time.

python -m sglang.launch_server   --model nvidia/Qwen3-32B-NVFP4   --reasoning-parser qwen3   --tool-call-parser qwen25 --quantization modelopt_fp4 --disable-piecewise-cuda-graph

SGLANG_TORCH_PROFILER_DIR="./" \
SGLANG_PROFILE_RECORD_SHAPES=true \
SGLANG_PROFILE_WITH_STACK=true \
python3 -m sglang.bench_one_batch_server \
  --model baseten-admin/glm-4.7-fp8-attn-fp4-mlp \
  --base-url http://localhost:30000 \
  --batch-size 16 \
  --input-len 1024 \
  --output-len 1024 \
  --profile \
  --profile-steps 10 \
  --show-report \
  --profile-by-stage

For a comparison of the strategy that cuDNN vs CUTLASS pick:

CUTLASS:

_ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalIN4cute5tupleIJiiiiEEENS1_10collective13CollectiveMmaINS1_42MainloopSm120TmaWarpSpecializedBlockScaledILi2ELi3ENS5_IJNS4_1CILi1EEESB_SB_EEENS1_51KernelTmaWarpSpecializedCooperativeBlockScaledSm120ILi3EEEEENS5_IJNSA_ILi128EEESG_NSA_ILi256EEEEEENS5_IJNS_12float_e2m1_tENS_13float_ue4m3_tEEEENS5_IJNS5_IJlSB_lEEENS4_6LayoutINS5_IJNS5_IJNS5_IJNSA_ILi32EEENSA_ILi4EEEEEEiEEENS5_IJNS5_IJNSA_ILi16EEESP_EEEiEEENS5_IJSB_iEEEEEENS5_IJSU_NS5_IJNS5_IJNSA_ILi0EEESB_EEENSA_ILi512EEEEEENS5_IJSX_iEEEEEEEEEEESL_S14_NS4_8TiledMMAINS4_8MMA_AtomIJNS4_5SM12011BLOCKSCALED19SM120_16x8x64_TN_VSISJ_SJ_fSK_Li16EEEEEENSN_INS5_IJSP_NSA_ILi2EEESB_EEENS5_IJSB_SP_SX_EEEEENS5_IJSG_NSN_INS5_IJNSA_ILi8EEES1C_S1C_EEENS5_IJSB_SS_S1G_EEEEENSA_ILi64EEEEEEEENS5_IJNS4_13SM90_TMA_LOADES1N_EEENS5_IJNS4_14ComposedLayoutINS4_7SwizzleILi3ELi4ELi3EEENS4_18smem_ptr_flag_bitsILi4EEENSN_INS5_IJS1G_SH_EEENS5_IJSH_SB_EEEEEEENSN_INS5_IJNS5_IJSQ_SB_EEENS5_IJST_SB_SP_EEEEEENS5_IJNS5_IJST_SZ_EEENS5_IJSY_SP_SZ_EEEEEEEEEEENS5_IJNS4_9Copy_AtomIJNS4_17SM75_U32x4_LDSM_NENS_15integer_subbyteILi4ELb0EEEEEENS26_IJNS4_13UniversalCopyISK_SK_EESK_EEEEEENS4_8identityES1O_S25_S2E_S2F_EENS_8epilogue10collective18CollectiveEpilogueINS2H_22Sm90TmaWarpSpecializedILi3ELi2ELi4ELb1ELb0EEEJSI_NS5_IJS1K_SO_EEENS_10bfloat16_tESM_S2N_SM_NS2H_6fusion15FusionCallbacksINS2H_23Sm120TmaWarpSpecializedILi3ELi2ELi4ELb1ELb0EEENS2O_17LinearCombinationIS2N_fS2N_fLNS_15FloatRoundStyleE2EEESI_S2M_JEEES1N_NS1P_INS1Q_ILi2ELi4ELi3EEENS1S_ILi16EEENSN_INS5_IJS1G_SO_EEENS5_IJSO_SB_EEEEEEENS4_17SM75_U32x2_LDSM_NENS4_14SM90_TMA_STOREES31_NS4_17SM90_U32x2_STSM_NENS26_IJS34_NS_6half_tEEEEvEEEvvEEEEvNT_6ParamsE -> NS5_IJNSA_ILi128EEESG_NSA_ILi256EEEE (so 128x128x256)

Accuracy:

python -m sglang.test.run_eval --base-url http://localhost:30000/ --eval-name gsm8k --num-examples 200 --max-tokens 16000 --repeat 5 --num-threads 100 --num-shots 5 --temperature 0.6 --top-p 0.95 --top-k 20 --min-p 0.0 --chat-template-kwargs '{"enable_thinking": true}'

Repeat: 5, mean: 0.969
Scores: ['0.970', '0.960', '0.975', '0.975', '0.965']

CuDNN:

cutlass3x_sm120_bstensorop_s16864gemm_block_scaled_ue4m3xe2m1_ue4m3xe2m1_f32_bf16_bf16_128x128x256_1x1x1_0_tnn_align32_o_vs16_bias_bf16_relu_stream_k (so now aligned on the strategy).

Accuracy:

Repeat: 5, mean: 0.976█████████████████████████████████████████████████████████████████████████████████████████████████▉                                                                                     | 115/200 [07:23<04:52,  3.44s/it]
Scores: ['0.970', '0.980', '0.970', '0.975', '0.985']

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the performance of NVFP4 GEMM operations on SM120 GPUs by introducing architecture-specific optimizations and restructuring the underlying kernel code. It also integrates a new CUTLASS backend, offering greater flexibility and potentially better performance for FP4 GEMM workloads, particularly for different matrix dimensions.

Highlights

• Performance Improvement for NVFP4 GEMM on SM120: Optimized NVFP4 GEMM operations for SM120 GPU architecture, achieving a reported 20% speedup for specific matrix sizes (e.g., M=16, N=6144, K=5120) through exhaustive tile size profiling and heuristic updates.
• Code Restructuring for SM100 and SM120 Architectures: Refactored the NVFP4 GEMM kernel code to separate implementations for SM100 and SM120 architectures, allowing for independent tuning and optimization based on their distinct features and capabilities.
• New CUTLASS Backend for FP4 GEMM: Introduced a dedicated 'cutlass' backend option for FP4 GEMM operations, providing an alternative to existing FlashInfer backends and enabling direct utilization of CUTLASS optimizations.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

The pull request significantly refactors the NVFP4 scaled matrix multiplication (GEMM) kernels by modularizing the CUDA C++ code into common utilities, SM100-specific, and SM120-specific implementations, including a new memory caching mechanism for CUDA workspaces. The Python codebase is updated to support a new 'cutlass' backend option for FP4 GEMM, which involves handling specific scale factor data types and adjusting weight transposition logic to align with the new CUTLASS JIT kernel's input requirements. Additionally, the import logic for the CUTLASS GEMM kernel is made more robust, allowing its use independently of FlashInfer.

[b8zhong]

/rerun-failed-ci j

[b8zhong]

@HydraQYH @BBuf, By the way, I just update the PR with the E2E acc, and the profiles to confirm the change.

[BBuf]

Suggested change

	if fp4_backend.is_cutlass():
	if fp4_backend.is_cutlass() and cutlass_fp4_gemm is not None:

[b8zhong]

Applied.

[BBuf]

The CSV schema looks inconsistent here. We now write [provider, M, N, K, ms, bandwidth_gbs] below, but the header still only has 5 columns. This will misalign downstream parsing.

[b8zhong]

Good point. I fix it w: writer.writerow(["provider", "m", "n", "k", "time_ms", "bandwidth_gbs"])

[BBuf]

sglang/python/sglang/srt/server_args.py

Line 4615 in 3867259

parser.add_argument(

should also be updated?

[HydraQYH]

@DarkSharpness It appears that the workspace required by the CUTLASS Kernel is requested by itself via cudaMallocAsync.
@b8zhong As far as I know, we may need to avoid such operations as much as possible; memory allocation should be done through PyTorch.

[DarkSharpness]

I would recommend allocating the workspace from Python Side. As an alternative, you may also try some utility functions here for C++ allocation https://github.com/sgl-project/sglang/blob/f5c225eeba3c6f027e5e7691784249d6a667026f/python/sglang/jit_kernel/include/sgl_kernel/ffi.h

[b8zhong]

I see. Could I use ffi::empty? Because, otherwise we need to bring the workspace size up to Python level (it could be useful when doing autotuning infra I think but is a bit more refactor required), this way it should still be using Torch allocator underneath.

[b8zhong]

@HydraQYH @DarkSharpness I do it in https://github.com/sgl-project/sglang/pull/21314/changes/2a354896f510e08658007523153c7fdfc9daf7e4..b38f8c60048eb0297be5336ebea79a48d115b68e. Let me know if it looks alright.

[HydraQYH]

Have you tried the Swap A/B approach? If Swap A/B works, I believe it can also improve performance on the Sm100.

[b8zhong]

Yes, I also agree. By the way, I'm only familiar with it's usage on SM90 with respect to WGMMA efficiency, but I assume it may or may not be compatible. I'll give it a try with CUTLASS profiler to be. a bit more comprehensive and add it to the generator, but is it OK to do it in the next PR (I plan to tune for large problem size too then) (ETA early next month), my current access to SM120 is a bit sporadic.

[b8zhong]

Latest UT result locally. I'm not sure if there is a 5090 machine running for jit-kernel tests (I don't think so).

[BBuf]

@HydraQYH @DarkSharpness Is it ready now? The author said they’ll try the Swap A/B optimization in the next PR, and I think that’s reasonable too.