From 4598aecdfa763846068b314b312b4b1d329356a2 Mon Sep 17 00:00:00 2001 From: jiacao-amd Date: Mon, 19 Oct 2026 20:52:57 -0700 Subject: [PATCH] [ROCm][Perf] Split-K BF16x3 router GEMM for decode batches The ROCm BF16x3 router tier (fp32 gate weights, e.g. MiniMax-M2/M3) only served batches of 2048+ tokens; decode batches between the gfx950 fp32 low-M kernel (M<=32) and 2048 fell back to F.linear on the fp32 pipe. Add a split-K Triton kernel that keeps the three bf16 weight terms in separate accumulators, then sums the K partitions in a fixed order, and route every batch above the fp32 low-M kernel and below 2048 tokens to it. On MI355X it is 1.6-2.1x faster than F.linear and 1.1-1.8x faster than the stacked hipBLASLt GEMM over that range at the MiniMax-M2/M3 router shapes, and more accurate than F.linear (rel. L2 vs fp64 1.1e-7 to 2.4e-7, vs 2e-7 to 1.4e-6). Co-Authored-By: Claude Signed-off-by: jiacao-amd --- tests/kernels/test_bf16x3_router_gemm_rocm.py | 88 ++++++++-- .../kernels/test_gate_linear_rocm_dispatch.py | 27 ++++ .../router/bf16x3_router_gemm_rocm.py | 151 ++++++++++++++++-- .../layers/fused_moe/router/gate_linear.py | 12 +- 4 files changed, 248 insertions(+), 30 deletions(-) diff --git a/tests/kernels/test_bf16x3_router_gemm_rocm.py b/tests/kernels/test_bf16x3_router_gemm_rocm.py index 69aa195b55f3..3dfed8db07f2 100644 --- a/tests/kernels/test_bf16x3_router_gemm_rocm.py +++ b/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 diff --git a/tests/kernels/test_gate_linear_rocm_dispatch.py b/tests/kernels/test_gate_linear_rocm_dispatch.py index 22fc355415f2..55c057c2a12c 100644 --- a/tests/kernels/test_gate_linear_rocm_dispatch.py +++ b/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) diff --git a/vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_rocm.py b/vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_rocm.py index 8d1fa84c8622..e735669b3324 100644 --- a/vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_rocm.py +++ b/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. diff --git a/vllm/model_executor/layers/fused_moe/router/gate_linear.py b/vllm/model_executor/layers/fused_moe/router/gate_linear.py index 3237783eceec..d6c1a646c6f1 100644 --- a/vllm/model_executor/layers/fused_moe/router/gate_linear.py +++ b/vllm/model_executor/layers/fused_moe/router/gate_linear.py @@ -30,7 +30,7 @@ class GateLinear(ReplicatedLinear): any shape) 4. bf16x3 CuteDSL kernel (SM100, bf16 in, fp32 weight) 5. ROCm bf16x3 router GEMM (gfx950, bf16 in, fp32 weight, fp32 out, - M>=2048) + M>32; split-K Triton below M=2048, hipBLASLt from there) 6. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 weight + fp32 out_dtype); gfx950 routes low M through the tier 2 kernel 7. F.linear via ReplicatedLinear (ultimate fallback) @@ -411,7 +411,7 @@ def rocm_bf16x3_router_gemm_dispatch_impl( weight: torch.Tensor, weight_split: torch.Tensor, ) -> torch.Tensor: - """Run the ROCm bf16x3 router GEMM, or fall back for small batches. + """Run the ROCm bf16x3 router GEMM, or the fp32 kernel for tiny batches. This must be a custom op because our torch.compile integration does not support runtime dispatching on num_tokens: models carrying fp32 router @@ -423,10 +423,8 @@ def rocm_bf16x3_router_gemm_dispatch_impl( bf16x3_router_gemm_rocm, ) - if bf16x3_router_gemm_rocm.is_supported(x, weight): - return bf16x3_router_gemm_rocm.bf16x3_router_gemm(x, weight_split) - # Below MIN_TOKENS, hand back to the low-M gfx950 kernel this tier is - # checked ahead of. + # The low-M gfx950 kernel this tier is checked ahead of is faster for the + # smallest batches. if x.shape[0] <= _FP32_ROUTER_GEMM_MAX_TOKENS: from vllm.model_executor.layers.fused_moe.router.rocm_fp32_router_gemm import ( # noqa: E501 can_use_rocm_fp32_router_gemm, @@ -436,6 +434,8 @@ def rocm_bf16x3_router_gemm_dispatch_impl( x = x.contiguous() if can_use_rocm_fp32_router_gemm(x, weight): return rocm_fp32_router_gemm(x, weight) + if bf16x3_router_gemm_rocm.is_supported(x, weight): + return bf16x3_router_gemm_rocm.bf16x3_router_gemm(x, weight_split) return torch.nn.functional.linear(x.float(), weight)