iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Kimi Perf] Tune hopper low latency gemm kernel, 4%~97% performance improvement (#54088)

Browse Source 
Signed-off-by: yewentao256 <zhyanwentao@126.com>
This commit is contained in:
Wentao Ye
2026-08-27 16:48:53 -04:00
committed by GitHub
parent 79378fe6e8
commit 9818bb3db8
2 changed files with 175 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/test_bf16_skinny_gemm.py
+77 -28
View File
@@ -1,6 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for BF16 skinny GEMMs and the Kimi-K3 SM103/SM100 selectors."""
"""Tests for BF16 skinny GEMMs and the Kimi-K3 SM90/SM100/SM103 selectors."""
from pathlib import Path
from types import SimpleNamespace
@@ -52,9 +52,16 @@ EXPECTED_SELECTIONS = {
(10240, 7168): (set(range(1, 5)), set()),
}
K3ProjectionTable = dict[tuple[int, int], k3_gemm.ProjectionSpec]
K3_GPU_TABLES = (
((10, 3), k3_gemm.KIMI_K3_PROJECTIONS),
((9, 0), k3_gemm.KIMI_K3_PROJECTIONS_SM90),
)
CUTE_CASES = [
(spec.n, spec.k, num_tokens)
for spec in k3_gemm.KIMI_K3_PROJECTIONS.values()
(capability, spec.n, spec.k, num_tokens)
for capability, table in K3_GPU_TABLES
for spec in table.values()
for num_tokens, _ in spec.cute_configs
]
@@ -186,6 +193,7 @@ def test_every_dsv3_routed_shape_is_instantiated() -> None:
specs = [
*KIMI_K3_PROJECTIONS.values(),
*k3_gemm.KIMI_K3_PROJECTIONS_SM100.values(),
*k3_gemm.KIMI_K3_PROJECTIONS_SM90.values(),
glm52_gemm.GLM52_QKV_A_PROJECTION,
glm52_gemm.GLM52_Q_B_PROJECTION,
]
@@ -448,8 +456,12 @@ EXPECTED_SELECTIONS_SM100 = {
}
def test_sm100_table_is_keyed_by_shape() -> None:
for (n, k), spec in k3_gemm.KIMI_K3_PROJECTIONS_SM100.items():
@pytest.mark.parametrize(
"table",
[k3_gemm.KIMI_K3_PROJECTIONS_SM100, k3_gemm.KIMI_K3_PROJECTIONS_SM90],
)
def test_device_table_is_keyed_by_shape(table: K3ProjectionTable) -> None:
for (n, k), spec in table.items():
assert (spec.n, spec.k) == (n, k)
@@ -468,8 +480,12 @@ def test_sm100_selector_table(key: tuple[int, int]) -> None:
assert backend is None
def test_sm100_build_plan_matches_table() -> None:
for spec in k3_gemm.KIMI_K3_PROJECTIONS_SM100.values():
@pytest.mark.parametrize(
"table",
[k3_gemm.KIMI_K3_PROJECTIONS_SM100, k3_gemm.KIMI_K3_PROJECTIONS_SM90],
)
def test_device_build_plan_matches_table(table: K3ProjectionTable) -> None:
for spec in table.values():
plan = k3_gemm._build_plan(spec)
for num_tokens in range(1, 17):
backend = k3_gemm._backend_for(spec, num_tokens, has_residual=False)
@@ -479,6 +495,12 @@ def test_sm100_build_plan_matches_table() -> None:
assert plan[num_tokens][0] == backend
def test_sm90_table_covers_measured_points() -> None:
specs = k3_gemm.KIMI_K3_PROJECTIONS_SM90.values()
assert len(k3_gemm.KIMI_K3_PROJECTIONS_SM90) == 18
assert sum(len(spec.cute_configs) + len(spec.dsv3_tokens) for spec in specs) == 90
def test_low_latency_table_capability_routing(
monkeypatch: pytest.MonkeyPatch,
) -> None:
@@ -493,6 +515,13 @@ def test_low_latency_table_capability_routing(
)
assert k3_gemm._low_latency_table() is k3_gemm.KIMI_K3_PROJECTIONS_SM100
monkeypatch.setattr(
k3_gemm.current_platform,
"is_device_capability",
lambda cc: cc == (9, 0),
)
assert k3_gemm._low_latency_table() is k3_gemm.KIMI_K3_PROJECTIONS_SM90
monkeypatch.setattr(
k3_gemm.current_platform, "is_device_capability", lambda cc: False
)
@@ -599,20 +628,34 @@ def test_installation_requires_bf16_sm103(
assert type(root.projection.quant_method) is k3_gemm.UnquantizedLinearMethod
def _require_sm103_and_dsv3() -> None:
if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 3):
pytest.skip("Kimi-K3 production selection requires SM103")
def _require_capability_and_dsv3(capability: tuple[int, int]) -> None:
if (
not torch.cuda.is_available()
or torch.cuda.get_device_capability() != capability
):
pytest.skip(f"Kimi-K3 selection requires SM{capability[0]}{capability[1]}")
if not hasattr(torch.ops._C, "dsv3_fused_a_gemm"):
pytest.skip("dsv3_fused_a_gemm was not built")
def _require_sm103_and_cute() -> None:
if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 3):
pytest.skip("Kimi-K3 production selection requires SM103")
def _require_capability_and_cute(capability: tuple[int, int]) -> None:
if (
not torch.cuda.is_available()
or torch.cuda.get_device_capability() != capability
):
pytest.skip(f"Kimi-K3 selection requires SM{capability[0]}{capability[1]}")
if not k3_gemm.shape_dynamic_skinny_gemm.is_available():
pytest.skip("CuTe DSL is not available")
def _require_sm103_and_dsv3() -> None:
_require_capability_and_dsv3((10, 3))
def _require_sm103_and_cute() -> None:
_require_capability_and_cute((10, 3))
@pytest.mark.parametrize("spec,config", GLM_CUTE_CASES)
def test_glm_cute_selected_shapes(
spec: glm52_gemm.GLM52ProjectionSpec,
@@ -683,9 +726,11 @@ def test_glm_cute_selected_shapes_cuda_graph_capture(
torch.testing.assert_close(output.float(), reference, rtol=2e-2, atol=2e-1)
@pytest.mark.parametrize("n,k,num_tokens", CUTE_CASES)
def test_cute_selected_shapes(n: int, k: int, num_tokens: int) -> None:
_require_sm103_and_cute()
@pytest.mark.parametrize("capability,n,k,num_tokens", CUTE_CASES)
def test_cute_selected_shapes(
capability: tuple[int, int], n: int, k: int, num_tokens: int
) -> None:
_require_capability_and_cute(capability)
torch.manual_seed(42)
x = torch.randn(num_tokens, k, dtype=torch.bfloat16, device="cuda")
weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda")
@@ -710,8 +755,9 @@ def _dsv3_probe_tokens(tokens: frozenset[int]) -> set[int]:
# Derived from the table rather than hand-listed, so a shape routed to dsv3
# cannot be added without being exercised here.
DSV3_CASES = sorted(
(num_tokens, spec.n, spec.k)
for spec in KIMI_K3_PROJECTIONS.values()
(capability, num_tokens, spec.n, spec.k)
for capability, table in K3_GPU_TABLES
for spec in table.values()
for num_tokens in _dsv3_probe_tokens(spec.dsv3_tokens)
)
@@ -725,11 +771,11 @@ GLM_DSV3_CASES = [
]
@pytest.mark.parametrize("num_tokens,n,k", DSV3_CASES)
def test_dsv3_selected_shapes(num_tokens: int, n: int, k: int) -> None:
_require_sm103_and_dsv3()
spec = k3_gemm.KIMI_K3_PROJECTIONS[(n, k)]
assert num_tokens in spec.dsv3_tokens
@pytest.mark.parametrize("capability,num_tokens,n,k", DSV3_CASES)
def test_dsv3_selected_shapes(
capability: tuple[int, int], num_tokens: int, n: int, k: int
) -> None:
_require_capability_and_dsv3(capability)
torch.manual_seed(42)
x = torch.randn(num_tokens, k, dtype=torch.bfloat16, device="cuda")
weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda")
@@ -784,11 +830,14 @@ def test_nonpacked_single_token_dsv3_falls_back() -> None:
torch.testing.assert_close(output, reference)
def test_selected_kernels_cuda_graph_capture() -> None:
_require_sm103_and_cute()
_require_sm103_and_dsv3()
cute_spec = k3_gemm.KIMI_K3_PROJECTIONS[(6288, 7168)]
dsv3_spec = k3_gemm.KIMI_K3_PROJECTIONS[(1536, 128)]
@pytest.mark.parametrize("capability,table", K3_GPU_TABLES)
def test_selected_kernels_cuda_graph_capture(
capability: tuple[int, int], table: K3ProjectionTable
) -> None:
_require_capability_and_cute(capability)
_require_capability_and_dsv3(capability)
cute_spec = table[(6288, 7168)]
dsv3_spec = table[(1536, 128)]
cute_x = torch.randn(1, cute_spec.k, dtype=torch.bfloat16, device="cuda")
cute_weight = torch.randn(
cute_spec.n, cute_spec.k, dtype=torch.bfloat16, device="cuda"
vllm/models/kimi_k3/nvidia/low_latency_gemm.py
+98 -7
View File
@@ -1,6 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Kimi-K3 decode GEMM selection for unquantized BF16 on SM103 and SM100.
"""Kimi-K3 decode GEMM selection for unquantized BF16 on SM90/SM100/SM103.
Dispatch is purely by local ``(N, K)`` shape and token count ``M`` — the module
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
resolved once per module at install time into a small ``{M: call}`` plan, so the
per-forward path is a single dict lookup.
The two supported capabilities carry separate measured tables:
The supported capabilities carry separate measured tables:
:data:`KIMI_K3_PROJECTIONS` was tuned on B300 (SM103),
:data:`KIMI_K3_PROJECTIONS_SM100` on B200 (SM100). The per-(shape, M) winners
genuinely differ between the two parts (e.g. 3584x7168 favors dsv3 at M2..8 on
B300 but only wins with CuTe at M1..3 on B200), so the tables must not be
merged.
:data:`KIMI_K3_PROJECTIONS_SM100` on B200 (SM100), and
:data:`KIMI_K3_PROJECTIONS_SM90` on H200 (SM90). The per-(shape, M) winners
genuinely differ between the parts, so the tables must not be merged.
"""
from __future__ import annotations
@@ -493,6 +492,96 @@ KIMI_K3_PROJECTIONS_SM100: dict[tuple[int, int], ProjectionSpec] = {
),
}
_SM90CuteConfig = tuple[int, int, int, int]
def _sm90_spec(
n: int,
k: int,
dsv3_tokens: frozenset[int] = frozenset(),
configs: tuple[_SM90CuteConfig, ...] = (),
) -> ProjectionSpec:
cute_configs = tuple(
(m, _cute(m, *config)) for m, config in enumerate(configs, start=1)
)
return ProjectionSpec(n, k, dsv3_tokens, cute_configs)
KIMI_K3_PROJECTIONS_SM90: dict[tuple[int, int], ProjectionSpec] = {
(1536, 128): _sm90_spec(1536, 128, frozenset(range(1, 9))),
(3072, 128): _sm90_spec(3072, 128, frozenset({1, 2, 5, 6, 7, 8, 9})),
(1536, 7168): _sm90_spec(
1536,
7168,
frozenset(range(7, 17)),
(
(224, 2, 4, 8),
(128, 3, 2, 8),
(128, 2, 1, 8),
(128, 2, 1, 8),
(128, 3, 1, 8),
(128, 3, 1, 8),
),
),
(3072, 7168): _sm90_spec(
3072,
7168,
configs=((224, 2, 4, 8), (128, 2, 2, 8), (64, 4, 1, 8), (128, 6, 1, 8)),
),
(2112, 7168): _sm90_spec(
2112,
7168,
frozenset(range(5, 17)),
((224, 2, 4, 8), (128, 4, 2, 8), (128, 2, 1, 8), (128, 2, 1, 8)),
),
(2304, 1536): _sm90_spec(
2304,
1536,
frozenset(range(3, 9)),
((96, 4, 2, 8), (96, 4, 1, 8)),
),
(4608, 1536): _sm90_spec(
4608,
1536,
configs=((96, 4, 2, 4), (96, 4, 1, 8), (64, 4, 1, 8)),
),
(3584, 7168): _sm90_spec(
3584,
7168,
configs=((224, 2, 4, 8), (128, 4, 2, 8), (128, 2, 1, 8)),
),
(6288, 7168): _sm90_spec(
6288,
7168,
configs=((224, 2, 4, 8), (128, 4, 2, 8), (64, 4, 1, 8)),
),
(12448, 7168): _sm90_spec(
12448,
7168,
configs=((224, 2, 4, 8), (224, 4, 2, 8), (128, 2, 1, 8)),
),
(7168, 768): _sm90_spec(7168, 768, configs=((96, 4, 2, 4), (96, 4, 1, 8))),
(7168, 1536): _sm90_spec(7168, 1536, configs=((96, 4, 2, 8), (96, 4, 1, 8))),
(7168, 3072): _sm90_spec(
7168,
3072,
configs=((96, 2, 4, 8), (64, 4, 2, 8), (64, 2, 2, 8)),
),
(7168, 3584): _sm90_spec(
7168,
3584,
configs=((224, 4, 2, 8), (64, 4, 2, 8), (64, 4, 2, 8)),
),
(7168, 4224): _sm90_spec(7168, 4224, configs=((96, 4, 2, 4), (32, 4, 2, 4))),
(7168, 8448): _sm90_spec(7168, 8448, configs=((96, 2, 4, 8), (32, 4, 2, 8))),
(7168, 14336): _sm90_spec(7168, 14336, configs=((224, 2, 4, 8), (128, 2, 2, 8))),
(20480, 7168): _sm90_spec(
20480,
7168,
configs=((224, 4, 2, 8), (224, 4, 2, 8), (128, 2, 1, 8)),
),
}
def _backend_for(
spec: ProjectionSpec, num_tokens: int, has_residual: bool
@@ -543,6 +632,8 @@ def _low_latency_table() -> dict[tuple[int, int], ProjectionSpec] | None:
return KIMI_K3_PROJECTIONS
if current_platform.is_device_capability((10, 0)):
return KIMI_K3_PROJECTIONS_SM100
if current_platform.is_device_capability((9, 0)):
return KIMI_K3_PROJECTIONS_SM90
return None
@@ -694,7 +785,7 @@ def enable_kimi_k3_low_latency_gemm(
Modules are matched purely by type, an exactly-unquantized method, and a
local ``(N, K)`` present in the current device's measured table
(:data:`KIMI_K3_PROJECTIONS` on SM103, :data:`KIMI_K3_PROJECTIONS_SM100`
on SM100).
on SM100, :data:`KIMI_K3_PROJECTIONS_SM90` on SM90).
"""
if dtype != torch.bfloat16:
return