[WhoisZihan]

Motivation

CUTLASS and xFormers already provides DualGEMM for sm80, which fuses the swiglu elementwise computation into the gate/up matrix multiplication. But they need rewriting and tuning on sm90+. Recently there is dualgemm PR #41887 for Blackwell using CuTeDSL, so the pr sticks with cutedsl sm90 dualgemm instead of cutlass in the end.

Modifications

Add an SM90 CuTe DSL kernel that fuses gate/up projections with SwiGLU for BF16 and FP16 decode batches of up to 16 tokens. It uses TMA loads, warp specialization and WGMMA, with two configurations and dispatch calibrated on H100 PCIe.
Includes tests for numerical accuracy, pipeline reuse, input validation, CUDA Graph replay and torch.compile. The kernel is exposed as a standalone operator. Model integration will be added later.

Accuracy Tests

No end-to-end test yet, while the unit tests are available which compares the actual output with reference

• actual: the sm90 dual gemm
• reference: fp64 gate/up matmul + fp32 silu and mul, with element-wise allclose, and rms check.

Speed Tests and Profiling

Currently there are 2 sets of tactics T1/T2 available with auto dispatch, and I've also tested another two sets of tactics T3/T4, but found them worse than T1/T2, so they are abandoned for now.

Configuration	T1	T2	T3 (not used)	T4 (not used)
Output features / CTA	64	64	128	128
Token tile (M ≤ 8 / M > 8)	8 / 16	8 / 16	8 / 16	8 / 16
K tile	128	64	128	64
Pipeline stages	3	5	3	3
Shared memory / CTA	102–108 KiB	85–90 KiB	198–204 KiB	99–102 KiB
CTA / SM limit from shared memory	2	2	1	2

BF16/FP16 speedup

num_tokens	K	I	dtype	Baseline (µs)	Auto (µs)	Speedup	Auto tactic	T1 (µs)	T2 (µs)
1	2048	7168	bf16	35.286	33.109	1.0658×	2	34.596	33.136
8	2048	7168	bf16	36.077	33.493	1.0771×	2	35.104	33.468
16	2048	7168	bf16	36.615	34.188	1.0712×	2	35.544	34.161
1	2048	7168	fp16	35.272	33.123	1.0649×	2	34.603	33.147
8	2048	7168	fp16	36.045	33.458	1.0773×	2	35.108	33.458
16	2048	7168	fp16	36.598	34.156	1.0716×	2	35.490	34.118
1	4096	3584	bf16	35.516	33.481	1.0603×	2	33.802	33.516
8	4096	3584	bf16	36.458	33.886	1.0759×	2	34.473	33.919
16	4096	3584	bf16	36.763	34.504	1.0654×	2	34.861	34.464
1	4096	3584	fp16	35.514	33.500	1.0601×	2	33.831	33.538
8	4096	3584	fp16	36.433	33.859	1.0759×	2	34.449	33.888
16	4096	3584	fp16	36.723	34.425	1.0671×	2	34.820	34.433
1	4096	7168	bf16	65.558	63.277	1.0362×	2	64.867	63.294
8	4096	7168	bf16	66.750	63.782	1.0465×	2	65.346	63.809
16	4096	7168	bf16	67.410	64.266	1.0488×	2	65.921	64.263
1	4096	7168	fp16	68.021	63.241	1.0759×	2	64.844	63.253
8	4096	7168	fp16	66.765	63.787	1.0467×	2	65.342	63.832
16	4096	7168	fp16	67.406	64.214	1.0494×	2	65.914	64.233
1	4096	14336	bf16	133.116	124.051	1.0731×	1	124.064	125.833
8	4096	14336	bf16	134.696	125.051	1.0770×	1	125.068	126.571
16	4096	14336	bf16	135.100	126.034	1.0717×	1	126.017	128.094
1	4096	14336	fp16	133.158	124.043	1.0735×	1	124.053	125.833
8	4096	14336	fp16	134.598	125.025	1.0766×	1	125.006	126.515
16	4096	14336	fp16	135.224	125.990	1.0733×	1	125.976	128.024
1	4096	18944	bf16	175.222	163.517	1.0717×	1	163.482	165.632
8	4096	18944	bf16	176.437	164.906	1.0698×	1	164.963	167.143
16	4096	18944	bf16	175.843	166.414	1.0567×	1	166.288	168.552
1	4096	18944	fp16	175.207	163.466	1.0718×	1	163.480	165.605
8	4096	18944	fp16	176.331	164.854	1.0696×	1	164.950	167.132
16	4096	18944	fp16	176.873	166.320	1.0633×	1	166.258	168.553
1	8192	14336	bf16	248.376	242.807	1.0229×	1	242.843	246.743
8	8192	14336	bf16	251.302	243.730	1.0311×	1	243.728	247.697
16	8192	14336	bf16	260.115	245.066	1.0614×	1	245.103	249.039
1	8192	14336	fp16	248.377	242.807	1.0231×	1	242.824	246.712
8	8192	14336	fp16	251.169	243.845	1.0301×	1	243.874	248.309
16	8192	14336	fp16	260.036	245.031	1.0612×	1	244.975	249.147

Follow Up

1. More tactics on more Hopper cards
   The statistics are run on H100-PCIe(sm90a) with only two sets of tactics, due to limited quota of cards. The percentage might vary a little bit on H100-SXM or H200. I'll consider adding more statistics in the future if possible.

2. FP8 support

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ❌ Run #38238620915
Latest PR Test (Extra): ❌ Run #38238620755
Latest PR Test (AMD ROCm 10): ❌ Run #38238620864

[github-actions]

Thanks @WhoisZihan. Closing this because it is still a draft and has not been updated in 24 days.

Reopen it if the work is still relevant.

Some directories moved recently, so an older branch may need retargeting:
sgl-kernel/ -> python/sglang/kernels/aot/, python/sglang/jit_kernel/
-> python/sglang/kernels/jit/, docs/ -> docs/docs/ (.mdx),
bench_serving.py -> benchmark/serving.py, test/srt/ -> test/registered/.