[Perf][Qwen4Exp] Add SM100 low-latency decode GEMM plans (#59214)
Signed-off-by: namgyu-youn <namgyu.dev@gmail.com> Co-authored-by: Thien Tran <gau.nernst@yahoo.com.sg>
namgyu-youn committed
Oct 1, 2026 at 12:40 UTC
4ed438b79ceaff5fe5d6d4a98b09bfb8718cceaa
2 files changed
+99
tests/kernels/test_bf16_skinny_gemm.py
+30
@@ -85,6 +85,12 @@ QWEN4_EXP_SM90_CASES = [
85
for num_tokens, config in plans.items()
86
]
87
88
+QWEN4_EXP_SM100_CASES = [
89
+ (n, k, num_tokens, config)
90
+ for (n, k), plans in qwen4_exp_gemm.QWEN4_EXP_SM100_GEMM_PLANS.items()
91
+ for num_tokens, config in plans.items()
92
+]
93
+
94
EXPECTED_CUTE_CONFIGS = {
95
(3072, 7168, 1): (224, 3, 4, 8),
96
(3072, 7168, 2): (128, 3, 2, 8),
@@ -717,6 +723,7 @@ def test_qwen4_exp_hopper_plans_are_valid() -> None:
723
"capability,expected_plans",
724
[
725
((10, 3), qwen4_exp_gemm.QWEN4_EXP_GEMM_PLANS),
726
+ ((10, 0), qwen4_exp_gemm.QWEN4_EXP_SM100_GEMM_PLANS),
727
((9, 0), qwen4_exp_gemm.QWEN4_EXP_SM90_GEMM_PLANS),
728
((8, 0), {}),
729
],
@@ -909,6 +916,29 @@ def test_qwen4_exp_sm90_selected_shapes(
916
assert cosine > 0.999
917
918
919
+@pytest.mark.parametrize("n,k,num_tokens,config", QWEN4_EXP_SM100_CASES)
920
+def test_qwen4_exp_sm100_selected_shapes(
921
+ n: int,
922
+ k: int,
923
+ num_tokens: int,
924
+ config: SkinnyGemmConfig,
925
+) -> None:
926
+ _require_capability_and_cute((10, 0))
927
+ torch.manual_seed(42 + num_tokens)
928
+ x = torch.randn(num_tokens, k, dtype=torch.bfloat16, device="cuda")
929
+ weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda")
930
+
931
+ selected = qwen4_exp_gemm.QWEN4_EXP_SM100_GEMM_PLANS[(n, k)][num_tokens]
932
+ assert selected == config
933
+ output = qwen4_exp_gemm._qwen4_exp_low_latency_gemm(x, weight)
934
+
935
+ reference = torch.nn.functional.linear(x, weight)
936
+ cosine = torch.nn.functional.cosine_similarity(
937
+ output.float().flatten(), reference.float().flatten(), dim=0
938
+ ).item()
939
+ assert cosine > 0.999
940
+
941
+
942
def test_glm52_q_b_nonpacked_single_row_falls_back() -> None:
943
_require_sm103_and_cute()
944
spec = glm52_gemm.GLM52_Q_B_PROJECTION
vllm/models/qwen4_exp/nvidia/low_latency_gemm.py
+69
@@ -144,6 +144,73 @@ QWEN4_EXP_SM90_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] =
144
}
145
146
147
+# B200 plans; the winning configs differ from the B300 table.
148
+QWEN4_EXP_SM100_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] = {
149
+ # GDN fused B/A projection, TP=4.
150
+ (24, 2560): {
151
+ 1: SkinnyGemmConfig(1, 128, 1, k_unroll=4, vector_width=4, static_k=2560),
152
+ 2: SkinnyGemmConfig(2, 128, 2, k_unroll=4, vector_width=4, static_k=2560),
153
+ 4: SkinnyGemmConfig(4, 128, 1, k_unroll=4, vector_width=4, static_k=2560),
154
+ 8: SkinnyGemmConfig(8, 128, 1, k_unroll=4, vector_width=4, static_k=2560),
155
+ 16: SkinnyGemmConfig(16, 128, 1, k_unroll=2, vector_width=4, static_k=2560),
156
+ },
157
+ # Shared-expert fused gate/up projection, TP=4.
158
+ (320, 2560): {
159
+ 1: SkinnyGemmConfig(1, 128, 1, k_unroll=4, vector_width=4, static_k=2560),
160
+ 2: SkinnyGemmConfig(2, 128, 1, k_unroll=4, vector_width=4, static_k=2560),
161
+ 4: SkinnyGemmConfig(4, 160, 1, k_unroll=2, static_k=2560),
162
+ 8: SkinnyGemmConfig(8, 128, 1, k_unroll=4, vector_width=4, static_k=2560),
163
+ 16: SkinnyGemmConfig(16, 128, 1, vector_width=4, static_k=2560),
164
+ },
165
+ # Final HC down projection, replicated in a TP=4 deployment.
166
+ (320, 10240): {
167
+ 1: SkinnyGemmConfig(1, 256, 1, k_unroll=4, static_k=10240),
168
+ 2: SkinnyGemmConfig(2, 256, 1, k_unroll=4, static_k=10240),
169
+ 4: SkinnyGemmConfig(4, 256, 1, k_unroll=2),
170
+ 8: SkinnyGemmConfig(8, 256, 1),
171
+ },
172
+ # HC merged down/injection projection, replicated in a TP=4 deployment.
173
+ (336, 10240): {
174
+ 1: SkinnyGemmConfig(1, 256, 1, k_unroll=4, static_k=10240),
175
+ 2: SkinnyGemmConfig(2, 256, 1, k_unroll=4, static_k=10240),
176
+ 8: SkinnyGemmConfig(8, 256, 3, static_k=10240),
177
+ },
178
+ # QSA indexer Q/K projection, replicated in a TP=4 deployment.
179
+ (640, 2560): {
180
+ 1: SkinnyGemmConfig(1, 128, 2, k_unroll=4, vector_width=4, static_k=2560),
181
+ 2: SkinnyGemmConfig(2, 128, 1, k_unroll=4, vector_width=4, static_k=2560),
182
+ 4: SkinnyGemmConfig(4, 128, 1, k_unroll=4, vector_width=4, static_k=2560),
183
+ 8: SkinnyGemmConfig(8, 128, 1, k_unroll=4, vector_width=4, static_k=2560),
184
+ },
185
+ # GDN and QSA output projections, TP=4.
186
+ (2560, 1536): {
187
+ 1: SkinnyGemmConfig(1, 128, 2, k_unroll=2, vector_width=4, static_k=1536),
188
+ 2: SkinnyGemmConfig(2, 64, 2, k_unroll=2, static_k=1536),
189
+ 4: SkinnyGemmConfig(4, 64, 2, k_unroll=2, static_k=1536),
190
+ },
191
+ # QSA fused QKV/gate projection, TP=4.
192
+ (3584, 2560): {
193
+ 1: SkinnyGemmConfig(1, 128, 4, k_unroll=2, vector_width=4, static_k=2560),
194
+ 2: SkinnyGemmConfig(2, 128, 4, k_unroll=2, vector_width=4, static_k=2560),
195
+ 4: SkinnyGemmConfig(4, 64, 2, k_unroll=4, vector_width=4, static_k=2560),
196
+ },
197
+ # GDN fused QKVZ projection, TP=4.
198
+ (4096, 2560): {
199
+ 1: SkinnyGemmConfig(1, 128, 2, vector_width=4, static_k=2560),
200
+ 2: SkinnyGemmConfig(2, 64, 2, k_unroll=2, static_k=2560),
201
+ 4: SkinnyGemmConfig(4, 64, 2, k_unroll=4, vector_width=4, static_k=2560),
202
+ },
203
+ # LM head, TP=4.
204
+ (62080, 2560): {
205
+ 1: SkinnyGemmConfig(1, 128, 2, k_unroll=4, vector_width=4),
206
+ },
207
+}
208
+
209
+
210
+def _is_sm100() -> bool:
211
+ return current_platform.is_device_capability((10, 0))
212
+
213
+
214
def _is_sm103() -> bool:
215
return current_platform.is_device_capability((10, 3))
216
@@ -155,6 +222,8 @@ def _is_sm90() -> bool:
222
def _gemm_plans() -> dict[tuple[int, int], dict[int, SkinnyGemmConfig]]:
223
if _is_sm103():
224
return QWEN4_EXP_GEMM_PLANS
225
+ if _is_sm100():
226
+ return QWEN4_EXP_SM100_GEMM_PLANS
227
if _is_sm90():
228
return QWEN4_EXP_SM90_GEMM_PLANS
229
return {}