From 182247257fb030fe0b41d342ae9e2a4486ae787e Mon Sep 17 00:00:00 2001 From: akii96 Date: Mon, 5 Oct 2026 20:28:07 +0300 Subject: [PATCH] [ROCm][Perf] Use a stacked hipBLASLt GEMM for the bf16x3 router (#52668) Signed-off-by: Aakif Nawaz --- .buildkite/test_areas/kernels.yaml | 2 + tests/kernels/test_bf16x3_router_gemm_rocm.py | 201 ++++++++++++++++++ .../kernels/test_gate_linear_rocm_dispatch.py | 90 ++++++++ .../router/bf16x3_router_gemm_rocm.py | 163 ++++++++++++++ .../layers/fused_moe/router/gate_linear.py | 142 ++++++++++++- 5 files changed, 594 insertions(+), 4 deletions(-) create mode 100644 tests/kernels/test_bf16x3_router_gemm_rocm.py create mode 100644 vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_rocm.py diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 515c3966f5..54e7215967 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -255,6 +255,7 @@ steps: --ignore=kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py --ignore=kernels/test_top_k_per_row.py --ignore=kernels/test_bf16x3_router_gemm_cutedsl.py + --ignore=kernels/test_bf16x3_router_gemm_rocm.py --ignore=kernels/test_ll_bf16_gemm.py --ignore=kernels/test_shuffle_rows.py # BROKEN on main, pending kernel fixes (B200): @@ -280,6 +281,7 @@ steps: kernels/test_apply_repetition_penalties.py kernels/test_awq_int4_to_int8.py kernels/test_bf16_skinny_gemm.py + kernels/test_bf16x3_router_gemm_rocm.py kernels/test_cache_kernels.py kernels/test_compressor_kv_cache.py kernels/test_cp_gather_fp8.py diff --git a/tests/kernels/test_bf16x3_router_gemm_rocm.py b/tests/kernels/test_bf16x3_router_gemm_rocm.py new file mode 100644 index 0000000000..69aa195b55 --- /dev/null +++ b/tests/kernels/test_bf16x3_router_gemm_rocm.py @@ -0,0 +1,201 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the ROCm BF16x3 router GEMM. + +The weight split and the eligibility policy are pure functions, so most of +this runs device-free. Only the GEMM correctness tests need a gfx950 runner. +""" + +import pytest +import torch + +from vllm.model_executor.layers.fused_moe.router import ( + bf16x3_router_gemm_rocm as rocm_bf16x3, +) +from vllm.platforms import current_platform + +pytestmark = pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm only") + +# (hidden_size, num_experts) for the models that ship fp32 router weights. +# (6144, 128) -> MiniMax-M3, (3072, 256) -> MiniMax-M2/M2.5 +ROUTER_SHAPES = [(6144, 128), (3072, 256)] +ROUTER_WEIGHT_SCALE = 0.053 + + +def _requires_gfx950(): + if not rocm_bf16x3.platform_supported(): + pytest.skip("ROCm bf16x3 router GEMM requires gfx950") + + +def _inputs(num_tokens, hidden_size, num_experts, seed=42, device="cuda"): + torch.manual_seed(seed) + x = torch.randn(num_tokens, hidden_size, device=device).to(torch.bfloat16) + w = torch.randn(num_experts, hidden_size, dtype=torch.float32, device=device) + return x, w * ROUTER_WEIGHT_SCALE + + +def _top_set(logits, k=4): + """Sorted top-k indices; topk's order among near-ties is not meaningful.""" + return logits.topk(k, dim=-1).indices.sort(dim=-1).values + + +@pytest.mark.parametrize(("hidden_size", "num_experts"), ROUTER_SHAPES) +def test_split_reconstructs_exactly(hidden_size: int, num_experts: int): + torch.manual_seed(42) + w = torch.randn(num_experts, hidden_size) * ROUTER_WEIGHT_SCALE + split = rocm_bf16x3.split_bf16x3(w) + + assert split.shape == (rocm_bf16x3.BF16X3_TERMS, num_experts, hidden_size) + assert split.dtype == torch.bfloat16 + # router_gemm reinterprets this as (TERMS * E, K) without a copy. + assert split.is_contiguous() + assert torch.equal(split.float().sum(0), w) + + +def test_split_is_exact_across_normal_exponents(): + exps = torch.arange(-100, 100, dtype=torch.float32) + mantissas = torch.linspace(1.0, 1.999, 64, dtype=torch.float32) + w = (mantissas[:, None] * torch.pow(2.0, exps)[None, :]).contiguous() + w = torch.cat([w, -w], dim=0) + assert torch.equal(rocm_bf16x3.split_bf16x3(w).float().sum(0), w) + + +def test_split_rejects_non_finite(): + """ValueError rather than assert: the caller downgrades to the fp32 path + instead of failing model load, and this must survive ``python -O``.""" + for bad in (float("inf"), float("-inf"), float("nan")): + with pytest.raises(ValueError): + rocm_bf16x3.split_bf16x3(torch.full((2, 4), bad, dtype=torch.float32)) + + +def test_split_rejects_non_fp32(): + with pytest.raises(ValueError): + rocm_bf16x3.split_bf16x3(torch.zeros(2, 4, dtype=torch.bfloat16)) + + +def test_split_does_not_mutate_input(): + w = torch.randn(8, 32, dtype=torch.float32) + original = w.clone() + rocm_bf16x3.split_bf16x3(w) + assert torch.equal(w, original) + + +def test_is_supported_rejects_ineligible_inputs(): + meta = torch.device("meta") + x = torch.zeros(8192, 6144, dtype=torch.bfloat16, device=meta) + 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) + 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]) + assert not rocm_bf16x3.is_supported(x.T.contiguous().T, w) + + +def test_is_supported_rejects_mismatched_devices(): + """The weight can be offloaded while the activations are not.""" + x = torch.zeros(8192, 6144, dtype=torch.bfloat16, device="meta") + w = torch.zeros(128, 6144, dtype=torch.float32) + assert not rocm_bf16x3.is_supported(x, w) + + +def test_is_supported_accepts_odd_expert_counts(): + """Unlike a tl.arange-indexed kernel, torch.mm needs no power-of-two E.""" + meta = torch.device("meta") + x = torch.zeros(8192, 6144, dtype=torch.bfloat16, device=meta) + for num_experts in (17, 384, 512): + w = torch.zeros(num_experts, 6144, dtype=torch.float32, device=meta) + assert rocm_bf16x3.is_supported(x, w) + + +# Measured on gfx950: this path lands at 3.4e-7 to 4.9e-7 and varies <0.2% +# across token counts. Fixed bounds rather than a live comparison, because the +# fp32 fallback swings 5.7e-7 to 1.4e-6 as torch picks different hipBLASLt +# kernels at different M. _FP32_FALLBACK_BEST is the best the fallback was ever +# observed to do, so beating it proves the claim without racing a moving target. +_MAX_REL_L2 = 2e-6 +_FP32_FALLBACK_BEST = 5.7e-7 + + +@pytest.mark.parametrize(("hidden_size", "num_experts"), ROUTER_SHAPES) +@pytest.mark.parametrize("num_tokens", [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) + out = rocm_bf16x3.bf16x3_router_gemm(x, rocm_bf16x3.split_bf16x3(w)) + ref = x.double() @ w.double().T + + assert out.shape == (num_tokens, num_experts) + assert out.dtype == torch.float32 + assert (out.double() - ref).norm() / ref.norm() < _MAX_REL_L2 + + +@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.""" + _requires_gfx950() + x, w = _inputs(8192, 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.double() - ref).norm() / ref.norm() < _MAX_REL_L2 + + +@pytest.mark.parametrize(("hidden_size", "num_experts"), ROUTER_SHAPES) +def test_more_accurate_than_fp32_fallback(hidden_size: int, num_experts: int): + """The claim is not just accuracy but beating the fp32 GEMM it replaces.""" + _requires_gfx950() + x, w = _inputs(16384, 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): + """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)) + + misrouted = int((_top_set(out) != ref_set).any(-1).sum()) + assert misrouted / num_tokens < 1e-4 + + +def test_writes_into_supplied_out(): + _requires_gfx950() + x, w = _inputs(8192, 6144, 128) + split = rocm_bf16x3.split_bf16x3(w) + out = torch.empty(8192, 128, dtype=torch.float32, device="cuda") + + returned = rocm_bf16x3.bf16x3_router_gemm(x, split, out=out) + assert returned is out + assert torch.equal(out, rocm_bf16x3.bf16x3_router_gemm(x, split)) + + +def test_validates_layout_assumptions(): + """The reduction assumes a contiguous (TERMS, E, K) split and a contiguous + out; nothing in the kernel enforces either.""" + _requires_gfx950() + x, w = _inputs(8192, 6144, 128) + split = rocm_bf16x3.split_bf16x3(w) + + with pytest.raises(AssertionError): + rocm_bf16x3.bf16x3_router_gemm(x, split.transpose(1, 2)) + with pytest.raises(AssertionError): + rocm_bf16x3.bf16x3_router_gemm(x.float(), split) + with pytest.raises(AssertionError): + rocm_bf16x3.bf16x3_router_gemm(x, split[:2]) + with pytest.raises(AssertionError): + rocm_bf16x3.bf16x3_router_gemm(x[:, :512], split) + wide = torch.empty(8192, 256, dtype=torch.float32, device="cuda") + with pytest.raises(AssertionError): + rocm_bf16x3.bf16x3_router_gemm(x, split, out=wide[:, :128]) diff --git a/tests/kernels/test_gate_linear_rocm_dispatch.py b/tests/kernels/test_gate_linear_rocm_dispatch.py index afa2e57f6a..06d69fc985 100644 --- a/tests/kernels/test_gate_linear_rocm_dispatch.py +++ b/tests/kernels/test_gate_linear_rocm_dispatch.py @@ -11,6 +11,8 @@ The ROCm branch is guarded on ``not bias`` because ``torch.mm`` has no bias term; a biased gate must fall back so the bias is not silently dropped. """ +import inspect + import pytest import torch @@ -168,3 +170,91 @@ def test_rocm_fp32_router_gemm_rejects_unsupported_configs( on_gfx950=on_gfx950, ) assert not gate.allow_fp32_router_gemm + + +# --------------------------------------------------------------------------- +# ROCm BF16x3 tier (fp32 router weights) +# --------------------------------------------------------------------------- +def _make_bf16x3_gate(monkeypatch, *, supported=True, **kwargs): + """Build a gate with the ROCm bf16x3 platform check mocked.""" + from vllm.model_executor.layers.fused_moe.router import ( + bf16x3_router_gemm_rocm as rocm_bf16x3, + ) + + monkeypatch.setattr(rocm_bf16x3, "platform_supported", lambda: supported) + kwargs.setdefault("params_dtype", torch.float32) + return _make_gate(monkeypatch, is_rocm=True, **kwargs) + + +def test_rocm_bf16x3_enabled_for_fp32_weight(monkeypatch): + gate = _make_bf16x3_gate(monkeypatch) + assert gate.allow_rocm_bf16x3_router_gemm + # Built by process_weights_after_loading, not lazily in forward(). + # Registered up front but still None, so named_buffers() skips it. + assert gate._bf16x3_weight is None + assert "_bf16x3_weight" in gate._buffers + + +def test_rocm_bf16x3_disabled_by_bias(monkeypatch): + # torch.mm has no bias term. + gate = _make_bf16x3_gate(monkeypatch, bias=True) + assert not gate.allow_rocm_bf16x3_router_gemm + + +def test_rocm_bf16x3_disabled_for_bf16_weight(monkeypatch): + # A bf16 weight belongs to the cuBLAS tier; there is nothing to split. + gate = _make_bf16x3_gate(monkeypatch, params_dtype=torch.bfloat16) + assert not gate.allow_rocm_bf16x3_router_gemm + + +def test_rocm_bf16x3_disabled_for_non_fp32_out_dtype(monkeypatch): + gate = _make_bf16x3_gate(monkeypatch, out_dtype=torch.bfloat16) + assert not gate.allow_rocm_bf16x3_router_gemm + + +def test_rocm_bf16x3_disabled_on_unsupported_arch(monkeypatch): + # Only gfx950 is benchmarked; gfx942 has a different fp32:bf16 MFMA ratio. + gate = _make_bf16x3_gate(monkeypatch, supported=False) + assert not gate.allow_rocm_bf16x3_router_gemm + assert not gate._rocm_bf16x3_weight_eligible + + +def test_rocm_bf16x3_no_split_built_when_tier_disabled(monkeypatch): + """An ineligible out_dtype must not allocate a split that is never read; + it is 1.5x the router weight on every layer.""" + gate = _make_bf16x3_gate(monkeypatch, out_dtype=torch.bfloat16) + gate.quant_method.process_weights_after_loading(gate) + assert gate._bf16x3_weight is None + + +def test_rocm_bf16x3_set_out_dtype_enables_tier(monkeypatch): + gate = _make_bf16x3_gate(monkeypatch, out_dtype=None) + assert not gate.allow_rocm_bf16x3_router_gemm + gate.set_out_dtype(torch.float32) + assert gate.allow_rocm_bf16x3_router_gemm + + +def test_rocm_bf16x3_dispatch_is_a_custom_op(): + """The num_tokens branch must stay opaque to Dynamo; see the op's impl.""" + assert hasattr(torch.ops.vllm, "rocm_bf16x3_router_gemm_dispatch") + x = torch.zeros(8192, 6144, dtype=torch.bfloat16, device="meta") + w = torch.zeros(128, 6144, dtype=torch.float32, device="meta") + split = torch.zeros(3, 128, 6144, dtype=torch.bfloat16, device="meta") + out = torch.ops.vllm.rocm_bf16x3_router_gemm_dispatch(x, w, split) + assert out.shape == (8192, 128) + assert out.dtype == torch.float32 + + +def test_rocm_bf16x3_is_reachable_for_low_m_shapes(monkeypatch): + """The gfx950 low-M fp32 tier claims these shapes too and returns + unconditionally, so bf16x3 is only reachable if checked first.""" + gate = _make_bf16x3_gate( + monkeypatch, input_size=6144, output_size=128, on_gfx950=True + ) + assert gate.allow_fp32_router_gemm + assert gate.allow_rocm_bf16x3_router_gemm + + source = inspect.getsource(type(gate).forward) + assert source.index("allow_rocm_bf16x3_router_gemm") < source.index( + "allow_fp32_router_gemm" + ) 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 new file mode 100644 index 0000000000..8d1fa84c86 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_rocm.py @@ -0,0 +1,163 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""ROCm BF16x3 MoE router GEMM for fp32 gate weights. + +Gates whose weights ship in fp32 (MiniMax-M2/M3, HunYuan-V3) otherwise land in +the ``F.linear`` fallback, which casts the bf16 activations up to fp32 and runs +the GEMM on the fp32 MFMA pipe. Splitting the weight into three bf16 terms +whose sum reproduces it exactly moves the GEMM onto the bf16 pipe and is more +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. +""" + +import torch + +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton + +__all__ = [ + "BF16X3_TERMS", + "MIN_TOKENS", + "bf16x3_router_gemm", + "is_supported", + "platform_supported", + "split_bf16x3", +] + +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 + +_MAX_INT32 = (1 << 31) - 1 + + +def split_bf16x3(weight: torch.Tensor) -> torch.Tensor: + """Split an fp32 weight into bf16 terms whose sum reproduces it exactly. + + Returns a contiguous ``(BF16X3_TERMS, E, K)`` tensor. Three + round-to-nearest bf16 terms carry 24 mantissa bits, matching fp32's 24, so + the reconstruction is bit-exact for finite normals down to ``2**-108``. + + Raises ValueError on non-finite input: ``inf`` rounds to bf16 ``inf`` and + ``inf - inf`` is NaN, which would poison every logit rather than just the + affected experts. + """ + if weight.dtype != torch.float32: + raise ValueError(f"expected fp32 weight, got {weight.dtype}") + if not torch.isfinite(weight).all(): + raise ValueError("router weight contains inf/NaN") + terms = [] + resid = weight.detach().clone() + for _ in range(BF16X3_TERMS): + hi = resid.to(torch.bfloat16) + terms.append(hi) + resid -= hi.to(torch.float32) + return torch.stack(terms).contiguous() + + +def platform_supported() -> bool: + """Whether this platform has a benchmarked bf16x3 router path. + + gfx942 has a different fp32:bf16 MFMA ratio and is not yet evaluated. + """ + if not current_platform.is_rocm(): + return False + from vllm.platforms.rocm import on_gfx950 + + return on_gfx950() + + +def is_supported(x: torch.Tensor, weight: torch.Tensor) -> bool: + """Whether this path can serve ``x @ weight.T`` -> fp32 profitably.""" + if x.dim() != 2 or weight.dim() != 2: + return False + n, k = weight.shape + m = x.shape[0] + return ( + m >= MIN_TOKENS + and x.dtype == torch.bfloat16 + and weight.dtype == torch.float32 + and x.shape[1] == k + # The weight may be offloaded while the activations are not; the + # fallback this guards would then fail on a device mismatch. + and x.device == weight.device + # torch.mm accepts any stride, but a non-K-contiguous activation costs + # a hidden copy that gives back the win over the fallback. + and x.stride(1) == 1 + and 0 < m * BF16X3_TERMS * n <= _MAX_INT32 + ) + + +@triton.jit +def _reduce_terms_kernel( + y_ptr, + out_ptr, + M, + stride_ym, + stride_om, + N: tl.constexpr, + TERMS: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + pid = tl.program_id(0) + offs_m = pid * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = tl.arange(0, BLOCK_N) + mask = (offs_m < M)[:, None] & (offs_n < N)[None, :] + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + for t in tl.static_range(TERMS): + acc += tl.load( + y_ptr + offs_m[:, None] * stride_ym + t * N + offs_n[None, :], + mask=mask, + other=0.0, + ) + tl.store(out_ptr + offs_m[:, None] * stride_om + offs_n[None, :], acc, mask=mask) + + +def bf16x3_router_gemm( + x: torch.Tensor, + weight_split: torch.Tensor, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Router logits in fp32 from bf16 activations and a bf16x3 split weight.""" + m, k = x.shape + terms, n, k_w = weight_split.shape + assert terms == BF16X3_TERMS, f"expected {BF16X3_TERMS} terms, got {terms}" + assert k == k_w, f"K mismatch: x has {k}, weight_split has {k_w}" + assert weight_split.dtype == torch.bfloat16, weight_split.dtype + assert weight_split.is_contiguous(), "weight_split must be contiguous" + assert x.dtype == torch.bfloat16, x.dtype + if out is None: + out = torch.empty((m, n), dtype=torch.float32, device=x.device) + else: + assert out.shape == (m, n) and out.is_contiguous(), "bad out layout" + + # 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. + y = torch.empty((m, terms * n), dtype=torch.float32, device=x.device) + torch.mm(x, weight_split.view(terms * n, k).T, out_dtype=torch.float32, out=y) + + # 64 rows/program keeps the accumulator small enough to stay in registers + # at the largest supported expert count. + block_m = 64 + _reduce_terms_kernel[(triton.cdiv(m, block_m),)]( + y, + out, + m, + y.stride(0), + out.stride(0), + N=n, + TERMS=terms, + BLOCK_M=block_m, + BLOCK_N=triton.next_power_of_2(n), + num_warps=4, + ) + return out 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 f91e9c449e..8c4961f269 100644 --- a/vllm/model_executor/layers/fused_moe/router/gate_linear.py +++ b/vllm/model_executor/layers/fused_moe/router/gate_linear.py @@ -4,6 +4,7 @@ import torch from torch.nn.parameter import Parameter import vllm._custom_ops as ops +from vllm.logger import init_logger from vllm.model_executor.custom_op import PluggableLayer from vllm.model_executor.layers.linear import ( ReplicatedLinear, @@ -13,6 +14,8 @@ from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.platforms import current_platform from vllm.utils.torch_utils import direct_register_custom_op +logger = init_logger(__name__) + @PluggableLayer.register("gate_linear") class GateLinear(ReplicatedLinear): @@ -23,8 +26,10 @@ class GateLinear(ReplicatedLinear): 2. fp32 specialized kernel (SM90+ or gfx950, bf16/fp32 in, fp32 out, M<=32, model-specific shapes) 3. bf16x3 CuteDSL kernel (SM100, bf16 in, fp32 weight) - 4. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 weight + fp32 out_dtype) - 5. F.linear via ReplicatedLinear (ultimate fallback) + 4. ROCm bf16x3 router GEMM (gfx950, bf16 in, fp32 weight, fp32 out, + M>=2048) + 5. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 weight + fp32 out_dtype) + 6. 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 @@ -140,6 +145,33 @@ class GateLinear(ReplicatedLinear): and self.out_dtype == torch.float32 ) + # ROCm bf16x3 router GEMM eligibility. out_dtype may still be None here + # and is folded in by set_out_dtype; the split itself is only built + # once the tier is known to be on (see _GateLinearMethod). + from vllm.model_executor.layers.fused_moe.router import ( + bf16x3_router_gemm_rocm, + ) + + self._rocm_bf16x3_weight_eligible = ( + self._router_gemm_no_bias + and self.weight.dtype == torch.float32 + and bf16x3_router_gemm_rocm.platform_supported() + ) + self.allow_rocm_bf16x3_router_gemm = ( + self._rocm_bf16x3_weight_eligible and self.out_dtype == torch.float32 + ) + # Filled in by process_weights_after_loading. Not persistent: it is + # derived from `weight` rather than loaded from a checkpoint. + self.register_buffer("_bf16x3_weight", None, persistent=False) + if self._rocm_bf16x3_weight_eligible: + # Safe to swap after create_weights() because + # UnquantizedLinearMethod holds no per-layer state. Asserted + # because the swap would silently drop quantization if a + # quant_config were ever plumbed into this layer. + assert isinstance(self.quant_method, UnquantizedLinearMethod) + self.quant_method = _GateLinearMethod() + logger.info_once("Enabled ROCm BF16x3 router GEMM.") + # 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. @@ -173,6 +205,10 @@ class GateLinear(ReplicatedLinear): self.allow_cublas_router_gemm = self.weight.dtype == torch.bfloat16 # out_dtype may start as None -> recompute eligibility here + self.allow_rocm_bf16x3_router_gemm = ( + self._rocm_bf16x3_weight_eligible and out_dtype == torch.float32 + ) + if self.allow_specialized_router_gemm: from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import ( is_available, @@ -205,6 +241,23 @@ class GateLinear(ReplicatedLinear): output = ll_bf16_gemm(x, self.weight) return self._return(output) + # Tier 4: ROCm bf16x3 router GEMM. Checked before tier 2 because both + # tiers claim the same gfx950 shapes and tier 2 returns + # unconditionally; the custom op below picks between them on the + # runtime num_tokens. + if self.allow_rocm_bf16x3_router_gemm and x.dtype == torch.bfloat16: + if self._bf16x3_weight is None: + logger.warning_once( + "ROCm BF16x3 router GEMM is enabled for %s but its weight " + "split is missing; falling back to fp32.", + self.prefix, + ) + else: + output = torch.ops.vllm.rocm_bf16x3_router_gemm_dispatch( + x, self.weight, self._bf16x3_weight + ) + return self._return(output) + # Tier 2: fp32 specialized kernel (model-specific shapes, M<=32) # Dispatch is wrapped in a custom op so that torch.compile/CUDA-graph # capture does not freeze the runtime num_tokens branch. @@ -226,12 +279,12 @@ class GateLinear(ReplicatedLinear): output = bf16x3_router_gemm(x, self.weight) return self._return(output) - # Tier 4: cuBLAS bf16→fp32 + # Tier 5: 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 self._return(output) - # Tier 5: F.linear (ReplicatedLinear) + # Tier 6: F.linear (ReplicatedLinear) if ( self.out_dtype is not None and self.is_unquantized @@ -249,9 +302,90 @@ class GateLinear(ReplicatedLinear): return output if not self.return_bias else (output, output_bias) +class _GateLinearMethod(UnquantizedLinearMethod): + """UnquantizedLinearMethod plus the ROCm bf16x3 weight split. + + Building the split here rather than lazily in forward() charges it to the + weights memory pool ahead of KV cache profiling, and rebuilds it on weight + reload. The latter matters because an in-place ``param.data.copy_()`` + leaves both ``data_ptr()`` and ``_version`` untouched, so a stale split + cannot be detected from inside forward(). + + ``allow_rocm_bf16x3_router_gemm`` is final by the time this runs: the only + caller of ``set_out_dtype`` does so during model construction. + """ + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + # The base CPU branch may delete layer.weight, so it must run first; + # nothing breaks today only because platform_supported() is False on + # CPU. + super().process_weights_after_loading(layer) + if not layer.allow_rocm_bf16x3_router_gemm: + return + try: + from vllm.model_executor.layers.fused_moe.router import ( + bf16x3_router_gemm_rocm, + ) + + layer._bf16x3_weight = bf16x3_router_gemm_rocm.split_bf16x3(layer.weight) + except ValueError as exc: + layer.allow_rocm_bf16x3_router_gemm = False + logger.warning_once( + "Disabling ROCm BF16x3 router GEMM for %s: %s", layer.prefix, exc + ) + + _FP32_ROUTER_GEMM_MAX_TOKENS = GateLinear.FP32_MAX_TOKENS +def rocm_bf16x3_router_gemm_dispatch_impl( + x: torch.Tensor, + weight: torch.Tensor, + weight_split: torch.Tensor, +) -> torch.Tensor: + """Run the ROCm bf16x3 router GEMM, or fall back for small batches. + + This must be a custom op because our torch.compile integration does not + support runtime dispatching on num_tokens: models carrying fp32 router + weights (MiniMax-M2, HunYuan-V3) call the gate from + ``@support_torch_compile`` model code and vLLM drops all Dynamo guards, so + a plain Python branch on ``x.shape[0]`` would be frozen at first trace. + """ + from vllm.model_executor.layers.fused_moe.router import ( + 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. + 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, + rocm_fp32_router_gemm, + ) + + x = x.contiguous() + if can_use_rocm_fp32_router_gemm(x, weight): + return rocm_fp32_router_gemm(x, weight) + return torch.nn.functional.linear(x.float(), weight) + + +def rocm_bf16x3_router_gemm_dispatch_fake( + x: torch.Tensor, + weight: torch.Tensor, + weight_split: torch.Tensor, +) -> torch.Tensor: + return x.new_empty((x.shape[0], weight.shape[0]), dtype=torch.float32) + + +direct_register_custom_op( + op_name="rocm_bf16x3_router_gemm_dispatch", + op_func=rocm_bf16x3_router_gemm_dispatch_impl, + fake_impl=rocm_bf16x3_router_gemm_dispatch_fake, +) + + def fp32_router_gemm_dispatch_impl( x: torch.Tensor, weight: torch.Tensor,