[Kimi K3][Kernel] Enable low-latency decode GEMM dispatch on SM100 (#53534)
Signed-off-by: Canlin Guo <canlinguosdu@gmail.com>
Canlin Guo committed
Aug 26, 2026 at 07:32 UTC
bc2d63e650f698aab24fd4d91889e0dcf16876f0
2 files changed
+323
-7
tests/kernels/test_bf16_skinny_gemm.py
+95
-1
@@ -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 SM103 selector."""
3
+"""Tests for BF16 skinny GEMMs and the Kimi-K3 SM103/SM100 selectors."""
4
5
from pathlib import Path
6
from types import SimpleNamespace
@@ -182,6 +182,7 @@ def test_every_dsv3_routed_shape_is_instantiated() -> None:
182
183
specs = [
184
*KIMI_K3_PROJECTIONS.values(),
185
+ *k3_gemm.KIMI_K3_PROJECTIONS_SM100.values(),
186
glm52_gemm.GLM52_QKV_A_PROJECTION,
187
glm52_gemm.GLM52_Q_B_PROJECTION,
188
]
@@ -404,6 +405,92 @@ def test_build_plan_matches_selector() -> None:
405
assert plan[num_tokens][0] == backend
406
407
408
+# Keyed by local (N, K): (cute token counts, dsv3 token counts). Mirrors
409
+# KIMI_K3_PROJECTIONS_SM100, measured on B200; intentionally different from
410
+# EXPECTED_SELECTIONS (e.g. 3584x7168 routes dsv3 at M2..8 on SM103 but only
411
+# wins with CuTe at M1..3 on SM100).
412
+EXPECTED_SELECTIONS_SM100 = {
413
+ (1536, 128): (set(), {1, 16}),
414
+ (3072, 128): (set(), {8}),
415
+ (1536, 7168): ({1, 2}, {4, 8}),
416
+ (3072, 7168): ({1, 2}, set()),
417
+ (2112, 7168): ({1, 2}, {4, 16}),
418
+ (2304, 1536): (set(), set(range(1, 17))),
419
+ (4608, 1536): (set(), {1, 2, 4}),
420
+ (6288, 7168): ({1}, set()),
421
+ (12448, 7168): ({1, 2, 3, 4}, set()),
422
+ (7168, 768): (set(), {1}),
423
+ # SM103-tuned config wins on B200 where the heuristic config tied.
424
+ (7168, 1536): ({1}, set()),
425
+ # M2 config lifted from the SM103 table (wins on B200).
426
+ (7168, 3072): ({1, 2}, set()),
427
+ (7168, 3584): ({1, 2}, set()),
428
+ (7168, 8448): ({1, 2, 3}, set()),
429
+ (8448, 7168): ({1, 2}, set()),
430
+ (16896, 7168): ({1}, set()),
431
+ (20480, 7168): ({1, 2, 3, 4}, set()),
432
+ (40960, 7168): ({1, 2, 3, 4}, set()),
433
+ (10240, 7168): ({1, 2, 3, 4}, set()),
434
+ (3584, 7168): ({1, 2, 3}, set()),
435
+ # TP16 shapes measured on B200.
436
+ (768, 7168): ({1, 2, 3, 4}, set()),
437
+ (1152, 1536): ({1}, set()),
438
+ (3216, 7168): ({1, 2}, set()),
439
+ (4224, 7168): ({1, 2, 3}, set()),
440
+}
441
+
442
+
443
+def test_sm100_table_is_keyed_by_shape() -> None:
444
+ for (n, k), spec in k3_gemm.KIMI_K3_PROJECTIONS_SM100.items():
445
+ assert (spec.n, spec.k) == (n, k)
446
+
447
+
448
+@pytest.mark.parametrize("key", EXPECTED_SELECTIONS_SM100)
449
+def test_sm100_selector_table(key: tuple[int, int]) -> None:
450
+ n, k = key
451
+ spec = k3_gemm.KIMI_K3_PROJECTIONS_SM100[key]
452
+ cute_tokens, dsv3_tokens = EXPECTED_SELECTIONS_SM100[key]
453
+ for num_tokens in range(1, 17):
454
+ backend = k3_gemm._backend_for(spec, num_tokens, has_residual=False)
455
+ if num_tokens in cute_tokens:
456
+ assert backend == "cute"
457
+ elif num_tokens in dsv3_tokens:
458
+ assert backend == "dsv3_fused_a"
459
+ else:
460
+ assert backend is None
461
+
462
+
463
+def test_sm100_build_plan_matches_table() -> None:
464
+ for spec in k3_gemm.KIMI_K3_PROJECTIONS_SM100.values():
465
+ plan = k3_gemm._build_plan(spec)
466
+ for num_tokens in range(1, 17):
467
+ backend = k3_gemm._backend_for(spec, num_tokens, has_residual=False)
468
+ if backend is None:
469
+ assert num_tokens not in plan
470
+ else:
471
+ assert plan[num_tokens][0] == backend
472
+
473
+
474
+def test_low_latency_table_capability_routing(
475
+ monkeypatch: pytest.MonkeyPatch,
476
+) -> None:
477
+ monkeypatch.setattr(k3_gemm, "_is_sm103", lambda: True)
478
+ assert k3_gemm._low_latency_table() is k3_gemm.KIMI_K3_PROJECTIONS
479
+
480
+ monkeypatch.setattr(k3_gemm, "_is_sm103", lambda: False)
481
+ monkeypatch.setattr(
482
+ k3_gemm.current_platform,
483
+ "is_device_capability",
484
+ lambda cc: cc == (10, 0),
485
+ )
486
+ assert k3_gemm._low_latency_table() is k3_gemm.KIMI_K3_PROJECTIONS_SM100
487
+
488
+ monkeypatch.setattr(
489
+ k3_gemm.current_platform, "is_device_capability", lambda cc: False
490
+ )
491
+ assert k3_gemm._low_latency_table() is None
492
+
493
+
494
def test_installation_is_shape_specific_and_unquantized(
495
monkeypatch: pytest.MonkeyPatch,
496
) -> None:
@@ -491,6 +578,13 @@ def test_installation_requires_bf16_sm103(
578
root.projection = FakeLinear()
579
monkeypatch.setattr(k3_gemm, "LinearBase", FakeLinear)
580
monkeypatch.setattr(k3_gemm, "_is_sm103", lambda: platform_enabled)
581
+ # Pin the rest of the capability probe so platform_enabled=False stays a
582
+ # no-op even when the test itself runs on SM100 hardware.
583
+ monkeypatch.setattr(
584
+ k3_gemm.current_platform,
585
+ "is_device_capability",
586
+ lambda cc: platform_enabled and cc == (10, 3),
587
+ )
588
589
k3_gemm.enable_kimi_k3_low_latency_gemm(root, dtype)
590
vllm/models/kimi_k3/nvidia/low_latency_gemm.py
+228
-6
@@ -1,12 +1,19 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Kimi-K3 decode GEMM selection for unquantized BF16 on SM103.
3
+"""Kimi-K3 decode GEMM selection for unquantized BF16 on SM103 and SM100.
4
5
Dispatch is purely by local ``(N, K)`` shape and token count ``M`` — the module
6
name plays no role. Each measured shape maps to a :class:`ProjectionSpec`
7
holding the winning backend per token count. The static part of the decision is
8
resolved once per module at install time into a small ``{M: call}`` plan, so the
9
per-forward path is a single dict lookup.
10
+
11
+The two supported capabilities carry separate measured tables:
12
+:data:`KIMI_K3_PROJECTIONS` was tuned on B300 (SM103),
13
+:data:`KIMI_K3_PROJECTIONS_SM100` on B200 (SM100). The per-(shape, M) winners
14
+genuinely differ between the two parts (e.g. 3584x7168 favors dsv3 at M2..8 on
15
+B300 but only wins with CuTe at M1..3 on B200), so the tables must not be
16
+merged.
17
"""
18
19
from __future__ import annotations
@@ -280,6 +287,204 @@ KIMI_K3_PROJECTIONS: dict[tuple[int, int], ProjectionSpec] = {
287
# vector_width=2, which measured slower than cuBLAS.
288
}
289
290
+# Keyed by local (N, K), B200 (SM100) entries. Measured on B200 at TP8/TP16
291
+# shapes over M=1..16 with the same >=5%-over-cuBLAS threshold as the SM103
292
+# table. Per-(shape, M) entries take the best of the SM103-tuned config and a
293
+# fresh SM100 measurement: several SM103 configs (7168x8448, 20480x7168,
294
+# 40960x7168, 10240x7168, 7168x3072 M2, 7168x1536) beat the heuristic config on
295
+# B200 outright, while a few (6288x7168 M2..4, 3072x7168 M3..5) lose to cuBLAS
296
+# on B200 and are intentionally not inherited. The 7168x3584 residual configs
297
+# are intentionally absent: the B200 residual path was not measured.
298
+KIMI_K3_PROJECTIONS_SM100: dict[tuple[int, int], ProjectionSpec] = {
299
+ (1536, 128): ProjectionSpec(1536, 128, frozenset({1, 16}), name="f_b_proj"),
300
+ (3072, 128): ProjectionSpec(3072, 128, frozenset({8}), name="f_b_proj"),
301
+ # Same mla_g_proj aux-stream/PDL capture caveat as the SM103 entry applies;
302
+ # only M4/M8 measured a dsv3 win, so this is the narrow set either way.
303
+ (1536, 7168): ProjectionSpec(
304
+ 1536,
305
+ 7168,
306
+ frozenset({4, 8}),
307
+ (
308
+ (1, _cute(1, 224, 3, 2)),
309
+ (2, _cute(2, 128, 2, 2)),
310
+ ),
311
+ name="shared_gate_up_proj/mla_g_proj",
312
+ ),
313
+ (3072, 7168): ProjectionSpec(
314
+ 3072,
315
+ 7168,
316
+ cute_configs=(
317
+ (1, _cute(1, 224, 3, 4)),
318
+ (2, _cute(2, 128, 2, 1)),
319
+ ),
320
+ name="shared_gate_up_proj",
321
+ ),
322
+ (2112, 7168): ProjectionSpec(
323
+ 2112,
324
+ 7168,
325
+ frozenset({4, 16}),
326
+ (
327
+ (1, _cute(1, 224, 3, 2)),
328
+ (2, _cute(2, 128, 2, 2)),
329
+ ),
330
+ name="fused_qkv_a_proj",
331
+ ),
332
+ (2304, 1536): ProjectionSpec(2304, 1536, _M1_TO_16, name="q_b_proj"),
333
+ (4608, 1536): ProjectionSpec(4608, 1536, frozenset({1, 2, 4}), name="q_b_proj"),
334
+ # M2..4 of the SM103 configs lose to cuBLAS on B200; only M1 carries over.
335
+ (6288, 7168): ProjectionSpec(
336
+ 6288,
337
+ 7168,
338
+ cute_configs=((1, _cute(1, 224, 3, 4)),),
339
+ name="in_proj_qkvgfab",
340
+ ),
341
+ (12448, 7168): ProjectionSpec(
342
+ 12448,
343
+ 7168,
344
+ cute_configs=(
345
+ (1, _cute(1, 224, 4, 2)),
346
+ (2, _cute(2, 64, 4, 2)),
347
+ (3, _cute(3, 64, 2, 2)),
348
+ (4, _cute(4, 128, 2, 1)),
349
+ ),
350
+ name="in_proj_qkvgfab",
351
+ ),
352
+ (7168, 768): ProjectionSpec(7168, 768, frozenset({1}), name="shared_down_proj"),
353
+ # SM103-tuned config wins on B200 (+25%) where the heuristic config tied.
354
+ (7168, 1536): ProjectionSpec(
355
+ 7168,
356
+ 1536,
357
+ cute_configs=((1, _cute(1, 96, 4, 2)),),
358
+ name="o_proj",
359
+ ),
360
+ (7168, 3072): ProjectionSpec(
361
+ 7168,
362
+ 3072,
363
+ cute_configs=(
364
+ (1, _cute(1, 96, 2, 4)),
365
+ (2, _cute(2, 32, 4, 4)),
366
+ ),
367
+ name="o_proj",
368
+ ),
369
+ (7168, 3584): ProjectionSpec(
370
+ 7168,
371
+ 3584,
372
+ cute_configs=(
373
+ (1, _cute(1, 64, 4, 2)),
374
+ (2, _cute(2, 64, 4, 2)),
375
+ ),
376
+ name="routed_expert_up_proj",
377
+ ),
378
+ (7168, 8448): ProjectionSpec(
379
+ 7168,
380
+ 8448,
381
+ cute_configs=(
382
+ (1, _cute(1, 32, 4, 4)),
383
+ (2, _cute(2, 96, 4, 1)),
384
+ (3, _cute(3, 96, 4, 1)),
385
+ ),
386
+ name="dense_down_proj",
387
+ ),
388
+ (8448, 7168): ProjectionSpec(
389
+ 8448,
390
+ 7168,
391
+ cute_configs=(
392
+ (1, _cute(1, 224, 3, 4)),
393
+ (2, _cute(2, 32, 4, 4)),
394
+ ),
395
+ name="dense_gate_up_proj",
396
+ ),
397
+ (16896, 7168): ProjectionSpec(
398
+ 16896,
399
+ 7168,
400
+ cute_configs=((1, _cute(1, 224, 3, 4)),),
401
+ name="dense_gate_up_proj",
402
+ ),
403
+ (20480, 7168): ProjectionSpec(
404
+ 20480,
405
+ 7168,
406
+ cute_configs=(
407
+ (1, _cute(1, 224, 4, 2)),
408
+ (2, _cute(2, 64, 4, 2)),
409
+ (3, _cute(3, 64, 2, 2)),
410
+ (4, _cute(4, 64, 4, 1)),
411
+ ),
412
+ name="lm_head",
413
+ ),
414
+ (40960, 7168): ProjectionSpec(
415
+ 40960,
416
+ 7168,
417
+ cute_configs=(
418
+ (1, _cute(1, 128, 4, 2)),
419
+ (2, _cute(2, 64, 4, 2)),
420
+ (3, _cute(3, 64, 4, 1)),
421
+ (4, _cute(4, 64, 4, 1)),
422
+ ),
423
+ name="lm_head",
424
+ ),
425
+ (10240, 7168): ProjectionSpec(
426
+ 10240,
427
+ 7168,
428
+ cute_configs=(
429
+ (1, _cute(1, 224, 4, 2)),
430
+ (2, _cute(2, 32, 2, 4)),
431
+ (3, _cute(3, 64, 4, 1)),
432
+ (4, _cute(4, 64, 4, 1)),
433
+ ),
434
+ name="lm_head",
435
+ ),
436
+ # Latent-MoE replicated projections. dsv3 measured M2..8 wins on B300 for
437
+ # 3584x7168 but loses to cuBLAS there on B200; only CuTe wins land here.
438
+ (3584, 7168): ProjectionSpec(
439
+ 3584,
440
+ 7168,
441
+ cute_configs=(
442
+ (1, _cute(1, 224, 2, 4)),
443
+ (2, _cute(2, 128, 2, 1)),
444
+ (3, _cute(3, 128, 2, 1)),
445
+ ),
446
+ name="routed_expert_down_proj",
447
+ ),
448
+ # TP16 shapes, newly measured on B200 (they only appear in TP16 deployments;
449
+ # on TP8 these table entries simply never match).
450
+ (768, 7168): ProjectionSpec(
451
+ 768,
452
+ 7168,
453
+ cute_configs=(
454
+ (1, _cute(1, 224, 2, 4)),
455
+ (2, _cute(2, 224, 2, 2)),
456
+ (3, _cute(3, 224, 2, 2)),
457
+ (4, _cute(4, 224, 2, 2)),
458
+ ),
459
+ name="mla_g_proj/shared_gate_up_proj",
460
+ ),
461
+ (1152, 1536): ProjectionSpec(
462
+ 1152,
463
+ 1536,
464
+ cute_configs=((1, _cute(1, 192, 3, 4)),),
465
+ name="q_b_proj",
466
+ ),
467
+ (3216, 7168): ProjectionSpec(
468
+ 3216,
469
+ 7168,
470
+ cute_configs=(
471
+ (1, _cute(1, 224, 3, 4)),
472
+ (2, _cute(2, 128, 4, 2)),
473
+ ),
474
+ name="in_proj_qkvgfab",
475
+ ),
476
+ (4224, 7168): ProjectionSpec(
477
+ 4224,
478
+ 7168,
479
+ cute_configs=(
480
+ (1, _cute(1, 224, 3, 4)),
481
+ (2, _cute(2, 128, 2, 1)),
482
+ (3, _cute(3, 64, 2, 2)),
483
+ ),
484
+ name="dense_gate_up_proj",
485
+ ),
486
+}
487
+
488
489
def _backend_for(
490
spec: ProjectionSpec, num_tokens: int, has_residual: bool
@@ -324,6 +529,15 @@ def _is_sm103() -> bool:
529
return current_platform.is_device_capability((10, 3))
530
531
532
+def _low_latency_table() -> dict[tuple[int, int], ProjectionSpec] | None:
533
+ """Measured dispatch table for the current device, or None if unsupported."""
534
+ if _is_sm103():
535
+ return KIMI_K3_PROJECTIONS
536
+ if current_platform.is_device_capability((10, 0)):
537
+ return KIMI_K3_PROJECTIONS_SM100
538
+ return None
539
+
540
+
541
def _is_packed_row_major(tensor: torch.Tensor) -> bool:
542
return tensor.dim() == 2 and tensor.stride() == (tensor.shape[1], 1)
543
@@ -393,9 +607,12 @@ def try_low_latency_gemm(
607
precomputed plan (see :func:`enable_kimi_k3_low_latency_gemm`) and does not
608
use this path.
609
"""
396
- if envs.VLLM_BATCH_INVARIANT or not _is_sm103() or not _runtime_ok(x, weight):
610
+ if envs.VLLM_BATCH_INVARIANT or not _runtime_ok(x, weight):
611
return None
398
- spec = KIMI_K3_PROJECTIONS.get((weight.shape[0], weight.shape[1]))
612
+ table = _low_latency_table()
613
+ if table is None:
614
+ return None
615
+ spec = table.get((weight.shape[0], weight.shape[1]))
616
if spec is None:
617
return None
618
if residual is None:
@@ -467,9 +684,14 @@ def enable_kimi_k3_low_latency_gemm(
684
"""Install shape-selected low-latency GEMMs and register CuTe warmups.
685
686
Modules are matched purely by type, an exactly-unquantized method, and a
470
- local ``(N, K)`` present in :data:`KIMI_K3_PROJECTIONS`.
687
+ local ``(N, K)`` present in the current device's measured table
688
+ (:data:`KIMI_K3_PROJECTIONS` on SM103, :data:`KIMI_K3_PROJECTIONS_SM100`
689
+ on SM100).
690
"""
472
- if dtype != torch.bfloat16 or not _is_sm103():
691
+ if dtype != torch.bfloat16:
692
+ return
693
+ table = _low_latency_table()
694
+ if table is None:
695
return
696
697
warmup_configs: set[SkinnyGemmConfig] = set()
@@ -490,7 +712,7 @@ def enable_kimi_k3_low_latency_gemm(
712
weight = getattr(child, "weight", None)
713
if weight is None or weight.dim() != 2:
714
continue
493
- spec = KIMI_K3_PROJECTIONS.get((weight.shape[0], weight.shape[1]))
715
+ spec = table.get((weight.shape[0], weight.shape[1]))
716
if spec is None:
717
continue
718
if is_linear: