[Kimi Perf] Tune hopper low latency gemm kernel, 4%~97% performance improvement (#54088)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Aug 27, 2026 at 16:48 UTC
9818bb3db8fd6fdb7742dfa7f8b5007f67a736fd
2 files changed
+175
-35
tests/kernels/test_bf16_skinny_gemm.py
+77
-28
@@ -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/SM100 selectors."""
3
+"""Tests for BF16 skinny GEMMs and the Kimi-K3 SM90/SM100/SM103 selectors."""
4
5
from pathlib import Path
6
from types import SimpleNamespace
@@ -52,9 +52,16 @@ EXPECTED_SELECTIONS = {
52
(10240, 7168): (set(range(1, 5)), set()),
53
}
54
55
+K3ProjectionTable = dict[tuple[int, int], k3_gemm.ProjectionSpec]
56
+K3_GPU_TABLES = (
57
+ ((10, 3), k3_gemm.KIMI_K3_PROJECTIONS),
58
+ ((9, 0), k3_gemm.KIMI_K3_PROJECTIONS_SM90),
59
+)
60
+
61
CUTE_CASES = [
56
- (spec.n, spec.k, num_tokens)
57
- for spec in k3_gemm.KIMI_K3_PROJECTIONS.values()
62
+ (capability, spec.n, spec.k, num_tokens)
63
+ for capability, table in K3_GPU_TABLES
64
+ for spec in table.values()
65
for num_tokens, _ in spec.cute_configs
66
]
67
@@ -186,6 +193,7 @@ def test_every_dsv3_routed_shape_is_instantiated() -> None:
193
specs = [
194
*KIMI_K3_PROJECTIONS.values(),
195
*k3_gemm.KIMI_K3_PROJECTIONS_SM100.values(),
196
+ *k3_gemm.KIMI_K3_PROJECTIONS_SM90.values(),
197
glm52_gemm.GLM52_QKV_A_PROJECTION,
198
glm52_gemm.GLM52_Q_B_PROJECTION,
199
]
@@ -448,8 +456,12 @@ EXPECTED_SELECTIONS_SM100 = {
456
}
457
458
451
-def test_sm100_table_is_keyed_by_shape() -> None:
452
- for (n, k), spec in k3_gemm.KIMI_K3_PROJECTIONS_SM100.items():
459
+@pytest.mark.parametrize(
460
+ "table",
461
+ [k3_gemm.KIMI_K3_PROJECTIONS_SM100, k3_gemm.KIMI_K3_PROJECTIONS_SM90],
462
+)
463
+def test_device_table_is_keyed_by_shape(table: K3ProjectionTable) -> None:
464
+ for (n, k), spec in table.items():
465
assert (spec.n, spec.k) == (n, k)
466
467
@@ -468,8 +480,12 @@ def test_sm100_selector_table(key: tuple[int, int]) -> None:
480
assert backend is None
481
482
471
-def test_sm100_build_plan_matches_table() -> None:
472
- for spec in k3_gemm.KIMI_K3_PROJECTIONS_SM100.values():
483
+@pytest.mark.parametrize(
484
+ "table",
485
+ [k3_gemm.KIMI_K3_PROJECTIONS_SM100, k3_gemm.KIMI_K3_PROJECTIONS_SM90],
486
+)
487
+def test_device_build_plan_matches_table(table: K3ProjectionTable) -> None:
488
+ for spec in table.values():
489
plan = k3_gemm._build_plan(spec)
490
for num_tokens in range(1, 17):
491
backend = k3_gemm._backend_for(spec, num_tokens, has_residual=False)
@@ -479,6 +495,12 @@ def test_sm100_build_plan_matches_table() -> None:
495
assert plan[num_tokens][0] == backend
496
497
498
+def test_sm90_table_covers_measured_points() -> None:
499
+ specs = k3_gemm.KIMI_K3_PROJECTIONS_SM90.values()
500
+ assert len(k3_gemm.KIMI_K3_PROJECTIONS_SM90) == 18
501
+ assert sum(len(spec.cute_configs) + len(spec.dsv3_tokens) for spec in specs) == 90
502
+
503
+
504
def test_low_latency_table_capability_routing(
505
monkeypatch: pytest.MonkeyPatch,
506
) -> None:
@@ -493,6 +515,13 @@ def test_low_latency_table_capability_routing(
515
)
516
assert k3_gemm._low_latency_table() is k3_gemm.KIMI_K3_PROJECTIONS_SM100
517
518
+ monkeypatch.setattr(
519
+ k3_gemm.current_platform,
520
+ "is_device_capability",
521
+ lambda cc: cc == (9, 0),
522
+ )
523
+ assert k3_gemm._low_latency_table() is k3_gemm.KIMI_K3_PROJECTIONS_SM90
524
+
525
monkeypatch.setattr(
526
k3_gemm.current_platform, "is_device_capability", lambda cc: False
527
)
@@ -599,20 +628,34 @@ def test_installation_requires_bf16_sm103(
628
assert type(root.projection.quant_method) is k3_gemm.UnquantizedLinearMethod
629
630
602
-def _require_sm103_and_dsv3() -> None:
603
- if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 3):
604
- pytest.skip("Kimi-K3 production selection requires SM103")
631
+def _require_capability_and_dsv3(capability: tuple[int, int]) -> None:
632
+ if (
633
+ not torch.cuda.is_available()
634
+ or torch.cuda.get_device_capability() != capability
635
+ ):
636
+ pytest.skip(f"Kimi-K3 selection requires SM{capability[0]}{capability[1]}")
637
if not hasattr(torch.ops._C, "dsv3_fused_a_gemm"):
638
pytest.skip("dsv3_fused_a_gemm was not built")
639
640
609
-def _require_sm103_and_cute() -> None:
610
- if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 3):
611
- pytest.skip("Kimi-K3 production selection requires SM103")
641
+def _require_capability_and_cute(capability: tuple[int, int]) -> None:
642
+ if (
643
+ not torch.cuda.is_available()
644
+ or torch.cuda.get_device_capability() != capability
645
+ ):
646
+ pytest.skip(f"Kimi-K3 selection requires SM{capability[0]}{capability[1]}")
647
if not k3_gemm.shape_dynamic_skinny_gemm.is_available():
648
pytest.skip("CuTe DSL is not available")
649
650
651
+def _require_sm103_and_dsv3() -> None:
652
+ _require_capability_and_dsv3((10, 3))
653
+
654
+
655
+def _require_sm103_and_cute() -> None:
656
+ _require_capability_and_cute((10, 3))
657
+
658
+
659
@pytest.mark.parametrize("spec,config", GLM_CUTE_CASES)
660
def test_glm_cute_selected_shapes(
661
spec: glm52_gemm.GLM52ProjectionSpec,
@@ -683,9 +726,11 @@ def test_glm_cute_selected_shapes_cuda_graph_capture(
726
torch.testing.assert_close(output.float(), reference, rtol=2e-2, atol=2e-1)
727
728
686
-@pytest.mark.parametrize("n,k,num_tokens", CUTE_CASES)
687
-def test_cute_selected_shapes(n: int, k: int, num_tokens: int) -> None:
688
- _require_sm103_and_cute()
729
+@pytest.mark.parametrize("capability,n,k,num_tokens", CUTE_CASES)
730
+def test_cute_selected_shapes(
731
+ capability: tuple[int, int], n: int, k: int, num_tokens: int
732
+) -> None:
733
+ _require_capability_and_cute(capability)
734
torch.manual_seed(42)
735
x = torch.randn(num_tokens, k, dtype=torch.bfloat16, device="cuda")
736
weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda")
@@ -710,8 +755,9 @@ def _dsv3_probe_tokens(tokens: frozenset[int]) -> set[int]:
755
# Derived from the table rather than hand-listed, so a shape routed to dsv3
756
# cannot be added without being exercised here.
757
DSV3_CASES = sorted(
713
- (num_tokens, spec.n, spec.k)
714
- for spec in KIMI_K3_PROJECTIONS.values()
758
+ (capability, num_tokens, spec.n, spec.k)
759
+ for capability, table in K3_GPU_TABLES
760
+ for spec in table.values()
761
for num_tokens in _dsv3_probe_tokens(spec.dsv3_tokens)
762
)
763
@@ -725,11 +771,11 @@ GLM_DSV3_CASES = [
771
]
772
773
728
-@pytest.mark.parametrize("num_tokens,n,k", DSV3_CASES)
729
-def test_dsv3_selected_shapes(num_tokens: int, n: int, k: int) -> None:
730
- _require_sm103_and_dsv3()
731
- spec = k3_gemm.KIMI_K3_PROJECTIONS[(n, k)]
732
- assert num_tokens in spec.dsv3_tokens
774
+@pytest.mark.parametrize("capability,num_tokens,n,k", DSV3_CASES)
775
+def test_dsv3_selected_shapes(
776
+ capability: tuple[int, int], num_tokens: int, n: int, k: int
777
+) -> None:
778
+ _require_capability_and_dsv3(capability)
779
torch.manual_seed(42)
780
x = torch.randn(num_tokens, k, dtype=torch.bfloat16, device="cuda")
781
weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda")
@@ -784,11 +830,14 @@ def test_nonpacked_single_token_dsv3_falls_back() -> None:
830
torch.testing.assert_close(output, reference)
831
832
787
-def test_selected_kernels_cuda_graph_capture() -> None:
788
- _require_sm103_and_cute()
789
- _require_sm103_and_dsv3()
790
- cute_spec = k3_gemm.KIMI_K3_PROJECTIONS[(6288, 7168)]
791
- dsv3_spec = k3_gemm.KIMI_K3_PROJECTIONS[(1536, 128)]
833
+@pytest.mark.parametrize("capability,table", K3_GPU_TABLES)
834
+def test_selected_kernels_cuda_graph_capture(
835
+ capability: tuple[int, int], table: K3ProjectionTable
836
+) -> None:
837
+ _require_capability_and_cute(capability)
838
+ _require_capability_and_dsv3(capability)
839
+ cute_spec = table[(6288, 7168)]
840
+ dsv3_spec = table[(1536, 128)]
841
cute_x = torch.randn(1, cute_spec.k, dtype=torch.bfloat16, device="cuda")
842
cute_weight = torch.randn(
843
cute_spec.n, cute_spec.k, dtype=torch.bfloat16, device="cuda"
vllm/models/kimi_k3/nvidia/low_latency_gemm.py
+98
-7
@@ -1,6 +1,6 @@
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 and SM100.
3
+"""Kimi-K3 decode GEMM selection for unquantized BF16 on SM90/SM100/SM103.
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`
@@ -8,12 +8,11 @@ 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:
11
+The 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.
13
+:data:`KIMI_K3_PROJECTIONS_SM100` on B200 (SM100), and
14
+:data:`KIMI_K3_PROJECTIONS_SM90` on H200 (SM90). The per-(shape, M) winners
15
+genuinely differ between the parts, so the tables must not be merged.
16
"""
17
18
from __future__ import annotations
@@ -493,6 +492,96 @@ KIMI_K3_PROJECTIONS_SM100: dict[tuple[int, int], ProjectionSpec] = {
492
),
493
}
494
495
+_SM90CuteConfig = tuple[int, int, int, int]
496
+
497
+
498
+def _sm90_spec(
499
+ n: int,
500
+ k: int,
501
+ dsv3_tokens: frozenset[int] = frozenset(),
502
+ configs: tuple[_SM90CuteConfig, ...] = (),
503
+) -> ProjectionSpec:
504
+ cute_configs = tuple(
505
+ (m, _cute(m, *config)) for m, config in enumerate(configs, start=1)
506
+ )
507
+ return ProjectionSpec(n, k, dsv3_tokens, cute_configs)
508
+
509
+
510
+KIMI_K3_PROJECTIONS_SM90: dict[tuple[int, int], ProjectionSpec] = {
511
+ (1536, 128): _sm90_spec(1536, 128, frozenset(range(1, 9))),
512
+ (3072, 128): _sm90_spec(3072, 128, frozenset({1, 2, 5, 6, 7, 8, 9})),
513
+ (1536, 7168): _sm90_spec(
514
+ 1536,
515
+ 7168,
516
+ frozenset(range(7, 17)),
517
+ (
518
+ (224, 2, 4, 8),
519
+ (128, 3, 2, 8),
520
+ (128, 2, 1, 8),
521
+ (128, 2, 1, 8),
522
+ (128, 3, 1, 8),
523
+ (128, 3, 1, 8),
524
+ ),
525
+ ),
526
+ (3072, 7168): _sm90_spec(
527
+ 3072,
528
+ 7168,
529
+ configs=((224, 2, 4, 8), (128, 2, 2, 8), (64, 4, 1, 8), (128, 6, 1, 8)),
530
+ ),
531
+ (2112, 7168): _sm90_spec(
532
+ 2112,
533
+ 7168,
534
+ frozenset(range(5, 17)),
535
+ ((224, 2, 4, 8), (128, 4, 2, 8), (128, 2, 1, 8), (128, 2, 1, 8)),
536
+ ),
537
+ (2304, 1536): _sm90_spec(
538
+ 2304,
539
+ 1536,
540
+ frozenset(range(3, 9)),
541
+ ((96, 4, 2, 8), (96, 4, 1, 8)),
542
+ ),
543
+ (4608, 1536): _sm90_spec(
544
+ 4608,
545
+ 1536,
546
+ configs=((96, 4, 2, 4), (96, 4, 1, 8), (64, 4, 1, 8)),
547
+ ),
548
+ (3584, 7168): _sm90_spec(
549
+ 3584,
550
+ 7168,
551
+ configs=((224, 2, 4, 8), (128, 4, 2, 8), (128, 2, 1, 8)),
552
+ ),
553
+ (6288, 7168): _sm90_spec(
554
+ 6288,
555
+ 7168,
556
+ configs=((224, 2, 4, 8), (128, 4, 2, 8), (64, 4, 1, 8)),
557
+ ),
558
+ (12448, 7168): _sm90_spec(
559
+ 12448,
560
+ 7168,
561
+ configs=((224, 2, 4, 8), (224, 4, 2, 8), (128, 2, 1, 8)),
562
+ ),
563
+ (7168, 768): _sm90_spec(7168, 768, configs=((96, 4, 2, 4), (96, 4, 1, 8))),
564
+ (7168, 1536): _sm90_spec(7168, 1536, configs=((96, 4, 2, 8), (96, 4, 1, 8))),
565
+ (7168, 3072): _sm90_spec(
566
+ 7168,
567
+ 3072,
568
+ configs=((96, 2, 4, 8), (64, 4, 2, 8), (64, 2, 2, 8)),
569
+ ),
570
+ (7168, 3584): _sm90_spec(
571
+ 7168,
572
+ 3584,
573
+ configs=((224, 4, 2, 8), (64, 4, 2, 8), (64, 4, 2, 8)),
574
+ ),
575
+ (7168, 4224): _sm90_spec(7168, 4224, configs=((96, 4, 2, 4), (32, 4, 2, 4))),
576
+ (7168, 8448): _sm90_spec(7168, 8448, configs=((96, 2, 4, 8), (32, 4, 2, 8))),
577
+ (7168, 14336): _sm90_spec(7168, 14336, configs=((224, 2, 4, 8), (128, 2, 2, 8))),
578
+ (20480, 7168): _sm90_spec(
579
+ 20480,
580
+ 7168,
581
+ configs=((224, 4, 2, 8), (224, 4, 2, 8), (128, 2, 1, 8)),
582
+ ),
583
+}
584
+
585
586
def _backend_for(
587
spec: ProjectionSpec, num_tokens: int, has_residual: bool
@@ -543,6 +632,8 @@ def _low_latency_table() -> dict[tuple[int, int], ProjectionSpec] | None:
632
return KIMI_K3_PROJECTIONS
633
if current_platform.is_device_capability((10, 0)):
634
return KIMI_K3_PROJECTIONS_SM100
635
+ if current_platform.is_device_capability((9, 0)):
636
+ return KIMI_K3_PROJECTIONS_SM90
637
return None
638
639
@@ -694,7 +785,7 @@ def enable_kimi_k3_low_latency_gemm(
785
Modules are matched purely by type, an exactly-unquantized method, and a
786
local ``(N, K)`` present in the current device's measured table
787
(:data:`KIMI_K3_PROJECTIONS` on SM103, :data:`KIMI_K3_PROJECTIONS_SM100`
697
- on SM100).
788
+ on SM100, :data:`KIMI_K3_PROJECTIONS_SM90` on SM90).
789
"""
790
if dtype != torch.bfloat16:
791
return