[WIP][Kernel] Generalized LL GEMMs with PDL

benchmarks/kernels/bench_ll_router_gemm.py

@@ -0,0 +1,105 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import argparse
+import os
+
+import torch
+
+from vllm import _custom_ops as ops
+from vllm.model_executor.layers.fused_moe.router.ll_router_gemm import (
+    ll_router_gemm,
+)
+from vllm.triton_utils import triton
+
+_HAS_DSV3 = hasattr(ops, "dsv3_router_gemm")
+
+_providers = ["ll-router-bf16", "ll-router-fp8", "cublas-bf16"]
+if _HAS_DSV3:
+    _providers.append("dsv3-trtllm")
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["M"],
+        x_vals=[1, 2, 4, 8, 16],
+        x_log=False,
+        line_arg="provider",
+        line_vals=_providers,
+        line_names=_providers,
+        ylabel="Latency (us, lower is better)",
+        plot_name="LL Router GEMM",
+        args={},
+    )
+)
+def benchmark(M, provider, N, K):
+    device = "cuda"
+    quantiles = [0.5, 0.2, 0.8]
+
+    if provider == "ll-router-bf16":
+        a = torch.randn(M, K, dtype=torch.bfloat16, device=device)
+        b = torch.randn(N, K, dtype=torch.bfloat16, device=device)
+        ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
+            lambda: ll_router_gemm(a, b), quantiles=quantiles
+        )
+
+    elif provider == "ll-router-fp8":
+        a = torch.randn(M, K, device=device).to(torch.float8_e4m3fn)
+        b = torch.randn(N, K, device=device).to(torch.float8_e4m3fn)
+        ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
+            lambda: ll_router_gemm(a, b), quantiles=quantiles
+        )
+
+    elif provider == "cublas-bf16":
+        a = torch.randn(M, K, dtype=torch.bfloat16, device=device)
+        b = torch.randn(N, K, dtype=torch.bfloat16, device=device)
+        out = torch.empty(M, N, dtype=torch.bfloat16, device=device)
+        ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
+            lambda: torch.mm(a, b.T, out=out), quantiles=quantiles
+        )
+
+    elif provider == "dsv3-trtllm":
+        # DSV3 only supports N∈{256,384}, K=7168
+        if N not in (256, 384) or K != 7168:
+            return float("nan"), float("nan"), float("nan")
+        from vllm import _custom_ops as ops
+
+        a = torch.randn(M, K, dtype=torch.bfloat16, device=device)
+        b = torch.randn(N, K, dtype=torch.bfloat16, device=device)
+        ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
+            lambda: ops.dsv3_router_gemm(a, b, torch.float32), quantiles=quantiles
+        )
+
+    # Return latency in us
+    return ms * 1000, min_ms * 1000, max_ms * 1000
+
+
+SHAPES = [
+    (256, 7168, "DSV3 router"),
+    (256, 2048, "Small K"),
+    (128, 5120, "DeepSeek V2"),
+    (8, 4096, "Mixtral-8x7B"),
+    (64, 2880, "Non-aligned K"),
+]
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--save-path", type=str, default=None)
+    args = parser.parse_args()
+
+    print(f"Device: {torch.cuda.get_device_name()}")
+    print()
+
+    for N, K, desc in SHAPES:
+        print(f"{desc}, N={N} K={K}:")
+        save_dir = args.save_path or f"bench_ll_router_n{N}_k{K}"
+        os.makedirs(save_dir, exist_ok=True)
+        benchmark.run(
+            print_data=True,
+            show_plots=False,
+            save_path=save_dir,
+            N=N,
+            K=K,
+        )
+        print()

tests/kernels/test_ll_router_gemm.py

@@ -0,0 +1,144 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+import pytest
+import torch
+
+from vllm.model_executor.layers.fused_moe.router.ll_router_gemm import (
+    is_available,
+    ll_router_gemm,
+)
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+
+@pytest.fixture(autouse=True)
+def _check_cutedsl():
+    if not is_available():
+        pytest.skip("cuteDSL (CUTLASS Python) not installed")
+
+
+# --- Test shapes ---
+
+# (N, K, description)
+SHAPES = [
+    (256, 7168, "DSV3 router"),
+    (256, 2048, "small K"),
+    (128, 5120, "DeepSeek V2"),
+    (8, 4096, "Mixtral-8x7B"),
+    (64, 2880, "non-aligned K"),
+]
+
+M_VALUES = [1, 2, 4, 8, 16]
+
+DTYPES_BF16 = [(torch.bfloat16, "bf16")]
+DTYPES_FP8 = [(torch.float8_e4m3fn, "fp8")]
+DTYPES_ALL = DTYPES_BF16 + DTYPES_FP8
+
+
+def _make_inputs(M, N, K, dtype):
+    if dtype == torch.bfloat16:
+        a = torch.randn(M, K, dtype=torch.bfloat16, device="cuda")
+        b = torch.randn(N, K, dtype=torch.bfloat16, device="cuda")
+    else:
+        a = torch.randn(M, K, device="cuda").to(dtype)
+        b = torch.randn(N, K, device="cuda").to(dtype)
+    return a, b
+
+
+def _reference(a, b):
+    """torch.mm fp32 reference."""
+    return torch.mm(a.float(), b.float().T)
+
+
+# --- Correctness tests ---
+
+
+@pytest.mark.parametrize("M", M_VALUES)
+@pytest.mark.parametrize("N,K,desc", SHAPES, ids=[s[2] for s in SHAPES])
+def test_bf16_correctness(M, N, K, desc):
+    a, b = _make_inputs(M, N, K, torch.bfloat16)
+    out = ll_router_gemm(a, b)
+    ref = _reference(a, b)
+    torch.testing.assert_close(out, ref, atol=1e-3, rtol=1e-3)
+
+
+@pytest.mark.parametrize("M", M_VALUES)
+@pytest.mark.parametrize(
+    "N,K,desc",
+    [(n, k, d) for n, k, d in SHAPES if k % 2 == 0],
+    ids=[s[2] for s in SHAPES if s[1] % 2 == 0],
+)
+def test_fp8_correctness(M, N, K, desc):
+    a, b = _make_inputs(M, N, K, torch.float8_e4m3fn)
+    out = ll_router_gemm(a, b)
+    ref = _reference(a, b)
+    torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2)
+
+
+@pytest.mark.parametrize("M", [1, 4, 16])
+def test_output_dtype_float32(M):
+    N, K = 64, 2048
+    a, b = _make_inputs(M, N, K, torch.bfloat16)
+    out = ll_router_gemm(a, b, output_dtype=torch.float32)
+    assert out.dtype == torch.float32
+    assert out.shape == (M, N)
+
+
+@pytest.mark.parametrize("M", [1, 16])
+def test_deterministic(M):
+    """Same inputs should produce identical outputs."""
+    N, K = 128, 4096
+    a, b = _make_inputs(M, N, K, torch.bfloat16)
+    out1 = ll_router_gemm(a, b)
+    out2 = ll_router_gemm(a, b)
+    torch.testing.assert_close(out1, out2, atol=0, rtol=0)
+
+
+@pytest.mark.parametrize("N", [1, 3, 7, 17, 64, 256])
+def test_arbitrary_N(N):
+    """N can be any positive integer."""
+    M, K = 4, 2048
+    a, b = _make_inputs(M, N, K, torch.bfloat16)
+    out = ll_router_gemm(a, b)
+    ref = _reference(a, b)
+    torch.testing.assert_close(out, ref, atol=1e-3, rtol=1e-3)
+
+
+@pytest.mark.parametrize("K", [128, 256, 512, 1024, 2048, 4096, 7168])
+def test_arbitrary_K(K):
+    """K can be any positive integer (multiple of 2 for fp8)."""
+    M, N = 4, 32
+    a, b = _make_inputs(M, N, K, torch.bfloat16)
+    out = ll_router_gemm(a, b)
+    ref = _reference(a, b)
+    torch.testing.assert_close(out, ref, atol=1e-3, rtol=1e-3)
+
+
+def test_small_K():
+    """Very small K that fits entirely in the scalar tail."""
+    M, N, K = 2, 8, 64
+    a, b = _make_inputs(M, N, K, torch.bfloat16)
+    out = ll_router_gemm(a, b)
+    ref = _reference(a, b)
+    torch.testing.assert_close(out, ref, atol=1e-3, rtol=1e-3)
+
+
+def test_cudagraph_compatible():
+    """Kernel must work inside a CUDA graph."""
+    M, N, K = 4, 64, 2048
+    a, b = _make_inputs(M, N, K, torch.bfloat16)
+
+    # Warm up compilation
+    ll_router_gemm(a, b)
+    torch.accelerator.synchronize()
+
+    g = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(g):
+        out = ll_router_gemm(a, b)
+    g.replay()
+    torch.accelerator.synchronize()
+
+    ref = _reference(a, b)
+    torch.testing.assert_close(out, ref, atol=1e-3, rtol=1e-3)