@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 1,223 lines 50.9 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from typing import NamedTuple
4
5 import torch
6
7 from vllm.platforms import current_platform
8 from vllm.platforms.interface import DeviceCapability
9
10
11 class _MatmulMConfig(NamedTuple):
12 block_m: int
13 block_n: int
14 num_warps: int
15 num_stages: int
16
17
18 class _MatmulShapeConfig(NamedTuple):
19 block_k: int
20 m_buckets: tuple[tuple[int, _MatmulMConfig], ...]
21
22
23 _BATCH_INVARIANT_MATMUL_TUNED_CONFIGS: dict[
24 str, dict[tuple[int, int], _MatmulShapeConfig]
25 ] = {
26 "ada": {
27 (12288, 2048): _MatmulShapeConfig(
28 block_k=64,
29 m_buckets=(
30 (1, _MatmulMConfig(16, 128, 8, 2)),
31 (4, _MatmulMConfig(16, 64, 8, 5)),
32 (8, _MatmulMConfig(16, 128, 8, 4)),
33 (16, _MatmulMConfig(16, 128, 4, 5)),
34 (32, _MatmulMConfig(32, 128, 4, 5)),
35 (64, _MatmulMConfig(32, 256, 4, 3)),
36 (256, _MatmulMConfig(128, 128, 8, 3)),
37 (512, _MatmulMConfig(128, 64, 4, 4)),
38 (1024, _MatmulMConfig(128, 128, 8, 3)),
39 (2048, _MatmulMConfig(128, 128, 4, 3)),
40 ),
41 ),
42 (2048, 6144): _MatmulShapeConfig(
43 block_k=128,
44 m_buckets=(
45 (1, _MatmulMConfig(16, 64, 4, 4)),
46 (4, _MatmulMConfig(16, 64, 8, 4)),
47 (8, _MatmulMConfig(16, 32, 8, 4)),
48 (16, _MatmulMConfig(16, 32, 8, 4)),
49 (32, _MatmulMConfig(16, 64, 8, 3)),
50 (64, _MatmulMConfig(32, 64, 4, 4)),
51 (256, _MatmulMConfig(32, 64, 4, 4)),
52 (512, _MatmulMConfig(64, 64, 4, 3)),
53 (1024, _MatmulMConfig(64, 64, 4, 3)),
54 (2048, _MatmulMConfig(64, 64, 4, 3)),
55 ),
56 ),
57 (4096, 2048): _MatmulShapeConfig(
58 block_k=64,
59 m_buckets=(
60 (1, _MatmulMConfig(16, 256, 8, 3)),
61 (4, _MatmulMConfig(16, 64, 4, 3)),
62 (8, _MatmulMConfig(16, 128, 8, 2)),
63 (16, _MatmulMConfig(16, 32, 4, 5)),
64 (32, _MatmulMConfig(16, 64, 4, 4)),
65 (64, _MatmulMConfig(64, 64, 8, 3)),
66 (256, _MatmulMConfig(64, 64, 4, 5)),
67 (512, _MatmulMConfig(64, 64, 4, 5)),
68 (1024, _MatmulMConfig(128, 64, 4, 4)),
69 (2048, _MatmulMConfig(128, 64, 8, 4)),
70 ),
71 ),
72 (151936, 2048): _MatmulShapeConfig(
73 block_k=64,
74 m_buckets=(
75 (1, _MatmulMConfig(16, 32, 8, 5)),
76 (4, _MatmulMConfig(16, 32, 8, 5)),
77 (8, _MatmulMConfig(16, 32, 8, 5)),
78 (16, _MatmulMConfig(16, 32, 4, 4)),
79 (32, _MatmulMConfig(16, 64, 4, 4)),
80 (64, _MatmulMConfig(32, 64, 4, 4)),
81 (256, _MatmulMConfig(128, 128, 4, 3)),
82 (512, _MatmulMConfig(128, 128, 4, 3)),
83 (1024, _MatmulMConfig(128, 128, 4, 3)),
84 (2048, _MatmulMConfig(128, 128, 4, 3)),
85 ),
86 ),
87 (2048, 2048): _MatmulShapeConfig(
88 block_k=64,
89 m_buckets=(
90 (1, _MatmulMConfig(32, 128, 8, 5)),
91 (4, _MatmulMConfig(16, 256, 8, 3)),
92 (8, _MatmulMConfig(32, 128, 8, 3)),
93 (16, _MatmulMConfig(32, 128, 8, 3)),
94 (32, _MatmulMConfig(16, 128, 8, 5)),
95 (64, _MatmulMConfig(16, 256, 4, 3)),
96 (256, _MatmulMConfig(64, 128, 8, 4)),
97 (512, _MatmulMConfig(64, 64, 4, 5)),
98 (1024, _MatmulMConfig(64, 64, 4, 5)),
99 (2048, _MatmulMConfig(64, 64, 4, 5)),
100 ),
101 ),
102 },
103 "hopper": {
104 (12288, 2048): _MatmulShapeConfig(
105 block_k=128,
106 m_buckets=(
107 (1, _MatmulMConfig(16, 256, 4, 3)),
108 (4, _MatmulMConfig(16, 256, 4, 4)),
109 (8, _MatmulMConfig(16, 64, 4, 4)),
110 (16, _MatmulMConfig(16, 64, 4, 4)),
111 (32, _MatmulMConfig(32, 64, 8, 5)),
112 (64, _MatmulMConfig(64, 64, 4, 4)),
113 (256, _MatmulMConfig(64, 128, 4, 3)),
114 (512, _MatmulMConfig(128, 128, 8, 3)),
115 (1024, _MatmulMConfig(128, 256, 8, 2)),
116 (2048, _MatmulMConfig(128, 256, 8, 2)),
117 ),
118 ),
119 (2048, 6144): _MatmulShapeConfig(
120 block_k=128,
121 m_buckets=(
122 (1, _MatmulMConfig(16, 32, 4, 4)),
123 (4, _MatmulMConfig(16, 64, 4, 4)),
124 (8, _MatmulMConfig(16, 64, 4, 4)),
125 (16, _MatmulMConfig(16, 32, 4, 3)),
126 (32, _MatmulMConfig(32, 32, 8, 3)),
127 (64, _MatmulMConfig(64, 32, 4, 5)),
128 (256, _MatmulMConfig(64, 128, 4, 3)),
129 (512, _MatmulMConfig(128, 128, 4, 3)),
130 (1024, _MatmulMConfig(64, 64, 4, 5)),
131 (2048, _MatmulMConfig(64, 128, 4, 4)),
132 ),
133 ),
134 (4096, 2048): _MatmulShapeConfig(
135 block_k=128,
136 m_buckets=(
137 (1, _MatmulMConfig(16, 64, 8, 5)),
138 (4, _MatmulMConfig(16, 64, 4, 5)),
139 (8, _MatmulMConfig(16, 128, 4, 3)),
140 (16, _MatmulMConfig(16, 64, 4, 5)),
141 (32, _MatmulMConfig(16, 64, 4, 4)),
142 (64, _MatmulMConfig(64, 32, 4, 3)),
143 (256, _MatmulMConfig(64, 256, 4, 2)),
144 (512, _MatmulMConfig(64, 64, 4, 3)),
145 (1024, _MatmulMConfig(64, 128, 4, 3)),
146 (2048, _MatmulMConfig(64, 64, 4, 4)),
147 ),
148 ),
149 (151936, 2048): _MatmulShapeConfig(
150 block_k=128,
151 m_buckets=(
152 (1, _MatmulMConfig(16, 256, 4, 4)),
153 (4, _MatmulMConfig(16, 256, 4, 4)),
154 (8, _MatmulMConfig(16, 256, 4, 3)),
155 (16, _MatmulMConfig(16, 256, 4, 3)),
156 (32, _MatmulMConfig(32, 256, 8, 3)),
157 (64, _MatmulMConfig(64, 128, 4, 3)),
158 (256, _MatmulMConfig(64, 128, 4, 3)),
159 (512, _MatmulMConfig(128, 128, 8, 3)),
160 (1024, _MatmulMConfig(128, 256, 8, 2)),
161 (2048, _MatmulMConfig(128, 256, 8, 2)),
162 ),
163 ),
164 (2048, 2048): _MatmulShapeConfig(
165 block_k=128,
166 m_buckets=(
167 (1, _MatmulMConfig(16, 64, 4, 5)),
168 (4, _MatmulMConfig(16, 128, 8, 3)),
169 (8, _MatmulMConfig(16, 32, 4, 3)),
170 (16, _MatmulMConfig(32, 64, 8, 5)),
171 (32, _MatmulMConfig(16, 128, 4, 5)),
172 (64, _MatmulMConfig(32, 64, 4, 2)),
173 (256, _MatmulMConfig(128, 64, 4, 4)),
174 (512, _MatmulMConfig(64, 256, 4, 2)),
175 (1024, _MatmulMConfig(64, 64, 4, 3)),
176 (2048, _MatmulMConfig(64, 128, 4, 3)),
177 ),
178 ),
179 },
180 "blackwell": {
181 (12288, 2048): _MatmulShapeConfig(
182 block_k=64,
183 m_buckets=(
184 (1, _MatmulMConfig(16, 128, 8, 5)),
185 (4, _MatmulMConfig(16, 128, 4, 5)),
186 (8, _MatmulMConfig(16, 128, 4, 5)),
187 (16, _MatmulMConfig(16, 128, 4, 5)),
188 (32, _MatmulMConfig(32, 128, 4, 5)),
189 (64, _MatmulMConfig(64, 128, 8, 5)),
190 (256, _MatmulMConfig(128, 256, 4, 4)),
191 (512, _MatmulMConfig(128, 256, 4, 4)),
192 (1024, _MatmulMConfig(128, 256, 4, 4)),
193 (2048, _MatmulMConfig(128, 256, 4, 4)),
194 ),
195 ),
196 (2048, 6144): _MatmulShapeConfig(
197 block_k=128,
198 m_buckets=(
199 (1, _MatmulMConfig(16, 32, 4, 5)),
200 (4, _MatmulMConfig(16, 32, 4, 5)),
201 (8, _MatmulMConfig(16, 32, 4, 5)),
202 (16, _MatmulMConfig(16, 32, 4, 5)),
203 (32, _MatmulMConfig(16, 32, 4, 5)),
204 (64, _MatmulMConfig(32, 32, 4, 5)),
205 (256, _MatmulMConfig(64, 64, 4, 5)),
206 (512, _MatmulMConfig(128, 64, 4, 4)),
207 (1024, _MatmulMConfig(128, 128, 8, 3)),
208 (2048, _MatmulMConfig(128, 128, 8, 3)),
209 ),
210 ),
211 (4096, 2048): _MatmulShapeConfig(
212 block_k=128,
213 m_buckets=(
214 (1, _MatmulMConfig(16, 64, 4, 5)),
215 (4, _MatmulMConfig(16, 64, 4, 5)),
216 (8, _MatmulMConfig(16, 32, 4, 5)),
217 (16, _MatmulMConfig(16, 32, 4, 5)),
218 (32, _MatmulMConfig(32, 32, 4, 5)),
219 (64, _MatmulMConfig(32, 64, 4, 4)),
220 (256, _MatmulMConfig(128, 64, 4, 4)),
221 (512, _MatmulMConfig(128, 128, 8, 3)),
222 (1024, _MatmulMConfig(128, 128, 8, 3)),
223 (2048, _MatmulMConfig(128, 128, 8, 3)),
224 ),
225 ),
226 (151936, 2048): _MatmulShapeConfig(
227 block_k=64,
228 m_buckets=(
229 (1, _MatmulMConfig(16, 256, 4, 5)),
230 (4, _MatmulMConfig(16, 256, 4, 5)),
231 (8, _MatmulMConfig(16, 256, 4, 5)),
232 (16, _MatmulMConfig(16, 256, 4, 5)),
233 (32, _MatmulMConfig(64, 256, 4, 5)),
234 (64, _MatmulMConfig(64, 256, 4, 5)),
235 (256, _MatmulMConfig(128, 256, 4, 4)),
236 (512, _MatmulMConfig(128, 256, 4, 4)),
237 (1024, _MatmulMConfig(128, 256, 4, 4)),
238 (2048, _MatmulMConfig(128, 256, 4, 4)),
239 ),
240 ),
241 (2048, 2048): _MatmulShapeConfig(
242 block_k=128,
243 m_buckets=(
244 (1, _MatmulMConfig(16, 32, 4, 4)),
245 (4, _MatmulMConfig(16, 32, 4, 4)),
246 (8, _MatmulMConfig(16, 32, 4, 5)),
247 (16, _MatmulMConfig(16, 32, 8, 5)),
248 (32, _MatmulMConfig(16, 32, 4, 4)),
249 (64, _MatmulMConfig(32, 32, 4, 5)),
250 (256, _MatmulMConfig(64, 64, 4, 4)),
251 (512, _MatmulMConfig(128, 64, 4, 4)),
252 (1024, _MatmulMConfig(128, 128, 4, 3)),
253 (2048, _MatmulMConfig(128, 128, 8, 3)),
254 ),
255 ),
256 },
257 "sm120": {
258 (12288, 2048): _MatmulShapeConfig( # Qwen3-1.7B gate_up_proj
259 block_k=64,
260 m_buckets=(
261 (1, _MatmulMConfig(16, 128, 8, 5)),
262 (4, _MatmulMConfig(16, 128, 8, 4)),
263 (8, _MatmulMConfig(16, 128, 8, 4)),
264 (16, _MatmulMConfig(16, 128, 4, 4)),
265 (32, _MatmulMConfig(32, 128, 8, 5)),
266 (64, _MatmulMConfig(64, 128, 4, 3)),
267 (256, _MatmulMConfig(64, 64, 4, 4)),
268 (512, _MatmulMConfig(64, 128, 4, 5)),
269 (1024, _MatmulMConfig(128, 128, 4, 3)),
270 (2048, _MatmulMConfig(128, 128, 4, 3)),
271 ),
272 ),
273 (2048, 6144): _MatmulShapeConfig( # Qwen3-1.7B down_proj
274 block_k=64,
275 m_buckets=(
276 (1, _MatmulMConfig(16, 32, 8, 5)),
277 (4, _MatmulMConfig(16, 32, 8, 3)),
278 (8, _MatmulMConfig(16, 32, 8, 4)),
279 (16, _MatmulMConfig(16, 32, 4, 5)),
280 (32, _MatmulMConfig(32, 32, 4, 5)),
281 (64, _MatmulMConfig(32, 32, 4, 4)),
282 (256, _MatmulMConfig(64, 64, 4, 5)),
283 (512, _MatmulMConfig(128, 64, 4, 3)),
284 (1024, _MatmulMConfig(128, 128, 4, 3)),
285 (2048, _MatmulMConfig(128, 64, 8, 4)),
286 ),
287 ),
288 (4096, 2048): _MatmulShapeConfig( # Qwen3-1.7B qkv_proj
289 block_k=128,
290 m_buckets=(
291 (1, _MatmulMConfig(16, 64, 4, 4)),
292 (4, _MatmulMConfig(16, 64, 8, 5)),
293 (8, _MatmulMConfig(16, 64, 8, 5)),
294 (16, _MatmulMConfig(16, 64, 8, 3)),
295 (32, _MatmulMConfig(16, 64, 8, 5)),
296 (64, _MatmulMConfig(32, 64, 4, 2)),
297 (256, _MatmulMConfig(64, 128, 4, 3)),
298 (512, _MatmulMConfig(64, 64, 4, 3)),
299 (1024, _MatmulMConfig(64, 128, 4, 3)),
300 (2048, _MatmulMConfig(64, 128, 4, 3)),
301 ),
302 ),
303 (151936, 2048): _MatmulShapeConfig( # Qwen3-1.7B lm_head
304 block_k=64,
305 m_buckets=(
306 (1, _MatmulMConfig(16, 64, 8, 5)),
307 (4, _MatmulMConfig(16, 128, 4, 4)),
308 (8, _MatmulMConfig(16, 128, 4, 5)),
309 (16, _MatmulMConfig(16, 128, 8, 3)),
310 (32, _MatmulMConfig(32, 128, 8, 3)),
311 (64, _MatmulMConfig(32, 128, 8, 5)),
312 (256, _MatmulMConfig(128, 128, 4, 3)),
313 (512, _MatmulMConfig(128, 128, 8, 3)),
314 (1024, _MatmulMConfig(128, 128, 4, 3)),
315 (2048, _MatmulMConfig(128, 128, 4, 3)),
316 ),
317 ),
318 (2048, 2048): _MatmulShapeConfig( # Qwen3-1.7B o_proj
319 block_k=32,
320 m_buckets=(
321 (1, _MatmulMConfig(16, 128, 4, 3)),
322 (4, _MatmulMConfig(16, 128, 4, 4)),
323 (8, _MatmulMConfig(16, 64, 4, 5)),
324 (16, _MatmulMConfig(32, 32, 8, 4)),
325 (32, _MatmulMConfig(16, 64, 4, 5)),
326 (64, _MatmulMConfig(32, 32, 4, 5)),
327 (256, _MatmulMConfig(128, 32, 4, 5)),
328 (512, _MatmulMConfig(128, 64, 4, 4)),
329 (1024, _MatmulMConfig(128, 128, 8, 5)),
330 (2048, _MatmulMConfig(128, 256, 8, 4)),
331 ),
332 ),
333 (19456, 2560): _MatmulShapeConfig( # Qwen3-4B gate_up_proj
334 block_k=64,
335 m_buckets=(
336 (1, _MatmulMConfig(16, 128, 8, 4)),
337 (4, _MatmulMConfig(16, 128, 8, 3)),
338 (8, _MatmulMConfig(16, 128, 8, 5)),
339 (16, _MatmulMConfig(16, 128, 8, 5)),
340 (32, _MatmulMConfig(32, 128, 4, 4)),
341 (64, _MatmulMConfig(64, 128, 8, 4)),
342 (256, _MatmulMConfig(128, 128, 4, 3)),
343 (512, _MatmulMConfig(64, 128, 4, 5)),
344 (1024, _MatmulMConfig(128, 128, 4, 3)),
345 (2048, _MatmulMConfig(128, 128, 4, 3)),
346 ),
347 ),
348 (2560, 9728): _MatmulShapeConfig( # Qwen3-4B down_proj
349 block_k=128,
350 m_buckets=(
351 (1, _MatmulMConfig(16, 64, 4, 5)),
352 (4, _MatmulMConfig(16, 64, 4, 4)),
353 (8, _MatmulMConfig(16, 64, 4, 4)),
354 (16, _MatmulMConfig(16, 32, 4, 5)),
355 (32, _MatmulMConfig(32, 32, 4, 3)),
356 (64, _MatmulMConfig(32, 64, 8, 4)),
357 (256, _MatmulMConfig(64, 64, 4, 3)),
358 (512, _MatmulMConfig(64, 128, 4, 3)),
359 (1024, _MatmulMConfig(64, 128, 4, 3)),
360 (2048, _MatmulMConfig(64, 128, 4, 3)),
361 ),
362 ),
363 (6144, 2560): _MatmulShapeConfig( # Qwen3-4B qkv_proj
364 block_k=64,
365 m_buckets=(
366 (1, _MatmulMConfig(16, 128, 8, 4)),
367 (4, _MatmulMConfig(16, 256, 4, 3)),
368 (8, _MatmulMConfig(16, 128, 4, 5)),
369 (16, _MatmulMConfig(16, 64, 8, 3)),
370 (32, _MatmulMConfig(16, 128, 8, 4)),
371 (64, _MatmulMConfig(64, 64, 4, 5)),
372 (256, _MatmulMConfig(128, 128, 8, 3)),
373 (512, _MatmulMConfig(64, 64, 4, 5)),
374 (1024, _MatmulMConfig(64, 128, 8, 4)),
375 (2048, _MatmulMConfig(128, 128, 4, 3)),
376 ),
377 ),
378 (151936, 2560): _MatmulShapeConfig( # Qwen3-4B lm_head
379 block_k=64,
380 m_buckets=(
381 (1, _MatmulMConfig(16, 64, 8, 5)),
382 (4, _MatmulMConfig(16, 64, 8, 5)),
383 (8, _MatmulMConfig(16, 64, 4, 5)),
384 (16, _MatmulMConfig(16, 64, 4, 5)),
385 (32, _MatmulMConfig(32, 64, 4, 5)),
386 (64, _MatmulMConfig(64, 64, 4, 5)),
387 (256, _MatmulMConfig(128, 128, 4, 3)),
388 (512, _MatmulMConfig(128, 128, 4, 3)),
389 (1024, _MatmulMConfig(128, 128, 8, 3)),
390 (2048, _MatmulMConfig(128, 128, 4, 3)),
391 ),
392 ),
393 (2560, 4096): _MatmulShapeConfig( # Qwen3-4B o_proj
394 block_k=128,
395 m_buckets=(
396 (1, _MatmulMConfig(16, 64, 8, 3)),
397 (4, _MatmulMConfig(16, 64, 8, 4)),
398 (8, _MatmulMConfig(16, 64, 8, 5)),
399 (16, _MatmulMConfig(16, 64, 4, 4)),
400 (32, _MatmulMConfig(16, 32, 4, 4)),
401 (64, _MatmulMConfig(16, 64, 8, 3)),
402 (256, _MatmulMConfig(64, 64, 8, 3)),
403 (512, _MatmulMConfig(64, 128, 4, 3)),
404 (1024, _MatmulMConfig(64, 128, 4, 3)),
405 (2048, _MatmulMConfig(64, 128, 4, 3)),
406 ),
407 ),
408 (24576, 4096): _MatmulShapeConfig( # Qwen3-8B gate_up_proj
409 block_k=64,
410 m_buckets=(
411 (1, _MatmulMConfig(16, 256, 8, 3)),
412 (4, _MatmulMConfig(16, 256, 8, 3)),
413 (8, _MatmulMConfig(16, 256, 8, 3)),
414 (16, _MatmulMConfig(16, 256, 8, 3)),
415 (32, _MatmulMConfig(32, 256, 8, 3)),
416 (64, _MatmulMConfig(64, 256, 4, 3)),
417 (256, _MatmulMConfig(64, 128, 4, 5)),
418 (512, _MatmulMConfig(128, 128, 4, 3)),
419 (1024, _MatmulMConfig(128, 128, 4, 3)),
420 (2048, _MatmulMConfig(128, 128, 4, 3)),
421 ),
422 ),
423 (4096, 12288): _MatmulShapeConfig( # Qwen3-8B down_proj
424 block_k=128,
425 m_buckets=(
426 (1, _MatmulMConfig(16, 64, 4, 5)),
427 (4, _MatmulMConfig(16, 64, 4, 5)),
428 (8, _MatmulMConfig(16, 64, 4, 4)),
429 (16, _MatmulMConfig(16, 32, 4, 4)),
430 (32, _MatmulMConfig(32, 32, 8, 5)),
431 (64, _MatmulMConfig(32, 64, 4, 5)),
432 (256, _MatmulMConfig(64, 128, 4, 3)),
433 (512, _MatmulMConfig(64, 64, 4, 3)),
434 (1024, _MatmulMConfig(64, 128, 4, 3)),
435 (2048, _MatmulMConfig(64, 128, 4, 3)),
436 ),
437 ),
438 (6144, 4096): _MatmulShapeConfig( # Qwen3-8B qkv_proj
439 block_k=64,
440 m_buckets=(
441 (1, _MatmulMConfig(16, 64, 8, 4)),
442 (4, _MatmulMConfig(16, 128, 8, 4)),
443 (8, _MatmulMConfig(16, 128, 4, 4)),
444 (16, _MatmulMConfig(16, 64, 4, 5)),
445 (32, _MatmulMConfig(16, 128, 8, 5)),
446 (64, _MatmulMConfig(64, 64, 4, 4)),
447 (256, _MatmulMConfig(128, 128, 8, 3)),
448 (512, _MatmulMConfig(64, 64, 4, 5)),
449 (1024, _MatmulMConfig(128, 64, 4, 3)),
450 (2048, _MatmulMConfig(128, 128, 4, 3)),
451 ),
452 ),
453 (151936, 4096): _MatmulShapeConfig( # Qwen3-8B lm_head
454 block_k=64,
455 m_buckets=(
456 (1, _MatmulMConfig(16, 64, 8, 5)),
457 (4, _MatmulMConfig(16, 64, 8, 5)),
458 (8, _MatmulMConfig(16, 64, 4, 5)),
459 (16, _MatmulMConfig(16, 64, 4, 5)),
460 (32, _MatmulMConfig(32, 64, 4, 5)),
461 (64, _MatmulMConfig(32, 128, 8, 5)),
462 (256, _MatmulMConfig(128, 128, 4, 3)),
463 (512, _MatmulMConfig(128, 128, 8, 3)),
464 (1024, _MatmulMConfig(128, 128, 4, 3)),
465 (2048, _MatmulMConfig(128, 128, 4, 3)),
466 ),
467 ),
468 (4096, 4096): _MatmulShapeConfig( # Qwen3-8B o_proj
469 block_k=128,
470 m_buckets=(
471 (1, _MatmulMConfig(16, 64, 8, 4)),
472 (4, _MatmulMConfig(16, 64, 4, 3)),
473 (8, _MatmulMConfig(16, 32, 4, 4)),
474 (16, _MatmulMConfig(16, 32, 4, 3)),
475 (32, _MatmulMConfig(16, 64, 4, 5)),
476 (64, _MatmulMConfig(64, 32, 8, 4)),
477 (256, _MatmulMConfig(64, 128, 4, 3)),
478 (512, _MatmulMConfig(64, 64, 4, 3)),
479 (1024, _MatmulMConfig(64, 128, 4, 3)),
480 (2048, _MatmulMConfig(64, 128, 4, 3)),
481 ),
482 ),
483 (6144, 2048): _MatmulShapeConfig( # Qwen3-1.7B tp2_gate_up_proj
484 block_k=64,
485 m_buckets=(
486 (1, _MatmulMConfig(16, 32, 8, 3)),
487 (4, _MatmulMConfig(16, 128, 8, 5)),
488 (8, _MatmulMConfig(16, 128, 4, 5)),
489 (16, _MatmulMConfig(16, 32, 4, 5)),
490 (32, _MatmulMConfig(32, 128, 8, 5)),
491 (64, _MatmulMConfig(32, 128, 4, 4)),
492 (256, _MatmulMConfig(128, 128, 4, 3)),
493 (512, _MatmulMConfig(64, 64, 4, 5)),
494 (1024, _MatmulMConfig(128, 64, 4, 3)),
495 (2048, _MatmulMConfig(128, 128, 4, 3)),
496 ),
497 ),
498 (2048, 3072): _MatmulShapeConfig( # Qwen3-1.7B tp2_down_proj
499 block_k=128,
500 m_buckets=(
501 (1, _MatmulMConfig(16, 32, 8, 5)),
502 (4, _MatmulMConfig(16, 64, 4, 3)),
503 (8, _MatmulMConfig(16, 64, 8, 4)),
504 (16, _MatmulMConfig(16, 64, 8, 3)),
505 (32, _MatmulMConfig(16, 64, 8, 3)),
506 (64, _MatmulMConfig(16, 64, 4, 5)),
507 (256, _MatmulMConfig(64, 64, 8, 3)),
508 (512, _MatmulMConfig(64, 128, 4, 3)),
509 (1024, _MatmulMConfig(64, 64, 4, 3)),
510 (2048, _MatmulMConfig(64, 128, 4, 3)),
511 ),
512 ),
513 (2048, 1024): _MatmulShapeConfig( # Qwen3-1.7B tp2_o_proj
514 block_k=64,
515 m_buckets=(
516 (1, _MatmulMConfig(16, 256, 8, 2)),
517 (4, _MatmulMConfig(16, 32, 8, 4)),
518 (8, _MatmulMConfig(32, 128, 8, 5)),
519 (16, _MatmulMConfig(32, 128, 4, 4)),
520 (32, _MatmulMConfig(64, 128, 4, 3)),
521 (64, _MatmulMConfig(16, 256, 4, 3)),
522 (256, _MatmulMConfig(32, 128, 4, 3)),
523 (512, _MatmulMConfig(64, 128, 4, 4)),
524 (1024, _MatmulMConfig(128, 128, 8, 3)),
525 (2048, _MatmulMConfig(64, 128, 4, 5)),
526 ),
527 ),
528 (75968, 2048): _MatmulShapeConfig( # Qwen3-1.7B tp2_lm_head
529 block_k=64,
530 m_buckets=(
531 (1, _MatmulMConfig(16, 256, 8, 3)),
532 (4, _MatmulMConfig(16, 256, 8, 3)),
533 (8, _MatmulMConfig(16, 256, 8, 3)),
534 (16, _MatmulMConfig(16, 256, 4, 3)),
535 (32, _MatmulMConfig(32, 256, 8, 3)),
536 (64, _MatmulMConfig(64, 256, 8, 3)),
537 (256, _MatmulMConfig(64, 128, 4, 5)),
538 (512, _MatmulMConfig(128, 128, 4, 3)),
539 (1024, _MatmulMConfig(128, 128, 4, 3)),
540 (2048, _MatmulMConfig(128, 128, 4, 3)),
541 ),
542 ),
543 (3072, 2048): _MatmulShapeConfig( # Qwen3-1.7B tp4_gate_up_proj
544 block_k=64,
545 m_buckets=(
546 (1, _MatmulMConfig(16, 256, 4, 3)),
547 (4, _MatmulMConfig(32, 128, 4, 5)),
548 (8, _MatmulMConfig(16, 128, 8, 3)),
549 (16, _MatmulMConfig(32, 128, 8, 4)),
550 (32, _MatmulMConfig(32, 128, 8, 5)),
551 (64, _MatmulMConfig(64, 32, 8, 3)),
552 (256, _MatmulMConfig(128, 64, 8, 3)),
553 (512, _MatmulMConfig(128, 128, 4, 3)),
554 (1024, _MatmulMConfig(64, 64, 4, 5)),
555 (2048, _MatmulMConfig(64, 128, 8, 4)),
556 ),
557 ),
558 (2048, 1536): _MatmulShapeConfig( # Qwen3-1.7B tp4_down_proj
559 block_k=64,
560 m_buckets=(
561 (1, _MatmulMConfig(16, 128, 8, 2)),
562 (4, _MatmulMConfig(16, 256, 4, 3)),
563 (8, _MatmulMConfig(32, 128, 4, 2)),
564 (16, _MatmulMConfig(32, 32, 8, 5)),
565 (32, _MatmulMConfig(32, 32, 8, 4)),
566 (64, _MatmulMConfig(32, 128, 4, 4)),
567 (256, _MatmulMConfig(64, 64, 4, 3)),
568 (512, _MatmulMConfig(128, 64, 8, 4)),
569 (1024, _MatmulMConfig(128, 128, 4, 3)),
570 (2048, _MatmulMConfig(128, 64, 8, 4)),
571 ),
572 ),
573 (1024, 2048): _MatmulShapeConfig( # Qwen3-1.7B tp4_qkv_proj
574 block_k=64,
575 m_buckets=(
576 (1, _MatmulMConfig(32, 128, 8, 4)),
577 (4, _MatmulMConfig(16, 32, 8, 4)),
578 (8, _MatmulMConfig(16, 32, 4, 4)),
579 (16, _MatmulMConfig(32, 128, 8, 3)),
580 (32, _MatmulMConfig(32, 128, 8, 4)),
581 (64, _MatmulMConfig(32, 64, 4, 5)),
582 (256, _MatmulMConfig(64, 32, 4, 3)),
583 (512, _MatmulMConfig(64, 64, 4, 3)),
584 (1024, _MatmulMConfig(64, 128, 4, 5)),
585 (2048, _MatmulMConfig(128, 128, 8, 3)),
586 ),
587 ),
588 (2048, 512): _MatmulShapeConfig( # Qwen3-1.7B tp4_o_proj
589 block_k=64,
590 m_buckets=(
591 (1, _MatmulMConfig(32, 32, 4, 5)),
592 (4, _MatmulMConfig(64, 128, 4, 2)),
593 (8, _MatmulMConfig(16, 128, 4, 2)),
594 (16, _MatmulMConfig(16, 256, 8, 3)),
595 (32, _MatmulMConfig(64, 32, 8, 4)),
596 (64, _MatmulMConfig(16, 128, 8, 4)),
597 (256, _MatmulMConfig(64, 64, 8, 5)),
598 (512, _MatmulMConfig(32, 128, 8, 4)),
599 (1024, _MatmulMConfig(64, 64, 8, 3)),
600 (2048, _MatmulMConfig(64, 128, 8, 3)),
601 ),
602 ),
603 (37984, 2048): _MatmulShapeConfig( # Qwen3-1.7B tp4_lm_head
604 block_k=64,
605 m_buckets=(
606 (1, _MatmulMConfig(16, 128, 8, 4)),
607 (4, _MatmulMConfig(16, 128, 8, 5)),
608 (8, _MatmulMConfig(16, 128, 4, 5)),
609 (16, _MatmulMConfig(16, 128, 8, 4)),
610 (32, _MatmulMConfig(32, 128, 4, 5)),
611 (64, _MatmulMConfig(64, 128, 4, 5)),
612 (256, _MatmulMConfig(64, 128, 4, 5)),
613 (512, _MatmulMConfig(128, 128, 4, 3)),
614 (1024, _MatmulMConfig(128, 128, 4, 3)),
615 (2048, _MatmulMConfig(128, 128, 4, 3)),
616 ),
617 ),
618 (1536, 2048): _MatmulShapeConfig( # Qwen3-1.7B tp8_gate_up_proj
619 block_k=64,
620 m_buckets=(
621 (1, _MatmulMConfig(32, 128, 8, 3)),
622 (4, _MatmulMConfig(16, 256, 4, 3)),
623 (8, _MatmulMConfig(16, 32, 8, 5)),
624 (16, _MatmulMConfig(32, 32, 8, 3)),
625 (32, _MatmulMConfig(32, 128, 8, 4)),
626 (64, _MatmulMConfig(16, 128, 8, 4)),
627 (256, _MatmulMConfig(32, 128, 8, 3)),
628 (512, _MatmulMConfig(64, 128, 8, 4)),
629 (1024, _MatmulMConfig(128, 128, 4, 3)),
630 (2048, _MatmulMConfig(64, 64, 4, 4)),
631 ),
632 ),
633 (2048, 768): _MatmulShapeConfig( # Qwen3-1.7B tp8_down_proj
634 block_k=128,
635 m_buckets=(
636 (1, _MatmulMConfig(32, 64, 4, 3)),
637 (4, _MatmulMConfig(32, 64, 8, 3)),
638 (8, _MatmulMConfig(16, 32, 4, 5)),
639 (16, _MatmulMConfig(32, 64, 4, 2)),
640 (32, _MatmulMConfig(16, 64, 8, 3)),
641 (64, _MatmulMConfig(16, 32, 8, 2)),
642 (256, _MatmulMConfig(32, 32, 4, 4)),
643 (512, _MatmulMConfig(32, 64, 4, 5)),
644 (1024, _MatmulMConfig(64, 128, 4, 3)),
645 (2048, _MatmulMConfig(64, 128, 4, 3)),
646 ),
647 ),
648 (512, 2048): _MatmulShapeConfig( # Qwen3-1.7B tp8_qkv_proj
649 block_k=64,
650 m_buckets=(
651 (1, _MatmulMConfig(64, 64, 8, 5)),
652 (4, _MatmulMConfig(16, 32, 4, 4)),
653 (8, _MatmulMConfig(16, 32, 4, 5)),
654 (16, _MatmulMConfig(32, 128, 8, 3)),
655 (32, _MatmulMConfig(64, 64, 8, 5)),
656 (64, _MatmulMConfig(32, 128, 8, 4)),
657 (256, _MatmulMConfig(16, 128, 4, 4)),
658 (512, _MatmulMConfig(32, 64, 8, 3)),
659 (1024, _MatmulMConfig(64, 64, 4, 4)),
660 (2048, _MatmulMConfig(64, 128, 4, 4)),
661 ),
662 ),
663 (2048, 256): _MatmulShapeConfig( # Qwen3-1.7B tp8_o_proj
664 block_k=64,
665 m_buckets=(
666 (1, _MatmulMConfig(32, 128, 8, 2)),
667 (4, _MatmulMConfig(16, 128, 4, 2)),
668 (8, _MatmulMConfig(32, 256, 4, 2)),
669 (16, _MatmulMConfig(32, 32, 8, 4)),
670 (32, _MatmulMConfig(16, 256, 4, 3)),
671 (64, _MatmulMConfig(128, 32, 8, 3)),
672 (256, _MatmulMConfig(16, 32, 8, 2)),
673 (512, _MatmulMConfig(16, 32, 4, 3)),
674 (1024, _MatmulMConfig(64, 256, 4, 3)),
675 (2048, _MatmulMConfig(64, 128, 8, 4)),
676 ),
677 ),
678 (18992, 2048): _MatmulShapeConfig( # Qwen3-1.7B tp8_lm_head
679 block_k=64,
680 m_buckets=(
681 (1, _MatmulMConfig(16, 64, 8, 5)),
682 (4, _MatmulMConfig(16, 128, 8, 5)),
683 (8, _MatmulMConfig(16, 128, 8, 3)),
684 (16, _MatmulMConfig(16, 128, 4, 5)),
685 (32, _MatmulMConfig(32, 128, 8, 5)),
686 (64, _MatmulMConfig(64, 128, 8, 4)),
687 (256, _MatmulMConfig(128, 128, 4, 3)),
688 (512, _MatmulMConfig(128, 128, 4, 3)),
689 (1024, _MatmulMConfig(128, 128, 8, 3)),
690 (2048, _MatmulMConfig(128, 128, 4, 3)),
691 ),
692 ),
693 (9728, 2560): _MatmulShapeConfig( # Qwen3-4B tp2_gate_up_proj
694 block_k=64,
695 m_buckets=(
696 (1, _MatmulMConfig(16, 32, 8, 4)),
697 (4, _MatmulMConfig(16, 64, 8, 3)),
698 (8, _MatmulMConfig(16, 128, 4, 5)),
699 (16, _MatmulMConfig(16, 32, 4, 5)),
700 (32, _MatmulMConfig(32, 64, 8, 4)),
701 (64, _MatmulMConfig(32, 128, 4, 3)),
702 (256, _MatmulMConfig(128, 128, 4, 3)),
703 (512, _MatmulMConfig(128, 128, 4, 3)),
704 (1024, _MatmulMConfig(64, 128, 4, 5)),
705 (2048, _MatmulMConfig(128, 128, 4, 3)),
706 ),
707 ),
708 (2560, 4864): _MatmulShapeConfig( # Qwen3-4B tp2_down_proj
709 block_k=128,
710 m_buckets=(
711 (1, _MatmulMConfig(16, 32, 4, 3)),
712 (4, _MatmulMConfig(16, 128, 8, 3)),
713 (8, _MatmulMConfig(16, 64, 8, 4)),
714 (16, _MatmulMConfig(32, 64, 8, 4)),
715 (32, _MatmulMConfig(16, 64, 4, 3)),
716 (64, _MatmulMConfig(32, 32, 8, 4)),
717 (256, _MatmulMConfig(64, 64, 4, 3)),
718 (512, _MatmulMConfig(64, 128, 4, 3)),
719 (1024, _MatmulMConfig(64, 128, 4, 3)),
720 (2048, _MatmulMConfig(64, 128, 4, 3)),
721 ),
722 ),
723 (3072, 2560): _MatmulShapeConfig( # Qwen3-4B tp2_qkv_proj
724 block_k=128,
725 m_buckets=(
726 (1, _MatmulMConfig(16, 128, 8, 3)),
727 (4, _MatmulMConfig(16, 32, 8, 2)),
728 (8, _MatmulMConfig(16, 32, 8, 4)),
729 (16, _MatmulMConfig(16, 64, 4, 3)),
730 (32, _MatmulMConfig(32, 64, 4, 4)),
731 (64, _MatmulMConfig(16, 128, 4, 3)),
732 (256, _MatmulMConfig(64, 128, 4, 3)),
733 (512, _MatmulMConfig(64, 64, 4, 3)),
734 (1024, _MatmulMConfig(64, 128, 4, 3)),
735 (2048, _MatmulMConfig(64, 128, 4, 3)),
736 ),
737 ),
738 (2560, 2048): _MatmulShapeConfig( # Qwen3-4B tp2_o_proj
739 block_k=64,
740 m_buckets=(
741 (1, _MatmulMConfig(32, 128, 8, 4)),
742 (4, _MatmulMConfig(32, 128, 4, 4)),
743 (8, _MatmulMConfig(16, 128, 4, 3)),
744 (16, _MatmulMConfig(32, 128, 4, 5)),
745 (32, _MatmulMConfig(16, 128, 8, 4)),
746 (64, _MatmulMConfig(16, 128, 8, 4)),
747 (256, _MatmulMConfig(128, 32, 8, 5)),
748 (512, _MatmulMConfig(128, 64, 8, 3)),
749 (1024, _MatmulMConfig(128, 128, 4, 3)),
750 (2048, _MatmulMConfig(128, 128, 8, 3)),
751 ),
752 ),
753 (75968, 2560): _MatmulShapeConfig( # Qwen3-4B tp2_lm_head
754 block_k=64,
755 m_buckets=(
756 (1, _MatmulMConfig(16, 64, 8, 5)),
757 (4, _MatmulMConfig(16, 64, 8, 5)),
758 (8, _MatmulMConfig(16, 64, 4, 5)),
759 (16, _MatmulMConfig(16, 64, 4, 5)),
760 (32, _MatmulMConfig(32, 64, 4, 5)),
761 (64, _MatmulMConfig(64, 256, 4, 3)),
762 (256, _MatmulMConfig(64, 128, 4, 5)),
763 (512, _MatmulMConfig(128, 128, 4, 3)),
764 (1024, _MatmulMConfig(128, 128, 4, 3)),
765 (2048, _MatmulMConfig(128, 128, 8, 3)),
766 ),
767 ),
768 (4864, 2560): _MatmulShapeConfig( # Qwen3-4B tp4_gate_up_proj
769 block_k=64,
770 m_buckets=(
771 (1, _MatmulMConfig(16, 32, 8, 5)),
772 (4, _MatmulMConfig(16, 128, 4, 4)),
773 (8, _MatmulMConfig(16, 32, 8, 4)),
774 (16, _MatmulMConfig(16, 128, 8, 4)),
775 (32, _MatmulMConfig(16, 128, 4, 5)),
776 (64, _MatmulMConfig(16, 128, 4, 3)),
777 (256, _MatmulMConfig(64, 128, 4, 5)),
778 (512, _MatmulMConfig(128, 128, 4, 3)),
779 (1024, _MatmulMConfig(128, 128, 4, 3)),
780 (2048, _MatmulMConfig(128, 128, 4, 3)),
781 ),
782 ),
783 (2560, 2432): _MatmulShapeConfig( # Qwen3-4B tp4_down_proj
784 block_k=64,
785 m_buckets=(
786 (1, _MatmulMConfig(16, 32, 8, 4)),
787 (4, _MatmulMConfig(16, 32, 4, 5)),
788 (8, _MatmulMConfig(32, 64, 8, 4)),
789 (16, _MatmulMConfig(32, 128, 8, 5)),
790 (32, _MatmulMConfig(16, 128, 4, 4)),
791 (64, _MatmulMConfig(64, 32, 4, 3)),
792 (256, _MatmulMConfig(32, 128, 8, 4)),
793 (512, _MatmulMConfig(64, 128, 4, 5)),
794 (1024, _MatmulMConfig(128, 128, 4, 3)),
795 (2048, _MatmulMConfig(128, 128, 4, 3)),
796 ),
797 ),
798 (1536, 2560): _MatmulShapeConfig( # Qwen3-4B tp4_qkv_proj
799 block_k=64,
800 m_buckets=(
801 (1, _MatmulMConfig(16, 32, 8, 4)),
802 (4, _MatmulMConfig(16, 32, 8, 3)),
803 (8, _MatmulMConfig(16, 32, 4, 4)),
804 (16, _MatmulMConfig(32, 128, 8, 3)),
805 (32, _MatmulMConfig(16, 128, 4, 3)),
806 (64, _MatmulMConfig(16, 64, 8, 4)),
807 (256, _MatmulMConfig(128, 32, 4, 5)),
808 (512, _MatmulMConfig(64, 128, 4, 5)),
809 (1024, _MatmulMConfig(64, 256, 4, 3)),
810 (2048, _MatmulMConfig(64, 128, 4, 4)),
811 ),
812 ),
813 (2560, 1024): _MatmulShapeConfig( # Qwen3-4B tp4_o_proj
814 block_k=64,
815 m_buckets=(
816 (1, _MatmulMConfig(32, 128, 8, 5)),
817 (4, _MatmulMConfig(16, 32, 8, 3)),
818 (8, _MatmulMConfig(16, 128, 4, 3)),
819 (16, _MatmulMConfig(32, 256, 4, 3)),
820 (32, _MatmulMConfig(16, 128, 4, 2)),
821 (64, _MatmulMConfig(32, 256, 8, 3)),
822 (256, _MatmulMConfig(32, 128, 4, 5)),
823 (512, _MatmulMConfig(64, 128, 8, 3)),
824 (1024, _MatmulMConfig(128, 128, 4, 3)),
825 (2048, _MatmulMConfig(128, 128, 4, 3)),
826 ),
827 ),
828 (37984, 2560): _MatmulShapeConfig( # Qwen3-4B tp4_lm_head
829 block_k=128,
830 m_buckets=(
831 (1, _MatmulMConfig(16, 128, 8, 3)),
832 (4, _MatmulMConfig(16, 128, 8, 3)),
833 (8, _MatmulMConfig(16, 128, 4, 3)),
834 (16, _MatmulMConfig(16, 32, 4, 5)),
835 (32, _MatmulMConfig(16, 32, 4, 5)),
836 (64, _MatmulMConfig(32, 32, 4, 5)),
837 (256, _MatmulMConfig(64, 128, 4, 3)),
838 (512, _MatmulMConfig(64, 128, 4, 3)),
839 (1024, _MatmulMConfig(64, 128, 4, 3)),
840 (2048, _MatmulMConfig(64, 128, 4, 3)),
841 ),
842 ),
843 (2432, 2560): _MatmulShapeConfig( # Qwen3-4B tp8_gate_up_proj
844 block_k=64,
845 m_buckets=(
846 (1, _MatmulMConfig(16, 32, 4, 5)),
847 (4, _MatmulMConfig(32, 128, 4, 3)),
848 (8, _MatmulMConfig(16, 128, 8, 3)),
849 (16, _MatmulMConfig(16, 128, 4, 4)),
850 (32, _MatmulMConfig(16, 128, 8, 4)),
851 (64, _MatmulMConfig(16, 128, 4, 3)),
852 (256, _MatmulMConfig(64, 64, 4, 5)),
853 (512, _MatmulMConfig(64, 128, 4, 4)),
854 (1024, _MatmulMConfig(128, 128, 8, 3)),
855 (2048, _MatmulMConfig(128, 128, 4, 3)),
856 ),
857 ),
858 (2560, 1216): _MatmulShapeConfig( # Qwen3-4B tp8_down_proj
859 block_k=64,
860 m_buckets=(
861 (1, _MatmulMConfig(16, 256, 8, 2)),
862 (4, _MatmulMConfig(16, 128, 8, 4)),
863 (8, _MatmulMConfig(16, 32, 4, 3)),
864 (16, _MatmulMConfig(32, 32, 4, 3)),
865 (32, _MatmulMConfig(16, 128, 8, 4)),
866 (64, _MatmulMConfig(32, 128, 8, 4)),
867 (256, _MatmulMConfig(32, 128, 4, 5)),
868 (512, _MatmulMConfig(64, 128, 8, 4)),
869 (1024, _MatmulMConfig(64, 256, 4, 3)),
870 (2048, _MatmulMConfig(128, 128, 4, 3)),
871 ),
872 ),
873 (768, 2560): _MatmulShapeConfig( # Qwen3-4B tp8_qkv_proj
874 block_k=128,
875 m_buckets=(
876 (1, _MatmulMConfig(32, 32, 4, 2)),
877 (4, _MatmulMConfig(16, 64, 8, 3)),
878 (8, _MatmulMConfig(16, 64, 8, 4)),
879 (16, _MatmulMConfig(16, 64, 4, 5)),
880 (32, _MatmulMConfig(32, 64, 8, 4)),
881 (64, _MatmulMConfig(16, 64, 8, 4)),
882 (256, _MatmulMConfig(64, 32, 4, 5)),
883 (512, _MatmulMConfig(64, 64, 8, 3)),
884 (1024, _MatmulMConfig(64, 128, 4, 3)),
885 (2048, _MatmulMConfig(64, 64, 4, 3)),
886 ),
887 ),
888 (2560, 512): _MatmulShapeConfig( # Qwen3-4B tp8_o_proj
889 block_k=64,
890 m_buckets=(
891 (1, _MatmulMConfig(16, 128, 8, 4)),
892 (4, _MatmulMConfig(16, 32, 4, 2)),
893 (8, _MatmulMConfig(16, 128, 4, 5)),
894 (16, _MatmulMConfig(32, 32, 4, 5)),
895 (32, _MatmulMConfig(32, 256, 4, 3)),
896 (64, _MatmulMConfig(32, 128, 8, 5)),
897 (256, _MatmulMConfig(16, 256, 4, 2)),
898 (512, _MatmulMConfig(32, 128, 8, 4)),
899 (1024, _MatmulMConfig(128, 128, 8, 3)),
900 (2048, _MatmulMConfig(128, 128, 4, 3)),
901 ),
902 ),
903 (18992, 2560): _MatmulShapeConfig( # Qwen3-4B tp8_lm_head
904 block_k=64,
905 m_buckets=(
906 (1, _MatmulMConfig(16, 64, 8, 4)),
907 (4, _MatmulMConfig(16, 64, 8, 5)),
908 (8, _MatmulMConfig(16, 128, 8, 3)),
909 (16, _MatmulMConfig(16, 128, 8, 5)),
910 (32, _MatmulMConfig(32, 128, 8, 5)),
911 (64, _MatmulMConfig(64, 128, 8, 4)),
912 (256, _MatmulMConfig(128, 128, 4, 3)),
913 (512, _MatmulMConfig(128, 128, 4, 3)),
914 (1024, _MatmulMConfig(128, 128, 4, 3)),
915 (2048, _MatmulMConfig(128, 128, 4, 3)),
916 ),
917 ),
918 (12288, 4096): _MatmulShapeConfig( # Qwen3-8B tp2_gate_up_proj
919 block_k=64,
920 m_buckets=(
921 (1, _MatmulMConfig(16, 128, 8, 3)),
922 (4, _MatmulMConfig(16, 128, 8, 5)),
923 (8, _MatmulMConfig(16, 128, 8, 4)),
924 (16, _MatmulMConfig(16, 128, 4, 4)),
925 (32, _MatmulMConfig(32, 128, 4, 4)),
926 (64, _MatmulMConfig(64, 128, 4, 5)),
927 (256, _MatmulMConfig(64, 64, 4, 5)),
928 (512, _MatmulMConfig(128, 64, 4, 3)),
929 (1024, _MatmulMConfig(128, 128, 4, 3)),
930 (2048, _MatmulMConfig(128, 128, 4, 3)),
931 ),
932 ),
933 (4096, 6144): _MatmulShapeConfig( # Qwen3-8B tp2_down_proj
934 block_k=64,
935 m_buckets=(
936 (1, _MatmulMConfig(16, 32, 8, 5)),
937 (4, _MatmulMConfig(16, 32, 8, 5)),
938 (8, _MatmulMConfig(16, 32, 8, 3)),
939 (16, _MatmulMConfig(16, 32, 4, 5)),
940 (32, _MatmulMConfig(16, 64, 8, 4)),
941 (64, _MatmulMConfig(32, 64, 4, 5)),
942 (256, _MatmulMConfig(64, 128, 4, 5)),
943 (512, _MatmulMConfig(128, 128, 4, 3)),
944 (1024, _MatmulMConfig(128, 64, 8, 4)),
945 (2048, _MatmulMConfig(128, 128, 4, 3)),
946 ),
947 ),
948 (3072, 4096): _MatmulShapeConfig( # Qwen3-8B tp2_qkv_proj
949 block_k=64,
950 m_buckets=(
951 (1, _MatmulMConfig(16, 32, 8, 4)),
952 (4, _MatmulMConfig(16, 32, 8, 3)),
953 (8, _MatmulMConfig(16, 128, 4, 3)),
954 (16, _MatmulMConfig(16, 32, 4, 4)),
955 (32, _MatmulMConfig(32, 32, 8, 5)),
956 (64, _MatmulMConfig(16, 128, 8, 4)),
957 (256, _MatmulMConfig(64, 128, 4, 5)),
958 (512, _MatmulMConfig(128, 128, 8, 3)),
959 (1024, _MatmulMConfig(64, 64, 4, 5)),
960 (2048, _MatmulMConfig(128, 64, 4, 3)),
961 ),
962 ),
963 (75968, 4096): _MatmulShapeConfig( # Qwen3-8B tp2_lm_head
964 block_k=64,
965 m_buckets=(
966 (1, _MatmulMConfig(16, 256, 8, 3)),
967 (4, _MatmulMConfig(16, 256, 8, 3)),
968 (8, _MatmulMConfig(16, 256, 8, 3)),
969 (16, _MatmulMConfig(16, 256, 8, 3)),
970 (32, _MatmulMConfig(32, 256, 8, 3)),
971 (64, _MatmulMConfig(64, 256, 8, 3)),
972 (256, _MatmulMConfig(128, 128, 4, 3)),
973 (512, _MatmulMConfig(128, 128, 8, 3)),
974 (1024, _MatmulMConfig(128, 128, 4, 3)),
975 (2048, _MatmulMConfig(128, 128, 4, 3)),
976 ),
977 ),
978 (4096, 3072): _MatmulShapeConfig( # Qwen3-8B tp4_down_proj
979 block_k=64,
980 m_buckets=(
981 (1, _MatmulMConfig(16, 256, 8, 3)),
982 (4, _MatmulMConfig(16, 128, 4, 5)),
983 (8, _MatmulMConfig(16, 128, 4, 4)),
984 (16, _MatmulMConfig(16, 64, 4, 5)),
985 (32, _MatmulMConfig(32, 128, 4, 3)),
986 (64, _MatmulMConfig(16, 128, 4, 4)),
987 (256, _MatmulMConfig(64, 128, 4, 5)),
988 (512, _MatmulMConfig(128, 128, 4, 3)),
989 (1024, _MatmulMConfig(128, 64, 8, 4)),
990 (2048, _MatmulMConfig(128, 128, 4, 3)),
991 ),
992 ),
993 (1536, 4096): _MatmulShapeConfig( # Qwen3-8B tp4_qkv_proj
994 block_k=128,
995 m_buckets=(
996 (1, _MatmulMConfig(16, 64, 4, 5)),
997 (4, _MatmulMConfig(16, 64, 8, 3)),
998 (8, _MatmulMConfig(32, 64, 4, 5)),
999 (16, _MatmulMConfig(16, 64, 8, 4)),
1000 (32, _MatmulMConfig(16, 32, 8, 3)),
1001 (64, _MatmulMConfig(16, 64, 8, 5)),
1002 (256, _MatmulMConfig(64, 64, 4, 3)),
1003 (512, _MatmulMConfig(64, 128, 4, 3)),
1004 (1024, _MatmulMConfig(64, 64, 4, 3)),
1005 (2048, _MatmulMConfig(64, 128, 4, 3)),
1006 ),
1007 ),
1008 (4096, 1024): _MatmulShapeConfig( # Qwen3-8B tp4_o_proj
1009 block_k=64,
1010 m_buckets=(
1011 (1, _MatmulMConfig(16, 128, 8, 4)),
1012 (4, _MatmulMConfig(32, 128, 8, 4)),
1013 (8, _MatmulMConfig(16, 256, 8, 3)),
1014 (16, _MatmulMConfig(16, 128, 8, 3)),
1015 (32, _MatmulMConfig(16, 128, 4, 4)),
1016 (64, _MatmulMConfig(64, 64, 4, 3)),
1017 (256, _MatmulMConfig(128, 64, 4, 3)),
1018 (512, _MatmulMConfig(64, 64, 4, 5)),
1019 (1024, _MatmulMConfig(128, 64, 4, 4)),
1020 (2048, _MatmulMConfig(128, 128, 4, 3)),
1021 ),
1022 ),
1023 (37984, 4096): _MatmulShapeConfig( # Qwen3-8B tp4_lm_head
1024 block_k=64,
1025 m_buckets=(
1026 (1, _MatmulMConfig(16, 128, 8, 3)),
1027 (4, _MatmulMConfig(16, 128, 4, 5)),
1028 (8, _MatmulMConfig(16, 128, 4, 5)),
1029 (16, _MatmulMConfig(16, 128, 8, 3)),
1030 (32, _MatmulMConfig(32, 128, 8, 3)),
1031 (64, _MatmulMConfig(64, 128, 8, 3)),
1032 (256, _MatmulMConfig(64, 128, 4, 5)),
1033 (512, _MatmulMConfig(128, 128, 4, 3)),
1034 (1024, _MatmulMConfig(128, 128, 4, 3)),
1035 (2048, _MatmulMConfig(128, 128, 4, 3)),
1036 ),
1037 ),
1038 (4096, 1536): _MatmulShapeConfig( # Qwen3-8B tp8_down_proj
1039 block_k=32,
1040 m_buckets=(
1041 (1, _MatmulMConfig(32, 128, 8, 5)),
1042 (4, _MatmulMConfig(16, 256, 4, 5)),
1043 (8, _MatmulMConfig(16, 64, 4, 4)),
1044 (16, _MatmulMConfig(16, 128, 8, 3)),
1045 (32, _MatmulMConfig(16, 64, 4, 3)),
1046 (64, _MatmulMConfig(16, 128, 4, 4)),
1047 (256, _MatmulMConfig(64, 128, 4, 4)),
1048 (512, _MatmulMConfig(128, 128, 8, 5)),
1049 (1024, _MatmulMConfig(128, 256, 8, 4)),
1050 (2048, _MatmulMConfig(128, 128, 8, 4)),
1051 ),
1052 ),
1053 (768, 4096): _MatmulShapeConfig( # Qwen3-8B tp8_qkv_proj
1054 block_k=128,
1055 m_buckets=(
1056 (1, _MatmulMConfig(16, 32, 8, 5)),
1057 (4, _MatmulMConfig(16, 32, 8, 3)),
1058 (8, _MatmulMConfig(32, 64, 4, 5)),
1059 (16, _MatmulMConfig(32, 64, 8, 3)),
1060 (32, _MatmulMConfig(32, 32, 8, 5)),
1061 (64, _MatmulMConfig(64, 32, 8, 3)),
1062 (256, _MatmulMConfig(32, 64, 4, 4)),
1063 (512, _MatmulMConfig(64, 64, 8, 3)),
1064 (1024, _MatmulMConfig(64, 128, 4, 3)),
1065 (2048, _MatmulMConfig(64, 64, 4, 3)),
1066 ),
1067 ),
1068 (4096, 512): _MatmulShapeConfig( # Qwen3-8B tp8_o_proj
1069 block_k=64,
1070 m_buckets=(
1071 (1, _MatmulMConfig(16, 128, 8, 2)),
1072 (4, _MatmulMConfig(32, 128, 4, 2)),
1073 (8, _MatmulMConfig(16, 256, 4, 2)),
1074 (16, _MatmulMConfig(32, 256, 4, 3)),
1075 (32, _MatmulMConfig(32, 256, 8, 3)),
1076 (64, _MatmulMConfig(32, 128, 4, 5)),
1077 (256, _MatmulMConfig(64, 32, 8, 4)),
1078 (512, _MatmulMConfig(128, 128, 8, 3)),
1079 (1024, _MatmulMConfig(128, 64, 4, 4)),
1080 (2048, _MatmulMConfig(128, 128, 8, 3)),
1081 ),
1082 ),
1083 (18992, 4096): _MatmulShapeConfig( # Qwen3-8B tp8_lm_head
1084 block_k=64,
1085 m_buckets=(
1086 (1, _MatmulMConfig(16, 64, 8, 5)),
1087 (4, _MatmulMConfig(16, 64, 8, 4)),
1088 (8, _MatmulMConfig(16, 64, 8, 5)),
1089 (16, _MatmulMConfig(16, 128, 4, 3)),
1090 (32, _MatmulMConfig(32, 64, 8, 5)),
1091 (64, _MatmulMConfig(64, 128, 8, 3)),
1092 (256, _MatmulMConfig(128, 128, 4, 3)),
1093 (512, _MatmulMConfig(128, 128, 4, 3)),
1094 (1024, _MatmulMConfig(128, 128, 4, 3)),
1095 (2048, _MatmulMConfig(128, 128, 4, 3)),
1096 ),
1097 ),
1098 },
1099 }
1100
1101 _TUNED_MATMUL_CONFIGS_FOR_DEVICE: dict[tuple[int, int], _MatmulShapeConfig] | None = (
1102 None
1103 )
1104 _TUNED_MATMUL_CONFIGS_RESOLVED = False
1105
1106
1107 def _get_tuned_matmul_arch_family(capability: DeviceCapability | None) -> str | None:
1108 if capability is None:
1109 return None
1110 if capability.major == 12:
1111 return "sm120"
1112 if capability.major == 10:
1113 return "blackwell"
1114 if capability.major == 9:
1115 return "hopper"
1116 if capability.major == 8 and capability.minor == 9:
1117 return "ada"
1118 return None
1119
1120
1121 def resolve_tuned_matmul_configs() -> None:
1122 global _TUNED_MATMUL_CONFIGS_FOR_DEVICE
1123 global _TUNED_MATMUL_CONFIGS_RESOLVED
1124
1125 if _TUNED_MATMUL_CONFIGS_RESOLVED:
1126 return
1127
1128 # Normal vLLM initialization resolves this before torch.compile tracing.
1129 capability = (
1130 current_platform.get_device_capability() if current_platform.is_cuda() else None
1131 )
1132 arch_family = _get_tuned_matmul_arch_family(capability)
1133 if arch_family is None:
1134 _TUNED_MATMUL_CONFIGS_FOR_DEVICE = None
1135 else:
1136 _TUNED_MATMUL_CONFIGS_FOR_DEVICE = _BATCH_INVARIANT_MATMUL_TUNED_CONFIGS.get(
1137 arch_family
1138 )
1139 _TUNED_MATMUL_CONFIGS_RESOLVED = True
1140
1141
1142 def _get_matmul_config(
1143 M: int,
1144 N: int,
1145 K: int,
1146 dtype: torch.dtype,
1147 default: dict[str, int],
1148 ) -> dict[str, int]:
1149 if not _TUNED_MATMUL_CONFIGS_RESOLVED:
1150 resolve_tuned_matmul_configs()
1151 if dtype != torch.bfloat16:
1152 return default
1153 device_configs = _TUNED_MATMUL_CONFIGS_FOR_DEVICE
1154 if device_configs is None:
1155 return default
1156 shape_config = device_configs.get((N, K))
1157 if shape_config is None:
1158 return default
1159 # Values above the tuned range reuse the largest bucket;
1160 # shape-wide BLOCK_K keeps this batch-invariant.
1161 m_config = shape_config.m_buckets[-1][1]
1162 for max_m, bucket_config in shape_config.m_buckets:
1163 if max_m >= M:
1164 m_config = bucket_config
1165 break
1166 return {
1167 "BLOCK_SIZE_M": m_config.block_m,
1168 "BLOCK_SIZE_N": m_config.block_n,
1169 "BLOCK_SIZE_K": shape_config.block_k,
1170 "GROUP_SIZE_M": 8,
1171 "num_warps": m_config.num_warps,
1172 "num_stages": m_config.num_stages,
1173 }
1174
1175
1176 def _get_descriptor_matmul_config(
1177 M: int, N: int, K: int, dtype: torch.dtype
1178 ) -> dict[str, int]:
1179 """Select a config for the XPU descriptor matmul from shape and dtype.
1180
1181 The default 128×128×64 config is optimized for large square-ish GEMMs but
1182 wastes resources on skinny shapes (small M during decode, or very tall/thin
1183 matrices during prefill). Shape-dependent tuning closes the gap with oneDNN.
1184
1185 Only BLOCK_SIZE_M/N and the launch parameters vary with M; dtype alone
1186 decides BLOCK_SIZE_K, so the K-reduction order never depends on batch size.
1187 """
1188 # fp32 uses smaller BLOCK_SIZE_K due to register pressure
1189 block_k = 32 if dtype == torch.float32 else 64
1190
1191 if M <= 16:
1192 # Decode: M=1-16. Tiny M means most of a 128-row tile is wasted.
1193 # Use small M-block, wide N-block to maximize useful work per tile.
1194 return {
1195 "BLOCK_SIZE_M": 16,
1196 "BLOCK_SIZE_N": 256,
1197 "BLOCK_SIZE_K": block_k,
1198 "GROUP_SIZE_M": 1,
1199 "num_stages": 4,
1200 "num_warps": 8,
1201 }
1202 elif M <= 64:
1203 # Small batch decode or very short prefill.
1204 return {
1205 "BLOCK_SIZE_M": 32,
1206 "BLOCK_SIZE_N": 128,
1207 "BLOCK_SIZE_K": block_k,
1208 "GROUP_SIZE_M": 4,
1209 "num_stages": 4,
1210 "num_warps": 8,
1211 }
1212 else:
1213 # Medium and large prefill (M > 64). 64×128 tiles provide the best
1214 # balance of register pressure vs parallelism on Intel XPU.
1215 # M=2048, N=4096 → 32×32 = 1024 tiles, well above ~160 compute units.
1216 return {
1217 "BLOCK_SIZE_M": 64,
1218 "BLOCK_SIZE_N": 128,
1219 "BLOCK_SIZE_K": block_k,
1220 "GROUP_SIZE_M": 8,
1221 "num_stages": 3,
1222 "num_warps": 8,
1223 }