[Kernel][Perf] Add sm_120 (RTX PRO 6000 / RTX 50) tuned configs for batch-invariant persistent matmul (#57456)
Signed-off-by: LioEinaudi <zhao3024667639@gmail.com>
Lio Einaudi committed
Sep 19, 2026 at 21:30 UTC
41c4a3ed4e45792739bd8a5bbcd01269b06eae2c
1 file changed
+229
vllm/model_executor/determinism/batch_invariant_configs.py
+229
@@ -254,6 +254,233 @@ _BATCH_INVARIANT_MATMUL_TUNED_CONFIGS: dict[
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
+ },
484
}
485
486
_TUNED_MATMUL_CONFIGS_FOR_DEVICE: dict[tuple[int, int], _MatmulShapeConfig] | None = (
@@ -265,6 +492,8 @@ _TUNED_MATMUL_CONFIGS_RESOLVED = False
492
def _get_tuned_matmul_arch_family(capability: DeviceCapability | None) -> str | None:
493
if capability is None:
494
return None
495
+ if capability.major == 12:
496
+ return "sm120"
497
if capability.major == 10:
498
return "blackwell"
499
if capability.major == 9: