[Perf][Feat] Add generic cuteDSL LL BF16 router (GEMM) (#42562)

Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Signed-off-by: Roberto L. Castro <38211239+LopezCastroRoberto@users.noreply.github.com>
Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com>

.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
 

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"])

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)

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()

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)

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)

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