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

[Kimi K3][Kernel] Enable low-latency decode GEMM dispatch on SM100 (#53534)

Browse Source 
Signed-off-by: Canlin Guo <canlinguosdu@gmail.com>
This commit is contained in:
Canlin Guo
2026-08-25 16:32:28 -07:00
committed by GitHub
parent bc39ded3c4
commit bc2d63e650
2 changed files with 323 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/test_bf16_skinny_gemm.py
+95 -1
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 selector."""
"""Tests for BF16 skinny GEMMs and the Kimi-K3 SM103/SM100 selectors."""
from pathlib import Path
from types import SimpleNamespace
@@ -182,6 +182,7 @@ def test_every_dsv3_routed_shape_is_instantiated() -> None:
specs = [
*KIMI_K3_PROJECTIONS.values(),
*k3_gemm.KIMI_K3_PROJECTIONS_SM100.values(),
glm52_gemm.GLM52_QKV_A_PROJECTION,
glm52_gemm.GLM52_Q_B_PROJECTION,
]
@@ -404,6 +405,92 @@ def test_build_plan_matches_selector() -> None:
assert plan[num_tokens][0] == backend
# Keyed by local (N, K): (cute token counts, dsv3 token counts). Mirrors
# KIMI_K3_PROJECTIONS_SM100, measured on B200; intentionally different from
# EXPECTED_SELECTIONS (e.g. 3584x7168 routes dsv3 at M2..8 on SM103 but only
# wins with CuTe at M1..3 on SM100).
EXPECTED_SELECTIONS_SM100 = {
(1536, 128): (set(), {1, 16}),
(3072, 128): (set(), {8}),
(1536, 7168): ({1, 2}, {4, 8}),
(3072, 7168): ({1, 2}, set()),
(2112, 7168): ({1, 2}, {4, 16}),
(2304, 1536): (set(), set(range(1, 17))),
(4608, 1536): (set(), {1, 2, 4}),
(6288, 7168): ({1}, set()),
(12448, 7168): ({1, 2, 3, 4}, set()),
(7168, 768): (set(), {1}),
# SM103-tuned config wins on B200 where the heuristic config tied.
(7168, 1536): ({1}, set()),
# M2 config lifted from the SM103 table (wins on B200).
(7168, 3072): ({1, 2}, set()),
(7168, 3584): ({1, 2}, set()),
(7168, 8448): ({1, 2, 3}, set()),
(8448, 7168): ({1, 2}, set()),
(16896, 7168): ({1}, set()),
(20480, 7168): ({1, 2, 3, 4}, set()),
(40960, 7168): ({1, 2, 3, 4}, set()),
(10240, 7168): ({1, 2, 3, 4}, set()),
(3584, 7168): ({1, 2, 3}, set()),
# TP16 shapes measured on B200.
(768, 7168): ({1, 2, 3, 4}, set()),
(1152, 1536): ({1}, set()),
(3216, 7168): ({1, 2}, set()),
(4224, 7168): ({1, 2, 3}, set()),
}
def test_sm100_table_is_keyed_by_shape() -> None:
for (n, k), spec in k3_gemm.KIMI_K3_PROJECTIONS_SM100.items():
assert (spec.n, spec.k) == (n, k)
@pytest.mark.parametrize("key", EXPECTED_SELECTIONS_SM100)
def test_sm100_selector_table(key: tuple[int, int]) -> None:
n, k = key
spec = k3_gemm.KIMI_K3_PROJECTIONS_SM100[key]
cute_tokens, dsv3_tokens = EXPECTED_SELECTIONS_SM100[key]
for num_tokens in range(1, 17):
backend = k3_gemm._backend_for(spec, num_tokens, has_residual=False)
if num_tokens in cute_tokens:
assert backend == "cute"
elif num_tokens in dsv3_tokens:
assert backend == "dsv3_fused_a"
else:
assert backend is None
def test_sm100_build_plan_matches_table() -> None:
for spec in k3_gemm.KIMI_K3_PROJECTIONS_SM100.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)
if backend is None:
assert num_tokens not in plan
else:
assert plan[num_tokens][0] == backend
def test_low_latency_table_capability_routing(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(k3_gemm, "_is_sm103", lambda: True)
assert k3_gemm._low_latency_table() is k3_gemm.KIMI_K3_PROJECTIONS
monkeypatch.setattr(k3_gemm, "_is_sm103", lambda: False)
monkeypatch.setattr(
k3_gemm.current_platform,
"is_device_capability",
lambda cc: cc == (10, 0),
)
assert k3_gemm._low_latency_table() is k3_gemm.KIMI_K3_PROJECTIONS_SM100
monkeypatch.setattr(
k3_gemm.current_platform, "is_device_capability", lambda cc: False
)
assert k3_gemm._low_latency_table() is None
def test_installation_is_shape_specific_and_unquantized(
monkeypatch: pytest.MonkeyPatch,
) -> None:
@@ -491,6 +578,13 @@ def test_installation_requires_bf16_sm103(
root.projection = FakeLinear()
monkeypatch.setattr(k3_gemm, "LinearBase", FakeLinear)
monkeypatch.setattr(k3_gemm, "_is_sm103", lambda: platform_enabled)
# Pin the rest of the capability probe so platform_enabled=False stays a
# no-op even when the test itself runs on SM100 hardware.
monkeypatch.setattr(
k3_gemm.current_platform,
"is_device_capability",
lambda cc: platform_enabled and cc == (10, 3),
)
k3_gemm.enable_kimi_k3_low_latency_gemm(root, dtype)
vllm/models/kimi_k3/nvidia/low_latency_gemm.py
+228 -6
View File
@@ -1,12 +1,19 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Kimi-K3 decode GEMM selection for unquantized BF16 on SM103.
"""Kimi-K3 decode GEMM selection for unquantized BF16 on SM103 and SM100.
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`
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:
: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.
"""
from __future__ import annotations
@@ -280,6 +287,204 @@ KIMI_K3_PROJECTIONS: dict[tuple[int, int], ProjectionSpec] = {
# vector_width=2, which measured slower than cuBLAS.
}
# Keyed by local (N, K), B200 (SM100) entries. Measured on B200 at TP8/TP16
# shapes over M=1..16 with the same >=5%-over-cuBLAS threshold as the SM103
# table. Per-(shape, M) entries take the best of the SM103-tuned config and a
# fresh SM100 measurement: several SM103 configs (7168x8448, 20480x7168,
# 40960x7168, 10240x7168, 7168x3072 M2, 7168x1536) beat the heuristic config on
# B200 outright, while a few (6288x7168 M2..4, 3072x7168 M3..5) lose to cuBLAS
# on B200 and are intentionally not inherited. The 7168x3584 residual configs
# are intentionally absent: the B200 residual path was not measured.
KIMI_K3_PROJECTIONS_SM100: dict[tuple[int, int], ProjectionSpec] = {
(1536, 128): ProjectionSpec(1536, 128, frozenset({1, 16}), name="f_b_proj"),
(3072, 128): ProjectionSpec(3072, 128, frozenset({8}), name="f_b_proj"),
# Same mla_g_proj aux-stream/PDL capture caveat as the SM103 entry applies;
# only M4/M8 measured a dsv3 win, so this is the narrow set either way.
(1536, 7168): ProjectionSpec(
1536,
7168,
frozenset({4, 8}),
(
(1, _cute(1, 224, 3, 2)),
(2, _cute(2, 128, 2, 2)),
),
name="shared_gate_up_proj/mla_g_proj",
),
(3072, 7168): ProjectionSpec(
3072,
7168,
cute_configs=(
(1, _cute(1, 224, 3, 4)),
(2, _cute(2, 128, 2, 1)),
),
name="shared_gate_up_proj",
),
(2112, 7168): ProjectionSpec(
2112,
7168,
frozenset({4, 16}),
(
(1, _cute(1, 224, 3, 2)),
(2, _cute(2, 128, 2, 2)),
),
name="fused_qkv_a_proj",
),
(2304, 1536): ProjectionSpec(2304, 1536, _M1_TO_16, name="q_b_proj"),
(4608, 1536): ProjectionSpec(4608, 1536, frozenset({1, 2, 4}), name="q_b_proj"),
# M2..4 of the SM103 configs lose to cuBLAS on B200; only M1 carries over.
(6288, 7168): ProjectionSpec(
6288,
7168,
cute_configs=((1, _cute(1, 224, 3, 4)),),
name="in_proj_qkvgfab",
),
(12448, 7168): ProjectionSpec(
12448,
7168,
cute_configs=(
(1, _cute(1, 224, 4, 2)),
(2, _cute(2, 64, 4, 2)),
(3, _cute(3, 64, 2, 2)),
(4, _cute(4, 128, 2, 1)),
),
name="in_proj_qkvgfab",
),
(7168, 768): ProjectionSpec(7168, 768, frozenset({1}), name="shared_down_proj"),
# SM103-tuned config wins on B200 (+25%) where the heuristic config tied.
(7168, 1536): ProjectionSpec(
7168,
1536,
cute_configs=((1, _cute(1, 96, 4, 2)),),
name="o_proj",
),
(7168, 3072): ProjectionSpec(
7168,
3072,
cute_configs=(
(1, _cute(1, 96, 2, 4)),
(2, _cute(2, 32, 4, 4)),
),
name="o_proj",
),
(7168, 3584): ProjectionSpec(
7168,
3584,
cute_configs=(
(1, _cute(1, 64, 4, 2)),
(2, _cute(2, 64, 4, 2)),
),
name="routed_expert_up_proj",
),
(7168, 8448): ProjectionSpec(
7168,
8448,
cute_configs=(
(1, _cute(1, 32, 4, 4)),
(2, _cute(2, 96, 4, 1)),
(3, _cute(3, 96, 4, 1)),
),
name="dense_down_proj",
),
(8448, 7168): ProjectionSpec(
8448,
7168,
cute_configs=(
(1, _cute(1, 224, 3, 4)),
(2, _cute(2, 32, 4, 4)),
),
name="dense_gate_up_proj",
),
(16896, 7168): ProjectionSpec(
16896,
7168,
cute_configs=((1, _cute(1, 224, 3, 4)),),
name="dense_gate_up_proj",
),
(20480, 7168): ProjectionSpec(
20480,
7168,
cute_configs=(
(1, _cute(1, 224, 4, 2)),
(2, _cute(2, 64, 4, 2)),
(3, _cute(3, 64, 2, 2)),
(4, _cute(4, 64, 4, 1)),
),
name="lm_head",
),
(40960, 7168): ProjectionSpec(
40960,
7168,
cute_configs=(
(1, _cute(1, 128, 4, 2)),
(2, _cute(2, 64, 4, 2)),
(3, _cute(3, 64, 4, 1)),
(4, _cute(4, 64, 4, 1)),
),
name="lm_head",
),
(10240, 7168): ProjectionSpec(
10240,
7168,
cute_configs=(
(1, _cute(1, 224, 4, 2)),
(2, _cute(2, 32, 2, 4)),
(3, _cute(3, 64, 4, 1)),
(4, _cute(4, 64, 4, 1)),
),
name="lm_head",
),
# Latent-MoE replicated projections. dsv3 measured M2..8 wins on B300 for
# 3584x7168 but loses to cuBLAS there on B200; only CuTe wins land here.
(3584, 7168): ProjectionSpec(
3584,
7168,
cute_configs=(
(1, _cute(1, 224, 2, 4)),
(2, _cute(2, 128, 2, 1)),
(3, _cute(3, 128, 2, 1)),
),
name="routed_expert_down_proj",
),
# TP16 shapes, newly measured on B200 (they only appear in TP16 deployments;
# on TP8 these table entries simply never match).
(768, 7168): ProjectionSpec(
768,
7168,
cute_configs=(
(1, _cute(1, 224, 2, 4)),
(2, _cute(2, 224, 2, 2)),
(3, _cute(3, 224, 2, 2)),
(4, _cute(4, 224, 2, 2)),
),
name="mla_g_proj/shared_gate_up_proj",
),
(1152, 1536): ProjectionSpec(
1152,
1536,
cute_configs=((1, _cute(1, 192, 3, 4)),),
name="q_b_proj",
),
(3216, 7168): ProjectionSpec(
3216,
7168,
cute_configs=(
(1, _cute(1, 224, 3, 4)),
(2, _cute(2, 128, 4, 2)),
),
name="in_proj_qkvgfab",
),
(4224, 7168): ProjectionSpec(
4224,
7168,
cute_configs=(
(1, _cute(1, 224, 3, 4)),
(2, _cute(2, 128, 2, 1)),
(3, _cute(3, 64, 2, 2)),
),
name="dense_gate_up_proj",
),
}
def _backend_for(
spec: ProjectionSpec, num_tokens: int, has_residual: bool
@@ -324,6 +529,15 @@ def _is_sm103() -> bool:
return current_platform.is_device_capability((10, 3))
def _low_latency_table() -> dict[tuple[int, int], ProjectionSpec] | None:
"""Measured dispatch table for the current device, or None if unsupported."""
if _is_sm103():
return KIMI_K3_PROJECTIONS
if current_platform.is_device_capability((10, 0)):
return KIMI_K3_PROJECTIONS_SM100
return None
def _is_packed_row_major(tensor: torch.Tensor) -> bool:
return tensor.dim() == 2 and tensor.stride() == (tensor.shape[1], 1)
@@ -393,9 +607,12 @@ def try_low_latency_gemm(
precomputed plan (see :func:`enable_kimi_k3_low_latency_gemm`) and does not
use this path.
"""
if envs.VLLM_BATCH_INVARIANT or not _is_sm103() or not _runtime_ok(x, weight):
if envs.VLLM_BATCH_INVARIANT or not _runtime_ok(x, weight):
return None
spec = KIMI_K3_PROJECTIONS.get((weight.shape[0], weight.shape[1]))
table = _low_latency_table()
if table is None:
return None
spec = table.get((weight.shape[0], weight.shape[1]))
if spec is None:
return None
if residual is None:
@@ -467,9 +684,14 @@ def enable_kimi_k3_low_latency_gemm(
"""Install shape-selected low-latency GEMMs and register CuTe warmups.
Modules are matched purely by type, an exactly-unquantized method, and a
local ``(N, K)`` present in :data:`KIMI_K3_PROJECTIONS`.
local ``(N, K)`` present in the current device's measured table
(:data:`KIMI_K3_PROJECTIONS` on SM103, :data:`KIMI_K3_PROJECTIONS_SM100`
on SM100).
"""
if dtype != torch.bfloat16 or not _is_sm103():
if dtype != torch.bfloat16:
return
table = _low_latency_table()
if table is None:
return
warmup_configs: set[SkinnyGemmConfig] = set()
@@ -490,7 +712,7 @@ def enable_kimi_k3_low_latency_gemm(
weight = getattr(child, "weight", None)
if weight is None or weight.dim() != 2:
continue
spec = KIMI_K3_PROJECTIONS.get((weight.shape[0], weight.shape[1]))
spec = table.get((weight.shape[0], weight.shape[1]))
if spec is None:
continue
if is_linear: