iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[ROCm][Perf] Use a stacked hipBLASLt GEMM for the bf16x3 router (#52668)

Browse Source 
Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>
This commit is contained in:
akii96
2026-10-05 17:28:07 +00:00
committed by GitHub
parent 54d93af9fb
commit 182247257f
5 changed files with 594 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.buildkite/test_areas/kernels.yaml
+2
View File
@@ -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
tests/kernels/test_bf16x3_router_gemm_rocm.py
+201
View File
@@ -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])
tests/kernels/test_gate_linear_rocm_dispatch.py
+90
View File
@@ -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"
)
vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_rocm.py
+163
View File
@@ -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
vllm/model_executor/layers/fused_moe/router/gate_linear.py
+138 -4
View File
@@ -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,