[ROCm][Perf] Split-K BF16x3 router GEMM for decode batches

tests/kernels/test_bf16x3_router_gemm_rocm.py

@@ -86,7 +86,9 @@ def test_is_supported_rejects_ineligible_inputs():
     w = torch.zeros(128, 6144, dtype=torch.float32, device=meta)
     assert rocm_bf16x3.is_supported(x, w)
 
-    assert not rocm_bf16x3.is_supported(x[: rocm_bf16x3.MIN_TOKENS - 1], w)
+    # The split-K kernel serves every batch below MM_MIN_TOKENS.
+    assert rocm_bf16x3.is_supported(x[:1], w)
+    assert not rocm_bf16x3.is_supported(x[:0], w)
     assert not rocm_bf16x3.is_supported(x.float(), w)
     assert not rocm_bf16x3.is_supported(x, w.bfloat16())
     assert not rocm_bf16x3.is_supported(x, w[:, :128])
@@ -119,7 +121,12 @@ def test_is_supported_accepts_odd_expert_counts():
 
 
 @pytest.mark.parametrize(("hidden_size", "num_experts"), ROUTER_SHAPES)
-@pytest.mark.parametrize("num_tokens", [2048, 8192, 16385])
+@pytest.mark.parametrize(
+    "num_tokens",
+    # Split-K below MM_MIN_TOKENS, including partial BLOCK_M tiles; the
+    # stacked GEMM from there on.
+    [1, 33, 96, 257, 1000, 2047, 2048, 8192, 16385],
+)
 def test_matches_fp64_reference(hidden_size: int, num_experts: int, num_tokens: int):
     _requires_gfx950()
     x, w = _inputs(num_tokens, hidden_size, num_experts)
@@ -131,37 +138,91 @@ def test_matches_fp64_reference(hidden_size: int, num_experts: int, num_tokens:
     assert (out.double() - ref).norm() / ref.norm() < _MAX_REL_L2
 
 
+@pytest.mark.parametrize("num_tokens", [96, 8192])
 @pytest.mark.parametrize("num_experts", [17, 384])
-def test_matches_fp64_for_non_power_of_two_experts(num_experts: int):
-    """Exercises the reduction's N mask, which the model shapes never hit."""
+def test_matches_fp64_for_non_power_of_two_experts(num_tokens: int, num_experts: int):
+    """Exercises the N masks, which the model shapes never hit."""
     _requires_gfx950()
-    x, w = _inputs(8192, 6144, num_experts)
+    x, w = _inputs(num_tokens, 6144, num_experts)
     out = rocm_bf16x3.bf16x3_router_gemm(x, rocm_bf16x3.split_bf16x3(w))
     ref = x.double() @ w.double().T
 
-    assert out.shape == (8192, num_experts)
+    assert out.shape == (num_tokens, num_experts)
     assert (out.double() - ref).norm() / ref.norm() < _MAX_REL_L2
 
 
+@pytest.mark.parametrize("hidden_size", [1000, 2880])
+def test_split_k_matches_fp64_for_ragged_hidden_size(hidden_size: int):
+    """K that no split or BLOCK_K divides exercises the K mask."""
+    _requires_gfx950()
+    x, w = _inputs(96, hidden_size, 128)
+    out = rocm_bf16x3.bf16x3_router_gemm(x, rocm_bf16x3.split_bf16x3(w))
+    ref = x.double() @ w.double().T
+    assert (out.double() - ref).norm() / ref.norm() < _MAX_REL_L2
+
+
+def test_split_k_honors_activation_row_stride():
+    """is_supported admits activations with padded rows; the kernel must read
+    them through their stride rather than assume contiguity."""
+    _requires_gfx950()
+    x, w = _inputs(96, 6144 + 64, 128)
+    x = x[:, :6144]
+    w = w[:, :6144].contiguous()
+    assert rocm_bf16x3.is_supported(x, w)
+    out = rocm_bf16x3.bf16x3_router_gemm(x, rocm_bf16x3.split_bf16x3(w))
+    expected = rocm_bf16x3.bf16x3_router_gemm(
+        x.contiguous(), rocm_bf16x3.split_bf16x3(w)
+    )
+    assert torch.equal(out, expected)
+
+
+@pytest.mark.parametrize("num_tokens", [33, 96, 1000])
+def test_split_k_is_deterministic(num_tokens: int):
+    """The K partitions are summed in a fixed order, not with atomics."""
+    _requires_gfx950()
+    x, w = _inputs(num_tokens, 6144, 128)
+    split = rocm_bf16x3.split_bf16x3(w)
+    first = rocm_bf16x3.bf16x3_router_gemm(x, split)
+    for _ in range(3):
+        assert torch.equal(rocm_bf16x3.bf16x3_router_gemm(x, split), first)
+
+
 @pytest.mark.parametrize(("hidden_size", "num_experts"), ROUTER_SHAPES)
-def test_more_accurate_than_fp32_fallback(hidden_size: int, num_experts: int):
+@pytest.mark.parametrize("num_tokens", list(range(1, 65)) + [96, 192, 1024, 2047])
+def test_split_k_config_bounds(hidden_size: int, num_experts: int, num_tokens: int):
+    block_m, block_n, split_k, _ = rocm_bf16x3._split_k_config(
+        num_tokens, num_experts, hidden_size
+    )
+    tiles = -(-num_tokens // block_m) * -(-num_experts // block_n)
+    assert split_k >= 1 and split_k & (split_k - 1) == 0
+    assert split_k == 1 or split_k * tiles <= 768
+    assert split_k == 1 or hidden_size // split_k >= 192
+
+
+@pytest.mark.parametrize(("hidden_size", "num_experts"), ROUTER_SHAPES)
+@pytest.mark.parametrize("num_tokens", [96, 1536, 16384])
+def test_more_accurate_than_fp32_fallback(
+    hidden_size: int, num_experts: int, num_tokens: int
+):
     """The claim is not just accuracy but beating the fp32 GEMM it replaces."""
     _requires_gfx950()
-    x, w = _inputs(16384, hidden_size, num_experts)
+    x, w = _inputs(num_tokens, hidden_size, num_experts)
     ref = x.double() @ w.double().T
     out = rocm_bf16x3.bf16x3_router_gemm(x, rocm_bf16x3.split_bf16x3(w))
     assert (out.double() - ref).norm() / ref.norm() < _FP32_FALLBACK_BEST
 
 
 @pytest.mark.parametrize(("hidden_size", "num_experts"), ROUTER_SHAPES)
-def test_expert_selection_matches_fp64(hidden_size: int, num_experts: int):
+@pytest.mark.parametrize("num_tokens", [2047, 16384])
+def test_expert_selection_matches_fp64(
+    hidden_size: int, num_experts: int, num_tokens: int
+):
     """Relative L2 is a proxy; the expert set is what reaches the model.
 
     Measures zero misroutes on both shapes, where the fp32 fallback misroutes a
     nonzero number at some token counts.
     """
     _requires_gfx950()
-    num_tokens = 16384
     x, w = _inputs(num_tokens, hidden_size, num_experts)
     ref_set = _top_set(x.double() @ w.double().T)
     out = rocm_bf16x3.bf16x3_router_gemm(x, rocm_bf16x3.split_bf16x3(w))
@@ -170,11 +231,12 @@ def test_expert_selection_matches_fp64(hidden_size: int, num_experts: int):
     assert misrouted / num_tokens < 1e-4
 
 
-def test_writes_into_supplied_out():
+@pytest.mark.parametrize("num_tokens", [96, 8192])
+def test_writes_into_supplied_out(num_tokens: int):
     _requires_gfx950()
-    x, w = _inputs(8192, 6144, 128)
+    x, w = _inputs(num_tokens, 6144, 128)
     split = rocm_bf16x3.split_bf16x3(w)
-    out = torch.empty(8192, 128, dtype=torch.float32, device="cuda")
+    out = torch.empty(num_tokens, 128, dtype=torch.float32, device="cuda")
 
     returned = rocm_bf16x3.bf16x3_router_gemm(x, split, out=out)
     assert returned is out

tests/kernels/test_gate_linear_rocm_dispatch.py

@@ -257,3 +257,30 @@ def test_rocm_bf16x3_is_reachable_for_low_m_shapes(monkeypatch):
     assert source.index("allow_rocm_bf16x3_router_gemm") < source.index(
         "allow_fp32_router_gemm"
     )
+
+
+@pytest.mark.parametrize("num_tokens", [1, 32, 33, 96, 2047, 2048])
+def test_rocm_bf16x3_dispatch_routes_by_num_tokens(num_tokens):
+    """The fp32 low-M kernel keeps the smallest batches; everything above it
+    takes the bf16x3 path, split-K below MM_MIN_TOKENS and hipBLASLt from
+    there."""
+    from vllm.model_executor.layers.fused_moe.router import (
+        bf16x3_router_gemm_rocm as rocm_bf16x3,
+    )
+    from vllm.model_executor.layers.fused_moe.router.rocm_fp32_router_gemm import (
+        rocm_fp32_router_gemm,
+    )
+
+    if not rocm_bf16x3.platform_supported():
+        pytest.skip("ROCm bf16x3 router GEMM requires gfx950")
+    torch.manual_seed(0)
+    x = torch.randn(num_tokens, 6144, device="cuda").to(torch.bfloat16)
+    w = torch.randn(128, 6144, device="cuda") * 0.053
+    split = rocm_bf16x3.split_bf16x3(w)
+
+    out = torch.ops.vllm.rocm_bf16x3_router_gemm_dispatch(x, w, split)
+    if num_tokens <= GateLinear.FP32_MAX_TOKENS:
+        expected = rocm_fp32_router_gemm(x, w)
+    else:
+        expected = rocm_bf16x3.bf16x3_router_gemm(x, split)
+    assert torch.equal(out, expected)

vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_rocm.py

@@ -9,9 +9,12 @@
 accurate than the path it replaces, which rounds the activations on the way to
 fp32 while this one leaves them untouched.
 
-The terms are stacked along the expert dimension into one ``(TERMS * E, K)``
-weight, so the router costs a single hipBLASLt GEMM plus a reduction over the
-``TERMS`` column blocks of its output.
+Large batches stack the terms along the expert dimension into one
+``(TERMS * E, K)`` weight, so the router costs a single hipBLASLt GEMM plus a
+reduction over the ``TERMS`` column blocks of its output. That pair has a
+launch and allocation floor that decode batches cannot amortize, so smaller
+batches instead run a split-K Triton kernel that accumulates the three terms
+in registers and sums the K partitions in a fixed order.
 """
 
 import torch
@@ -21,7 +24,7 @@
 
 __all__ = [
     "BF16X3_TERMS",
-    "MIN_TOKENS",
+    "MM_MIN_TOKENS",
     "bf16x3_router_gemm",
     "is_supported",
     "platform_supported",
@@ -30,11 +33,12 @@
 
 BF16X3_TERMS = 3
 
-# Below this the fp32 fallback is faster: the GEMM plus reduction has a ~32us
-# floor from launch and allocation that small batches cannot amortize.
-# Measured on gfx950; the crossover is 1536 at (K=6144, E=128) and (K=3072,
-# E=256), and this leaves margin for untested shapes.
-MIN_TOKENS = 2048
+# From this many tokens on the stacked hipBLASLt GEMM serves the router; below
+# it the split-K kernel does. The GEMM plus reduction has a ~10us floor from
+# launch and allocation that small batches cannot amortize, where the split-K
+# kernel is 1.1-1.8x faster than the GEMM and 1.6-2.1x faster than the fp32
+# F.linear fallback on gfx950 at (K=6144, E=128) and (K=3072, E=256).
+MM_MIN_TOKENS = 2048
 
 _MAX_INT32 = (1 << 31) - 1
 
@@ -76,13 +80,17 @@ def platform_supported() -> bool:
 
 
 def is_supported(x: torch.Tensor, weight: torch.Tensor) -> bool:
-    """Whether this path can serve ``x @ weight.T`` -> fp32 profitably."""
+    """Whether this path can serve ``x @ weight.T`` -> fp32 profitably.
+
+    Batches small enough for the gfx950 fp32 low-M kernel are faster there;
+    the caller is expected to try that kernel first.
+    """
     if x.dim() != 2 or weight.dim() != 2:
         return False
     n, k = weight.shape
     m = x.shape[0]
     return (
-        m >= MIN_TOKENS
+        m > 0
         and x.dtype == torch.bfloat16
         and weight.dtype == torch.float32
         and x.shape[1] == k
@@ -93,7 +101,126 @@ def is_supported(x: torch.Tensor, weight: torch.Tensor) -> bool:
         # a hidden copy that gives back the win over the fallback.
         and x.stride(1) == 1
         and 0 < m * BF16X3_TERMS * n <= _MAX_INT32
+        # The split-K kernel indexes both operands with int32 offsets.
+        and BF16X3_TERMS * n * k <= _MAX_INT32
+        and m * x.stride(0) <= _MAX_INT32
+    )
+
+
+@triton.jit
+def _split_k_partials_kernel(
+    x_ptr,
+    w_ptr,
+    partials_ptr,
+    M,
+    N,
+    K,
+    stride_xm,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    K_PER_SPLIT: tl.constexpr,
+):
+    pid_m = tl.program_id(0)
+    pid_n = tl.program_id(1)
+    split = tl.program_id(2)
+    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    offs_k = tl.arange(0, BLOCK_K)
+    mask_m = offs_m < M
+    mask_n = offs_n < N
+    # One accumulator per term: the small terms are summed among themselves
+    # before they meet the large one.
+    acc0 = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    acc1 = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    acc2 = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    for k_start in range(0, K_PER_SPLIT, BLOCK_K):
+        k = split * K_PER_SPLIT + k_start + offs_k
+        mask_k = k < K
+        x = tl.load(
+            x_ptr + offs_m[:, None] * stride_xm + k[None, :],
+            mask=mask_m[:, None] & mask_k[None, :],
+            other=0.0,
+        )
+        w_offs = offs_n[None, :] * K + k[:, None]
+        w_mask = mask_n[None, :] & mask_k[:, None]
+        w0 = tl.load(w_ptr + w_offs, mask=w_mask, other=0.0)
+        w1 = tl.load(w_ptr + N * K + w_offs, mask=w_mask, other=0.0)
+        w2 = tl.load(w_ptr + 2 * N * K + w_offs, mask=w_mask, other=0.0)
+        acc0 = tl.dot(x, w0, acc0)
+        acc1 = tl.dot(x, w1, acc1)
+        acc2 = tl.dot(x, w2, acc2)
+    tl.store(
+        partials_ptr + (split * M + offs_m[:, None]) * N + offs_n[None, :],
+        acc0 + (acc1 + acc2),
+        mask=mask_m[:, None] & mask_n[None, :],
+    )
+
+
+@triton.jit
+def _split_k_reduce_kernel(
+    partials_ptr,
+    out_ptr,
+    numel,
+    SPLIT_K: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < numel
+    acc = tl.zeros((BLOCK,), dtype=tl.float32)
+    # Fixed split order, so the result does not depend on scheduling.
+    for split in tl.static_range(SPLIT_K):
+        acc += tl.load(partials_ptr + split * numel + offs, mask=mask, other=0.0)
+    tl.store(out_ptr + offs, acc, mask=mask)
+
+
+def _split_k_config(m: int, n: int, k: int) -> tuple[int, int, int, int]:
+    """Return (BLOCK_M, BLOCK_N, SPLIT_K, num_warps), swept on gfx950.
+
+    Split K while the grid stays within 768 programs and each split keeps at
+    least 192 of K. Over M in [64, 2047] and six router shapes this lands
+    within 3% of the best split on average and 16% in the worst case.
+    """
+    block_m = 16 if m <= 16 else 32 if m <= 64 else 64
+    block_n = 32
+    tiles = triton.cdiv(m, block_m) * triton.cdiv(n, block_n)
+    split_k = 1
+    while split_k * 2 * tiles <= 768 and k // (split_k * 2) >= 192:
+        split_k *= 2
+    return block_m, block_n, split_k, 4
+
+
+def _bf16x3_router_gemm_split_k(
+    x: torch.Tensor, weight_split: torch.Tensor, out: torch.Tensor
+) -> torch.Tensor:
+    m, k = x.shape
+    n = weight_split.shape[1]
+    block_m, block_n, split_k, num_warps = _split_k_config(m, n, k)
+    block_k = 64
+    k_per_split = triton.cdiv(triton.cdiv(k, split_k), block_k) * block_k
+    partials = torch.empty((split_k, m, n), dtype=torch.float32, device=x.device)
+    grid = (triton.cdiv(m, block_m), triton.cdiv(n, block_n), split_k)
+    _split_k_partials_kernel[grid](
+        x,
+        weight_split,
+        partials,
+        m,
+        n,
+        k,
+        x.stride(0),
+        BLOCK_M=block_m,
+        BLOCK_N=block_n,
+        BLOCK_K=block_k,
+        K_PER_SPLIT=k_per_split,
+        num_warps=num_warps,
+        num_stages=2,
+        matrix_instr_nonkdim=16,
+    )
+    block = 1024
+    _split_k_reduce_kernel[(triton.cdiv(m * n, block),)](
+        partials, out, m * n, SPLIT_K=split_k, BLOCK=block, num_warps=4
     )
+    return out
 
 
 @triton.jit
@@ -139,6 +266,8 @@ def bf16x3_router_gemm(
         out = torch.empty((m, n), dtype=torch.float32, device=x.device)
     else:
         assert out.shape == (m, n) and out.is_contiguous(), "bad out layout"
+    if m < MM_MIN_TOKENS:
+        return _bf16x3_router_gemm_split_k(x, weight_split, out)
 
     # Rows [t*E, (t+1)*E) of the (TERMS * E, K) view hold term t, so the GEMM
     # emits one column block per term and the reduction sums them.