From bc2d63e650f698aab24fd4d91889e0dcf16876f0 Mon Sep 17 00:00:00 2001 From: Canlin Guo Date: Wed, 26 Aug 2026 07:32:28 +0800 Subject: [PATCH] [Kimi K3][Kernel] Enable low-latency decode GEMM dispatch on SM100 (#53534) Signed-off-by: Canlin Guo --- tests/kernels/test_bf16_skinny_gemm.py | 96 ++++++- .../models/kimi_k3/nvidia/low_latency_gemm.py | 234 +++++++++++++++++- 2 files changed, 323 insertions(+), 7 deletions(-) diff --git a/tests/kernels/test_bf16_skinny_gemm.py b/tests/kernels/test_bf16_skinny_gemm.py index 09fb87d8d8..d542f9f483 100644 --- a/tests/kernels/test_bf16_skinny_gemm.py +++ b/tests/kernels/test_bf16_skinny_gemm.py @@ -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) diff --git a/vllm/models/kimi_k3/nvidia/low_latency_gemm.py b/vllm/models/kimi_k3/nvidia/low_latency_gemm.py index 64ea39b487..40dd51cb6d 100644 --- a/vllm/models/kimi_k3/nvidia/low_latency_gemm.py +++ b/vllm/models/kimi_k3/nvidia/low_latency_gemm.py @@ -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: