| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for BF16 skinny GEMMs and model-specific selectors.""" |
| 4 | |
| 5 | from pathlib import Path |
| 6 | from types import SimpleNamespace |
| 7 | from typing import Any |
| 8 | |
| 9 | import pytest |
| 10 | import regex as re |
| 11 | import torch |
| 12 | from torch import nn |
| 13 | |
| 14 | from vllm.model_executor.kernels.linear.cute_dsl.skinny_gemm import ( |
| 15 | SkinnyGemmConfig, |
| 16 | ) |
| 17 | from vllm.models.deepseek_v32.nvidia import glm52_low_latency_gemm as glm52_gemm |
| 18 | from vllm.models.kimi_k3.nvidia import low_latency_gemm as k3_gemm |
| 19 | from vllm.models.kimi_k3.nvidia.low_latency_gemm import KIMI_K3_PROJECTIONS |
| 20 | from vllm.models.qwen4_exp.nvidia import low_latency_gemm as qwen4_exp_gemm |
| 21 | from vllm.platforms import current_platform |
| 22 | |
| 23 | # Keyed by local (N, K): (cute token counts, dsv3 token counts). 1536x7168 is |
| 24 | # the unified shared_gate_up_proj/mla_g_proj entry (dsv3 M1..16). |
| 25 | EXPECTED_SELECTIONS = { |
| 26 | (1536, 128): (set(), set(range(1, 17))), |
| 27 | (3072, 128): (set(), set(range(1, 17))), |
| 28 | (1536, 7168): (set(), set(range(1, 17))), |
| 29 | (3072, 7168): (set(range(1, 6)), set()), |
| 30 | (2112, 7168): (set(), set(range(1, 17))), |
| 31 | (2304, 1536): (set(), set(range(1, 17))), |
| 32 | (4608, 1536): (set(), set(range(1, 17))), |
| 33 | (3584, 7168): ({1}, set(range(2, 9))), |
| 34 | (6288, 7168): (set(range(1, 5)), set()), |
| 35 | (12448, 7168): (set(range(1, 4)), set()), |
| 36 | (7168, 768): (set(), set(range(1, 17))), |
| 37 | (7168, 1536): ({1}, set()), |
| 38 | (7168, 3072): ({1, 2}, set()), |
| 39 | (7168, 3584): ({1, 2}, set()), |
| 40 | (7168, 4224): ({1}, set()), |
| 41 | (7168, 8448): (set(range(1, 4)), set()), |
| 42 | (8448, 7168): ({1, 2}, set()), |
| 43 | (16896, 7168): ({1, 2}, set()), |
| 44 | (7168, 14336): ({1, 2}, set()), |
| 45 | (20480, 7168): (set(range(1, 5)), set()), |
| 46 | (40960, 7168): (set(range(1, 5)), set()), |
| 47 | # TP16. |
| 48 | (3216, 7168): (set(range(1, 6)), set(range(9, 16))), |
| 49 | (768, 7168): (set(range(1, 5)), set(range(5, 17))), |
| 50 | (1152, 1536): ({1}, set(range(2, 17))), |
| 51 | (768, 128): (set(), set(range(1, 17))), |
| 52 | (7168, 384): (set(), set(range(1, 9))), |
| 53 | (4224, 7168): (set(range(1, 4)), set(range(4, 9))), |
| 54 | (10240, 7168): (set(range(1, 5)), set()), |
| 55 | } |
| 56 | |
| 57 | K3ProjectionTable = dict[tuple[int, int], k3_gemm.ProjectionSpec] |
| 58 | K3_GPU_TABLES = ( |
| 59 | ((10, 3), k3_gemm.KIMI_K3_PROJECTIONS), |
| 60 | ((9, 0), k3_gemm.KIMI_K3_PROJECTIONS_SM90), |
| 61 | ) |
| 62 | |
| 63 | CUTE_CASES = [ |
| 64 | (capability, spec.n, spec.k, num_tokens) |
| 65 | for capability, table in K3_GPU_TABLES |
| 66 | for spec in table.values() |
| 67 | for num_tokens, _ in spec.cute_configs |
| 68 | ] |
| 69 | |
| 70 | RESIDUAL_CUTE_CASES = [ |
| 71 | (spec.n, spec.k, num_tokens) |
| 72 | for spec in k3_gemm.KIMI_K3_PROJECTIONS.values() |
| 73 | for num_tokens, _ in spec.residual_configs |
| 74 | ] |
| 75 | |
| 76 | GLM_CUTE_CASES = [ |
| 77 | (spec, config) |
| 78 | for spec in glm52_gemm.GLM52_PROJECTIONS.values() |
| 79 | for _, config in spec.cute_configs |
| 80 | ] |
| 81 | |
| 82 | QWEN4_EXP_SM90_CASES = [ |
| 83 | (n, k, num_tokens, config) |
| 84 | for (n, k), plans in qwen4_exp_gemm.QWEN4_EXP_SM90_GEMM_PLANS.items() |
| 85 | for num_tokens, config in plans.items() |
| 86 | ] |
| 87 | |
| 88 | QWEN4_EXP_SM100_CASES = [ |
| 89 | (n, k, num_tokens, config) |
| 90 | for (n, k), plans in qwen4_exp_gemm.QWEN4_EXP_SM100_GEMM_PLANS.items() |
| 91 | for num_tokens, config in plans.items() |
| 92 | ] |
| 93 | |
| 94 | EXPECTED_CUTE_CONFIGS = { |
| 95 | (3072, 7168, 1): (224, 3, 4, 8), |
| 96 | (3072, 7168, 2): (128, 3, 2, 8), |
| 97 | (3072, 7168, 3): (128, 2, 1, 8), |
| 98 | (3072, 7168, 4): (64, 2, 2, 8), |
| 99 | (3072, 7168, 5): (128, 3, 1, 8), |
| 100 | (3584, 7168, 1): (224, 2, 4, 8), |
| 101 | (6288, 7168, 1): (224, 3, 4, 8), |
| 102 | (6288, 7168, 2): (64, 3, 2, 8), |
| 103 | (6288, 7168, 3): (32, 3, 4, 8), |
| 104 | (6288, 7168, 4): (128, 6, 1, 8), |
| 105 | (12448, 7168, 1): (224, 4, 2, 8), |
| 106 | (12448, 7168, 2): (64, 4, 2, 8), |
| 107 | (12448, 7168, 3): (64, 2, 2, 8), |
| 108 | (7168, 1536, 1): (96, 4, 2, 8), |
| 109 | (7168, 3072, 1): (96, 2, 4, 8), |
| 110 | (7168, 3072, 2): (32, 4, 4, 8), |
| 111 | (7168, 3584, 1): (224, 4, 2, 8), |
| 112 | (7168, 3584, 2): (64, 4, 2, 8), |
| 113 | (7168, 4224, 1): (96, 4, 2, 4), |
| 114 | (7168, 8448, 1): (32, 4, 4, 8), |
| 115 | (7168, 8448, 2): (96, 4, 1, 8), |
| 116 | (7168, 8448, 3): (96, 4, 1, 8), |
| 117 | (8448, 7168, 1): (224, 3, 4, 8), |
| 118 | (8448, 7168, 2): (32, 4, 4, 8), |
| 119 | (16896, 7168, 1): (224, 6, 4, 8), |
| 120 | (16896, 7168, 2): (32, 4, 4, 8), |
| 121 | (7168, 14336, 1): (256, 2, 1, 4), |
| 122 | (7168, 14336, 2): (224, 4, 2, 8), |
| 123 | (20480, 7168, 1): (224, 4, 2, 8), |
| 124 | (20480, 7168, 2): (64, 4, 2, 8), |
| 125 | (20480, 7168, 3): (64, 2, 2, 8), |
| 126 | (20480, 7168, 4): (64, 4, 1, 8), |
| 127 | (40960, 7168, 1): (128, 4, 2, 8), |
| 128 | (40960, 7168, 2): (64, 4, 2, 8), |
| 129 | (40960, 7168, 3): (64, 2, 2, 8), |
| 130 | (40960, 7168, 4): (64, 4, 1, 8), |
| 131 | # TP16. |
| 132 | (3216, 7168, 1): (224, 3, 4, 8), |
| 133 | (3216, 7168, 2): (128, 4, 2, 8), |
| 134 | (3216, 7168, 3): (128, 2, 1, 8), |
| 135 | (3216, 7168, 4): (64, 2, 2, 8), |
| 136 | (3216, 7168, 5): (128, 3, 1, 8), |
| 137 | (768, 7168, 1): (224, 2, 4, 8), |
| 138 | (768, 7168, 2): (224, 2, 2, 8), |
| 139 | (768, 7168, 3): (224, 2, 2, 8), |
| 140 | (768, 7168, 4): (224, 2, 2, 8), |
| 141 | (1152, 1536, 1): (192, 3, 4, 8), |
| 142 | (4224, 7168, 1): (224, 3, 4, 8), |
| 143 | (4224, 7168, 2): (128, 2, 1, 8), |
| 144 | (4224, 7168, 3): (64, 2, 2, 8), |
| 145 | (10240, 7168, 1): (224, 4, 2, 8), |
| 146 | (10240, 7168, 2): (32, 2, 4, 8), |
| 147 | (10240, 7168, 3): (64, 4, 1, 8), |
| 148 | (10240, 7168, 4): (64, 4, 1, 8), |
| 149 | } |
| 150 | |
| 151 | EXPECTED_RESIDUAL_CUTE_CONFIGS = { |
| 152 | (7168, 3584, 1): (64, 4, 2, 8), |
| 153 | (7168, 3584, 2): (64, 7, 2, 8), |
| 154 | (7168, 3584, 3): (64, 2, 1, 8), |
| 155 | (7168, 3584, 4): (64, 2, 1, 8), |
| 156 | } |
| 157 | |
| 158 | |
| 159 | def _config_tuple(config) -> tuple[int, int, int, int]: |
| 160 | return ( |
| 161 | config.block_size, |
| 162 | config.outputs_per_block, |
| 163 | config.k_unroll, |
| 164 | config.vector_width, |
| 165 | ) |
| 166 | |
| 167 | |
| 168 | def test_table_is_keyed_by_shape() -> None: |
| 169 | for (n, k), spec in k3_gemm.KIMI_K3_PROJECTIONS.items(): |
| 170 | assert (spec.n, spec.k) == (n, k) |
| 171 | |
| 172 | |
| 173 | def test_every_dsv3_routed_shape_is_instantiated() -> None: |
| 174 | """dsv3_fused_a_gemm specializes on (K, N); an unlisted shape raises. |
| 175 | |
| 176 | The table routes by shape while the kernel is built per shape, so a missing |
| 177 | instantiation only shows up at the token counts that route to dsv3. Checking |
| 178 | it here needs no GPU, which is the point -- a GPU-only check is exactly what |
| 179 | let (3216, 7168) ship without its DISPATCH_DSV3_SHAPE(7168, 3216). |
| 180 | """ |
| 181 | source = ( |
| 182 | Path(__file__).resolve().parents[2] |
| 183 | / "csrc" |
| 184 | / "libtorch_stable" |
| 185 | / "dsv3_fused_a_gemm.cu" |
| 186 | ).read_text(encoding="utf-8") |
| 187 | # Benchmark-only shapes live behind VLLM_K3_BENCH_SHAPES and are not built |
| 188 | # by default, so they must not count as available. |
| 189 | production_macros = source.split("#ifdef VLLM_K3_BENCH_SHAPES")[0] |
| 190 | explicit = source.split("#undef DISPATCH_DSV3_SHAPE")[1].split( |
| 191 | "#ifdef VLLM_K3_BENCH_SHAPES" |
| 192 | )[0] |
| 193 | compiled = { |
| 194 | (int(hd_in), int(hd_out)) |
| 195 | for hd_in, hd_out in re.findall( |
| 196 | r"DISPATCH_DSV3_SHAPE\((\d+),\s*(\d+)\)", production_macros |
| 197 | ) |
| 198 | } | { |
| 199 | (int(hd_in), int(hd_out)) |
| 200 | for hd_in, hd_out in re.findall( |
| 201 | r"hd_in == (\d+) && hd_out == (\d+)", |
| 202 | production_macros + explicit, |
| 203 | ) |
| 204 | } |
| 205 | assert compiled, "failed to parse the dispatch list" |
| 206 | |
| 207 | specs = [ |
| 208 | *KIMI_K3_PROJECTIONS.values(), |
| 209 | *k3_gemm.KIMI_K3_PROJECTIONS_SM100.values(), |
| 210 | *k3_gemm.KIMI_K3_PROJECTIONS_SM90.values(), |
| 211 | glm52_gemm.GLM52_QKV_A_PROJECTION, |
| 212 | glm52_gemm.GLM52_Q_B_PROJECTION, |
| 213 | ] |
| 214 | missing = sorted( |
| 215 | (spec.n, spec.k) |
| 216 | for spec in specs |
| 217 | if spec.dsv3_tokens and (spec.k, spec.n) not in compiled |
| 218 | ) |
| 219 | assert not missing, ( |
| 220 | f"routed to dsv3 with no instantiation: {missing}; add " |
| 221 | "DISPATCH_DSV3_SHAPE(K, N) for each" |
| 222 | ) |
| 223 | |
| 224 | |
| 225 | def test_packed_row_major_rejects_single_row_slice() -> None: |
| 226 | packed = torch.empty(1, 128) |
| 227 | sliced = torch.empty(1, 144)[:, :128] |
| 228 | |
| 229 | assert packed.is_contiguous() |
| 230 | assert sliced.is_contiguous() |
| 231 | assert k3_gemm._is_packed_row_major(packed) |
| 232 | assert not k3_gemm._is_packed_row_major(sliced) |
| 233 | |
| 234 | |
| 235 | def test_cute_configs_match_measured_table() -> None: |
| 236 | configs = [ |
| 237 | (spec.n, spec.k, num_tokens, config) |
| 238 | for spec in k3_gemm.KIMI_K3_PROJECTIONS.values() |
| 239 | for num_tokens, config in spec.cute_configs |
| 240 | ] |
| 241 | actual = { |
| 242 | (n, k, num_tokens): _config_tuple(config) |
| 243 | for n, k, num_tokens, config in configs |
| 244 | } |
| 245 | assert actual == EXPECTED_CUTE_CONFIGS |
| 246 | static_k_configs = { |
| 247 | (n, k, num_tokens, config.static_k) |
| 248 | for n, k, num_tokens, config in configs |
| 249 | if config.static_k is not None |
| 250 | } |
| 251 | assert static_k_configs == {(7168, 14336, 1, 14336)} |
| 252 | |
| 253 | |
| 254 | def test_residual_cute_configs_match_measured_table() -> None: |
| 255 | actual = { |
| 256 | (spec.n, spec.k, num_tokens): _config_tuple(config) |
| 257 | for spec in k3_gemm.KIMI_K3_PROJECTIONS.values() |
| 258 | for num_tokens, config in spec.residual_configs |
| 259 | } |
| 260 | assert actual == EXPECTED_RESIDUAL_CUTE_CONFIGS |
| 261 | |
| 262 | |
| 263 | def test_kda_overlap_configs_match_measured_table() -> None: |
| 264 | assert k3_gemm.KDA_PROJECTION_OVERLAP_MAX_TOKENS == 14 |
| 265 | assert k3_gemm.KDA_SKINNY_N_MAX_TOKENS == 14 |
| 266 | assert k3_gemm.KDA_SKINNY_K_MAX_TOKENS == 14 |
| 267 | assert _config_tuple(k3_gemm.KDA_M1_QKVG_CONFIG) == (64, 4, 2, 8) |
| 268 | assert _config_tuple(k3_gemm.KDA_M1_FAB_CONFIG) == (224, 1, 2, 8) |
| 269 | assert { |
| 270 | num_tokens: _config_tuple(config) |
| 271 | for num_tokens, config in k3_gemm.KDA_QKVG_CONFIGS.items() |
| 272 | } == { |
| 273 | 1: (64, 4, 2, 8), |
| 274 | 2: (64, 3, 2, 8), |
| 275 | } |
| 276 | |
| 277 | |
| 278 | @pytest.mark.skipif( |
| 279 | not current_platform.is_cuda(), reason="KDA CuTeDSL helper requires CUDA" |
| 280 | ) |
| 281 | @pytest.mark.parametrize( |
| 282 | "compute_capability,supported", |
| 283 | [((9, 0), False), ((10, 0), True), ((10, 3), True), ((12, 0), True)], |
| 284 | ) |
| 285 | def test_kda_mixed_precision_bf16_fma_capability( |
| 286 | monkeypatch: pytest.MonkeyPatch, |
| 287 | compute_capability: tuple[int, int], |
| 288 | supported: bool, |
| 289 | ) -> None: |
| 290 | from vllm.models.kimi_k3.nvidia.ops.cute_dsl import kda_skinny_gemm |
| 291 | |
| 292 | capability = compute_capability[0] * 10 + compute_capability[1] |
| 293 | monkeypatch.setattr( |
| 294 | kda_skinny_gemm.current_platform, |
| 295 | "has_device_capability", |
| 296 | lambda target, device_id=0: capability >= target, |
| 297 | ) |
| 298 | |
| 299 | assert kda_skinny_gemm._has_mixed_precision_bf16_fma() is supported |
| 300 | |
| 301 | |
| 302 | @pytest.mark.parametrize("num_tokens", [2, 8, 14]) |
| 303 | def test_kda_skinny_gemms_sm90_cuda_graph(num_tokens: int) -> None: |
| 304 | """Hopper must compile and capture the pre-SM100 FMA specialization.""" |
| 305 | _require_capability_and_cute((9, 0)) |
| 306 | from vllm.models.kimi_k3.nvidia.ops.cute_dsl.kda_skinny_gemm import ( |
| 307 | KdaSkinnyGemm, |
| 308 | ) |
| 309 | |
| 310 | torch.manual_seed(43 + num_tokens) |
| 311 | hidden_states = torch.randn(num_tokens, 7168, dtype=torch.bfloat16, device="cuda") |
| 312 | f_ab_weight = torch.randn(144, 7168, dtype=torch.bfloat16, device="cuda") |
| 313 | f_b_weight = torch.randn(1536, 128, dtype=torch.bfloat16, device="cuda") |
| 314 | gemm = KdaSkinnyGemm() |
| 315 | |
| 316 | gemm.run_k(gemm.run_n(hidden_states, f_ab_weight), f_b_weight) |
| 317 | torch.accelerator.synchronize() |
| 318 | |
| 319 | capture_stream = torch.cuda.Stream() |
| 320 | capture_stream.wait_stream(torch.cuda.current_stream()) |
| 321 | graph = torch.cuda.CUDAGraph() |
| 322 | with torch.cuda.graph(graph, stream=capture_stream): |
| 323 | projected_fab = gemm.run_n(hidden_states, f_ab_weight) |
| 324 | projected_fb = gemm.run_k(projected_fab, f_b_weight) |
| 325 | torch.cuda.current_stream().wait_stream(capture_stream) |
| 326 | graph.replay() |
| 327 | torch.accelerator.synchronize() |
| 328 | |
| 329 | expected_fab = torch.nn.functional.linear( |
| 330 | hidden_states.float(), f_ab_weight.float() |
| 331 | ) |
| 332 | expected_fb = torch.nn.functional.linear( |
| 333 | expected_fab[:, :128].to(torch.bfloat16).float(), f_b_weight.float() |
| 334 | ) |
| 335 | for actual, expected in ( |
| 336 | (projected_fab, expected_fab), |
| 337 | (projected_fb, expected_fb), |
| 338 | ): |
| 339 | cosine = torch.nn.functional.cosine_similarity( |
| 340 | actual.float().flatten(), expected.flatten(), dim=0 |
| 341 | ).item() |
| 342 | assert cosine > 0.999 |
| 343 | |
| 344 | |
| 345 | @pytest.mark.parametrize("tp_size", [1, 2, 4, 16]) |
| 346 | def test_kda_projection_overlap_is_tp8_only(tp_size: int) -> None: |
| 347 | from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention |
| 348 | |
| 349 | kda = KimiK3DeltaAttention.__new__(KimiK3DeltaAttention) |
| 350 | nn.Module.__init__(kda) |
| 351 | kda.tp_size = tp_size |
| 352 | kda._projection_overlap_max_tokens = 0 |
| 353 | |
| 354 | assert not k3_gemm._enable_kda_projection_overlap(kda) |
| 355 | assert kda._projection_overlap_max_tokens == 0 |
| 356 | |
| 357 | |
| 358 | @pytest.mark.parametrize( |
| 359 | ("capability", "expected"), |
| 360 | [ |
| 361 | ((9, 0), None), |
| 362 | ((10, 0), "cute-dsl"), |
| 363 | ((10, 3), "cute-dsl"), |
| 364 | ((10, 7), None), |
| 365 | ((12, 0), None), |
| 366 | ], |
| 367 | ) |
| 368 | def test_kda_qkvg_flashinfer_backend( |
| 369 | monkeypatch: pytest.MonkeyPatch, |
| 370 | capability: tuple[int, int], |
| 371 | expected: str | None, |
| 372 | ) -> None: |
| 373 | monkeypatch.setattr( |
| 374 | current_platform, |
| 375 | "is_device_capability", |
| 376 | lambda candidate: candidate == capability, |
| 377 | ) |
| 378 | |
| 379 | assert k3_gemm._kda_qkvg_flashinfer_backend() == expected |
| 380 | |
| 381 | |
| 382 | @pytest.mark.parametrize("backend", ["cute-dsl", None]) |
| 383 | def test_kda_qkvg_autotune_enables_full_overlap( |
| 384 | monkeypatch: pytest.MonkeyPatch, |
| 385 | backend: str | None, |
| 386 | ) -> None: |
| 387 | flashinfer_gemm = pytest.importorskip("flashinfer.gemm") |
| 388 | |
| 389 | from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention |
| 390 | |
| 391 | kda = KimiK3DeltaAttention.__new__(KimiK3DeltaAttention) |
| 392 | nn.Module.__init__(kda) |
| 393 | kda._projection_overlap_max_tokens = max(k3_gemm.KDA_QKVG_CONFIGS) |
| 394 | kda.in_proj_qkvgfab = nn.Module() |
| 395 | kda.in_proj_qkvgfab.weight = nn.Parameter( |
| 396 | torch.empty(6144, 8, dtype=torch.bfloat16), |
| 397 | requires_grad=False, |
| 398 | ) |
| 399 | calls = [] |
| 400 | |
| 401 | def fake_mm_bf16(a, b, *, pdl, backend): |
| 402 | calls.append((a.shape, b.shape, pdl, backend)) |
| 403 | return torch.empty(a.shape[0], b.shape[1], dtype=a.dtype) |
| 404 | |
| 405 | monkeypatch.setattr(flashinfer_gemm, "mm_bf16", fake_mm_bf16) |
| 406 | monkeypatch.setattr( |
| 407 | k3_gemm, |
| 408 | "_kda_qkvg_flashinfer_backend", |
| 409 | lambda: backend, |
| 410 | ) |
| 411 | |
| 412 | k3_gemm.autotune_kda_qkvg(kda) |
| 413 | |
| 414 | expected_calls = [] |
| 415 | if backend is not None: |
| 416 | expected_calls.append( |
| 417 | (torch.Size([14, 8]), torch.Size([8, 6144]), True, "cute-dsl") |
| 418 | ) |
| 419 | assert calls == expected_calls |
| 420 | assert ( |
| 421 | kda._projection_overlap_max_tokens == k3_gemm.KDA_PROJECTION_OVERLAP_MAX_TOKENS |
| 422 | ) |
| 423 | |
| 424 | |
| 425 | def test_glm52_projection_plans_are_separate() -> None: |
| 426 | qkv_a = glm52_gemm.GLM52_QKV_A_PROJECTION |
| 427 | q_b = glm52_gemm.GLM52_Q_B_PROJECTION |
| 428 | |
| 429 | qkv_a_plan = qkv_a.build_plan() |
| 430 | q_b_plan = q_b.build_plan() |
| 431 | |
| 432 | assert (qkv_a.n, qkv_a.k, set(qkv_a_plan)) == ( |
| 433 | 2624, |
| 434 | 6144, |
| 435 | set(range(1, 17)), |
| 436 | ) |
| 437 | assert (q_b.n, q_b.k, set(q_b_plan)) == ( |
| 438 | 2048, |
| 439 | 2048, |
| 440 | set(range(1, 17)), |
| 441 | ) |
| 442 | assert { |
| 443 | num_tokens |
| 444 | for num_tokens, (backend, _) in qkv_a_plan.items() |
| 445 | if backend == "cute" |
| 446 | } == {1, 2} |
| 447 | assert { |
| 448 | num_tokens |
| 449 | for num_tokens, (backend, _) in qkv_a_plan.items() |
| 450 | if backend == "dsv3_fused_a" |
| 451 | } == set(range(3, 17)) |
| 452 | assert { |
| 453 | num_tokens for num_tokens, (backend, _) in q_b_plan.items() if backend == "cute" |
| 454 | } == {1, 2} |
| 455 | assert { |
| 456 | num_tokens |
| 457 | for num_tokens, (backend, _) in q_b_plan.items() |
| 458 | if backend == "dsv3_fused_a" |
| 459 | } == set(range(3, 17)) |
| 460 | |
| 461 | eh = glm52_gemm.GLM52_EH_PROJECTION |
| 462 | eh_plan = eh.build_plan() |
| 463 | assert (eh.n, eh.k) == (6144, 12288) |
| 464 | # The MTP eh_proj has no dsv3 winners; M >= 4 falls back to cuBLAS. |
| 465 | assert set(eh_plan) == {1, 2, 3} |
| 466 | assert all(backend == "cute" for backend, _ in eh_plan.values()) |
| 467 | |
| 468 | |
| 469 | def test_glm52_layout_rejects_nonpacked_single_row_view() -> None: |
| 470 | single_row = torch.empty(1, 144)[:, :128] |
| 471 | multiple_rows = torch.empty(2, 144)[:, :128] |
| 472 | |
| 473 | assert single_row.stride() == multiple_rows.stride() == (144, 1) |
| 474 | assert not glm52_gemm._is_supported_row_major(single_row) |
| 475 | assert not glm52_gemm._is_supported_row_major(multiple_rows) |
| 476 | |
| 477 | |
| 478 | def test_glm52_installer_maps_only_selected_unquantized_shapes( |
| 479 | monkeypatch: pytest.MonkeyPatch, |
| 480 | ) -> None: |
| 481 | class FakeLinearBase(nn.Module): |
| 482 | def __init__( |
| 483 | self, |
| 484 | n: int, |
| 485 | k: int, |
| 486 | quant_method: object, |
| 487 | ) -> None: |
| 488 | super().__init__() |
| 489 | self.weight = nn.Parameter( |
| 490 | torch.empty( |
| 491 | n, |
| 492 | k, |
| 493 | dtype=torch.bfloat16, |
| 494 | device="meta", |
| 495 | ) |
| 496 | ) |
| 497 | self.quant_method = quant_method |
| 498 | |
| 499 | qkv_a = glm52_gemm.GLM52_QKV_A_PROJECTION |
| 500 | q_b = glm52_gemm.GLM52_Q_B_PROJECTION |
| 501 | root = nn.Module() |
| 502 | root.attn = nn.Module() |
| 503 | root.attn.fused_qkv_a_proj = FakeLinearBase( |
| 504 | qkv_a.n, qkv_a.k, glm52_gemm.UnquantizedLinearMethod() |
| 505 | ) |
| 506 | root.attn.q_b_proj = FakeLinearBase( |
| 507 | q_b.n, q_b.k, glm52_gemm.UnquantizedLinearMethod() |
| 508 | ) |
| 509 | root.same_shape_other_name = FakeLinearBase( |
| 510 | qkv_a.n, qkv_a.k, glm52_gemm.UnquantizedLinearMethod() |
| 511 | ) |
| 512 | root.quantized = nn.Module() |
| 513 | quantized_method = object() |
| 514 | root.quantized.q_b_proj = FakeLinearBase(q_b.n, q_b.k, quantized_method) |
| 515 | root.wrong_shape = nn.Module() |
| 516 | root.wrong_shape.fused_qkv_a_proj = FakeLinearBase( |
| 517 | qkv_a.n + 1, qkv_a.k, glm52_gemm.UnquantizedLinearMethod() |
| 518 | ) |
| 519 | monkeypatch.setattr(glm52_gemm, "LinearBase", FakeLinearBase) |
| 520 | monkeypatch.setattr(glm52_gemm, "_is_sm103", lambda: True) |
| 521 | monkeypatch.setattr( |
| 522 | glm52_gemm.shape_dynamic_skinny_gemm, |
| 523 | "is_available", |
| 524 | lambda: False, |
| 525 | ) |
| 526 | |
| 527 | glm52_gemm.enable_glm52_low_latency_gemm( |
| 528 | root, |
| 529 | torch.bfloat16, |
| 530 | ) |
| 531 | |
| 532 | assert isinstance( |
| 533 | root.attn.fused_qkv_a_proj.quant_method, |
| 534 | glm52_gemm.GLM52LowLatencyLinearMethod, |
| 535 | ) |
| 536 | assert isinstance( |
| 537 | root.attn.q_b_proj.quant_method, |
| 538 | glm52_gemm.GLM52LowLatencyLinearMethod, |
| 539 | ) |
| 540 | assert root.attn.fused_qkv_a_proj.quant_method._plan == qkv_a.build_plan() |
| 541 | assert root.attn.q_b_proj.quant_method._plan == q_b.build_plan() |
| 542 | assert isinstance( |
| 543 | root.same_shape_other_name.quant_method, |
| 544 | glm52_gemm.GLM52LowLatencyLinearMethod, |
| 545 | ) |
| 546 | assert root.quantized.q_b_proj.quant_method is quantized_method |
| 547 | assert ( |
| 548 | type(root.wrong_shape.fused_qkv_a_proj.quant_method) |
| 549 | is glm52_gemm.UnquantizedLinearMethod |
| 550 | ) |
| 551 | |
| 552 | |
| 553 | @pytest.mark.parametrize("key", EXPECTED_SELECTIONS) |
| 554 | def test_sm103_selector_table(key: tuple[int, int]) -> None: |
| 555 | n, k = key |
| 556 | cute_tokens, dsv3_tokens = EXPECTED_SELECTIONS[key] |
| 557 | for num_tokens in range(1, 17): |
| 558 | backend = k3_gemm.select_kimi_k3_backend(num_tokens, n, k) |
| 559 | if num_tokens in cute_tokens: |
| 560 | assert backend == "cute" |
| 561 | elif num_tokens in dsv3_tokens: |
| 562 | assert backend == "dsv3_fused_a" |
| 563 | else: |
| 564 | assert backend is None |
| 565 | |
| 566 | |
| 567 | @pytest.mark.parametrize("key", EXPECTED_SELECTIONS) |
| 568 | def test_selector_requires_supported_shape_and_tokens(key: tuple[int, int]) -> None: |
| 569 | n, k = key |
| 570 | assert k3_gemm.select_kimi_k3_backend(0, n, k) is None |
| 571 | assert k3_gemm.select_kimi_k3_backend(17, n, k) is None |
| 572 | assert k3_gemm.select_kimi_k3_backend(1, n + 1, k) is None |
| 573 | assert k3_gemm.select_kimi_k3_backend(1, n, k + 1) is None |
| 574 | |
| 575 | |
| 576 | def test_unlisted_shape_and_unselected_tokens_fall_back() -> None: |
| 577 | # Shape absent from the table. |
| 578 | assert k3_gemm.select_kimi_k3_backend(1, 1000, 1000) is None |
| 579 | # o_proj (7168,1536) is CuTe M1 only; M2+ falls back. |
| 580 | assert k3_gemm.select_kimi_k3_backend(2, 7168, 1536) is None |
| 581 | |
| 582 | |
| 583 | @pytest.mark.parametrize("num_tokens", range(1, 17)) |
| 584 | def test_sm103_residual_selector_table(num_tokens: int) -> None: |
| 585 | backend = k3_gemm.select_kimi_k3_backend(num_tokens, 7168, 3584, has_residual=True) |
| 586 | assert backend == ("cute" if num_tokens <= 4 else None) |
| 587 | |
| 588 | |
| 589 | def test_build_plan_matches_selector() -> None: |
| 590 | for spec in k3_gemm.KIMI_K3_PROJECTIONS.values(): |
| 591 | plan = k3_gemm._build_plan(spec) |
| 592 | for num_tokens in range(1, 17): |
| 593 | backend = k3_gemm.select_kimi_k3_backend(num_tokens, spec.n, spec.k) |
| 594 | if backend is None: |
| 595 | assert num_tokens not in plan |
| 596 | else: |
| 597 | assert plan[num_tokens][0] == backend |
| 598 | |
| 599 | |
| 600 | # Keyed by local (N, K): (cute token counts, dsv3 token counts). Mirrors |
| 601 | # KIMI_K3_PROJECTIONS_SM100, measured on B200; intentionally different from |
| 602 | # EXPECTED_SELECTIONS (e.g. 3584x7168 routes dsv3 at M2..8 on SM103 but only |
| 603 | # wins with CuTe at M1..3 on SM100). |
| 604 | EXPECTED_SELECTIONS_SM100 = { |
| 605 | (1536, 128): (set(), {1, 16}), |
| 606 | (3072, 128): (set(), {8}), |
| 607 | (1536, 7168): ({1, 2}, {4, 8}), |
| 608 | (3072, 7168): ({1, 2}, set()), |
| 609 | (2112, 7168): ({1, 2}, {4, 16}), |
| 610 | (2304, 1536): (set(), set(range(1, 17))), |
| 611 | (4608, 1536): (set(), {1, 2, 4}), |
| 612 | (6288, 7168): ({1}, set()), |
| 613 | (12448, 7168): ({1, 2, 3, 4}, set()), |
| 614 | (7168, 768): (set(), {1}), |
| 615 | # SM103-tuned config wins on B200 where the heuristic config tied. |
| 616 | (7168, 1536): ({1}, set()), |
| 617 | # M2 config lifted from the SM103 table (wins on B200). |
| 618 | (7168, 3072): ({1, 2}, set()), |
| 619 | (7168, 3584): ({1, 2}, set()), |
| 620 | (7168, 8448): ({1, 2, 3}, set()), |
| 621 | (8448, 7168): ({1, 2}, set()), |
| 622 | (16896, 7168): ({1}, set()), |
| 623 | (20480, 7168): ({1, 2, 3, 4}, set()), |
| 624 | (40960, 7168): ({1, 2, 3, 4}, set()), |
| 625 | (10240, 7168): ({1, 2, 3, 4}, set()), |
| 626 | (3584, 7168): ({1, 2, 3}, set()), |
| 627 | # TP16 shapes measured on B200. |
| 628 | (768, 7168): ({1, 2, 3, 4}, set()), |
| 629 | (1152, 1536): ({1}, set()), |
| 630 | (3216, 7168): ({1, 2}, set()), |
| 631 | (4224, 7168): ({1, 2, 3}, set()), |
| 632 | } |
| 633 | |
| 634 | |
| 635 | @pytest.mark.parametrize( |
| 636 | "table", |
| 637 | [k3_gemm.KIMI_K3_PROJECTIONS_SM100, k3_gemm.KIMI_K3_PROJECTIONS_SM90], |
| 638 | ) |
| 639 | def test_device_table_is_keyed_by_shape(table: K3ProjectionTable) -> None: |
| 640 | for (n, k), spec in table.items(): |
| 641 | assert (spec.n, spec.k) == (n, k) |
| 642 | |
| 643 | |
| 644 | @pytest.mark.parametrize("key", EXPECTED_SELECTIONS_SM100) |
| 645 | def test_sm100_selector_table(key: tuple[int, int]) -> None: |
| 646 | n, k = key |
| 647 | spec = k3_gemm.KIMI_K3_PROJECTIONS_SM100[key] |
| 648 | cute_tokens, dsv3_tokens = EXPECTED_SELECTIONS_SM100[key] |
| 649 | for num_tokens in range(1, 17): |
| 650 | backend = k3_gemm._backend_for(spec, num_tokens, has_residual=False) |
| 651 | if num_tokens in cute_tokens: |
| 652 | assert backend == "cute" |
| 653 | elif num_tokens in dsv3_tokens: |
| 654 | assert backend == "dsv3_fused_a" |
| 655 | else: |
| 656 | assert backend is None |
| 657 | |
| 658 | |
| 659 | @pytest.mark.parametrize( |
| 660 | "table", |
| 661 | [k3_gemm.KIMI_K3_PROJECTIONS_SM100, k3_gemm.KIMI_K3_PROJECTIONS_SM90], |
| 662 | ) |
| 663 | def test_device_build_plan_matches_table(table: K3ProjectionTable) -> None: |
| 664 | for spec in table.values(): |
| 665 | plan = k3_gemm._build_plan(spec) |
| 666 | for num_tokens in range(1, 17): |
| 667 | backend = k3_gemm._backend_for(spec, num_tokens, has_residual=False) |
| 668 | if backend is None: |
| 669 | assert num_tokens not in plan |
| 670 | else: |
| 671 | assert plan[num_tokens][0] == backend |
| 672 | |
| 673 | |
| 674 | def test_sm90_table_covers_measured_points() -> None: |
| 675 | specs = k3_gemm.KIMI_K3_PROJECTIONS_SM90.values() |
| 676 | assert len(k3_gemm.KIMI_K3_PROJECTIONS_SM90) == 18 |
| 677 | assert sum(len(spec.cute_configs) + len(spec.dsv3_tokens) for spec in specs) == 90 |
| 678 | |
| 679 | |
| 680 | def test_low_latency_table_capability_routing( |
| 681 | monkeypatch: pytest.MonkeyPatch, |
| 682 | ) -> None: |
| 683 | monkeypatch.setattr(k3_gemm, "_is_sm103", lambda: True) |
| 684 | assert k3_gemm._low_latency_table() is k3_gemm.KIMI_K3_PROJECTIONS |
| 685 | |
| 686 | monkeypatch.setattr(k3_gemm, "_is_sm103", lambda: False) |
| 687 | monkeypatch.setattr( |
| 688 | k3_gemm.current_platform, |
| 689 | "is_device_capability", |
| 690 | lambda cc: cc == (10, 0), |
| 691 | ) |
| 692 | assert k3_gemm._low_latency_table() is k3_gemm.KIMI_K3_PROJECTIONS_SM100 |
| 693 | |
| 694 | monkeypatch.setattr( |
| 695 | k3_gemm.current_platform, |
| 696 | "is_device_capability", |
| 697 | lambda cc: cc == (9, 0), |
| 698 | ) |
| 699 | assert k3_gemm._low_latency_table() is k3_gemm.KIMI_K3_PROJECTIONS_SM90 |
| 700 | |
| 701 | monkeypatch.setattr( |
| 702 | k3_gemm.current_platform, "is_device_capability", lambda cc: False |
| 703 | ) |
| 704 | assert k3_gemm._low_latency_table() is None |
| 705 | |
| 706 | |
| 707 | def test_qwen4_exp_hopper_plans_are_valid() -> None: |
| 708 | plans = qwen4_exp_gemm.QWEN4_EXP_SM90_GEMM_PLANS |
| 709 | |
| 710 | assert len(plans) == 9 |
| 711 | assert sum(map(len, plans.values())) == 31 |
| 712 | assert (320, 10240) in plans |
| 713 | assert (10240, 320) not in plans |
| 714 | for (n, k), shape_plans in plans.items(): |
| 715 | for num_tokens, config in shape_plans.items(): |
| 716 | assert config.num_rows == num_tokens |
| 717 | assert n % config.outputs_per_block == 0 |
| 718 | assert k % (config.block_size * config.vector_width) == 0 |
| 719 | assert config.static_k in (None, k) |
| 720 | |
| 721 | |
| 722 | @pytest.mark.parametrize( |
| 723 | "capability,expected_plans", |
| 724 | [ |
| 725 | ((10, 3), qwen4_exp_gemm.QWEN4_EXP_GEMM_PLANS), |
| 726 | ((10, 0), qwen4_exp_gemm.QWEN4_EXP_SM100_GEMM_PLANS), |
| 727 | ((9, 0), qwen4_exp_gemm.QWEN4_EXP_SM90_GEMM_PLANS), |
| 728 | ((8, 0), {}), |
| 729 | ], |
| 730 | ) |
| 731 | def test_qwen4_exp_gemm_capability_routing( |
| 732 | monkeypatch: pytest.MonkeyPatch, |
| 733 | capability: tuple[int, int], |
| 734 | expected_plans: dict[tuple[int, int], dict[int, SkinnyGemmConfig]], |
| 735 | ) -> None: |
| 736 | monkeypatch.setattr( |
| 737 | qwen4_exp_gemm.current_platform, |
| 738 | "is_device_capability", |
| 739 | lambda target: capability == target, |
| 740 | ) |
| 741 | |
| 742 | assert qwen4_exp_gemm._gemm_plans() == expected_plans |
| 743 | |
| 744 | |
| 745 | def test_installation_is_shape_specific_and_unquantized( |
| 746 | monkeypatch: pytest.MonkeyPatch, |
| 747 | ) -> None: |
| 748 | class FakeLinear(nn.Module): |
| 749 | def __init__(self, quant_method: object, n: int, k: int) -> None: |
| 750 | super().__init__() |
| 751 | self.quant_method = quant_method |
| 752 | self.weight = torch.empty(n, k) |
| 753 | |
| 754 | class FakeHead(nn.Module): |
| 755 | def __init__(self, n: int, k: int) -> None: |
| 756 | super().__init__() |
| 757 | self.quant_method = k3_gemm.UnquantizedEmbeddingMethod() |
| 758 | self.weight = torch.empty(n, k) |
| 759 | |
| 760 | root = nn.Module() |
| 761 | # dsv3-only shape (no cute warmup contribution). |
| 762 | root.dsv3_only = FakeLinear(k3_gemm.UnquantizedLinearMethod(), 2304, 1536) |
| 763 | # quantized: must be left untouched. |
| 764 | quantized_method = object() |
| 765 | root.quantized = FakeLinear(quantized_method, 6288, 7168) |
| 766 | # cute shape. |
| 767 | root.cute = FakeLinear(k3_gemm.UnquantizedLinearMethod(), 6288, 7168) |
| 768 | # cute + residual shape. |
| 769 | root.residual = FakeLinear(k3_gemm.UnquantizedLinearMethod(), 7168, 3584) |
| 770 | # shape absent from the table: must be left untouched. |
| 771 | root.unlisted = FakeLinear(k3_gemm.UnquantizedLinearMethod(), 1234, 5678) |
| 772 | root.lm_head = FakeHead(20480, 7168) |
| 773 | |
| 774 | monkeypatch.setattr(k3_gemm, "LinearBase", FakeLinear) |
| 775 | monkeypatch.setattr(k3_gemm, "ParallelLMHead", FakeHead) |
| 776 | monkeypatch.setattr(k3_gemm, "_is_sm103", lambda: True) |
| 777 | warmup_configs: set[SkinnyGemmConfig] = set() |
| 778 | residual_warmup_configs: set[SkinnyGemmConfig] = set() |
| 779 | monkeypatch.setattr(k3_gemm.shape_dynamic_skinny_gemm, "is_available", lambda: True) |
| 780 | |
| 781 | def request_warmup_configs(dtype, configs, *, has_residual=False): |
| 782 | target = residual_warmup_configs if has_residual else warmup_configs |
| 783 | target.update(configs) |
| 784 | |
| 785 | monkeypatch.setattr( |
| 786 | k3_gemm.shape_dynamic_skinny_gemm, |
| 787 | "request_warmup_configs", |
| 788 | request_warmup_configs, |
| 789 | ) |
| 790 | |
| 791 | k3_gemm.enable_kimi_k3_low_latency_gemm(root, torch.bfloat16) |
| 792 | |
| 793 | assert isinstance(root.dsv3_only.quant_method, k3_gemm.KimiK3LowLatencyLinearMethod) |
| 794 | assert isinstance(root.cute.quant_method, k3_gemm.KimiK3LowLatencyLinearMethod) |
| 795 | assert isinstance(root.residual.quant_method, k3_gemm.KimiK3LowLatencyLinearMethod) |
| 796 | assert root.quantized.quant_method is quantized_method |
| 797 | assert type(root.unlisted.quant_method) is k3_gemm.UnquantizedLinearMethod |
| 798 | assert isinstance( |
| 799 | root.lm_head.quant_method, k3_gemm.KimiK3LowLatencyEmbeddingMethod |
| 800 | ) |
| 801 | # Warmup covers only the installed modules' local (N, K). |
| 802 | assert warmup_configs == { |
| 803 | config |
| 804 | for key in ((6288, 7168), (7168, 3584), (20480, 7168)) |
| 805 | for _, config in k3_gemm.KIMI_K3_PROJECTIONS[key].cute_configs |
| 806 | } |
| 807 | assert residual_warmup_configs == { |
| 808 | config |
| 809 | for _, config in k3_gemm.KIMI_K3_PROJECTIONS[(7168, 3584)].residual_configs |
| 810 | } |
| 811 | |
| 812 | |
| 813 | @pytest.mark.parametrize( |
| 814 | "dtype,platform_enabled", |
| 815 | [(torch.float16, True), (torch.bfloat16, False)], |
| 816 | ) |
| 817 | def test_installation_requires_bf16_sm103( |
| 818 | monkeypatch: pytest.MonkeyPatch, |
| 819 | dtype: torch.dtype, |
| 820 | platform_enabled: bool, |
| 821 | ) -> None: |
| 822 | class FakeLinear(nn.Module): |
| 823 | def __init__(self) -> None: |
| 824 | super().__init__() |
| 825 | self.quant_method = k3_gemm.UnquantizedLinearMethod() |
| 826 | self.weight = torch.empty(2304, 1536) |
| 827 | |
| 828 | root = nn.Module() |
| 829 | root.projection = FakeLinear() |
| 830 | monkeypatch.setattr(k3_gemm, "LinearBase", FakeLinear) |
| 831 | monkeypatch.setattr(k3_gemm, "_is_sm103", lambda: platform_enabled) |
| 832 | # Pin the rest of the capability probe so platform_enabled=False stays a |
| 833 | # no-op even when the test itself runs on SM100 hardware. |
| 834 | monkeypatch.setattr( |
| 835 | k3_gemm.current_platform, |
| 836 | "is_device_capability", |
| 837 | lambda cc: platform_enabled and cc == (10, 3), |
| 838 | ) |
| 839 | |
| 840 | k3_gemm.enable_kimi_k3_low_latency_gemm(root, dtype) |
| 841 | |
| 842 | assert type(root.projection.quant_method) is k3_gemm.UnquantizedLinearMethod |
| 843 | |
| 844 | |
| 845 | def _require_capability_and_dsv3(capability: tuple[int, int]) -> None: |
| 846 | if ( |
| 847 | not torch.cuda.is_available() |
| 848 | or torch.cuda.get_device_capability() != capability |
| 849 | ): |
| 850 | pytest.skip(f"Kimi-K3 selection requires SM{capability[0]}{capability[1]}") |
| 851 | if not hasattr(torch.ops._C, "dsv3_fused_a_gemm"): |
| 852 | pytest.skip("dsv3_fused_a_gemm was not built") |
| 853 | |
| 854 | |
| 855 | def _require_capability_and_cute(capability: tuple[int, int]) -> None: |
| 856 | if ( |
| 857 | not torch.cuda.is_available() |
| 858 | or torch.cuda.get_device_capability() != capability |
| 859 | ): |
| 860 | pytest.skip(f"CuTe DSL selection requires SM{capability[0]}{capability[1]}") |
| 861 | if not k3_gemm.shape_dynamic_skinny_gemm.is_available(): |
| 862 | pytest.skip("CuTe DSL is not available") |
| 863 | |
| 864 | |
| 865 | def _require_sm103_and_dsv3() -> None: |
| 866 | _require_capability_and_dsv3((10, 3)) |
| 867 | |
| 868 | |
| 869 | def _require_sm103_and_cute() -> None: |
| 870 | _require_capability_and_cute((10, 3)) |
| 871 | |
| 872 | |
| 873 | @pytest.mark.parametrize("spec,config", GLM_CUTE_CASES) |
| 874 | def test_glm_cute_selected_shapes( |
| 875 | spec: glm52_gemm.GLM52ProjectionSpec, |
| 876 | config: SkinnyGemmConfig, |
| 877 | ) -> None: |
| 878 | _require_sm103_and_cute() |
| 879 | torch.manual_seed(42) |
| 880 | x = torch.randn( |
| 881 | config.num_rows, |
| 882 | spec.k, |
| 883 | dtype=torch.bfloat16, |
| 884 | device="cuda", |
| 885 | ) |
| 886 | weight = torch.randn(spec.n, spec.k, dtype=torch.bfloat16, device="cuda") |
| 887 | plan = spec.build_plan() |
| 888 | |
| 889 | output = glm52_gemm.run_glm52_plan(plan, x, weight) |
| 890 | |
| 891 | assert output is not None |
| 892 | reference = x.float() @ weight.float().t() |
| 893 | torch.testing.assert_close(output.float(), reference, rtol=2e-2, atol=2e-1) |
| 894 | |
| 895 | |
| 896 | @pytest.mark.parametrize("n,k,num_tokens,config", QWEN4_EXP_SM90_CASES) |
| 897 | def test_qwen4_exp_sm90_selected_shapes( |
| 898 | n: int, |
| 899 | k: int, |
| 900 | num_tokens: int, |
| 901 | config: SkinnyGemmConfig, |
| 902 | ) -> None: |
| 903 | _require_capability_and_cute((9, 0)) |
| 904 | torch.manual_seed(42 + num_tokens) |
| 905 | x = torch.randn(num_tokens, k, dtype=torch.bfloat16, device="cuda") |
| 906 | weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") |
| 907 | |
| 908 | selected = qwen4_exp_gemm.QWEN4_EXP_SM90_GEMM_PLANS[(n, k)][num_tokens] |
| 909 | assert selected == config |
| 910 | output = qwen4_exp_gemm._qwen4_exp_low_latency_gemm(x, weight) |
| 911 | |
| 912 | reference = torch.nn.functional.linear(x, weight) |
| 913 | cosine = torch.nn.functional.cosine_similarity( |
| 914 | output.float().flatten(), reference.float().flatten(), dim=0 |
| 915 | ).item() |
| 916 | assert cosine > 0.999 |
| 917 | |
| 918 | |
| 919 | @pytest.mark.parametrize("n,k,num_tokens,config", QWEN4_EXP_SM100_CASES) |
| 920 | def test_qwen4_exp_sm100_selected_shapes( |
| 921 | n: int, |
| 922 | k: int, |
| 923 | num_tokens: int, |
| 924 | config: SkinnyGemmConfig, |
| 925 | ) -> None: |
| 926 | _require_capability_and_cute((10, 0)) |
| 927 | torch.manual_seed(42 + num_tokens) |
| 928 | x = torch.randn(num_tokens, k, dtype=torch.bfloat16, device="cuda") |
| 929 | weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") |
| 930 | |
| 931 | selected = qwen4_exp_gemm.QWEN4_EXP_SM100_GEMM_PLANS[(n, k)][num_tokens] |
| 932 | assert selected == config |
| 933 | output = qwen4_exp_gemm._qwen4_exp_low_latency_gemm(x, weight) |
| 934 | |
| 935 | reference = torch.nn.functional.linear(x, weight) |
| 936 | cosine = torch.nn.functional.cosine_similarity( |
| 937 | output.float().flatten(), reference.float().flatten(), dim=0 |
| 938 | ).item() |
| 939 | assert cosine > 0.999 |
| 940 | |
| 941 | |
| 942 | def test_glm52_q_b_nonpacked_single_row_falls_back() -> None: |
| 943 | _require_sm103_and_cute() |
| 944 | spec = glm52_gemm.GLM52_Q_B_PROJECTION |
| 945 | storage = torch.randn( |
| 946 | 1, |
| 947 | glm52_gemm.GLM52_QKV_A_PROJECTION.n, |
| 948 | dtype=torch.bfloat16, |
| 949 | device="cuda", |
| 950 | ) |
| 951 | x = storage[:, : spec.k] |
| 952 | weight = torch.randn(spec.n, spec.k, dtype=torch.bfloat16, device="cuda") |
| 953 | |
| 954 | assert x.stride() == (glm52_gemm.GLM52_QKV_A_PROJECTION.n, 1) |
| 955 | assert not glm52_gemm._runtime_ok(x, weight) |
| 956 | output = glm52_gemm.run_glm52_plan(spec.build_plan(), x, weight) |
| 957 | |
| 958 | assert output is None |
| 959 | |
| 960 | |
| 961 | @pytest.mark.parametrize("spec,config", GLM_CUTE_CASES) |
| 962 | def test_glm_cute_selected_shapes_cuda_graph_capture( |
| 963 | spec: glm52_gemm.GLM52ProjectionSpec, |
| 964 | config: SkinnyGemmConfig, |
| 965 | ) -> None: |
| 966 | _require_sm103_and_cute() |
| 967 | x = torch.randn( |
| 968 | config.num_rows, |
| 969 | spec.k, |
| 970 | dtype=torch.bfloat16, |
| 971 | device="cuda", |
| 972 | ) |
| 973 | weight = torch.randn(spec.n, spec.k, dtype=torch.bfloat16, device="cuda") |
| 974 | plan = spec.build_plan() |
| 975 | glm52_gemm.run_glm52_plan(plan, x, weight) |
| 976 | torch.accelerator.synchronize() |
| 977 | |
| 978 | graph = torch.cuda.CUDAGraph() |
| 979 | with torch.cuda.graph(graph): |
| 980 | output = glm52_gemm.run_glm52_plan(plan, x, weight) |
| 981 | graph.replay() |
| 982 | torch.accelerator.synchronize() |
| 983 | |
| 984 | assert output is not None |
| 985 | reference = x.float() @ weight.float().t() |
| 986 | torch.testing.assert_close(output.float(), reference, rtol=2e-2, atol=2e-1) |
| 987 | |
| 988 | |
| 989 | @pytest.mark.parametrize("capability,n,k,num_tokens", CUTE_CASES) |
| 990 | def test_cute_selected_shapes( |
| 991 | capability: tuple[int, int], n: int, k: int, num_tokens: int |
| 992 | ) -> None: |
| 993 | _require_capability_and_cute(capability) |
| 994 | torch.manual_seed(42) |
| 995 | x = torch.randn(num_tokens, k, dtype=torch.bfloat16, device="cuda") |
| 996 | weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") |
| 997 | |
| 998 | output = k3_gemm.try_low_latency_gemm(x, weight) |
| 999 | |
| 1000 | assert output is not None |
| 1001 | reference = torch.nn.functional.linear(x, weight) |
| 1002 | cosine = torch.nn.functional.cosine_similarity( |
| 1003 | output.float().flatten(), reference.float().flatten(), dim=0 |
| 1004 | ).item() |
| 1005 | assert cosine > 0.999 |
| 1006 | |
| 1007 | |
| 1008 | @pytest.mark.parametrize("num_tokens", [1, 2, 3, 4, 8, 9, 10, 11, 12, 13, 14]) |
| 1009 | def test_kda_projection_overlap_cuda_graph(num_tokens: int) -> None: |
| 1010 | """The TP8 split must preserve projection results across both streams.""" |
| 1011 | _require_sm103_and_cute() |
| 1012 | _require_sm103_and_dsv3() |
| 1013 | torch.manual_seed(43) |
| 1014 | hidden_states = torch.randn(num_tokens, 7168, dtype=torch.bfloat16, device="cuda") |
| 1015 | packed_weight = torch.randn(6288, 7168, dtype=torch.bfloat16, device="cuda") |
| 1016 | f_b_weight = torch.randn(1536, 128, dtype=torch.bfloat16, device="cuda") |
| 1017 | aux_stream = torch.cuda.Stream() |
| 1018 | events = (torch.cuda.Event(), torch.cuda.Event()) |
| 1019 | |
| 1020 | k3_gemm.run_kda_projection_overlap( |
| 1021 | hidden_states, |
| 1022 | packed_weight, |
| 1023 | f_b_weight, |
| 1024 | aux_stream, |
| 1025 | events, |
| 1026 | ) |
| 1027 | torch.accelerator.synchronize() |
| 1028 | |
| 1029 | capture_stream = torch.cuda.Stream() |
| 1030 | capture_stream.wait_stream(torch.cuda.current_stream()) |
| 1031 | graph = torch.cuda.CUDAGraph() |
| 1032 | with torch.cuda.graph(graph, stream=capture_stream): |
| 1033 | projected_qkvg, g1, beta = k3_gemm.run_kda_projection_overlap( |
| 1034 | hidden_states, |
| 1035 | packed_weight, |
| 1036 | f_b_weight, |
| 1037 | aux_stream, |
| 1038 | events, |
| 1039 | ) |
| 1040 | torch.cuda.current_stream().wait_stream(capture_stream) |
| 1041 | graph.replay() |
| 1042 | torch.accelerator.synchronize() |
| 1043 | |
| 1044 | assert beta.stride() == ((140 if num_tokens == 1 else 144), 1) |
| 1045 | |
| 1046 | expected_qkvg = torch.nn.functional.linear( |
| 1047 | hidden_states.float(), packed_weight[:6144].float() |
| 1048 | ) |
| 1049 | expected_fab = torch.nn.functional.linear( |
| 1050 | hidden_states.float(), packed_weight[6144:6284].float() |
| 1051 | ) |
| 1052 | expected_g1 = torch.nn.functional.linear( |
| 1053 | expected_fab[:, :128].to(torch.bfloat16).float(), |
| 1054 | f_b_weight.float(), |
| 1055 | ) |
| 1056 | for actual, expected in ( |
| 1057 | (projected_qkvg, expected_qkvg), |
| 1058 | (g1, expected_g1), |
| 1059 | (beta, expected_fab[:, 128:]), |
| 1060 | ): |
| 1061 | cosine = torch.nn.functional.cosine_similarity( |
| 1062 | actual.float().flatten(), expected.flatten(), dim=0 |
| 1063 | ).item() |
| 1064 | assert cosine > 0.999 |
| 1065 | |
| 1066 | |
| 1067 | def _dsv3_probe_tokens(tokens: frozenset[int]) -> set[int]: |
| 1068 | """Extremes, plus both sides of the kernel's num_tokens<=8 tile_n branch.""" |
| 1069 | if not tokens: |
| 1070 | return set() |
| 1071 | return {min(tokens), max(tokens)} | ({8, 9} & set(tokens)) |
| 1072 | |
| 1073 | |
| 1074 | # Derived from the table rather than hand-listed, so a shape routed to dsv3 |
| 1075 | # cannot be added without being exercised here. |
| 1076 | DSV3_CASES = sorted( |
| 1077 | (capability, num_tokens, spec.n, spec.k) |
| 1078 | for capability, table in K3_GPU_TABLES |
| 1079 | for spec in table.values() |
| 1080 | for num_tokens in _dsv3_probe_tokens(spec.dsv3_tokens) |
| 1081 | ) |
| 1082 | |
| 1083 | GLM_DSV3_CASES = [ |
| 1084 | (num_tokens, spec) |
| 1085 | for spec in ( |
| 1086 | glm52_gemm.GLM52_QKV_A_PROJECTION, |
| 1087 | glm52_gemm.GLM52_Q_B_PROJECTION, |
| 1088 | ) |
| 1089 | for num_tokens in sorted(_dsv3_probe_tokens(spec.dsv3_tokens)) |
| 1090 | ] |
| 1091 | |
| 1092 | |
| 1093 | @pytest.mark.parametrize("capability,num_tokens,n,k", DSV3_CASES) |
| 1094 | def test_dsv3_selected_shapes( |
| 1095 | capability: tuple[int, int], num_tokens: int, n: int, k: int |
| 1096 | ) -> None: |
| 1097 | _require_capability_and_dsv3(capability) |
| 1098 | torch.manual_seed(42) |
| 1099 | x = torch.randn(num_tokens, k, dtype=torch.bfloat16, device="cuda") |
| 1100 | weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") |
| 1101 | |
| 1102 | output = k3_gemm.try_low_latency_gemm(x, weight) |
| 1103 | |
| 1104 | assert output is not None |
| 1105 | reference = torch.nn.functional.linear(x, weight) |
| 1106 | cosine = torch.nn.functional.cosine_similarity( |
| 1107 | output.float().flatten(), reference.float().flatten(), dim=0 |
| 1108 | ).item() |
| 1109 | assert cosine > 0.999 |
| 1110 | |
| 1111 | |
| 1112 | @pytest.mark.parametrize("num_tokens,spec", GLM_DSV3_CASES) |
| 1113 | def test_glm_dsv3_selected_shapes( |
| 1114 | num_tokens: int, |
| 1115 | spec: glm52_gemm.GLM52ProjectionSpec, |
| 1116 | ) -> None: |
| 1117 | _require_sm103_and_dsv3() |
| 1118 | torch.manual_seed(42) |
| 1119 | x = torch.randn(num_tokens, spec.k, dtype=torch.bfloat16, device="cuda") |
| 1120 | weight = torch.randn(spec.n, spec.k, dtype=torch.bfloat16, device="cuda") |
| 1121 | |
| 1122 | output = glm52_gemm.run_glm52_plan(spec.build_plan(), x, weight) |
| 1123 | |
| 1124 | assert output is not None |
| 1125 | reference = torch.nn.functional.linear(x, weight) |
| 1126 | cosine = torch.nn.functional.cosine_similarity( |
| 1127 | output.float().flatten(), reference.float().flatten(), dim=0 |
| 1128 | ).item() |
| 1129 | assert cosine > 0.999 |
| 1130 | |
| 1131 | |
| 1132 | @pytest.mark.parametrize("num_tokens", [1, 4]) |
| 1133 | def test_kda_f_b_nonpacked_dsv3_dispatches(num_tokens: int) -> None: |
| 1134 | _require_sm103_and_dsv3() |
| 1135 | n, k = 1536, 128 |
| 1136 | storage = torch.randn(num_tokens, 6288, dtype=torch.bfloat16, device="cuda") |
| 1137 | x = storage[:, 6144 : 6144 + k] |
| 1138 | weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") |
| 1139 | spec = k3_gemm.KIMI_K3_PROJECTIONS[(n, k)] |
| 1140 | method = k3_gemm.KimiK3LowLatencyLinearMethod( |
| 1141 | k3_gemm._build_plan(spec), k3_gemm._build_residual_plan(spec) |
| 1142 | ) |
| 1143 | |
| 1144 | assert x.stride() == (6288, 1) |
| 1145 | assert k3_gemm._runtime_ok(x, weight) |
| 1146 | output = method.apply(SimpleNamespace(weight=weight), x) |
| 1147 | |
| 1148 | reference = torch.nn.functional.linear(x, weight) |
| 1149 | torch.testing.assert_close(output, reference) |
| 1150 | |
| 1151 | |
| 1152 | @pytest.mark.parametrize("capability,table", K3_GPU_TABLES) |
| 1153 | def test_selected_kernels_cuda_graph_capture( |
| 1154 | capability: tuple[int, int], table: K3ProjectionTable |
| 1155 | ) -> None: |
| 1156 | _require_capability_and_cute(capability) |
| 1157 | _require_capability_and_dsv3(capability) |
| 1158 | cute_spec = table[(6288, 7168)] |
| 1159 | dsv3_spec = table[(1536, 128)] |
| 1160 | cute_x = torch.randn(1, cute_spec.k, dtype=torch.bfloat16, device="cuda") |
| 1161 | cute_weight = torch.randn( |
| 1162 | cute_spec.n, cute_spec.k, dtype=torch.bfloat16, device="cuda" |
| 1163 | ) |
| 1164 | dsv3_x = torch.randn(1, dsv3_spec.k, dtype=torch.bfloat16, device="cuda") |
| 1165 | dsv3_weight = torch.randn( |
| 1166 | dsv3_spec.n, dsv3_spec.k, dtype=torch.bfloat16, device="cuda" |
| 1167 | ) |
| 1168 | k3_gemm.try_low_latency_gemm(cute_x, cute_weight) |
| 1169 | k3_gemm.try_low_latency_gemm(dsv3_x, dsv3_weight) |
| 1170 | torch.accelerator.synchronize() |
| 1171 | |
| 1172 | graph = torch.cuda.CUDAGraph() |
| 1173 | with torch.cuda.graph(graph): |
| 1174 | cute_output = k3_gemm.try_low_latency_gemm(cute_x, cute_weight) |
| 1175 | dsv3_output = k3_gemm.try_low_latency_gemm(dsv3_x, dsv3_weight) |
| 1176 | graph.replay() |
| 1177 | torch.accelerator.synchronize() |
| 1178 | |
| 1179 | assert cute_output is not None |
| 1180 | assert dsv3_output is not None |
| 1181 | for output, activation, weight in ( |
| 1182 | (cute_output, cute_x, cute_weight), |
| 1183 | (dsv3_output, dsv3_x, dsv3_weight), |
| 1184 | ): |
| 1185 | reference = torch.nn.functional.linear(activation, weight) |
| 1186 | cosine = torch.nn.functional.cosine_similarity( |
| 1187 | output.float().flatten(), reference.float().flatten(), dim=0 |
| 1188 | ).item() |
| 1189 | assert cosine > 0.999 |
| 1190 | |
| 1191 | |
| 1192 | @pytest.mark.parametrize("num_tokens", [1, 8, 9, 16]) |
| 1193 | def test_dsv3_cuda_graph_capture_tile_branches(num_tokens: int) -> None: |
| 1194 | """Capture DSV3 across the num_tokens<=8 vs >8 tile_n branch.""" |
| 1195 | _require_sm103_and_dsv3() |
| 1196 | spec = k3_gemm.KIMI_K3_PROJECTIONS[(1536, 128)] |
| 1197 | x = torch.randn(num_tokens, spec.k, dtype=torch.bfloat16, device="cuda") |
| 1198 | weight = torch.randn(spec.n, spec.k, dtype=torch.bfloat16, device="cuda") |
| 1199 | k3_gemm.try_low_latency_gemm(x, weight) |
| 1200 | torch.accelerator.synchronize() |
| 1201 | |
| 1202 | graph = torch.cuda.CUDAGraph() |
| 1203 | with torch.cuda.graph(graph): |
| 1204 | output = k3_gemm.try_low_latency_gemm(x, weight) |
| 1205 | graph.replay() |
| 1206 | torch.accelerator.synchronize() |
| 1207 | |
| 1208 | assert output is not None |
| 1209 | reference = torch.nn.functional.linear(x, weight) |
| 1210 | cosine = torch.nn.functional.cosine_similarity( |
| 1211 | output.float().flatten(), reference.float().flatten(), dim=0 |
| 1212 | ).item() |
| 1213 | assert cosine > 0.999 |
| 1214 | |
| 1215 | |
| 1216 | @pytest.mark.parametrize("n,k,num_tokens", RESIDUAL_CUTE_CASES) |
| 1217 | def test_cute_residual_epilogue(n: int, k: int, num_tokens: int) -> None: |
| 1218 | _require_sm103_and_cute() |
| 1219 | torch.manual_seed(42 + num_tokens) |
| 1220 | x = torch.randn(num_tokens, k, dtype=torch.bfloat16, device="cuda") |
| 1221 | weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") |
| 1222 | residual = torch.randn(num_tokens, n, dtype=torch.bfloat16, device="cuda") |
| 1223 | spec = k3_gemm.KIMI_K3_PROJECTIONS[(n, k)] |
| 1224 | config = spec.residual_config(num_tokens) |
| 1225 | assert config is not None |
| 1226 | |
| 1227 | output = k3_gemm.shape_dynamic_skinny_gemm(x, weight, config, residual) |
| 1228 | |
| 1229 | reference = x.float() @ weight.float().t() + residual.float() |
| 1230 | cosine = torch.nn.functional.cosine_similarity( |
| 1231 | output.float().flatten(), reference.flatten(), dim=0 |
| 1232 | ).item() |
| 1233 | assert cosine > 0.999 |
| 1234 | |
| 1235 | |
| 1236 | @pytest.mark.parametrize("num_tokens", range(1, 17)) |
| 1237 | def test_cute_residual_epilogue_all_supported_token_counts(num_tokens: int) -> None: |
| 1238 | _require_sm103_and_cute() |
| 1239 | from vllm.model_executor.kernels.linear.cute_dsl.skinny_gemm import ( |
| 1240 | ShapeDynamicSkinnyGemm, |
| 1241 | ) |
| 1242 | |
| 1243 | n, k = 64, 512 |
| 1244 | x = torch.randn(num_tokens, k, dtype=torch.bfloat16, device="cuda") |
| 1245 | weight = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") |
| 1246 | residual = torch.randn(num_tokens, n, dtype=torch.bfloat16, device="cuda") |
| 1247 | config = ShapeDynamicSkinnyGemm._config(num_tokens, n, k) |
| 1248 | |
| 1249 | output = k3_gemm.shape_dynamic_skinny_gemm(x, weight, config, residual) |
| 1250 | |
| 1251 | reference = x.float() @ weight.float().t() + residual.float() |
| 1252 | torch.testing.assert_close(output.float(), reference, rtol=2e-2, atol=2e-1) |
| 1253 | |
| 1254 | |
| 1255 | @pytest.mark.parametrize("num_tokens", range(1, 5)) |
| 1256 | def test_cute_residual_epilogue_cuda_graph_capture(num_tokens: int) -> None: |
| 1257 | _require_sm103_and_cute() |
| 1258 | spec = k3_gemm.KIMI_K3_PROJECTIONS[(7168, 3584)] |
| 1259 | config = spec.residual_config(num_tokens) |
| 1260 | assert config is not None |
| 1261 | x = torch.randn(num_tokens, spec.k, dtype=torch.bfloat16, device="cuda") |
| 1262 | weight = torch.randn(spec.n, spec.k, dtype=torch.bfloat16, device="cuda") |
| 1263 | residual = torch.randn(num_tokens, spec.n, dtype=torch.bfloat16, device="cuda") |
| 1264 | k3_gemm.shape_dynamic_skinny_gemm(x, weight, config, residual) |
| 1265 | torch.accelerator.synchronize() |
| 1266 | |
| 1267 | graph = torch.cuda.CUDAGraph() |
| 1268 | with torch.cuda.graph(graph): |
| 1269 | output = k3_gemm.shape_dynamic_skinny_gemm(x, weight, config, residual) |
| 1270 | graph.replay() |
| 1271 | torch.accelerator.synchronize() |
| 1272 | |
| 1273 | reference = x.float() @ weight.float().t() + residual.float() |
| 1274 | cosine = torch.nn.functional.cosine_similarity( |
| 1275 | output.float().flatten(), reference.flatten(), dim=0 |
| 1276 | ).item() |
| 1277 | assert cosine > 0.999 |
| 1278 | |
| 1279 | |
| 1280 | class _SkinnyGemmSpy: |
| 1281 | """Wraps the skinny-GEMM singleton to record whether CuTe was invoked.""" |
| 1282 | |
| 1283 | def __init__(self, real: Any) -> None: |
| 1284 | self._real = real |
| 1285 | self.calls: list[int] = [] |
| 1286 | |
| 1287 | def __call__(self, a, b, config=None, residual=None): |
| 1288 | self.calls.append(a.shape[0]) |
| 1289 | return self._real(a, b, config, residual) |
| 1290 | |
| 1291 | def is_available(self) -> bool: |
| 1292 | return self._real.is_available() |
| 1293 | |
| 1294 | |
| 1295 | @pytest.mark.parametrize("num_tokens", [1, 2, 3, 4]) |
| 1296 | def test_latent_moe_production_layout_residual( |
| 1297 | monkeypatch: pytest.MonkeyPatch, |
| 1298 | num_tokens: int, |
| 1299 | ) -> None: |
| 1300 | """The real Latent-MoE residual is a non-packed slice of a cat buffer. |
| 1301 | |
| 1302 | The strict packed-row-major guard rejects such a slice at every token count |
| 1303 | (a size-1 leading dim reads as contiguous but its stride is not packed), so |
| 1304 | the CuTe residual epilogue never fires for this production layout and the |
| 1305 | method falls back to addmm. Output is correct regardless of the path. |
| 1306 | """ |
| 1307 | _require_sm103_and_cute() |
| 1308 | latent_dim, shared_dim = 3584, 7168 # routed_expert_up_proj K, N |
| 1309 | torch.manual_seed(7 + num_tokens) |
| 1310 | buf = torch.randn( |
| 1311 | num_tokens, latent_dim + shared_dim, dtype=torch.bfloat16, device="cuda" |
| 1312 | ) |
| 1313 | latent = buf[:, :latent_dim] # non-contiguous view (row stride = full width) |
| 1314 | residual = buf[:, latent_dim:] # non-contiguous view |
| 1315 | weight = torch.randn(shared_dim, latent_dim, dtype=torch.bfloat16, device="cuda") |
| 1316 | |
| 1317 | spec = k3_gemm.KIMI_K3_PROJECTIONS[(shared_dim, latent_dim)] |
| 1318 | method = k3_gemm.KimiK3LowLatencyLinearMethod( |
| 1319 | k3_gemm._build_plan(spec), k3_gemm._build_residual_plan(spec) |
| 1320 | ) |
| 1321 | spy = _SkinnyGemmSpy(k3_gemm.shape_dynamic_skinny_gemm) |
| 1322 | monkeypatch.setattr(k3_gemm, "shape_dynamic_skinny_gemm", spy) |
| 1323 | |
| 1324 | layer = SimpleNamespace(weight=weight) |
| 1325 | output = method.apply_with_residual(layer, latent, residual) |
| 1326 | |
| 1327 | reference = latent.float() @ weight.float().t() + residual.float() |
| 1328 | cosine = torch.nn.functional.cosine_similarity( |
| 1329 | output.float().flatten(), reference.flatten(), dim=0 |
| 1330 | ).item() |
| 1331 | assert cosine > 0.999 # correct regardless of the path taken |
| 1332 | assert not spy.calls, ( |
| 1333 | "non-packed buf-slice residual must fall back to addmm at every M" |
| 1334 | ) |
| 1335 | |
| 1336 | |
| 1337 | def test_residual_dispatch_falls_back_to_addmm( |
| 1338 | monkeypatch: pytest.MonkeyPatch, |
| 1339 | ) -> None: |
| 1340 | fallback = torch.randn(2, 3) |
| 1341 | residual = torch.randn(2, 3) |
| 1342 | x = torch.randn(2, 4) |
| 1343 | weight = torch.randn(3, 4) |
| 1344 | monkeypatch.setattr(torch, "addmm", lambda *args: fallback) |
| 1345 | # CPU tensors fail the runtime check, forcing the addmm fallback. |
| 1346 | method = k3_gemm.KimiK3LowLatencyLinearMethod({}, {}) |
| 1347 | |
| 1348 | output = method.apply_with_residual(SimpleNamespace(weight=weight), x, residual) |
| 1349 | |
| 1350 | assert output is fallback |
| 1351 | |
| 1352 | |
| 1353 | @pytest.mark.skipif( |
| 1354 | not current_platform.is_cuda(), |
| 1355 | reason="The base UnquantizedLinearMethod is platform-dispatched: on ROCm it " |
| 1356 | "goes through the vllm::rocm_unquantized_gemm custom op, which has no CPU " |
| 1357 | "kernel, so the CPU-tensor fallback cannot be exercised there.", |
| 1358 | ) |
| 1359 | def test_fallback_preserves_default_method() -> None: |
| 1360 | x = torch.randn(2, 4) |
| 1361 | weight = torch.randn(3, 4) |
| 1362 | # CPU tensors fail the runtime check, forcing the base-method fallback. |
| 1363 | method = k3_gemm.KimiK3LowLatencyLinearMethod({}, {}) |
| 1364 | |
| 1365 | output = method.apply(SimpleNamespace(weight=weight), x) |
| 1366 | |
| 1367 | torch.testing.assert_close(output, torch.nn.functional.linear(x, weight)) |