diff --git a/tests/kernels/test_bf16_skinny_gemm.py b/tests/kernels/test_bf16_skinny_gemm.py index 4988b4afd3..3fd03db49a 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/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" diff --git a/vllm/models/kimi_k3/nvidia/low_latency_gemm.py b/vllm/models/kimi_k3/nvidia/low_latency_gemm.py index 834dc880ed..ecab56dbe4 100644 --- a/vllm/models/kimi_k3/nvidia/low_latency_gemm.py +++ b/vllm/models/kimi_k3/nvidia/low_latency_gemm.py @@ -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