[Kernel][Qwen] Add Hopper LL-GEMM tuning table for Qwen4Exp (#54560)
Signed-off-by: Zheng Cai <8370601+zigzagcai@users.noreply.github.com>
Zheng Cai committed
Sep 1, 2026 at 19:21 UTC
191cecd51e250de77618dda924359214b4d88119
2 files changed
+151
-6
tests/kernels/test_bf16_skinny_gemm.py
+69
-2
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Tests for BF16 skinny GEMMs and the Kimi-K3 SM90/SM100/SM103 selectors."""
3
+"""Tests for BF16 skinny GEMMs and model-specific selectors."""
4
5
from pathlib import Path
6
from types import SimpleNamespace
@@ -17,6 +17,7 @@ from vllm.model_executor.kernels.linear.cute_dsl.skinny_gemm import (
17
from vllm.models.deepseek_v32.nvidia import glm52_low_latency_gemm as glm52_gemm
18
from vllm.models.kimi_k3.nvidia import low_latency_gemm as k3_gemm
19
from vllm.models.kimi_k3.nvidia.low_latency_gemm import KIMI_K3_PROJECTIONS
20
+from vllm.models.qwen4_exp.nvidia import low_latency_gemm as qwen4_exp_gemm
21
22
# Keyed by local (N, K): (cute token counts, dsv3 token counts). 1536x7168 is
23
# the unified shared_gate_up_proj/mla_g_proj entry (dsv3 M1..16).
@@ -77,6 +78,12 @@ GLM_CUTE_CASES = [
78
for _, config in spec.cute_configs
79
]
80
81
+QWEN4_EXP_SM90_CASES = [
82
+ (n, k, num_tokens, config)
83
+ for (n, k), plans in qwen4_exp_gemm.QWEN4_EXP_SM90_GEMM_PLANS.items()
84
+ for num_tokens, config in plans.items()
85
+]
86
+
87
EXPECTED_CUTE_CONFIGS = {
88
(3072, 7168, 1): (224, 3, 4, 8),
89
(3072, 7168, 2): (128, 3, 2, 8),
@@ -528,6 +535,43 @@ def test_low_latency_table_capability_routing(
535
assert k3_gemm._low_latency_table() is None
536
537
538
+def test_qwen4_exp_hopper_plans_are_valid() -> None:
539
+ plans = qwen4_exp_gemm.QWEN4_EXP_SM90_GEMM_PLANS
540
+
541
+ assert len(plans) == 9
542
+ assert sum(map(len, plans.values())) == 31
543
+ assert (320, 10240) in plans
544
+ assert (10240, 320) not in plans
545
+ for (n, k), shape_plans in plans.items():
546
+ for num_tokens, config in shape_plans.items():
547
+ assert config.num_rows == num_tokens
548
+ assert n % config.outputs_per_block == 0
549
+ assert k % (config.block_size * config.vector_width) == 0
550
+ assert config.static_k in (None, k)
551
+
552
+
553
+@pytest.mark.parametrize(
554
+ "capability,expected_plans",
555
+ [
556
+ ((10, 3), qwen4_exp_gemm.QWEN4_EXP_GEMM_PLANS),
557
+ ((9, 0), qwen4_exp_gemm.QWEN4_EXP_SM90_GEMM_PLANS),
558
+ ((8, 0), {}),
559
+ ],
560
+)
561
+def test_qwen4_exp_gemm_capability_routing(
562
+ monkeypatch: pytest.MonkeyPatch,
563
+ capability: tuple[int, int],
564
+ expected_plans: dict[tuple[int, int], dict[int, SkinnyGemmConfig]],
565
+) -> None:
566
+ monkeypatch.setattr(
567
+ qwen4_exp_gemm.current_platform,
568
+ "is_device_capability",
569
+ lambda target: capability == target,
570
+ )
571
+
572
+ assert qwen4_exp_gemm._gemm_plans() == expected_plans
573
+
574
+
575
def test_installation_is_shape_specific_and_unquantized(
576
monkeypatch: pytest.MonkeyPatch,
577
) -> None:
@@ -643,7 +687,7 @@ def _require_capability_and_cute(capability: tuple[int, int]) -> None:
687
not torch.cuda.is_available()
688
or torch.cuda.get_device_capability() != capability
689
):
646
- pytest.skip(f"Kimi-K3 selection requires SM{capability[0]}{capability[1]}")
690
+ pytest.skip(f"CuTe DSL selection requires SM{capability[0]}{capability[1]}")
691
if not k3_gemm.shape_dynamic_skinny_gemm.is_available():
692
pytest.skip("CuTe DSL is not available")
693
@@ -679,6 +723,29 @@ def test_glm_cute_selected_shapes(
723
torch.testing.assert_close(output.float(), reference, rtol=2e-2, atol=2e-1)
724
725
726
+@pytest.mark.parametrize("n,k,num_tokens,config", QWEN4_EXP_SM90_CASES)
727
+def test_qwen4_exp_sm90_selected_shapes(
728
+ n: int,
729
+ k: int,
730
+ num_tokens: int,
731
+ config: SkinnyGemmConfig,
732
+) -> None:
733
+ _require_capability_and_cute((9, 0))
734
+ torch.manual_seed(42 + num_tokens)
735
+ x = torch.randn(num_tokens, k, dtype=torch.bfloat16, device="cuda")
736
+ weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda")
737
+
738
+ selected = qwen4_exp_gemm.QWEN4_EXP_SM90_GEMM_PLANS[(n, k)][num_tokens]
739
+ assert selected == config
740
+ output = qwen4_exp_gemm._qwen4_exp_low_latency_gemm(x, weight)
741
+
742
+ reference = torch.nn.functional.linear(x, weight)
743
+ cosine = torch.nn.functional.cosine_similarity(
744
+ output.float().flatten(), reference.float().flatten(), dim=0
745
+ ).item()
746
+ assert cosine > 0.999
747
+
748
+
749
def test_glm52_q_b_nonpacked_single_row_falls_back() -> None:
750
_require_sm103_and_cute()
751
spec = glm52_gemm.GLM52_Q_B_PROJECTION
vllm/models/qwen4_exp/nvidia/low_latency_gemm.py
+82
-4
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Qwen4Exp decode GEMM selection on Blackwell.
3
+"""Qwen4Exp decode GEMM selection on Hopper and Blackwell.
4
5
Dispatch follows Kimi-K3 and uses the local ``(N, K)`` shape and token count.
6
Plans contain measured CUDA graph capture sizes; other token counts use the
@@ -78,11 +78,88 @@ QWEN4_EXP_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] = {
78
},
79
}
80
81
+# H200 plans selected by exhaustive CUDA graph replay measurements over
82
+# M={1, 2, 4, 8, 16}. Only points that beat the standard linear implementation
83
+# in both hot-cache and L2-flush measurements are retained; other token counts
84
+# keep the standard implementation and its GEMM heuristics.
85
+QWEN4_EXP_SM90_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] = {
86
+ # GDN fused QKVZ projection, TP=4.
87
+ (4096, 2560): {
88
+ 1: SkinnyGemmConfig(1, 128, 2, vector_width=4, static_k=2560),
89
+ 2: SkinnyGemmConfig(2, 64, 4, vector_width=4, static_k=2560),
90
+ },
91
+ # GDN and QSA output projections, TP=4.
92
+ (2560, 1536): {
93
+ 1: SkinnyGemmConfig(1, 128, 4, vector_width=2, static_k=1536),
94
+ 2: SkinnyGemmConfig(2, 128, 4, vector_width=4, static_k=1536),
95
+ 4: SkinnyGemmConfig(4, 64, 4, k_unroll=6, vector_width=4),
96
+ },
97
+ # GDN fused B/A projection, TP=4.
98
+ (24, 2560): {
99
+ 1: SkinnyGemmConfig(1, 128, 3, vector_width=4, static_k=2560),
100
+ 2: SkinnyGemmConfig(2, 64, 2, vector_width=4, static_k=2560),
101
+ 4: SkinnyGemmConfig(4, 64, 1, static_k=2560),
102
+ 8: SkinnyGemmConfig(8, 128, 1, vector_width=4, static_k=2560),
103
+ 16: SkinnyGemmConfig(16, 128, 1, vector_width=4, static_k=2560),
104
+ },
105
+ # QSA fused QKV/gate projection, TP=4.
106
+ (3584, 2560): {
107
+ 1: SkinnyGemmConfig(1, 128, 4, vector_width=2, static_k=2560),
108
+ 2: SkinnyGemmConfig(2, 64, 4, k_unroll=5),
109
+ },
110
+ # QSA indexer Q/K projection, replicated in a TP=4 deployment.
111
+ (640, 2560): {
112
+ 1: SkinnyGemmConfig(1, 256, 2, vector_width=2, static_k=2560),
113
+ 2: SkinnyGemmConfig(2, 128, 2, vector_width=4, static_k=2560),
114
+ 4: SkinnyGemmConfig(4, 128, 1, vector_width=2, static_k=2560),
115
+ 8: SkinnyGemmConfig(8, 128, 2, vector_width=4, static_k=2560),
116
+ },
117
+ # Shared-expert fused gate/up projection, TP=4.
118
+ (320, 2560): {
119
+ 1: SkinnyGemmConfig(1, 64, 2, vector_width=4, static_k=2560),
120
+ 2: SkinnyGemmConfig(2, 128, 4, k_unroll=5, vector_width=4),
121
+ 4: SkinnyGemmConfig(4, 160, 1, k_unroll=2),
122
+ 8: SkinnyGemmConfig(8, 128, 1, vector_width=4, static_k=2560),
123
+ 16: SkinnyGemmConfig(16, 128, 1, vector_width=4, static_k=2560),
124
+ },
125
+ # LM head, TP=4.
126
+ (62080, 2560): {
127
+ 1: SkinnyGemmConfig(1, 64, 2, vector_width=2, static_k=2560),
128
+ 2: SkinnyGemmConfig(2, 64, 2, vector_width=2, static_k=2560),
129
+ },
130
+ # HC merged down/injection projection, replicated in a TP=4 deployment.
131
+ (336, 10240): {
132
+ 1: SkinnyGemmConfig(1, 256, 1, k_unroll=5),
133
+ 2: SkinnyGemmConfig(2, 256, 3, static_k=10240),
134
+ 4: SkinnyGemmConfig(4, 256, 3, static_k=10240),
135
+ 8: SkinnyGemmConfig(8, 256, 3, static_k=10240),
136
+ },
137
+ # Final HC down projection, replicated in a TP=4 deployment.
138
+ (320, 10240): {
139
+ 1: SkinnyGemmConfig(1, 256, 1, static_k=10240),
140
+ 2: SkinnyGemmConfig(2, 128, 1, k_unroll=10),
141
+ 4: SkinnyGemmConfig(4, 128, 1, k_unroll=10),
142
+ 8: SkinnyGemmConfig(8, 128, 1, k_unroll=10),
143
+ },
144
+}
145
+
146
147
def _is_sm103() -> bool:
148
return current_platform.is_device_capability((10, 3))
149
150
151
+def _is_sm90() -> bool:
152
+ return current_platform.is_device_capability((9, 0))
153
+
154
+
155
+def _gemm_plans() -> dict[tuple[int, int], dict[int, SkinnyGemmConfig]]:
156
+ if _is_sm103():
157
+ return QWEN4_EXP_GEMM_PLANS
158
+ if _is_sm90():
159
+ return QWEN4_EXP_SM90_GEMM_PLANS
160
+ return {}
161
+
162
+
163
def _is_packed_row_major(tensor: torch.Tensor) -> bool:
164
return tensor.dim() == 2 and tensor.stride() == (tensor.shape[1], 1)
165
@@ -124,7 +201,7 @@ class Qwen4ExpLowLatencyEmbeddingMethod(
201
202
203
def _qwen4_exp_low_latency_gemm(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
127
- plan = QWEN4_EXP_GEMM_PLANS.get((weight.shape[0], weight.shape[1]))
204
+ plan = _gemm_plans().get((weight.shape[0], weight.shape[1]))
205
config = None if plan is None else plan.get(x.shape[0])
206
if (
207
config is not None
@@ -152,7 +229,8 @@ def enable_qwen4_exp_low_latency_gemm(
229
module: nn.Module,
230
dtype: torch.dtype,
231
) -> None:
155
- if dtype != torch.bfloat16 or not _is_sm103():
232
+ plans = _gemm_plans()
233
+ if dtype != torch.bfloat16 or not plans:
234
return
235
if not shape_dynamic_skinny_gemm.is_available():
236
return
@@ -172,7 +250,7 @@ def enable_qwen4_exp_low_latency_gemm(
250
weight = getattr(child, "weight", None)
251
if weight is None or weight.dim() != 2:
252
continue
175
- plan = QWEN4_EXP_GEMM_PLANS.get((weight.shape[0], weight.shape[1]))
253
+ plan = plans.get((weight.shape[0], weight.shape[1]))
254
if plan is None:
255
continue
256
if is_linear: