@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 332 lines 12.8 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
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
7 standard linear implementation.
8 """
9
10 import torch
11 from torch import nn
12
13 import vllm.envs as envs
14 from vllm.model_executor.kernels.linear.cute_dsl.skinny_gemm import (
15 SkinnyGemmConfig,
16 shape_dynamic_skinny_gemm,
17 )
18 from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
19 from vllm.model_executor.layers.vocab_parallel_embedding import (
20 ParallelLMHead,
21 UnquantizedEmbeddingMethod,
22 )
23 from vllm.platforms import current_platform
24 from vllm.utils.torch_utils import direct_register_custom_op
25
26 QWEN4_EXP_GEMM_PLANS: dict[tuple[int, int], dict[int, SkinnyGemmConfig]] = {
27 # GDN fused QKVZ projection, TP=4.
28 (4096, 2560): {
29 1: SkinnyGemmConfig(1, 64, 4, k_unroll=4),
30 2: SkinnyGemmConfig(2, 64, 4, k_unroll=4),
31 },
32 # GDN and QSA output projections, TP=4.
33 (2560, 1536): {
34 1: SkinnyGemmConfig(1, 128, 2, k_unroll=2, vector_width=4),
35 2: SkinnyGemmConfig(2, 128, 2, k_unroll=2, vector_width=4),
36 4: SkinnyGemmConfig(4, 64, 2, k_unroll=2),
37 },
38 # GDN fused B/A projection, TP=4.
39 (24, 2560): {
40 1: SkinnyGemmConfig(1, 128, 2, k_unroll=4, vector_width=4),
41 2: SkinnyGemmConfig(2, 128, 2, k_unroll=4, vector_width=4),
42 4: SkinnyGemmConfig(4, 128, 2, k_unroll=4, vector_width=4),
43 8: SkinnyGemmConfig(8, 128, 1, k_unroll=4, vector_width=4),
44 16: SkinnyGemmConfig(16, 128, 1, k_unroll=4, vector_width=4),
45 },
46 # QSA fused QKV/gate projection, TP=4.
47 (3584, 2560): {
48 1: SkinnyGemmConfig(1, 128, 4, k_unroll=4, vector_width=4),
49 2: SkinnyGemmConfig(2, 64, 2, k_unroll=2),
50 4: SkinnyGemmConfig(4, 64, 2, k_unroll=2),
51 },
52 # QSA indexer Q/K projection, replicated in a TP=4 deployment.
53 (640, 2560): {
54 1: SkinnyGemmConfig(1, 128, 1, k_unroll=4, vector_width=4),
55 2: SkinnyGemmConfig(2, 128, 1, k_unroll=4, vector_width=4),
56 4: SkinnyGemmConfig(4, 128, 1, k_unroll=4, vector_width=4),
57 8: SkinnyGemmConfig(8, 128, 1, k_unroll=4, vector_width=4),
58 },
59 # Shared-expert fused gate/up projection, TP=4.
60 (320, 2560): {
61 1: SkinnyGemmConfig(1, 128, 2, k_unroll=4, vector_width=4),
62 2: SkinnyGemmConfig(2, 128, 2, k_unroll=4, vector_width=4),
63 4: SkinnyGemmConfig(4, 128, 2, k_unroll=4, vector_width=4),
64 8: SkinnyGemmConfig(8, 64, 1, k_unroll=4),
65 16: SkinnyGemmConfig(16, 128, 2, k_unroll=4, vector_width=4),
66 },
67 # LM head, TP=4.
68 (62080, 2560): {
69 1: SkinnyGemmConfig(1, 64, 4, k_unroll=2),
70 2: SkinnyGemmConfig(2, 32, 4, k_unroll=2),
71 },
72 # HC merged down/injection projection, replicated in a TP=4 deployment.
73 (336, 10240): {
74 1: SkinnyGemmConfig(1, 128, 1, static_k=10240),
75 2: SkinnyGemmConfig(2, 128, 1, static_k=10240),
76 4: SkinnyGemmConfig(4, 128, 2, static_k=10240),
77 8: SkinnyGemmConfig(8, 128, 1, k_unroll=4),
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 # 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
217
218 def _is_sm90() -> bool:
219 return current_platform.is_device_capability((9, 0))
220
221
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 {}
230
231
232 def _is_packed_row_major(tensor: torch.Tensor) -> bool:
233 return tensor.dim() == 2 and tensor.stride() == (tensor.shape[1], 1)
234
235
236 def _runtime_ok(x: torch.Tensor, weight: torch.Tensor) -> bool:
237 return (
238 not envs.VLLM_BATCH_INVARIANT
239 and _is_packed_row_major(x)
240 and _is_packed_row_major(weight)
241 and x.dtype == torch.bfloat16
242 and weight.dtype == torch.bfloat16
243 and x.is_cuda
244 and weight.is_cuda
245 and x.device == weight.device
246 and x.shape[1] == weight.shape[1]
247 )
248
249
250 class _Qwen4ExpLowLatencyApply:
251 def apply(
252 self,
253 layer: nn.Module,
254 x: torch.Tensor,
255 bias: torch.Tensor | None = None,
256 ) -> torch.Tensor:
257 if bias is None and not envs.VLLM_BATCH_INVARIANT:
258 return torch.ops.vllm.qwen4_exp_low_latency_gemm(x, layer.weight)
259 return super().apply(layer, x, bias) # type: ignore[misc]
260
261
262 class Qwen4ExpLowLatencyLinearMethod(_Qwen4ExpLowLatencyApply, UnquantizedLinearMethod):
263 pass
264
265
266 class Qwen4ExpLowLatencyEmbeddingMethod(
267 _Qwen4ExpLowLatencyApply, UnquantizedEmbeddingMethod
268 ):
269 pass
270
271
272 def _qwen4_exp_low_latency_gemm(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
273 plan = _gemm_plans().get((weight.shape[0], weight.shape[1]))
274 config = None if plan is None else plan.get(x.shape[0])
275 if (
276 config is not None
277 and _runtime_ok(x, weight)
278 and shape_dynamic_skinny_gemm.is_available()
279 ):
280 return shape_dynamic_skinny_gemm(x, weight, config)
281 return torch.nn.functional.linear(x, weight)
282
283
284 def _qwen4_exp_low_latency_gemm_fake(
285 x: torch.Tensor, weight: torch.Tensor
286 ) -> torch.Tensor:
287 return x.new_empty((*x.shape[:-1], weight.shape[0]))
288
289
290 direct_register_custom_op(
291 op_name="qwen4_exp_low_latency_gemm",
292 op_func=_qwen4_exp_low_latency_gemm,
293 fake_impl=_qwen4_exp_low_latency_gemm_fake,
294 )
295
296
297 def enable_qwen4_exp_low_latency_gemm(
298 module: nn.Module,
299 dtype: torch.dtype,
300 ) -> None:
301 plans = _gemm_plans()
302 if dtype != torch.bfloat16 or not plans:
303 return
304 if not shape_dynamic_skinny_gemm.is_available():
305 return
306
307 warmup_configs: set[SkinnyGemmConfig] = set()
308 for child in module.modules():
309 is_linear = (
310 isinstance(child, LinearBase)
311 and type(child.quant_method) is UnquantizedLinearMethod
312 )
313 is_head = (
314 isinstance(child, ParallelLMHead)
315 and type(child.quant_method) is UnquantizedEmbeddingMethod
316 )
317 if not (is_linear or is_head):
318 continue
319 weight = getattr(child, "weight", None)
320 if weight is None or weight.dim() != 2:
321 continue
322 plan = plans.get((weight.shape[0], weight.shape[1]))
323 if plan is None:
324 continue
325 if is_linear:
326 child.quant_method = Qwen4ExpLowLatencyLinearMethod()
327 else:
328 child.quant_method = Qwen4ExpLowLatencyEmbeddingMethod()
329 warmup_configs.update(plan.values())
330
331 if warmup_configs:
332 shape_dynamic_skinny_gemm.request_warmup_configs(dtype, warmup_configs)