From 0762f2afeb74f790e8c5ebe2b95a012eb38d499e Mon Sep 17 00:00:00 2001 From: "Roberto L. Castro" <38211239+LopezCastroRoberto@users.noreply.github.com> Date: Tue, 14 Jul 2026 08:01:21 +0200 Subject: [PATCH] [Perf][Feat] Add generic cuteDSL LL BF16 router (GEMM) (#42562) Signed-off-by: LopezCastroRoberto Signed-off-by: Lucas Wilkinson Signed-off-by: Roberto L. Castro <38211239+LopezCastroRoberto@users.noreply.github.com> Co-authored-by: Lucas Wilkinson --- .buildkite/test_areas/kernels.yaml | 5 + tests/kernels/test_ll_bf16_gemm.py | 553 +++++++++++++++++ .../kernels/linear/cute_dsl/__init__.py | 0 .../linear/cute_dsl/_ll_bf16_dotprod.py | 313 ++++++++++ .../linear/cute_dsl/_ll_bf16_splitk.py | 585 ++++++++++++++++++ .../kernels/linear/cute_dsl/ll_bf16.py | 282 +++++++++ .../layers/fused_moe/router/gate_linear.py | 54 +- vllm/model_executor/warmup/kernel_warmup.py | 28 + 8 files changed, 1812 insertions(+), 8 deletions(-) create mode 100644 tests/kernels/test_ll_bf16_gemm.py create mode 100644 vllm/model_executor/kernels/linear/cute_dsl/__init__.py create mode 100644 vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py create mode 100644 vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_splitk.py create mode 100644 vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 894cda34119..d1f3efd6f07 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -232,6 +232,10 @@ steps: - vllm/v1/attention/backends/mla/flashinfer_mla.py - vllm/v1/attention/selector.py - vllm/platforms/cuda.py + - vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py + - vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py + - vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_splitk.py + - tests/kernels/test_ll_bf16_gemm.py - tests/kernels/test_top_k_per_row.py commands: - nvidia-smi @@ -260,6 +264,7 @@ steps: - pytest -v -s tests/kernels/moe/test_flashinfer_moe.py - pytest -v -s tests/kernels/moe/test_trtllm_nvfp4_moe.py - pytest -v -s tests/kernels/moe/test_cutedsl_moe.py + - pytest -v -s tests/kernels/test_ll_bf16_gemm.py # e2e - pytest -v -s tests/models/quantization/test_nvfp4.py diff --git a/tests/kernels/test_ll_bf16_gemm.py b/tests/kernels/test_ll_bf16_gemm.py new file mode 100644 index 00000000000..f530f6f7414 --- /dev/null +++ b/tests/kernels/test_ll_bf16_gemm.py @@ -0,0 +1,553 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for cuteDSL low-latency router GEMM (dot-product + split-K).""" + +import pytest +import torch +import torch.nn.functional as F + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.fixture(autouse=True, scope="module") +def _require_sm90_and_cutedsl(): + if torch.cuda.get_device_capability()[0] < 9: + pytest.skip("Requires SM90+ (Hopper/Blackwell)") + from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import ( + is_available, + ) + + if not is_available(): + pytest.skip("cuteDSL (CUTLASS Python) not installed") + + +# ===== Helpers ===== + + +def _ref(a, b): + return torch.mm(a.float(), b.float().T) + + +def _assert_close(out, ref, *, min_cos_sim=0.99, context=""): + assert out.device.type == "cuda", f"{context}: not on CUDA" + assert torch.isfinite(out).all(), f"{context}: NaN/Inf" + cos = F.cosine_similarity( + out.reshape(-1).float(), ref.reshape(-1).float(), dim=0 + ).item() + assert cos > min_cos_sim, ( + f"{context}: cos_sim {cos:.6f} < {min_cos_sim} " + f"(abs_err={(out.float() - ref.float()).abs().max().item():.2e})" + ) + + +def _can_precompile(a, b): + return ( + a.dim() == 2 + and b.dim() == 2 + and a.dtype == torch.bfloat16 + and b.dtype == torch.bfloat16 + and a.device.type == "cuda" + and b.device.type == "cuda" + and a.device == b.device + and a.shape[1] == b.shape[1] + and a.is_contiguous() + and b.is_contiguous() + ) + + +def _gemm(a, b): + from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import ( + ll_bf16_gemm, + ll_bf16_gemm_kernel, + ) + + if _can_precompile(a, b): + compile_key = ll_bf16_gemm_kernel.dispatch( + M=a.shape[0], K=a.shape[1], N=b.shape[0] + ) + ll_bf16_gemm_kernel.compile(compile_key) + return ll_bf16_gemm(a, b) + + +# ===== Shapes ===== + +SHAPES = [ + (256, 7168, "DSV3"), + (256, 14400, "DSV4-Flash"), + (128, 5120, "DeepSeek-V2"), + (8, 4096, "Mixtral-8x7B"), + (64, 2880, "non-tile-aligned-K"), + (256, 2048, "split-K-boundary"), +] + +SHAPES_SPLITK = [(n, k, d) for n, k, d in SHAPES if k >= 2048] + + +# ================================================================= +# Dot-product kernel (M<=4 or K<2048) +# ================================================================= + + +@pytest.mark.parametrize("M", [1, 2, 3, 4]) +@pytest.mark.parametrize("N,K,desc", SHAPES, ids=[s[2] for s in SHAPES]) +def test_dotprod(M, N, K, desc): + torch.manual_seed(42) + a = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") + b = torch.randn(N, K, dtype=torch.bfloat16, device="cuda") + out = _gemm(a, b) + assert out.dtype == torch.float32 + assert out.shape == (M, N) + _assert_close(out, _ref(a, b), context=f"dotprod {M}x{N}x{K}") + + +@pytest.mark.parametrize("M", [1, 4, 8, 16]) +def test_dotprod_small_K(M): + """K<2048 forces dot-product regardless of M.""" + torch.manual_seed(42) + a = torch.randn(M, 1024, dtype=torch.bfloat16, device="cuda") + b = torch.randn(64, 1024, dtype=torch.bfloat16, device="cuda") + _assert_close(_gemm(a, b), _ref(a, b), context=f"small_K M={M}") + + +@pytest.mark.parametrize("K", [16, 32, 64, 128, 256, 512, 1024, 1536]) +def test_dotprod_K_sweep(K): + torch.manual_seed(42) + a = torch.randn(4, K, dtype=torch.bfloat16, device="cuda") + b = torch.randn(32, K, dtype=torch.bfloat16, device="cuda") + _assert_close(_gemm(a, b), _ref(a, b), context=f"K={K}") + + +# ================================================================= +# Split-K kernel (M>4 and K>=2048) +# ================================================================= + + +@pytest.mark.parametrize("M", [5, 6, 8, 12, 16]) +@pytest.mark.parametrize("N,K,desc", SHAPES_SPLITK, ids=[s[2] for s in SHAPES_SPLITK]) +def test_splitk(M, N, K, desc): + torch.manual_seed(42) + a = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") + b = torch.randn(N, K, dtype=torch.bfloat16, device="cuda") + out = _gemm(a, b) + assert out.dtype == torch.float32 + assert out.shape == (M, N) + _assert_close(out, _ref(a, b), context=f"splitk {M}x{N}x{K}") + + +@pytest.mark.parametrize("K", [2048, 2304, 2880, 3072, 4096, 5120, 7168, 14400]) +def test_splitk_K_sweep(K): + """Includes non-tile-aligned K (2880) and uneven split (2304).""" + torch.manual_seed(42) + a = torch.randn(8, K, dtype=torch.bfloat16, device="cuda") + b = torch.randn(64, K, dtype=torch.bfloat16, device="cuda") + _assert_close(_gemm(a, b), _ref(a, b), context=f"splitk K={K}") + + +# ================================================================= +# Dispatch boundary (M=4/5, K=2032/2048) +# ================================================================= + + +@pytest.mark.parametrize("M", [4, 5]) +@pytest.mark.parametrize("K", [2032, 2048]) +def test_dispatch_boundary(M, K): + torch.manual_seed(42) + a = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") + b = torch.randn(64, K, dtype=torch.bfloat16, device="cuda") + path = "splitk" if M > 4 and K >= 2048 else "dotprod" + _assert_close(_gemm(a, b), _ref(a, b), context=f"M={M} K={K} ({path})") + + +# ================================================================= +# Arbitrary N +# ================================================================= + + +@pytest.mark.parametrize("N", [1, 3, 7, 16, 17, 64, 128, 256, 384]) +def test_arbitrary_N_dotprod(N): + torch.manual_seed(42) + a = torch.randn(4, 2048, dtype=torch.bfloat16, device="cuda") + b = torch.randn(N, 2048, dtype=torch.bfloat16, device="cuda") + out = _gemm(a, b) + assert out.shape == (4, N) + _assert_close(out, _ref(a, b), context=f"dotprod N={N}") + + +@pytest.mark.parametrize("N", [1, 8, 16, 17, 64, 128, 256]) +def test_arbitrary_N_splitk(N): + torch.manual_seed(42) + a = torch.randn(8, 4096, dtype=torch.bfloat16, device="cuda") + b = torch.randn(N, 4096, dtype=torch.bfloat16, device="cuda") + out = _gemm(a, b) + assert out.shape == (8, N) + _assert_close(out, _ref(a, b), context=f"splitk N={N}") + + +# ================================================================= +# Single token (M=1, decode path) +# ================================================================= + + +@pytest.mark.parametrize( + "N,K", + [(256, 7168), (256, 14400), (8, 4096), (384, 7168)], + ids=["DSV3", "DSV4-Flash", "Mixtral", "DSV4-Pro"], +) +def test_single_token(N, K): + torch.manual_seed(42) + a = torch.randn(1, K, dtype=torch.bfloat16, device="cuda") + b = torch.randn(N, K, dtype=torch.bfloat16, device="cuda") + out = _gemm(a, b) + assert out.shape == (1, N) + _assert_close(out, _ref(a, b), context=f"M=1 {N}x{K}") + + +# ================================================================= +# Numerical robustness +# ================================================================= + + +@pytest.mark.parametrize("M,K", [(4, 2048), (8, 4096)], ids=["dotprod", "splitk"]) +def test_large_values(M, K): + torch.manual_seed(42) + a = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") * 100 + b = torch.randn(64, K, dtype=torch.bfloat16, device="cuda") * 100 + out = _gemm(a, b) + assert torch.isfinite(out).all() + _assert_close(out, _ref(a, b), context=f"large M={M}") + + +def test_near_zero(): + torch.manual_seed(42) + a = torch.randn(4, 2048, dtype=torch.bfloat16, device="cuda") * 1e-4 + b = torch.randn(64, 2048, dtype=torch.bfloat16, device="cuda") * 1e-4 + out = _gemm(a, b) + assert torch.isfinite(out).all() + assert out.abs().max() < 1.0 + + +def test_zeros(): + a = torch.zeros(4, 2048, dtype=torch.bfloat16, device="cuda") + b = torch.randn(64, 2048, dtype=torch.bfloat16, device="cuda") + assert (_gemm(a, b) == 0).all() + + +def test_ones(): + a = torch.ones(1, 2048, dtype=torch.bfloat16, device="cuda") + b = torch.randn(32, 2048, dtype=torch.bfloat16, device="cuda") + _assert_close(_gemm(a, b), _ref(a, b), context="ones") + + +# ================================================================= +# Output dtype +# ================================================================= + + +@pytest.mark.parametrize("M,K", [(4, 2048), (8, 4096)], ids=["dotprod", "splitk"]) +def test_output_fp32(M, K): + torch.manual_seed(42) + a = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") + b = torch.randn(64, K, dtype=torch.bfloat16, device="cuda") + assert _gemm(a, b).dtype == torch.float32 + + +# ================================================================= +# Determinism +# ================================================================= + + +@pytest.mark.parametrize( + "M,K", [(1, 4096), (4, 4096), (5, 4096), (8, 4096), (16, 4096)] +) +def test_deterministic(M, K): + torch.manual_seed(42) + a = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") + b = torch.randn(128, K, dtype=torch.bfloat16, device="cuda") + torch.testing.assert_close(_gemm(a, b), _gemm(a, b), atol=0, rtol=0) + + +# ================================================================= +# Cross-kernel consistency +# ================================================================= + + +@pytest.mark.parametrize("K", [2048, 4096, 7168]) +def test_dotprod_vs_splitk(K): + """First 4 rows from split-K (M=5) match dot-product (M=4).""" + torch.manual_seed(42) + a = torch.randn(5, K, dtype=torch.bfloat16, device="cuda") + b = torch.randn(64, K, dtype=torch.bfloat16, device="cuda") + _assert_close( + _gemm(a, b)[:4], + _gemm(a[:4], b), + min_cos_sim=0.999, + context=f"cross-kernel K={K}", + ) + + +# ================================================================= +# CUDA graph +# ================================================================= + + +@pytest.mark.parametrize("M,K", [(4, 2048), (8, 4096)], ids=["dotprod", "splitk"]) +def test_cudagraph(M, K): + torch.manual_seed(42) + a = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") + b = torch.randn(64, K, dtype=torch.bfloat16, device="cuda") + _gemm(a, b) + torch.accelerator.synchronize() + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + out = _gemm(a, b) + for _ in range(5): + g.replay() + torch.accelerator.synchronize() + _assert_close(out, _ref(a, b), context=f"cudagraph M={M}") + + +def test_cudagraph_20x_replay(): + torch.manual_seed(42) + a = torch.randn(4, 4096, dtype=torch.bfloat16, device="cuda") + b = torch.randn(256, 4096, dtype=torch.bfloat16, device="cuda") + _gemm(a, b) + torch.accelerator.synchronize() + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + out = _gemm(a, b) + results = [] + for _ in range(20): + g.replay() + torch.accelerator.synchronize() + results.append(out.clone()) + for i in range(1, len(results)): + torch.testing.assert_close( + results[0], results[i], atol=0, rtol=0, msg=f"Replay {i} differs" + ) + + +def test_cudagraph_input_update(): + torch.manual_seed(42) + a = torch.randn(4, 2048, dtype=torch.bfloat16, device="cuda") + b = torch.randn(64, 2048, dtype=torch.bfloat16, device="cuda") + _gemm(a, b) + torch.accelerator.synchronize() + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + out = _gemm(a, b) + a.copy_(torch.randn_like(a)) + g.replay() + torch.accelerator.synchronize() + _assert_close(out, _ref(a, b), context="cudagraph input update") + + +# ================================================================= +# GateLinear dispatch integration +# ================================================================= + + +def _make_gate_linear(monkeypatch, *, params_dtype, out_dtype=torch.float32): + monkeypatch.setattr( + "vllm.model_executor.layers.linear.get_tensor_model_parallel_rank", + lambda: 0, + ) + monkeypatch.setattr( + "vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", + lambda: 1, + ) + monkeypatch.setattr( + "vllm.model_executor.parameter.get_tensor_model_parallel_rank", + lambda: 0, + ) + monkeypatch.setattr( + "vllm.model_executor.parameter.get_tensor_model_parallel_world_size", + lambda: 1, + ) + + from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear + + monkeypatch.setattr( + "vllm.model_executor.kernels.linear.cute_dsl.ll_bf16.is_available", + lambda: True, + ) + return GateLinear( + input_size=2048, + output_size=64, + bias=False, + out_dtype=out_dtype, + params_dtype=params_dtype, + ).cuda() + + +def test_gate_linear_uses_ll_bf16_for_bf16_fast_path(monkeypatch): + gate = _make_gate_linear(monkeypatch, params_dtype=torch.bfloat16) + x = torch.randn(4, 2048, dtype=torch.bfloat16, device="cuda") + calls = [] + + def fake_ll_bf16_gemm(hidden_states, router_weight): + calls.append((hidden_states, router_weight)) + return torch.full( + (hidden_states.shape[0], router_weight.shape[0]), + 1.0, + dtype=torch.float32, + device=hidden_states.device, + ) + + monkeypatch.setattr( + "vllm.model_executor.kernels.linear.cute_dsl.ll_bf16.ll_bf16_gemm", + fake_ll_bf16_gemm, + ) + out, bias = gate(x) + assert bias is None + assert out.shape == (4, 64) + assert out.dtype == torch.float32 + assert len(calls) == 1 + assert calls[0][0] is x + assert calls[0][1] is gate.weight + + +def test_gate_linear_fp32_weight_falls_back(monkeypatch): + gate = _make_gate_linear(monkeypatch, params_dtype=torch.float32) + assert not gate.allow_ll_bf16_gemm + x = torch.randn(4, 2048, dtype=torch.bfloat16, device="cuda") + + def fail_ll_bf16_gemm(hidden_states, router_weight): + raise AssertionError("ll_bf16_gemm should not run for fp32 weights") + + monkeypatch.setattr( + "vllm.model_executor.kernels.linear.cute_dsl.ll_bf16.ll_bf16_gemm", + fail_ll_bf16_gemm, + ) + out, _ = gate(x) + assert out.shape == (4, 64) + assert out.dtype == torch.float32 + + +def test_gate_linear_non_bf16_activation_falls_back(monkeypatch): + gate = _make_gate_linear(monkeypatch, params_dtype=torch.bfloat16) + x = torch.randn(4, 2048, dtype=torch.float16, device="cuda") + + def fail_ll_bf16_gemm(hidden_states, router_weight): + raise AssertionError("ll_bf16_gemm should not run for non-bf16 activations") + + monkeypatch.setattr( + "vllm.model_executor.kernels.linear.cute_dsl.ll_bf16.ll_bf16_gemm", + fail_ll_bf16_gemm, + ) + out, _ = gate(x) + assert out.shape == (4, 64) + assert out.dtype == torch.float32 + + +def test_gate_linear_set_out_dtype_enables_ll_bf16(monkeypatch): + gate = _make_gate_linear(monkeypatch, params_dtype=torch.bfloat16, out_dtype=None) + assert not gate.allow_ll_bf16_gemm + gate.set_out_dtype(torch.float32) + assert gate.allow_ll_bf16_gemm + + +def test_gate_linear_non_fp32_out_dtype_disables_ll_bf16(monkeypatch): + gate = _make_gate_linear(monkeypatch, params_dtype=torch.bfloat16, out_dtype=None) + gate.set_out_dtype(torch.bfloat16) + assert not gate.allow_ll_bf16_gemm + + +def test_gate_linear_m_gt_16_falls_back(monkeypatch): + gate = _make_gate_linear(monkeypatch, params_dtype=torch.bfloat16) + x = torch.randn(17, 2048, dtype=torch.bfloat16, device="cuda") + + def fail_ll_bf16_gemm(hidden_states, router_weight): + raise AssertionError("ll_bf16_gemm should not run for M > 16") + + monkeypatch.setattr( + "vllm.model_executor.kernels.linear.cute_dsl.ll_bf16.ll_bf16_gemm", + fail_ll_bf16_gemm, + ) + out, _ = gate(x) + assert out.shape == (17, 64) + assert out.dtype == torch.float32 + + +# ================================================================= +# Negative tests — invalid inputs +# ================================================================= + + +@pytest.mark.parametrize( + "M,K,N,dtype", + [ + pytest.param(4, 2048, 64, torch.float32, id="fp32_input"), + pytest.param(4, 2048, 64, torch.float16, id="fp16_input"), + pytest.param(8, 4096, 64, torch.float32, id="fp32_splitk_path"), + ], +) +def test_invalid_dtype(M, K, N, dtype): + a = torch.randn(M, K, device="cuda", dtype=dtype) + b = torch.randn(N, K, device="cuda", dtype=dtype) + with pytest.raises(ValueError, match="dtype=bfloat16"): + _gemm(a, b) + + +def test_invalid_device_cpu(): + a = torch.randn(4, 2048, dtype=torch.bfloat16) + b = torch.randn(64, 2048, dtype=torch.bfloat16) + with pytest.raises(ValueError, match="device_type=cuda"): + _gemm(a, b) + + +def test_invalid_1d_input(): + a = torch.randn(2048, dtype=torch.bfloat16, device="cuda") + b = torch.randn(64, 2048, dtype=torch.bfloat16, device="cuda") + with pytest.raises(ValueError, match="2D tensors"): + _gemm(a, b) + + +def test_mismatched_K(): + a = torch.randn(4, 2048, dtype=torch.bfloat16, device="cuda") + b = torch.randn(64, 1024, dtype=torch.bfloat16, device="cuda") + with pytest.raises(ValueError, match="matching K dimensions"): + _gemm(a, b) + + +def test_invalid_K_divisibility(): + a = torch.randn(4, 2049, dtype=torch.bfloat16, device="cuda") + b = torch.randn(64, 2049, dtype=torch.bfloat16, device="cuda") + + with pytest.raises(ValueError, match="K to be divisible by 8"): + _gemm(a, b) + + +def test_non_contiguous_input(): + a = torch.randn(2048, 4, dtype=torch.bfloat16, device="cuda").T + b = torch.randn(64, 2048, dtype=torch.bfloat16, device="cuda") + assert not a.is_contiguous() + with pytest.raises(ValueError, match="contiguous row-major"): + _gemm(a, b) + + +def test_invalid_output_dtype(): + a = torch.randn(4, 2048, dtype=torch.bfloat16, device="cuda") + b = torch.randn(64, 2048, dtype=torch.bfloat16, device="cuda") + with pytest.raises(ValueError, match="output_dtype=torch.float32"): + from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import ll_bf16_gemm + + ll_bf16_gemm(a, b, output_dtype=torch.bfloat16) + + +def test_cache_miss_compiles_dotprod(): + from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import LLBf16Gemm + + torch.manual_seed(42) + a = torch.randn(3, 64, dtype=torch.bfloat16, device="cuda") + b = torch.randn(17, 64, dtype=torch.bfloat16, device="cuda") + kernel = LLBf16Gemm() + out = kernel(a, b) + assert out.shape == (3, 17) + _assert_close(out, _ref(a, b), context="cache miss dotprod") + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/vllm/model_executor/kernels/linear/cute_dsl/__init__.py b/vllm/model_executor/kernels/linear/cute_dsl/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py b/vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py new file mode 100644 index 00000000000..f6931ceb1c8 --- /dev/null +++ b/vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py @@ -0,0 +1,313 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import cutlass +import cutlass.cute as cute +from cuda.bindings.driver import CUstream +from cutlass import const_expr + + +class LLBf16Dotprod: + """BF16 router GEMM kernel based on CTA-local dot products. + + This kernel computes C[M, N] = A[M, K] @ B[N, K]^T for bf16 inputs + and fp32 output. It launches one CTA per output column, distributes K + across CTA threads with vectorized loads, accumulates one fp32 dot product + per token, and reduces through warp shuffles plus shared memory. + + :param k: Compile-time K dimension specialized into the generated kernel. + :type k: int + :param bs: Threads per CTA and K-stripe width used by the reduction. + :type bs: int + :param main_vec_width: bf16 elements loaded per thread in the main loop. + :type main_vec_width: int + :param tail_vec_width: bf16 elements loaded per thread in the vector tail. + :type tail_vec_width: int + :param use_pdl: Whether to launch with Programmatic Dependent Launch. + :type use_pdl: bool + + :note: Supported A/B data types: + - BFloat16/BFloat16 + :note: Supported accumulator data types: + - Float32 + :note: Supported C data types: + - Float32 + :note: Constraints: + - K must preserve 16-byte row alignment for contiguous bf16 inputs. + + :note: K is handled as vectorized main/tail loops plus scalar remainder. + + :compile-key: ``(M, K, bs)`` selects the token count, hidden size, + and CTA thread/K-stripe width specialization. + """ + + def __init__( + self, + k: int, + bs: int = 128, + main_vec_width: int = 8, + tail_vec_width: int = 4, + use_pdl: bool = False, + ): + """Initialize the dot-product kernel configuration. + + This configuration fixes the CTA thread count, reduction warp count, + bf16 vector widths, and K-loop decomposition used by the generated + kernel. + + :param k: Hidden size K used to specialize the K-loop decomposition. + :type k: int + :param bs: Threads per CTA and K-stripe width for the reduction. + :type bs: int + :param main_vec_width: BF16 elements loaded per thread in the main loop. + :type main_vec_width: int + :param tail_vec_width: BF16 elements loaded per thread in the vector tail. + :type tail_vec_width: int + :param use_pdl: Whether to launch with Programmatic Dependent Launch. + :type use_pdl: bool + """ + self.bs = bs + self.main_vec_width = main_vec_width + self.tail_vec_width = tail_vec_width + self.use_pdl = use_pdl + self.num_warps = bs // cute.arch.WARP_SIZE + self._init_k_tiles(k) + + def _vectorized_elems(self, k_extent: int, vec_width: int) -> int: + vector_tile = vec_width * self.bs + return (k_extent // vector_tile) * vector_tile + + def _init_k_tiles(self, k: int) -> None: + """Split K into vector loops, scalar rounds, and ragged tail.""" + self.k_main_elems = self._vectorized_elems(k, self.main_vec_width) + self.k_after_main = k - self.k_main_elems + self.k_tail_elems = self._vectorized_elems( + self.k_after_main, self.tail_vec_width + ) + self.k_done_all = self.k_main_elems + self.k_tail_elems + self.scalar_rem = k - self.k_done_all + self.ks_full = self.scalar_rem // self.bs + self.ks_part = self.scalar_rem % self.bs + self.k_scalar_full = self.ks_full * self.bs + self.k_part_offset = self.k_done_all + self.k_scalar_full + self.main_tiles = self.k_main_elems // (self.main_vec_width * self.bs) + self.tail_tiles = self.k_tail_elems // (self.tail_vec_width * self.bs) + + @cute.jit + def _vector_dotprod( + self, + acc: cute.Tensor, + tA: cute.Tensor, + tB: cute.Tensor, + M: cutlass.Constexpr, + num_tiles: cutlass.Constexpr, + align_bytes: cutlass.Constexpr, + ): + for tile in cutlass.range_constexpr(num_tiles): + bt = tB[None, tile] + br = cute.make_rmem_tensor_like(bt) + cute.autovec_copy(bt, br) + br_f32 = br.load().to(cutlass.Float32) + + for m in cutlass.range_constexpr(M): + at = tA[m, None, tile] + ar = cute.make_rmem_tensor_like(at) + cute.autovec_copy(at, ar) + vec_width: cutlass.Constexpr = cute.size(ar) + for v in cutlass.range_constexpr(vec_width): + acc[m] = acc[m] + ar[v].to(cutlass.Float32) * br_f32[v] + + def _make_thread_vector_slice( + self, + gA_vec: cute.Tensor, + gB_vec: cute.Tensor, + tidx: cutlass.Int32, + n_idx: cutlass.Int32, + bs: cutlass.Constexpr, + ): + # (M/N, K_TILE, K_LANE, K_VEC); tidx selects K_LANE. + tA = cute.logical_divide(gA_vec, (None, (None, bs))) + tB = cute.logical_divide(gB_vec, (None, (None, bs))) + return tA[None, (None, (tidx, None))], tB[n_idx, (None, (tidx, None))] + + def _make_k_slice( + self, + gX: cute.Tensor, + k_offset: cutlass.Constexpr, + k_extent: cutlass.Constexpr, + ): + k_layout_extent: cutlass.Constexpr = ( + 1 if const_expr(k_extent == 0) else k_extent + ) + + if const_expr(k_offset == 0): + return cute.local_tile( + gX, (cute.size(gX, mode=[0]), k_layout_extent), (0, 0) + ) + return cute.local_tile( + cute.domain_offset((0, k_offset), gX), + (cute.size(gX, mode=[0]), k_layout_extent), + (0, 0), + ) + + @cute.jit + def __call__( + self, + gA: cute.Tensor, + gB: cute.Tensor, + gC: cute.Tensor, + M: cutlass.Constexpr, + K_dim: cutlass.Constexpr, + N_dim: cutlass.Int32, + stream: CUstream, + ): + """Execute the dot-product GEMM operation in steps: + - Launch one CTA per output column ``n`` with ``bs`` threads. + - Keep one FP32 accumulator per token ``m`` in each thread. + - Traverse K with vectorized 128-bit, vectorized 64-bit, scalar, and + ragged-tail loops from the precomputed K decomposition. + - Reduce each token accumulator first within the warp, then across + warps through shared memory, and store ``C[:, n]``. + + :param gA: Input tensor A with shape ``[M, K]``. + :type gA: cute.Tensor + :param gB: Input tensor B with shape ``[N, K]``. + :type gB: cute.Tensor + :param gC: Output tensor C with shape ``[M, N]``. + :type gC: cute.Tensor + :param M: Token count selected by the compile key. + :type M: cutlass.Constexpr + :param K_dim: Hidden size selected by the compile key. + :type K_dim: cutlass.Constexpr + :param N_dim: Output column count used for the launch grid. + :type N_dim: cutlass.Int32 + :param stream: CUDA stream for asynchronous execution. + :type stream: CUstream + """ + self.kernel( + gA, + gB, + gC, + M, + self.main_vec_width, + self.tail_vec_width, + self.bs, + self.num_warps, + self.k_main_elems, + self.k_tail_elems, + self.k_done_all, + self.ks_full, + self.ks_part, + self.k_scalar_full, + self.k_part_offset, + self.main_tiles, + self.tail_tiles, + ).launch( + grid=[N_dim, 1, 1], + block=[self.bs, 1, 1], + smem=M * 4 * self.num_warps, + stream=stream, + use_pdl=self.use_pdl, + min_blocks_per_mp=1, + ) + + @cute.kernel + def kernel( + self, + gA: cute.Tensor, + gB: cute.Tensor, + gC: cute.Tensor, + M: cutlass.Constexpr, + main_vec_width: cutlass.Constexpr, + tail_vec_width: cutlass.Constexpr, + bs: cutlass.Constexpr, + num_warps: cutlass.Constexpr, + k_main_elems: cutlass.Constexpr, + k_tail_elems: cutlass.Constexpr, + k_done_all: cutlass.Constexpr, + ks_full: cutlass.Constexpr, + ks_part: cutlass.Constexpr, + k_scalar_full: cutlass.Constexpr, + k_part_offset: cutlass.Constexpr, + main_tiles: cutlass.Constexpr, + tail_tiles: cutlass.Constexpr, + ): + tidx, _, _ = cute.arch.thread_idx() + n_idx, _, _ = cute.arch.block_idx() + wid = cute.arch.warp_idx() + + # One FP32 accumulator per token. + acc = cute.make_rmem_tensor((M,), cutlass.Float32) + acc.fill(0.0) + + if const_expr(self.use_pdl): + cute.arch.griddepcontrol_wait() + + # 128-bit vectorized main loop + if const_expr(k_main_elems > 0): + gA_main = self._make_k_slice(gA, 0, k_main_elems) + gB_main = self._make_k_slice(gB, 0, k_main_elems) + gA_vec = cute.logical_divide(gA_main, (None, main_vec_width)) + gB_vec = cute.logical_divide(gB_main, (None, main_vec_width)) + tA, tB = self._make_thread_vector_slice(gA_vec, gB_vec, tidx, n_idx, bs) + self._vector_dotprod(acc, tA, tB, M, main_tiles, 16) + + # 64-bit vectorized tail (K remainder after main loop) + if const_expr(k_tail_elems > 0): + gA_tail = self._make_k_slice(gA, k_main_elems, k_tail_elems) + gB_tail = self._make_k_slice(gB, k_main_elems, k_tail_elems) + gA_tail_vec = cute.logical_divide(gA_tail, (None, tail_vec_width)) + gB_tail_vec = cute.logical_divide(gB_tail, (None, tail_vec_width)) + tA_t, tB_t = self._make_thread_vector_slice( + gA_tail_vec, gB_tail_vec, tidx, n_idx, bs + ) + self._vector_dotprod(acc, tA_t, tB_t, M, tail_tiles, 8) + + # Full scalar rounds use CuTe width-1 tiles; KS_PART is the ragged tail. + if const_expr(ks_full > 0): + gA_scalar = self._make_k_slice(gA, k_done_all, k_scalar_full) + gB_scalar = self._make_k_slice(gB, k_done_all, k_scalar_full) + gA_scalar_vec = cute.logical_divide(gA_scalar, (None, 1)) + gB_scalar_vec = cute.logical_divide(gB_scalar, (None, 1)) + tA_s, tB_s = self._make_thread_vector_slice( + gA_scalar_vec, gB_scalar_vec, tidx, n_idx, bs + ) + self._vector_dotprod(acc, tA_s, tB_s, M, ks_full, 2) + + # Only threads below KS_PART load the ragged tail. + if const_expr(ks_part > 0): + gA_part = self._make_k_slice(gA, k_part_offset, ks_part) + gB_part = self._make_k_slice(gB, k_part_offset, ks_part) + if tidx < ks_part: + bv2 = gB_part[n_idx, tidx].to(cutlass.Float32) + for m in cutlass.range_constexpr(M): + acc[m] = acc[m] + gA_part[m, tidx].to(cutlass.Float32) * bv2 + + # Intra-warp shuffle reduction + for m in cutlass.range_constexpr(M): + acc[m] = cute.arch.warp_reduction_sum(acc[m]) + + # Cross-warp reduction via shared memory + smem_red_layout = cute.make_layout((M, num_warps), stride=(num_warps, 1)) + smem = cutlass.utils.SmemAllocator() + sm = smem.allocate_tensor(cutlass.Float32, smem_red_layout, byte_alignment=16) + with cute.arch.elect_one(): + for m in cutlass.range_constexpr(M): + sm[m, wid] = acc[m] + + # Final reduction and output + cute.arch.sync_threads() + if tidx == 0: + for m in cutlass.range_constexpr(M): + partials = sm[m, None].load() + gC[m, n_idx] = partials.reduce( + cute.ReductionOp.ADD, + init_val=cutlass.Float32(0.0), + reduction_profile=0, + ) + if const_expr(self.use_pdl): + cute.arch.griddepcontrol_launch_dependents() + + +def make_host_bf16(k_val: int, bs: int = 128): + return LLBf16Dotprod(k=k_val, bs=bs) diff --git a/vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_splitk.py b/vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_splitk.py new file mode 100644 index 00000000000..b44a5f6e2c9 --- /dev/null +++ b/vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_splitk.py @@ -0,0 +1,585 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import math + +import cutlass +import cutlass.cute as cute +from cuda.bindings.driver import CUstream +from cutlass import const_expr +from cutlass._mlir import ir as _ir +from cutlass._mlir.dialects import llvm as _llvm +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.pipeline import sm90 as pipeline + + +# Map a local smem address to the same peer CTA DSMEM offset. +@dsl_user_op +def set_block_rank(smem_ptr, peer_rank, *, loc=None, ip=None): + dsmem_ptr = cute.arch.map_dsmem_ptr(smem_ptr, peer_rank, loc=loc, ip=ip) + return cutlass.Int32(dsmem_ptr.toint(loc=loc, ip=ip)) + + +# Plain DSMEM store; mbarrier helpers do not model this reduction. +@dsl_user_op +def st_shared_remote_f32(remote_addr, val, *, loc=None, ip=None): + i32 = _ir.IntegerType.get_signless(32) + addr_ir = remote_addr.ir_value(loc=loc, ip=ip) + val_ir = val.ir_value(loc=loc, ip=ip) + _llvm.inline_asm( + i32, + [addr_ir, val_ir], + "st.shared::cluster.f32 [$0], $1; mov.u32 $2, 0;", + "r,f,=r", + has_side_effects=True, # keep the inline store ordered + loc=loc, + ip=ip, + ) + + +class LLBf16SplitK: + """BF16 router GEMM kernel based on clustered split-K MMA. + + This kernel computes C[M, N] = A[M, K] @ B[N, K]^T for bf16 inputs + and fp32 output. It partitions K across a CTA cluster, uses DMA warps to + stage A/B tiles with cp.async, uses MMA warps to accumulate fp32 partials, + and reduces split-K partials through DSMEM before storing C. + + :param ab_dtype: Element type for A and B operands. + :param acc_dtype: Accumulator type used by MMA and reductions. + :param out_dtype: Output element type. The public wrapper uses fp32. + :param tile_n: CTA tile size in N. M is fixed to 16 for router batches. + :type tile_n: int + :param tile_k: K tile size staged through shared memory. + :type tile_k: int + :param num_stages: Number of cp.async pipeline stages. + :type num_stages: int + :param num_dma_warps: Producer warps that issue GMEM to SMEM copies. + :type num_dma_warps: int + :param split_k: Number of CTAs in the cluster-level split-K reduction. + :type split_k: int + :param use_pdl: Whether to launch with Programmatic Dependent Launch. + :type use_pdl: bool + + :note: Supported A/B data types: + - BFloat16/BFloat16 + :note: Supported accumulator data types: + - Float32 + :note: Supported C data types: + - Float32 + :note: Constraints: + - K must preserve 16-byte row alignment for contiguous bf16 inputs. + + :compile-key: ``(split_k, num_stages)`` selects the cluster split + count and cp.async pipeline depth specialization. + """ + + def __init__( + self, + ab_dtype=cutlass.BFloat16, + acc_dtype=cutlass.Float32, + out_dtype=cutlass.Float32, + tile_n: int = 16, + tile_k: int = 256, + num_stages: int = 2, + num_dma_warps: int = 4, + split_k: int = 8, + use_pdl: bool = False, + ): + """Initialize the split-K kernel configuration. + + This configuration fixes the CTA tile shape, cp.async pipeline depth, + producer/consumer warp split, and CTA-cluster split count used by the + DSMEM reduction. + + :param ab_dtype: Element type for A and B operands. + :param acc_dtype: Accumulator type used by MMA and reductions. + :param out_dtype: Output element type. + :param tile_n: CTA tile size in N. M is fixed to 16. + :type tile_n: int + :param tile_k: K tile size staged through shared memory. + :type tile_k: int + :param num_stages: Number of cp.async pipeline stages. + :type num_stages: int + :param num_dma_warps: Producer warps that issue GMEM to SMEM copies. + :type num_dma_warps: int + :param split_k: CTAs in the cluster-level split-K reduction. + :type split_k: int + :param use_pdl: Whether to launch with Programmatic Dependent Launch. + :type use_pdl: bool + """ + self.ab_dtype = ab_dtype + self.acc_dtype = acc_dtype + self.out_dtype = out_dtype + self.tile_m = 16 + self.tile_n = tile_n + self.tile_k = tile_k + self.copy_bits = 128 + self.num_stages = num_stages + self.split_k = split_k + self.use_pdl = use_pdl + self.mma_shape = (16, 8, 16) # mma.sync.aligned.m16n8k16 + self.atom_layout = (1, 1, 1) # one MMA atom per warp + self.num_dma_warps = num_dma_warps + self.num_mma_warps = 4 + self.num_dma_threads = self.num_dma_warps * cute.arch.WARP_SIZE + self.num_mma_threads = self.num_mma_warps * cute.arch.WARP_SIZE + self.num_threads = self.num_dma_threads + self.num_mma_threads + self.num_epilogue_elems = self.tile_m * self.tile_n + self.epilogue_elems_per_thread = self.num_epilogue_elems // self.num_mma_threads + + def _make_smem_layout_AB(self, dtype, copy_bits, smem_tiler): + """Build the staged swizzled SMEM layout for A or B tiles.""" + major_size = min(smem_tiler[1], 64) + # Match swizzle span to contiguous K bytes, capped by CuTe 3-bit swizzle. + swizzle_bits = int(math.log2(major_size * dtype.width // copy_bits)) + swizzle_bits = min(swizzle_bits, 3) + # Tile the swizzled atom across (M_or_N, K, stages). + layout_atom_outer = cute.make_layout((8, major_size), stride=(major_size, 1)) + layout_atom = cute.make_composed_layout( + cute.make_swizzle(swizzle_bits, 3, 3), 0, layout_atom_outer + ) + return cute.tile_to_shape(layout_atom, smem_tiler, (0, 1, 2)) + + def _make_gmem_tiled_copy(self, atom_copy, dtype, copy_bits, num_threads): + """Build the per-thread cp.async vector-copy layout.""" + # Lay threads across K so each lane issues one vector copy. + copy_elems = copy_bits // dtype.width + k_threads = cute.size(self.tile_k) // copy_elems # threads along K + thread_layout = cute.make_layout( + (num_threads // k_threads, k_threads), stride=(k_threads, 1) + ) + value_layout = cute.make_layout((1, copy_elems)) + return cute.make_tiled_copy_tv(atom_copy, thread_layout, value_layout) + + @cute.jit + def _fill_pred(self, pred_flat, coord_tensor, k_tile, dim_limit, K_total): + # Predicate one K tile and keep the stage-broadcast view in sync. + coord_ktile = coord_tensor[None, None, 0, k_tile] + num_vec = pred_flat.shape[0] + num_mn = pred_flat.shape[1] + for v in cutlass.range_constexpr(num_vec): + # pred_flat is (K_VEC, M/N) for one K_TILE. + for j in cutlass.range_constexpr(num_mn): + pred_flat[v, j] = cute.elem_less( + coord_ktile[(0, v), j], (dim_limit, K_total) + ) + + def _make_pred(self, tXcX, k_tile, dim_limit, K_total): + # pred_flat is (K_VEC, M/N); pred is (K_VEC, M/N, STAGE). + num_vec = tXcX.shape[0][1] + num_mn = cute.size(tXcX, mode=[1]) + pred_flat = cute.make_rmem_tensor( + cute.make_layout((num_vec, num_mn), stride=(num_mn, 1)), + cutlass.Boolean, + ) + pred = cute.make_tensor( + pred_flat.iterator, + cute.make_layout( + (num_vec, num_mn, cute.size(tXcX, mode=[2])), + stride=(num_mn, 1, 0), + ), + ) + return pred_flat, pred + + @cute.jit + def __call__( + self, + mA: cute.Tensor, + mB: cute.Tensor, + mC: cute.Tensor, + stream: CUstream, + scale: float = 1.0, + ): + """Execute the split-K GEMM operation in steps: + - Build swizzled staged SMEM layouts for A ``[16, tile_k, stages]`` + and B ``[tile_n, tile_k, stages]``. + - Build 128-bit cp.async GMEM-to-SMEM tiled copies for DMA warps. + - Build an m16n8k16 BF16 MMA tiled across the CTA N tile. + - Launch grid ``ceil(M/16) x ceil(N/tile_n) x split_k`` with one + cluster along split-K, so each cluster rank owns K tiles + ``rank, rank + split_k, ...``. + - Reduce MMA-warp partials within the CTA, exchange split-K partials + through DSMEM, then reduce cluster partials and store C. + + :param mA: Input tensor A with shape ``[M, K]``. + :type mA: cute.Tensor + :param mB: Input tensor B with shape ``[N, K]``. + :type mB: cute.Tensor + :param mC: Output tensor C with shape ``[M, N]``. + :type mC: cute.Tensor + :param stream: CUDA stream for asynchronous execution. + :type stream: CUstream + :param scale: Epilogue scale applied before storing C. + :type scale: float + """ + bM, bN, bK = self.tile_m, self.tile_n, self.tile_k + copy_bits: cutlass.Constexpr = self.copy_bits + sA_layout = self._make_smem_layout_AB( + mA.element_type, copy_bits, (bM, bK, self.num_stages) + ) + sB_layout = self._make_smem_layout_AB( + mB.element_type, copy_bits, (bN, bK, self.num_stages) + ) + + @cute.struct + class SharedStorage: + a: cute.struct.Align[ + cute.struct.MemRange[mA.element_type, cute.cosize(sA_layout)], 16 + ] + b: cute.struct.Align[ + cute.struct.MemRange[mB.element_type, cute.cosize(sB_layout)], 16 + ] + mbar: cute.struct.Align[ + cute.struct.MemRange[cutlass.Int64, self.num_stages * 2], 8 + ] + + atom_g2s = cute.make_copy_atom( + cute.nvgpu.cpasync.CopyG2SOp( + cache_mode=cute.nvgpu.cpasync.LoadCacheMode.GLOBAL + ), + mA.element_type, + num_bits_per_copy=copy_bits, + ) # cp.async GMEM -> SMEM, bypassing L1 + tiled_copy_A = self._make_gmem_tiled_copy( + atom_g2s, mA.element_type, copy_bits, self.num_dma_threads + ) + tiled_copy_B = self._make_gmem_tiled_copy( + atom_g2s, mB.element_type, copy_bits, self.num_dma_threads + ) + op = cute.nvgpu.warp.MmaF16BF16Op(self.ab_dtype, self.acc_dtype, self.mma_shape) + # Repeat the m16n8k16 atom along N to cover the CTA output tile. + perm_mnk = ( + self.atom_layout[0] * self.mma_shape[0], + self.atom_layout[1] * self.mma_shape[1] * (self.tile_n // 8), + self.atom_layout[2] * self.mma_shape[2], + ) + tiled_mma = cute.make_tiled_mma( + op, cute.make_layout(self.atom_layout), permutation_mnk=perm_mnk + ) + tiler_mn = (bM, bN) + grid_m, grid_n = cute.ceil_div(mC.shape, tiler_mn) + self.kernel( + mA, + mB, + mC, + scale, + sA_layout, + sB_layout, + tiled_copy_A, + tiled_copy_B, + tiled_mma, + SharedStorage, + ).launch( + grid=[ + cute.size(grid_m), + cute.size(grid_n), + self.split_k, + ], + block=[self.num_threads, 1, 1], + cluster=[ + 1, + 1, + self.split_k, + ], # split-K CTAs form one cluster + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.kernel + def kernel( + self, + mA, + mB, + mC, + scale: cutlass.Float32, + sA_layout: cute.ComposedLayout, + sB_layout: cute.ComposedLayout, + tiled_copy_A: cute.TiledCopy, + tiled_copy_B: cute.TiledCopy, + tiled_mma: cute.TiledMma, + shared_storage: cutlass.Constexpr, + ): + bM, bN, bK = self.tile_m, self.tile_n, self.tile_k + num_stages = self.num_stages + tidx, _, _ = cute.arch.thread_idx() + bid_m, bid_n, bid_z = cute.arch.block_idx() + warp_idx = cute.arch.warp_idx() + lane_id = cute.arch.lane_idx() + num_dma_warps: cutlass.Constexpr = self.num_dma_warps + is_dma_warp = warp_idx < num_dma_warps + dma_tidx = tidx + mma_tidx = tidx - self.num_dma_threads + mma_warp_idx = warp_idx - num_dma_warps + + cta_tiler = (bM, bN, bK) + coord = (bid_m, bid_n, None) # all K tiles + # CTA-local tiles. + gA = cute.local_tile( + mA, tiler=cta_tiler, coord=coord, proj=(1, None, 1) + ) # skip N + gB = cute.local_tile( + mB, tiler=cta_tiler, coord=coord, proj=(None, 1, 1) + ) # skip M + gC = cute.local_tile( + mC, tiler=cta_tiler, coord=coord, proj=(1, 1, None) + ) # skip K + + mcA = cute.make_identity_tensor(mA.layout.shape) + mcB = cute.make_identity_tensor(mB.layout.shape) + mcC = cute.make_identity_tensor(mC.layout.shape) + # Coordinate modes: cA=(M,K,k_tile), cB=(N,K,k_tile), cC=(M,N). + cA = cute.local_tile(mcA, tiler=cta_tiler, coord=coord, proj=(1, None, 1)) + cB = cute.local_tile(mcB, tiler=cta_tiler, coord=coord, proj=(None, 1, 1)) + cC = cute.local_tile(mcC, tiler=cta_tiler, coord=coord, proj=(1, 1, None)) + + # 128-bit cp.async copies require 16-byte aligned GMEM views. + gA = cute.make_tensor(gA.iterator.align(16), gA.layout) + gB = cute.make_tensor(gB.iterator.align(16), gB.layout) + + smem = cutlass.utils.SmemAllocator() + storage_ptr = smem.allocate(shared_storage.size_in_bytes(), byte_alignment=16) # type: ignore[attr-defined] + storage = shared_storage(storage_ptr) # type: ignore[call-arg] + sA = storage.a.get_tensor(sA_layout) + sB = storage.b.get_tensor(sB_layout) + + # Pipeline cp.async producers into MMA consumers. + producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, self.num_dma_threads + ) + consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, self.num_mma_threads + ) + mainloop_pipeline = pipeline.PipelineCpAsync.create( + barrier_storage=storage.mbar.data_ptr(), + num_stages=num_stages, + producer_group=producer_group, + consumer_group=consumer_group, + ) + + # Round-robin split-K: split z handles tiles z, z+split_k, ... + K_total = cute.size(mA, mode=[1]) + k_tile_count = cute.size(gA, mode=[2]) + k_start = bid_z + num_k_tiles = cute.ceil_div(k_tile_count - k_start, self.split_k) + + if is_dma_warp: + # DMA warps trade registers for copy throughput. + cute.arch.setmaxregister_decrease(40) + thr_A = tiled_copy_A.get_slice(dma_tidx) + thr_B = tiled_copy_B.get_slice(dma_tidx) + tAgA = thr_A.partition_S(gA) + tAsA = thr_A.partition_D(sA) + tBgB = thr_B.partition_S(gB) + tBsB = thr_B.partition_D(sB) + tAcA = thr_A.partition_S(cA) + tBcB = thr_B.partition_S(cB) + + # Build M/K and N/K predicates, broadcast across K-tile copies. + tApA_flat, tApA = self._make_pred(tAcA, k_start, mA.shape[0], K_total) + tBpB_flat, tBpB = self._make_pred(tBcB, k_start, mB.shape[0], K_total) + self._fill_pred(tApA_flat, tAcA, k_start, mA.shape[0], K_total) + self._fill_pred(tBpB_flat, tBcB, k_start, mB.shape[0], K_total) + + producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, num_stages + ) + + # Prime the first pipeline stage. + mainloop_pipeline.producer_acquire(producer_state) + cute.copy( + tiled_copy_B, + tBgB[None, None, None, k_start], + tBsB[None, None, None, producer_state.index], + pred=tBpB, + ) + if const_expr(self.use_pdl): + cute.arch.griddepcontrol_wait() + cute.copy( + tiled_copy_A, + tAgA[None, None, None, k_start], + tAsA[None, None, None, producer_state.index], + pred=tApA, + ) + mainloop_pipeline.producer_commit(producer_state) + producer_state.advance() + + for k_tile in cutlass.range( + k_start + self.split_k, k_tile_count, self.split_k, unroll=1 + ): + self._fill_pred(tApA_flat, tAcA, k_tile, mA.shape[0], K_total) + self._fill_pred(tBpB_flat, tBcB, k_tile, mB.shape[0], K_total) + mainloop_pipeline.producer_acquire(producer_state) + cute.copy( + tiled_copy_A, + tAgA[None, None, None, k_tile], + tAsA[None, None, None, producer_state.index], + pred=tApA, + ) + cute.copy( + tiled_copy_B, + tBgB[None, None, None, k_tile], + tBsB[None, None, None, producer_state.index], + pred=tBpB, + ) + mainloop_pipeline.producer_commit(producer_state) + producer_state.advance() + + mainloop_pipeline.producer_tail(producer_state) + + else: + # MMA warps with k-phase interleaving + cute.arch.setmaxregister_increase(232) # large MMA fragments + + num_mma_warps: cutlass.Constexpr = self.num_mma_warps + + thr_mma = tiled_mma.get_slice(lane_id) + tCsA = thr_mma.partition_A(sA) + tCsB = thr_mma.partition_B(sB) + tCgC = thr_mma.partition_C(gC) + + tCrA = tiled_mma.make_fragment_A(tCsA[None, None, None, 0]) + tCrB = tiled_mma.make_fragment_B(tCsB[None, None, None, 0]) + tCrC = tiled_mma.make_fragment_C(tCgC) + tCrC.fill(0.0) + + # ldmatrix moves SMEM fragments into MMA registers. + atom_s2r_A = cute.make_copy_atom( + cute.nvgpu.warp.LdMatrix8x8x16bOp(False, 4), mA.element_type + ) # non-transposed, x4 + atom_s2r_B = cute.make_copy_atom( + cute.nvgpu.warp.LdMatrix8x8x16bOp(False, 4), mB.element_type + ) + + # SMEM -> register copy path. + tiled_s2r_A = cute.make_tiled_copy_A(atom_s2r_A, tiled_mma) + tiled_s2r_B = cute.make_tiled_copy_B(atom_s2r_B, tiled_mma) + thr_s2r_A = tiled_s2r_A.get_slice(lane_id) + thr_s2r_B = tiled_s2r_B.get_slice(lane_id) + tCsA_v = thr_s2r_A.partition_S(sA) # views, not copies + tCrA_v = thr_s2r_A.retile(tCrA) + tCsB_v = thr_s2r_B.partition_S(sB) + tCrB_v = thr_s2r_B.retile(tCrB) + # Split the MMA K-fragments across the MMA warps. + tCsA_warp_v = cute.logical_divide(tCsA_v, (None, None, num_mma_warps, None)) + tCsB_warp_v = cute.logical_divide(tCsB_v, (None, None, num_mma_warps, None)) + + num_k_blocks = cute.size(tCrA, mode=[2]) + k_blocks_per_warp: cutlass.Constexpr = num_k_blocks // num_mma_warps + + consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, num_stages + ) + + # Shape-dynamic split-K count, so this stays a runtime range. + for _ in cutlass.range(num_k_tiles, unroll_full=True): + mainloop_pipeline.consumer_wait(consumer_state) + for ki in cutlass.range_constexpr(k_blocks_per_warp): + cute.copy( + tiled_s2r_A, + tCsA_warp_v[ + None, None, (mma_warp_idx, ki), consumer_state.index + ], + tCrA_v[None, None, 0], + ) # ldmatrix + cute.copy( + tiled_s2r_B, + tCsB_warp_v[ + None, None, (mma_warp_idx, ki), consumer_state.index + ], + tCrB_v[None, None, 0], + ) + cute.gemm( + tiled_mma, tCrC, tCrA[None, None, 0], tCrB[None, None, 0], tCrC + ) # mma.sync + mainloop_pipeline.consumer_release(consumer_state) + consumer_state.advance() + + # Cluster reduction epilogue. + # Reduce per-warp accumulators within this CTA. + num_elems: cutlass.Constexpr = self.num_epilogue_elems + elems_per_thread: cutlass.Constexpr = self.epilogue_elems_per_thread + # Map MMA threads to linear CTA output elements. + epilogue_thread_layout = cute.make_layout( + (elems_per_thread, self.num_mma_threads), + stride=(self.num_mma_threads, 1), + ) + epilogue_slots = cute.make_tensor(0, epilogue_thread_layout) + epilogue_slot_coords = cute.make_identity_tensor((bN, bM)) + # Layout: (mma_warp, linear MN element). + smem_red = cute.make_tensor( + cute.arch.alloc_smem( + cutlass.Float32, num_elems * num_mma_warps, alignment=16 + ), + cute.make_layout((num_mma_warps, num_elems), stride=(num_elems, 1)), + ) + smem_warp = cute.make_tensor( + cute.domain_offset((mma_warp_idx, 0), smem_red).iterator, + cute.make_layout((bM, bN), stride=(bN, 1)), + ) + tCsC_partial = thr_mma.partition_C(smem_warp) + cute.autovec_copy(tCrC, tCsC_partial) + cute.arch.sync_threads() + + # Layout: (split-K rank, linear MN element). + partials = cute.make_tensor( + cute.arch.alloc_smem( + cutlass.Float32, num_elems * self.split_k, alignment=16 + ), + cute.make_layout((self.split_k, num_elems), stride=(num_elems, 1)), + ) + cta_rank = cute.arch.block_idx_in_cluster() + + for ei in cutlass.range_constexpr(elems_per_thread): + elem_idx = epilogue_slots[ei, mma_tidx] + local_coord = cute.select(epilogue_slot_coords[elem_idx], mode=[1, 0]) + total = cutlass.Float32(0.0) + if cute.elem_less(local_coord, gC.shape): + total = ( + smem_red[None, elem_idx] + .load() + .reduce( + cute.ReductionOp.ADD, + init_val=cutlass.Float32(0.0), + reduction_profile=0, + ) + ) + total = total * scale + partials[cta_rank, elem_idx] = total + + cute.arch.sync_threads() + + # Broadcast this CTA's partials to peer DSMEM. + for ei in cutlass.range_constexpr(elems_per_thread): + elem_idx = epilogue_slots[ei, mma_tidx] + my_slot = cute.domain_offset((cta_rank, elem_idx), partials).iterator + my_val = partials[cta_rank, elem_idx] + for peer in cutlass.range_constexpr(self.split_k): + remote = set_block_rank(my_slot, cutlass.Int32(peer)) + st_shared_remote_f32(remote, my_val) + + cute.arch.cluster_arrive() + cute.arch.cluster_wait() # peer DSMEM stores are now visible + + if const_expr(self.use_pdl) and mma_tidx == 0: + cute.arch.griddepcontrol_launch_dependents() + cute.arch.sync_threads() + + # Reduce split-K partials and write global output. + for ei in cutlass.range_constexpr(elems_per_thread): + elem_idx = epilogue_slots[ei, mma_tidx] + local_coord = cute.select(epilogue_slot_coords[elem_idx], mode=[1, 0]) + global_coord = cC[local_coord] + if cute.elem_less(global_coord, mC.shape): + acc = ( + partials[None, elem_idx] + .load() + .reduce( + cute.ReductionOp.ADD, + init_val=cutlass.Float32(0.0), + reduction_profile=0, + ) + ) + gC[local_coord] = acc + + cute.arch.sync_threads() diff --git a/vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py b/vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py new file mode 100644 index 00000000000..16bf9632965 --- /dev/null +++ b/vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py @@ -0,0 +1,282 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import logging +import math +from collections.abc import Iterable +from dataclasses import dataclass +from typing import Any, Literal + +import torch + +logger = logging.getLogger(__name__) + +_cutedsl_available: bool | None = None + + +def is_available() -> bool: + global _cutedsl_available + if _cutedsl_available is not None: + return _cutedsl_available + try: + import cutlass # noqa: F401 + import cutlass.cute # noqa: F401 + + _cutedsl_available = True + except ImportError: + _cutedsl_available = False + logger.info("cuteDSL (CUTLASS Python) not available, ll_bf16_gemm disabled") + return _cutedsl_available + + +_DEFAULT_DOTPROD_BS = 128 +_DEFAULT_DOTPROD_MAX_M = 4 +_DEFAULT_SPLITK_CONFIG = (6, 4) +_TUNED_DOTPROD_MAX_M: dict[tuple[int, int], int] = { + (7168, 256): 6, +} +_TUNED_CONFIGS: dict[tuple[int, int], dict[int, tuple[int, int]]] = { + (7168, 384): { + 5: (4, 4), + **{M: (5, 4) for M in range(6, 17)}, + }, +} + + +_cute_ctx = None + + +def _cute(): + global _cute_ctx + if _cute_ctx is not None: + return _cute_ctx + import cutlass.cute as cute + from cuda.bindings.driver import CUstream + + _cute_ctx = (cute, CUstream) + return _cute_ctx + + +def _stream(): + _, CUstream = _cute() + from vllm.utils.torch_utils import current_stream + + return CUstream(current_stream().cuda_stream) + + +def _use_pdl() -> bool: + from vllm.platforms import current_platform + + return current_platform.is_arch_support_pdl() + + +class LLBf16Gemm: + @dataclass(frozen=True, slots=True) + class CompileKey: + backend: Literal["dotprod", "splitk"] + M: int = 0 + K: int = 0 + bs: int = 0 + split_k: int = 0 + num_stages: int = 0 + + def __init__(self) -> None: + # Dot-prod: keyed on (M, K, bs), because M and K are Constexpr. + self._compiled_cache: dict[tuple[int, int, int], Any] = {} + # Split-K: keyed on (split_k, num_stages), fully shape-dynamic. + self._splitk_cache: dict[tuple[int, int], Any] = {} + + def dispatch(self, *, M: int, K: int, N: int) -> CompileKey: + dotprod_max_m = _TUNED_DOTPROD_MAX_M.get((K, N), _DEFAULT_DOTPROD_MAX_M) + if dotprod_max_m >= M or K < 2048: + return self.CompileKey(backend="dotprod", M=M, K=K, bs=_DEFAULT_DOTPROD_BS) + + split_k, num_stages = _TUNED_CONFIGS.get((K, N), {}).get( + M, _DEFAULT_SPLITK_CONFIG + ) + return self.CompileKey(backend="splitk", split_k=split_k, num_stages=num_stages) + + def get_warmup_keys( + self, + *, + shapes: Iterable[tuple[int, int]], + m_values: Iterable[int], + ) -> list[CompileKey]: + return list( + dict.fromkeys( + self.dispatch(M=M, K=K, N=N) for K, N in shapes for M in m_values + ) + ) + + @staticmethod + def _fake_gemm_tensors(*, M, K, N, divisibility: int): + from cutlass import BFloat16, Float32 + from quack.compile_utils import make_fake_tensor + + hidden_states = make_fake_tensor(BFloat16, (M, K), divisibility=divisibility) + router_weight = make_fake_tensor(BFloat16, (N, K), divisibility=divisibility) + output = make_fake_tensor(Float32, (M, N), divisibility=1) + return hidden_states, router_weight, output + + def _compile_splitk(self, compile_key: CompileKey) -> None: + cute, _ = _cute() + from ._ll_bf16_splitk import LLBf16SplitK + + hidden_states, router_weight, output = self._fake_gemm_tensors( + M=cute.sym_int(), + K=cute.sym_int(), + N=cute.sym_int(), + divisibility=8, + ) + gemm = LLBf16SplitK( + tile_n=16, + tile_k=256, + num_stages=compile_key.num_stages, + num_dma_warps=4, + split_k=compile_key.split_k, + use_pdl=_use_pdl(), + ) + compiled = cute.compile( + gemm, + hidden_states, + router_weight, + output, + _stream(), + options="--enable-tvm-ffi", + ) + self._splitk_cache[(compile_key.split_k, compile_key.num_stages)] = compiled + logger.debug( + "Compiled ll_bf16_splitk: sk=%d ns=%d", + compile_key.split_k, + compile_key.num_stages, + ) + + def _compile_dotprod(self, compile_key: CompileKey) -> None: + cute, _ = _cute() + from ._ll_bf16_dotprod import LLBf16Dotprod + + N = cute.sym_int() + stride_divisibility = math.gcd(8, compile_key.K) + hidden_states, router_weight, output = self._fake_gemm_tensors( + M=compile_key.M, + K=compile_key.K, + N=N, + divisibility=stride_divisibility, + ) + gemm = LLBf16Dotprod(k=compile_key.K, bs=compile_key.bs, use_pdl=_use_pdl()) + compiled = cute.compile( + gemm, + hidden_states, + router_weight, + output, + compile_key.M, + compile_key.K, + 1, # runtime N placeholder for fake-tensor compile + _stream(), + options="--enable-tvm-ffi --ptxas-options -maxrregcount=64", + ) + self._compiled_cache[(compile_key.M, compile_key.K, compile_key.bs)] = compiled + logger.debug( + "Compiled ll_bf16_dotprod: M=%d, K=%d, bs=%d", + compile_key.M, + compile_key.K, + compile_key.bs, + ) + + def compile(self, compile_key: CompileKey) -> None: + if compile_key.backend == "splitk": + splitk_cache_key = (compile_key.split_k, compile_key.num_stages) + if splitk_cache_key not in self._splitk_cache: + self._compile_splitk(compile_key) + return + + dotprod_cache_key = (compile_key.M, compile_key.K, compile_key.bs) + if dotprod_cache_key not in self._compiled_cache: + self._compile_dotprod(compile_key) + + def warmup( + self, + *, + shapes: Iterable[tuple[int, int]], + m_values: Iterable[int], + ) -> None: + for compile_key in self.get_warmup_keys(shapes=shapes, m_values=m_values): + self.compile(compile_key) + + @staticmethod + def _validate_inputs( + hidden_states: torch.Tensor, + router_weight: torch.Tensor, + output_dtype: torch.dtype, + ) -> None: + if hidden_states.dim() != 2 or router_weight.dim() != 2: + raise ValueError("hidden_states and router_weight must be 2D tensors") + if ( + hidden_states.dtype != torch.bfloat16 + or router_weight.dtype != torch.bfloat16 + ): + raise ValueError("hidden_states and router_weight must have dtype=bfloat16") + if hidden_states.device.type != "cuda" or router_weight.device.type != "cuda": + raise ValueError( + "hidden_states and router_weight must have device_type=cuda" + ) + if hidden_states.device != router_weight.device: + raise ValueError( + "hidden_states and router_weight must be on the same CUDA device" + ) + if output_dtype != torch.float32: + raise ValueError("ll_bf16_gemm only supports output_dtype=torch.float32") + if hidden_states.shape[1] != router_weight.shape[1]: + raise ValueError( + "hidden_states and router_weight must have matching K dimensions" + ) + # Kernels use vectorized bf16 loads and require 16-byte row alignment. + if hidden_states.shape[1] % 8 != 0: + raise ValueError("ll_bf16_gemm requires K to be divisible by 8") + if not hidden_states.is_contiguous() or not router_weight.is_contiguous(): + raise ValueError( + "hidden_states and router_weight must be contiguous row-major inputs" + ) + + def __call__( + self, + hidden_states: torch.Tensor, # [M, K] bf16 + router_weight: torch.Tensor, # [N, K] bf16 + output_dtype: torch.dtype = torch.float32, + ) -> torch.Tensor: # [M, N] fp32 + self._validate_inputs(hidden_states, router_weight, output_dtype) + + M, K = hidden_states.shape + N = router_weight.shape[0] + compile_key = self.dispatch(M=M, K=K, N=N) + if compile_key.backend == "splitk": + splitk_cache_key = (compile_key.split_k, compile_key.num_stages) + if splitk_cache_key not in self._splitk_cache: + self.compile(compile_key) + kernel = self._splitk_cache[splitk_cache_key] + else: + dotprod_cache_key = (compile_key.M, compile_key.K, compile_key.bs) + if dotprod_cache_key not in self._compiled_cache: + self.compile(compile_key) + kernel = self._compiled_cache[dotprod_cache_key] + + stream = _stream() + output = torch.empty(M, N, dtype=output_dtype, device=hidden_states.device) + if compile_key.backend == "splitk": + kernel(hidden_states, router_weight, output, stream, 1.0) + else: + kernel(hidden_states, router_weight, output, N, stream) + return output + + +ll_bf16_gemm_kernel = LLBf16Gemm() + + +def ll_bf16_gemm( + hidden_states: torch.Tensor, + router_weight: torch.Tensor, + output_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + return ll_bf16_gemm_kernel(hidden_states, router_weight, output_dtype) diff --git a/vllm/model_executor/layers/fused_moe/router/gate_linear.py b/vllm/model_executor/layers/fused_moe/router/gate_linear.py index 37e7018a867..e55cd5f4f2a 100644 --- a/vllm/model_executor/layers/fused_moe/router/gate_linear.py +++ b/vllm/model_executor/layers/fused_moe/router/gate_linear.py @@ -14,11 +14,13 @@ from vllm.utils.torch_utils import direct_register_custom_op class GateLinear(ReplicatedLinear): """MoE gate linear layer with multi-tier GEMM dispatch: - 1. DSV3 specialized kernel (SM90+, M<=16, H=7168 E=256/384, H=6144 E=256) - 2. fp32 specialized kernel (SM90+, bf16/fp32 in, fp32 out, M<=32, + 1. cuteDSL ll_bf16_gemm (SM90+, M<=16, bf16 in, fp32 out, + K divisible by 8) + 2. DSV3 specialized kernel (SM90+, M<=16, H=7168 E=256/384, H=6144 E=256) + 3. fp32 specialized kernel (SM90+, bf16/fp32 in, fp32 out, M<=32, (H, E) in {(3072, 256), (6144, 128), (6144, 256)}) - 3. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 weight + fp32 out_dtype) - 4. F.linear via ReplicatedLinear (ultimate fallback) + 4. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 weight + fp32 out_dtype) + 5. F.linear via ReplicatedLinear (ultimate fallback) The ``out_dtype`` attribute is mutable and can be set after init (e.g. when the required dtype depends on the expert quantization @@ -99,6 +101,21 @@ class GateLinear(ReplicatedLinear): and self.out_dtype == torch.float32 ) + # cuteDSL ll_bf16_gemm eligibility. Any dims supported, but SM90+ required bc: + # 1. PDL support. Both dot-product and split-K kernels. + # 2. Thread Block Clusters. Split-K kernel for cross-CTA reduction. + self.allow_ll_bf16_gemm = False + if can_use_specialized_kernels: + from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import ( + is_available, + ) + + self.allow_ll_bf16_gemm = ( + self.weight.dtype == torch.bfloat16 + and self.out_dtype == torch.float32 + and is_available() + ) + def set_out_dtype(self, out_dtype: torch.dtype) -> None: """Set output dtype for the router logits after init. @@ -116,10 +133,31 @@ class GateLinear(ReplicatedLinear): ): self.allow_cublas_router_gemm = self.weight.dtype == torch.bfloat16 + # out_dtype may start as None -> recompute eligibility here + if self.allow_specialized_router_gemm: + from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import ( + is_available, + ) + + self.allow_ll_bf16_gemm = ( + self.weight.dtype == torch.bfloat16 + and out_dtype == torch.float32 + and is_available() + ) + def forward( self, x: torch.Tensor ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]: - # Tier 1: DSV3 specialized kernel + # Tier 1: cuteDSL ll_bf16_gemm (SM90+, any dims) + if self.allow_ll_bf16_gemm and x.shape[0] <= 16 and x.dtype == torch.bfloat16: + from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import ( + ll_bf16_gemm, + ) + + output = ll_bf16_gemm(x, self.weight) + return output, None + + # Tier 2: DSV3 specialized kernel (fallback for when cuteDSL unavailable) if self.allow_dsv3_router_gemm and x.shape[0] <= self._dsv3_max_batch: output = ops.dsv3_router_gemm( hidden_states=x, @@ -128,7 +166,7 @@ class GateLinear(ReplicatedLinear): ) return output, None - # Tier 2: fp32 specialized kernel (H=3072, E=256, M<=32) + # Tier 3: fp32 specialized kernel (H=3072, E=256, M<=32) # Dispatch is wrapped in a custom op so that torch.compile/CUDA-graph # capture does not freeze the runtime num_tokens branch. if self.allow_fp32_router_gemm and x.dtype in ( @@ -138,12 +176,12 @@ class GateLinear(ReplicatedLinear): output = torch.ops.vllm.fp32_router_gemm_dispatch(x, self.weight) return output, None - # Tier 3: cuBLAS bf16→fp32 + # Tier 4: cuBLAS bf16→fp32 if self.allow_cublas_router_gemm and x.dtype == torch.bfloat16: output = torch.mm(x, self.weight.T, out_dtype=torch.float32) return output, None - # Tier 4: F.linear (ReplicatedLinear) + # Tier 5: F.linear (ReplicatedLinear) if self.out_dtype is not None and x.dtype != self.weight.dtype: x = x.to(self.weight.dtype) output, output_bias = super().forward(x) diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index d96a2e594bb..c55ad148b2d 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -42,6 +42,31 @@ if TYPE_CHECKING: logger = init_logger(__name__) +_LL_BF16_WARMUP_MODEL_SHAPES: tuple[tuple[int, int], ...] = ( + (7168, 256), # DSV3 + (7168, 384), # DSV4-Pro + (14400, 256), # DSV4-Flash +) +_LL_BF16_WARMUP_M_RANGE = range(1, 17) + + +def _warmup_ll_bf16_router_gemm() -> None: + from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import ( + is_available as is_ll_bf16_gemm_available, + ) + from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import ( + ll_bf16_gemm_kernel, + ) + + if not is_ll_bf16_gemm_available(): + return + + logger.info("Warming up ll_bf16 router GEMM kernels.") + ll_bf16_gemm_kernel.warmup( + shapes=_LL_BF16_WARMUP_MODEL_SHAPES, + m_values=_LL_BF16_WARMUP_M_RANGE, + ) + def kernel_warmup(worker: "Worker"): from vllm.model_executor.warmup.minimax_m3_msa_warmup import ( @@ -94,6 +119,9 @@ def kernel_warmup(worker: "Worker"): elif has_flashinfer() and current_platform.has_device_capability(90): flashinfer_autotune(worker.model_runner) + if current_platform.has_device_capability(90): + _warmup_ll_bf16_router_gemm() + # FlashInfer attention warmup # Only warmup if the model has FlashInfer attention groups # and is not a pooling model