| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for cuteDSL low-latency router GEMM (dot-product + split-K).""" |
| 4 | |
| 5 | import pytest |
| 6 | import torch |
| 7 | import torch.nn.functional as F |
| 8 | |
| 9 | pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") |
| 10 | |
| 11 | |
| 12 | @pytest.fixture(autouse=True, scope="module") |
| 13 | def _require_sm90_and_cutedsl(): |
| 14 | if torch.cuda.get_device_capability()[0] < 9: |
| 15 | pytest.skip("Requires SM90+ (Hopper/Blackwell)") |
| 16 | from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import ( |
| 17 | is_available, |
| 18 | ) |
| 19 | |
| 20 | if not is_available(): |
| 21 | pytest.skip("cuteDSL (CUTLASS Python) not installed") |
| 22 | |
| 23 | |
| 24 | # ===== Helpers ===== |
| 25 | |
| 26 | |
| 27 | def _ref(a, b): |
| 28 | return torch.mm(a.float(), b.float().T) |
| 29 | |
| 30 | |
| 31 | def _assert_close(out, ref, *, min_cos_sim=0.99, context=""): |
| 32 | assert out.device.type == "cuda", f"{context}: not on CUDA" |
| 33 | assert torch.isfinite(out).all(), f"{context}: NaN/Inf" |
| 34 | cos = F.cosine_similarity( |
| 35 | out.reshape(-1).float(), ref.reshape(-1).float(), dim=0 |
| 36 | ).item() |
| 37 | assert cos > min_cos_sim, ( |
| 38 | f"{context}: cos_sim {cos:.6f} < {min_cos_sim} " |
| 39 | f"(abs_err={(out.float() - ref.float()).abs().max().item():.2e})" |
| 40 | ) |
| 41 | |
| 42 | |
| 43 | def _can_precompile(a, b): |
| 44 | return ( |
| 45 | a.dim() == 2 |
| 46 | and b.dim() == 2 |
| 47 | and a.dtype == torch.bfloat16 |
| 48 | and b.dtype == torch.bfloat16 |
| 49 | and a.device.type == "cuda" |
| 50 | and b.device.type == "cuda" |
| 51 | and a.device == b.device |
| 52 | and a.shape[1] == b.shape[1] |
| 53 | and a.is_contiguous() |
| 54 | and b.is_contiguous() |
| 55 | ) |
| 56 | |
| 57 | |
| 58 | def _gemm(a, b): |
| 59 | from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import ( |
| 60 | ll_bf16_gemm, |
| 61 | ll_bf16_gemm_kernel, |
| 62 | ) |
| 63 | |
| 64 | if _can_precompile(a, b): |
| 65 | compile_key = ll_bf16_gemm_kernel.dispatch( |
| 66 | M=a.shape[0], K=a.shape[1], N=b.shape[0] |
| 67 | ) |
| 68 | ll_bf16_gemm_kernel.compile(compile_key) |
| 69 | return ll_bf16_gemm(a, b) |
| 70 | |
| 71 | |
| 72 | def test_c1_pdl_kernel_is_selected_automatically(monkeypatch): |
| 73 | from vllm.model_executor.kernels.linear.cute_dsl import ll_bf16 |
| 74 | |
| 75 | calls = [] |
| 76 | |
| 77 | def default_kernel(hidden_states, router_weight, output_dtype): |
| 78 | calls.append(("default", hidden_states.shape[0])) |
| 79 | return hidden_states |
| 80 | |
| 81 | def c1_pdl_kernel(hidden_states, router_weight, output_dtype): |
| 82 | calls.append(("c1_pdl", hidden_states.shape[0])) |
| 83 | return hidden_states |
| 84 | |
| 85 | monkeypatch.setattr(ll_bf16, "ll_bf16_gemm_kernel", default_kernel) |
| 86 | monkeypatch.setattr(ll_bf16, "ll_bf16_gemm_c1_pdl_kernel", c1_pdl_kernel) |
| 87 | weight = torch.empty(4, 8, device="cuda", dtype=torch.bfloat16) |
| 88 | |
| 89 | ll_bf16.ll_bf16_gemm(torch.empty(1, 8, device="cuda", dtype=torch.bfloat16), weight) |
| 90 | ll_bf16.ll_bf16_gemm(torch.empty(2, 8, device="cuda", dtype=torch.bfloat16), weight) |
| 91 | |
| 92 | assert calls == [("c1_pdl", 1), ("default", 2)] |
| 93 | |
| 94 | |
| 95 | # ===== Shapes ===== |
| 96 | |
| 97 | SHAPES = [ |
| 98 | (256, 7168, "DSV3"), |
| 99 | (256, 14400, "DSV4-Flash"), |
| 100 | (128, 5120, "DeepSeek-V2"), |
| 101 | (8, 4096, "Mixtral-8x7B"), |
| 102 | (64, 2880, "non-tile-aligned-K"), |
| 103 | (256, 2048, "split-K-boundary"), |
| 104 | ] |
| 105 | |
| 106 | SHAPES_SPLITK = [(n, k, d) for n, k, d in SHAPES if k >= 2048] |
| 107 | |
| 108 | |
| 109 | # ================================================================= |
| 110 | # Dot-product kernel (M<=4 or K<2048) |
| 111 | # ================================================================= |
| 112 | |
| 113 | |
| 114 | @pytest.mark.parametrize("M", [1, 2, 3, 4]) |
| 115 | @pytest.mark.parametrize("N,K,desc", SHAPES, ids=[s[2] for s in SHAPES]) |
| 116 | def test_dotprod(M, N, K, desc): |
| 117 | torch.manual_seed(42) |
| 118 | a = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") |
| 119 | b = torch.randn(N, K, dtype=torch.bfloat16, device="cuda") |
| 120 | out = _gemm(a, b) |
| 121 | assert out.dtype == torch.float32 |
| 122 | assert out.shape == (M, N) |
| 123 | _assert_close(out, _ref(a, b), context=f"dotprod {M}x{N}x{K}") |
| 124 | |
| 125 | |
| 126 | @pytest.mark.parametrize("M", [1, 4, 8, 16]) |
| 127 | def test_dotprod_small_K(M): |
| 128 | """K<2048 forces dot-product regardless of M.""" |
| 129 | torch.manual_seed(42) |
| 130 | a = torch.randn(M, 1024, dtype=torch.bfloat16, device="cuda") |
| 131 | b = torch.randn(64, 1024, dtype=torch.bfloat16, device="cuda") |
| 132 | _assert_close(_gemm(a, b), _ref(a, b), context=f"small_K M={M}") |
| 133 | |
| 134 | |
| 135 | @pytest.mark.parametrize("K", [16, 32, 64, 128, 256, 512, 1024, 1536]) |
| 136 | def test_dotprod_K_sweep(K): |
| 137 | torch.manual_seed(42) |
| 138 | a = torch.randn(4, K, dtype=torch.bfloat16, device="cuda") |
| 139 | b = torch.randn(32, K, dtype=torch.bfloat16, device="cuda") |
| 140 | _assert_close(_gemm(a, b), _ref(a, b), context=f"K={K}") |
| 141 | |
| 142 | |
| 143 | # ================================================================= |
| 144 | # Split-K kernel (M>4 and K>=2048) |
| 145 | # ================================================================= |
| 146 | |
| 147 | |
| 148 | @pytest.mark.parametrize("M", [5, 6, 8, 12, 16]) |
| 149 | @pytest.mark.parametrize("N,K,desc", SHAPES_SPLITK, ids=[s[2] for s in SHAPES_SPLITK]) |
| 150 | def test_splitk(M, N, K, desc): |
| 151 | torch.manual_seed(42) |
| 152 | a = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") |
| 153 | b = torch.randn(N, K, dtype=torch.bfloat16, device="cuda") |
| 154 | out = _gemm(a, b) |
| 155 | assert out.dtype == torch.float32 |
| 156 | assert out.shape == (M, N) |
| 157 | _assert_close(out, _ref(a, b), context=f"splitk {M}x{N}x{K}") |
| 158 | |
| 159 | |
| 160 | @pytest.mark.parametrize("K", [2048, 2304, 2880, 3072, 4096, 5120, 7168, 14400]) |
| 161 | def test_splitk_K_sweep(K): |
| 162 | """Includes non-tile-aligned K (2880) and uneven split (2304).""" |
| 163 | torch.manual_seed(42) |
| 164 | a = torch.randn(8, K, dtype=torch.bfloat16, device="cuda") |
| 165 | b = torch.randn(64, K, dtype=torch.bfloat16, device="cuda") |
| 166 | _assert_close(_gemm(a, b), _ref(a, b), context=f"splitk K={K}") |
| 167 | |
| 168 | |
| 169 | # ================================================================= |
| 170 | # Dispatch boundary (M=4/5, K=2032/2048) |
| 171 | # ================================================================= |
| 172 | |
| 173 | |
| 174 | @pytest.mark.parametrize("M", [4, 5]) |
| 175 | @pytest.mark.parametrize("K", [2032, 2048]) |
| 176 | def test_dispatch_boundary(M, K): |
| 177 | torch.manual_seed(42) |
| 178 | a = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") |
| 179 | b = torch.randn(64, K, dtype=torch.bfloat16, device="cuda") |
| 180 | path = "splitk" if M > 4 and K >= 2048 else "dotprod" |
| 181 | _assert_close(_gemm(a, b), _ref(a, b), context=f"M={M} K={K} ({path})") |
| 182 | |
| 183 | |
| 184 | # ================================================================= |
| 185 | # Arbitrary N |
| 186 | # ================================================================= |
| 187 | |
| 188 | |
| 189 | @pytest.mark.parametrize("N", [1, 3, 7, 16, 17, 64, 128, 256, 384]) |
| 190 | def test_arbitrary_N_dotprod(N): |
| 191 | torch.manual_seed(42) |
| 192 | a = torch.randn(4, 2048, dtype=torch.bfloat16, device="cuda") |
| 193 | b = torch.randn(N, 2048, dtype=torch.bfloat16, device="cuda") |
| 194 | out = _gemm(a, b) |
| 195 | assert out.shape == (4, N) |
| 196 | _assert_close(out, _ref(a, b), context=f"dotprod N={N}") |
| 197 | |
| 198 | |
| 199 | @pytest.mark.parametrize("N", [1, 8, 16, 17, 64, 128, 256]) |
| 200 | def test_arbitrary_N_splitk(N): |
| 201 | torch.manual_seed(42) |
| 202 | a = torch.randn(8, 4096, dtype=torch.bfloat16, device="cuda") |
| 203 | b = torch.randn(N, 4096, dtype=torch.bfloat16, device="cuda") |
| 204 | out = _gemm(a, b) |
| 205 | assert out.shape == (8, N) |
| 206 | _assert_close(out, _ref(a, b), context=f"splitk N={N}") |
| 207 | |
| 208 | |
| 209 | # ================================================================= |
| 210 | # Single token (M=1, decode path) |
| 211 | # ================================================================= |
| 212 | |
| 213 | |
| 214 | @pytest.mark.parametrize( |
| 215 | "N,K", |
| 216 | [(256, 7168), (256, 14400), (8, 4096), (384, 7168), (264, 6144)], |
| 217 | ids=["DSV3", "DSV4-Flash", "Mixtral", "DSV4-Pro", "Inkling"], |
| 218 | ) |
| 219 | def test_single_token(N, K): |
| 220 | torch.manual_seed(42) |
| 221 | a = torch.randn(1, K, dtype=torch.bfloat16, device="cuda") |
| 222 | b = torch.randn(N, K, dtype=torch.bfloat16, device="cuda") |
| 223 | out = _gemm(a, b) |
| 224 | assert out.shape == (1, N) |
| 225 | _assert_close(out, _ref(a, b), context=f"M=1 {N}x{K}") |
| 226 | |
| 227 | |
| 228 | def test_inkling_max_tokens(): |
| 229 | torch.manual_seed(42) |
| 230 | a = torch.randn(64, 6144, dtype=torch.bfloat16, device="cuda") |
| 231 | b = torch.randn(264, 6144, dtype=torch.bfloat16, device="cuda") |
| 232 | out = _gemm(a, b) |
| 233 | assert out.shape == (64, 264) |
| 234 | _assert_close(out, _ref(a, b), context="Inkling M=64") |
| 235 | |
| 236 | |
| 237 | # ================================================================= |
| 238 | # Numerical robustness |
| 239 | # ================================================================= |
| 240 | |
| 241 | |
| 242 | @pytest.mark.parametrize("M,K", [(4, 2048), (8, 4096)], ids=["dotprod", "splitk"]) |
| 243 | def test_large_values(M, K): |
| 244 | torch.manual_seed(42) |
| 245 | a = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") * 100 |
| 246 | b = torch.randn(64, K, dtype=torch.bfloat16, device="cuda") * 100 |
| 247 | out = _gemm(a, b) |
| 248 | assert torch.isfinite(out).all() |
| 249 | _assert_close(out, _ref(a, b), context=f"large M={M}") |
| 250 | |
| 251 | |
| 252 | def test_near_zero(): |
| 253 | torch.manual_seed(42) |
| 254 | a = torch.randn(4, 2048, dtype=torch.bfloat16, device="cuda") * 1e-4 |
| 255 | b = torch.randn(64, 2048, dtype=torch.bfloat16, device="cuda") * 1e-4 |
| 256 | out = _gemm(a, b) |
| 257 | assert torch.isfinite(out).all() |
| 258 | assert out.abs().max() < 1.0 |
| 259 | |
| 260 | |
| 261 | def test_zeros(): |
| 262 | a = torch.zeros(4, 2048, dtype=torch.bfloat16, device="cuda") |
| 263 | b = torch.randn(64, 2048, dtype=torch.bfloat16, device="cuda") |
| 264 | assert (_gemm(a, b) == 0).all() |
| 265 | |
| 266 | |
| 267 | def test_ones(): |
| 268 | a = torch.ones(1, 2048, dtype=torch.bfloat16, device="cuda") |
| 269 | b = torch.randn(32, 2048, dtype=torch.bfloat16, device="cuda") |
| 270 | _assert_close(_gemm(a, b), _ref(a, b), context="ones") |
| 271 | |
| 272 | |
| 273 | # ================================================================= |
| 274 | # Output dtype |
| 275 | # ================================================================= |
| 276 | |
| 277 | |
| 278 | @pytest.mark.parametrize("M,K", [(4, 2048), (8, 4096)], ids=["dotprod", "splitk"]) |
| 279 | def test_output_fp32(M, K): |
| 280 | torch.manual_seed(42) |
| 281 | a = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") |
| 282 | b = torch.randn(64, K, dtype=torch.bfloat16, device="cuda") |
| 283 | assert _gemm(a, b).dtype == torch.float32 |
| 284 | |
| 285 | |
| 286 | # ================================================================= |
| 287 | # Determinism |
| 288 | # ================================================================= |
| 289 | |
| 290 | |
| 291 | @pytest.mark.parametrize( |
| 292 | "M,K", [(1, 4096), (4, 4096), (5, 4096), (8, 4096), (16, 4096)] |
| 293 | ) |
| 294 | def test_deterministic(M, K): |
| 295 | torch.manual_seed(42) |
| 296 | a = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") |
| 297 | b = torch.randn(128, K, dtype=torch.bfloat16, device="cuda") |
| 298 | torch.testing.assert_close(_gemm(a, b), _gemm(a, b), atol=0, rtol=0) |
| 299 | |
| 300 | |
| 301 | # ================================================================= |
| 302 | # Cross-kernel consistency |
| 303 | # ================================================================= |
| 304 | |
| 305 | |
| 306 | @pytest.mark.parametrize("K", [2048, 4096, 7168]) |
| 307 | def test_dotprod_vs_splitk(K): |
| 308 | """First 4 rows from split-K (M=5) match dot-product (M=4).""" |
| 309 | torch.manual_seed(42) |
| 310 | a = torch.randn(5, K, dtype=torch.bfloat16, device="cuda") |
| 311 | b = torch.randn(64, K, dtype=torch.bfloat16, device="cuda") |
| 312 | _assert_close( |
| 313 | _gemm(a, b)[:4], |
| 314 | _gemm(a[:4], b), |
| 315 | min_cos_sim=0.999, |
| 316 | context=f"cross-kernel K={K}", |
| 317 | ) |
| 318 | |
| 319 | |
| 320 | # ================================================================= |
| 321 | # CUDA graph |
| 322 | # ================================================================= |
| 323 | |
| 324 | |
| 325 | @pytest.mark.parametrize("M,K", [(4, 2048), (8, 4096)], ids=["dotprod", "splitk"]) |
| 326 | def test_cudagraph(M, K): |
| 327 | torch.manual_seed(42) |
| 328 | a = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") |
| 329 | b = torch.randn(64, K, dtype=torch.bfloat16, device="cuda") |
| 330 | _gemm(a, b) |
| 331 | torch.accelerator.synchronize() |
| 332 | |
| 333 | g = torch.cuda.CUDAGraph() |
| 334 | with torch.cuda.graph(g): |
| 335 | out = _gemm(a, b) |
| 336 | for _ in range(5): |
| 337 | g.replay() |
| 338 | torch.accelerator.synchronize() |
| 339 | _assert_close(out, _ref(a, b), context=f"cudagraph M={M}") |
| 340 | |
| 341 | |
| 342 | def test_cudagraph_20x_replay(): |
| 343 | torch.manual_seed(42) |
| 344 | a = torch.randn(4, 4096, dtype=torch.bfloat16, device="cuda") |
| 345 | b = torch.randn(256, 4096, dtype=torch.bfloat16, device="cuda") |
| 346 | _gemm(a, b) |
| 347 | torch.accelerator.synchronize() |
| 348 | |
| 349 | g = torch.cuda.CUDAGraph() |
| 350 | with torch.cuda.graph(g): |
| 351 | out = _gemm(a, b) |
| 352 | results = [] |
| 353 | for _ in range(20): |
| 354 | g.replay() |
| 355 | torch.accelerator.synchronize() |
| 356 | results.append(out.clone()) |
| 357 | for i in range(1, len(results)): |
| 358 | torch.testing.assert_close( |
| 359 | results[0], results[i], atol=0, rtol=0, msg=f"Replay {i} differs" |
| 360 | ) |
| 361 | |
| 362 | |
| 363 | def test_cudagraph_input_update(): |
| 364 | torch.manual_seed(42) |
| 365 | a = torch.randn(4, 2048, dtype=torch.bfloat16, device="cuda") |
| 366 | b = torch.randn(64, 2048, dtype=torch.bfloat16, device="cuda") |
| 367 | _gemm(a, b) |
| 368 | torch.accelerator.synchronize() |
| 369 | |
| 370 | g = torch.cuda.CUDAGraph() |
| 371 | with torch.cuda.graph(g): |
| 372 | out = _gemm(a, b) |
| 373 | a.copy_(torch.randn_like(a)) |
| 374 | g.replay() |
| 375 | torch.accelerator.synchronize() |
| 376 | _assert_close(out, _ref(a, b), context="cudagraph input update") |
| 377 | |
| 378 | |
| 379 | # ================================================================= |
| 380 | # GateLinear dispatch integration |
| 381 | # ================================================================= |
| 382 | |
| 383 | |
| 384 | def _make_gate_linear(monkeypatch, *, params_dtype, out_dtype=torch.float32): |
| 385 | monkeypatch.setattr( |
| 386 | "vllm.model_executor.layers.linear.get_tensor_model_parallel_rank", |
| 387 | lambda: 0, |
| 388 | ) |
| 389 | monkeypatch.setattr( |
| 390 | "vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", |
| 391 | lambda: 1, |
| 392 | ) |
| 393 | monkeypatch.setattr( |
| 394 | "vllm.model_executor.parameter.get_tensor_model_parallel_rank", |
| 395 | lambda: 0, |
| 396 | ) |
| 397 | monkeypatch.setattr( |
| 398 | "vllm.model_executor.parameter.get_tensor_model_parallel_world_size", |
| 399 | lambda: 1, |
| 400 | ) |
| 401 | |
| 402 | from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear |
| 403 | |
| 404 | monkeypatch.setattr( |
| 405 | "vllm.model_executor.kernels.linear.cute_dsl.ll_bf16.is_available", |
| 406 | lambda: True, |
| 407 | ) |
| 408 | return GateLinear( |
| 409 | input_size=2048, |
| 410 | output_size=64, |
| 411 | bias=False, |
| 412 | out_dtype=out_dtype, |
| 413 | params_dtype=params_dtype, |
| 414 | ).cuda() |
| 415 | |
| 416 | |
| 417 | def test_gate_linear_uses_ll_bf16_for_bf16_fast_path(monkeypatch): |
| 418 | gate = _make_gate_linear(monkeypatch, params_dtype=torch.bfloat16) |
| 419 | x = torch.randn(4, 2048, dtype=torch.bfloat16, device="cuda") |
| 420 | calls = [] |
| 421 | |
| 422 | def fake_ll_bf16_gemm(hidden_states, router_weight): |
| 423 | calls.append((hidden_states, router_weight)) |
| 424 | return torch.full( |
| 425 | (hidden_states.shape[0], router_weight.shape[0]), |
| 426 | 1.0, |
| 427 | dtype=torch.float32, |
| 428 | device=hidden_states.device, |
| 429 | ) |
| 430 | |
| 431 | monkeypatch.setattr( |
| 432 | "vllm.model_executor.kernels.linear.cute_dsl.ll_bf16.ll_bf16_gemm", |
| 433 | fake_ll_bf16_gemm, |
| 434 | ) |
| 435 | out, bias = gate(x) |
| 436 | assert bias is None |
| 437 | assert out.shape == (4, 64) |
| 438 | assert out.dtype == torch.float32 |
| 439 | assert len(calls) == 1 |
| 440 | assert calls[0][0] is x |
| 441 | assert calls[0][1] is gate.weight |
| 442 | |
| 443 | |
| 444 | def test_gate_linear_fp32_weight_falls_back(monkeypatch): |
| 445 | gate = _make_gate_linear(monkeypatch, params_dtype=torch.float32) |
| 446 | assert not gate.allow_ll_bf16_gemm |
| 447 | x = torch.randn(4, 2048, dtype=torch.bfloat16, device="cuda") |
| 448 | |
| 449 | def fail_ll_bf16_gemm(hidden_states, router_weight): |
| 450 | raise AssertionError("ll_bf16_gemm should not run for fp32 weights") |
| 451 | |
| 452 | monkeypatch.setattr( |
| 453 | "vllm.model_executor.kernels.linear.cute_dsl.ll_bf16.ll_bf16_gemm", |
| 454 | fail_ll_bf16_gemm, |
| 455 | ) |
| 456 | out, _ = gate(x) |
| 457 | assert out.shape == (4, 64) |
| 458 | assert out.dtype == torch.float32 |
| 459 | |
| 460 | |
| 461 | def test_gate_linear_non_bf16_activation_falls_back(monkeypatch): |
| 462 | gate = _make_gate_linear(monkeypatch, params_dtype=torch.bfloat16) |
| 463 | x = torch.randn(4, 2048, dtype=torch.float16, device="cuda") |
| 464 | |
| 465 | def fail_ll_bf16_gemm(hidden_states, router_weight): |
| 466 | raise AssertionError("ll_bf16_gemm should not run for non-bf16 activations") |
| 467 | |
| 468 | monkeypatch.setattr( |
| 469 | "vllm.model_executor.kernels.linear.cute_dsl.ll_bf16.ll_bf16_gemm", |
| 470 | fail_ll_bf16_gemm, |
| 471 | ) |
| 472 | out, _ = gate(x) |
| 473 | assert out.shape == (4, 64) |
| 474 | assert out.dtype == torch.float32 |
| 475 | |
| 476 | |
| 477 | def test_gate_linear_set_out_dtype_enables_ll_bf16(monkeypatch): |
| 478 | gate = _make_gate_linear(monkeypatch, params_dtype=torch.bfloat16, out_dtype=None) |
| 479 | assert not gate.allow_ll_bf16_gemm |
| 480 | gate.set_out_dtype(torch.float32) |
| 481 | assert gate.allow_ll_bf16_gemm |
| 482 | |
| 483 | |
| 484 | def test_gate_linear_non_fp32_out_dtype_disables_ll_bf16(monkeypatch): |
| 485 | gate = _make_gate_linear(monkeypatch, params_dtype=torch.bfloat16, out_dtype=None) |
| 486 | gate.set_out_dtype(torch.bfloat16) |
| 487 | assert not gate.allow_ll_bf16_gemm |
| 488 | |
| 489 | |
| 490 | def test_gate_linear_m_gt_16_falls_back(monkeypatch): |
| 491 | gate = _make_gate_linear(monkeypatch, params_dtype=torch.bfloat16) |
| 492 | x = torch.randn(17, 2048, dtype=torch.bfloat16, device="cuda") |
| 493 | |
| 494 | def fail_ll_bf16_gemm(hidden_states, router_weight): |
| 495 | raise AssertionError("ll_bf16_gemm should not run for M > 16") |
| 496 | |
| 497 | monkeypatch.setattr( |
| 498 | "vllm.model_executor.kernels.linear.cute_dsl.ll_bf16.ll_bf16_gemm", |
| 499 | fail_ll_bf16_gemm, |
| 500 | ) |
| 501 | out, _ = gate(x) |
| 502 | assert out.shape == (17, 64) |
| 503 | assert out.dtype == torch.float32 |
| 504 | |
| 505 | |
| 506 | # ================================================================= |
| 507 | # Negative tests — invalid inputs |
| 508 | # ================================================================= |
| 509 | |
| 510 | |
| 511 | @pytest.mark.parametrize( |
| 512 | "M,K,N,dtype", |
| 513 | [ |
| 514 | pytest.param(4, 2048, 64, torch.float32, id="fp32_input"), |
| 515 | pytest.param(4, 2048, 64, torch.float16, id="fp16_input"), |
| 516 | pytest.param(8, 4096, 64, torch.float32, id="fp32_splitk_path"), |
| 517 | ], |
| 518 | ) |
| 519 | def test_invalid_dtype(M, K, N, dtype): |
| 520 | a = torch.randn(M, K, device="cuda", dtype=dtype) |
| 521 | b = torch.randn(N, K, device="cuda", dtype=dtype) |
| 522 | with pytest.raises(ValueError, match="dtype=bfloat16"): |
| 523 | _gemm(a, b) |
| 524 | |
| 525 | |
| 526 | def test_invalid_device_cpu(): |
| 527 | a = torch.randn(4, 2048, dtype=torch.bfloat16) |
| 528 | b = torch.randn(64, 2048, dtype=torch.bfloat16) |
| 529 | with pytest.raises(ValueError, match="device_type=cuda"): |
| 530 | _gemm(a, b) |
| 531 | |
| 532 | |
| 533 | def test_invalid_1d_input(): |
| 534 | a = torch.randn(2048, dtype=torch.bfloat16, device="cuda") |
| 535 | b = torch.randn(64, 2048, dtype=torch.bfloat16, device="cuda") |
| 536 | with pytest.raises(ValueError, match="2D tensors"): |
| 537 | _gemm(a, b) |
| 538 | |
| 539 | |
| 540 | def test_mismatched_K(): |
| 541 | a = torch.randn(4, 2048, dtype=torch.bfloat16, device="cuda") |
| 542 | b = torch.randn(64, 1024, dtype=torch.bfloat16, device="cuda") |
| 543 | with pytest.raises(ValueError, match="matching K dimensions"): |
| 544 | _gemm(a, b) |
| 545 | |
| 546 | |
| 547 | def test_invalid_K_divisibility(): |
| 548 | a = torch.randn(4, 2049, dtype=torch.bfloat16, device="cuda") |
| 549 | b = torch.randn(64, 2049, dtype=torch.bfloat16, device="cuda") |
| 550 | |
| 551 | with pytest.raises(ValueError, match="K to be divisible by 8"): |
| 552 | _gemm(a, b) |
| 553 | |
| 554 | |
| 555 | def test_non_contiguous_input(): |
| 556 | a = torch.randn(2048, 4, dtype=torch.bfloat16, device="cuda").T |
| 557 | b = torch.randn(64, 2048, dtype=torch.bfloat16, device="cuda") |
| 558 | assert not a.is_contiguous() |
| 559 | with pytest.raises(ValueError, match="contiguous row-major"): |
| 560 | _gemm(a, b) |
| 561 | |
| 562 | |
| 563 | def test_invalid_output_dtype(): |
| 564 | a = torch.randn(4, 2048, dtype=torch.bfloat16, device="cuda") |
| 565 | b = torch.randn(64, 2048, dtype=torch.bfloat16, device="cuda") |
| 566 | with pytest.raises(ValueError, match="output_dtype=torch.float32"): |
| 567 | from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import ll_bf16_gemm |
| 568 | |
| 569 | ll_bf16_gemm(a, b, output_dtype=torch.bfloat16) |
| 570 | |
| 571 | |
| 572 | def test_cache_miss_compiles_dotprod(): |
| 573 | from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import LLBf16Gemm |
| 574 | |
| 575 | torch.manual_seed(42) |
| 576 | a = torch.randn(3, 64, dtype=torch.bfloat16, device="cuda") |
| 577 | b = torch.randn(17, 64, dtype=torch.bfloat16, device="cuda") |
| 578 | kernel = LLBf16Gemm() |
| 579 | out = kernel(a, b) |
| 580 | assert out.shape == (3, 17) |
| 581 | _assert_close(out, _ref(a, b), context="cache miss dotprod") |
| 582 | |
| 583 | |
| 584 | if __name__ == "__main__": |
| 585 | pytest.main([__file__, "-v"]) |