[Hardware][AMD] Enable fused bf16→fp32 router GEMM on ROCm (#50268)
[Hardware][AMD] Enable fused bf16→fp32 router GEMM on ROCm (#50268) Signed-off-by: Matvei Pashkovskii <mpashkov@amd.com> Co-authored-by: Claude <noreply@anthropic.com>
Matvei Pashkovskii committed
Aug 12, 2026 at 05:46 UTC
20727e2841ed7c3194ad5d6410486fc72dac4556
2 files changed
+112
-3
tests/kernels/test_gate_linear_rocm_dispatch.py
new
+98
@@ -0,0 +1,98 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Platform-agnostic eligibility tests for GateLinear's fused router GEMM.
4
+
5
+These assert the ``allow_cublas_router_gemm`` dispatch flag directly, so they
6
+run device-free by mocking the platform predicates. The flag decides whether
7
+the bf16xbf16->fp32 router GEMM uses ``torch.mm``'s fused out_dtype epilogue
8
+(one kernel) or falls back to a bf16 matmul plus a standalone bf16->fp32 copy.
9
+
10
+The ROCm branch is guarded on ``not bias`` because ``torch.mm`` has no bias
11
+term; a biased gate must fall back so the bias is not silently dropped.
12
+"""
13
+
14
+import torch
15
+
16
+import vllm.model_executor.layers.fused_moe.router.gate_linear as gate_linear_mod
17
+from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear
18
+
19
+
20
+def _make_gate(
21
+ monkeypatch,
22
+ *,
23
+ is_rocm: bool,
24
+ is_cuda: bool = False,
25
+ bias: bool = False,
26
+ params_dtype: torch.dtype = torch.bfloat16,
27
+ out_dtype: torch.dtype | None = torch.float32,
28
+) -> GateLinear:
29
+ """Build a GateLinear with platform predicates mocked, no GPU needed."""
30
+ for target in (
31
+ "vllm.model_executor.layers.linear",
32
+ "vllm.model_executor.parameter",
33
+ ):
34
+ monkeypatch.setattr(
35
+ f"{target}.get_tensor_model_parallel_rank",
36
+ lambda: 0,
37
+ )
38
+ monkeypatch.setattr(
39
+ f"{target}.get_tensor_model_parallel_world_size",
40
+ lambda: 1,
41
+ )
42
+
43
+ platform = gate_linear_mod.current_platform
44
+ monkeypatch.setattr(platform, "is_cuda", lambda: is_cuda)
45
+ monkeypatch.setattr(platform, "is_rocm", lambda: is_rocm)
46
+ # Force the CUDA specialized-kernel gate off so ROCm eligibility is the
47
+ # only thing under test (these are the SM90/SM100 capability checks).
48
+ monkeypatch.setattr(platform, "is_device_capability", lambda *a, **k: False)
49
+ monkeypatch.setattr(platform, "is_device_capability_family", lambda *a, **k: False)
50
+
51
+ return GateLinear(
52
+ input_size=2048,
53
+ output_size=64,
54
+ bias=bias,
55
+ out_dtype=out_dtype,
56
+ params_dtype=params_dtype,
57
+ )
58
+
59
+
60
+def test_rocm_no_bias_bf16_fp32_enables_fused_gemm(monkeypatch):
61
+ gate = _make_gate(monkeypatch, is_rocm=True, bias=False)
62
+ assert not gate.allow_specialized_router_gemm
63
+ assert gate.allow_cublas_router_gemm
64
+
65
+
66
+def test_rocm_bias_disables_fused_gemm(monkeypatch):
67
+ # torch.mm cannot add a bias, so a biased gate must not take the fused path.
68
+ gate = _make_gate(monkeypatch, is_rocm=True, bias=True)
69
+ assert not gate.allow_cublas_router_gemm
70
+
71
+
72
+def test_rocm_fp32_weight_disables_fused_gemm(monkeypatch):
73
+ gate = _make_gate(monkeypatch, is_rocm=True, params_dtype=torch.float32)
74
+ assert not gate.allow_cublas_router_gemm
75
+
76
+
77
+def test_rocm_non_fp32_out_dtype_disables_fused_gemm(monkeypatch):
78
+ gate = _make_gate(monkeypatch, is_rocm=True, out_dtype=torch.bfloat16)
79
+ assert not gate.allow_cublas_router_gemm
80
+
81
+
82
+def test_non_rocm_non_cuda_disables_fused_gemm(monkeypatch):
83
+ # Neither the CUDA specialized path nor the ROCm branch applies.
84
+ gate = _make_gate(monkeypatch, is_rocm=False, is_cuda=False)
85
+ assert not gate.allow_cublas_router_gemm
86
+
87
+
88
+def test_rocm_set_out_dtype_enables_fused_gemm(monkeypatch):
89
+ gate = _make_gate(monkeypatch, is_rocm=True, bias=False, out_dtype=None)
90
+ assert not gate.allow_cublas_router_gemm
91
+ gate.set_out_dtype(torch.float32)
92
+ assert gate.allow_cublas_router_gemm
93
+
94
+
95
+def test_rocm_set_out_dtype_respects_bias_guard(monkeypatch):
96
+ gate = _make_gate(monkeypatch, is_rocm=True, bias=True, out_dtype=None)
97
+ gate.set_out_dtype(torch.float32)
98
+ assert not gate.allow_cublas_router_gemm
vllm/model_executor/layers/fused_moe/router/gate_linear.py
+14
-3
@@ -114,9 +114,17 @@ class GateLinear(ReplicatedLinear):
114
if self.allow_bf16x3_router_gemm:
115
logger.info_once("Enabled experimental SM100 BF16x3 router GEMM.")
116
117
- # cuBLAS bf16→fp32 eligibility
117
+ # Fused bf16 x bf16 -> fp32 GEMM eligibility. torch.mm's out_dtype
118
+ # epilogue folds the fp32 cast into the GEMM, removing the standalone
119
+ # bf16->fp32 copy kernel that otherwise runs before grouped_topk.
120
+ # cuBLAS on CUDA (SM90+, via allow_specialized_router_gemm); hipBLASLt on
121
+ # ROCm, which supports the same out_dtype epilogue.
122
+ self._router_gemm_no_bias = not bias
123
self.allow_cublas_router_gemm = (
119
- self.allow_specialized_router_gemm
124
+ (
125
+ self.allow_specialized_router_gemm
126
+ or (current_platform.is_rocm() and self._router_gemm_no_bias)
127
+ )
128
and self.weight.dtype == torch.bfloat16
129
and self.out_dtype == torch.float32
130
)
@@ -148,7 +156,10 @@ class GateLinear(ReplicatedLinear):
156
157
if (
158
not self.allow_cublas_router_gemm
151
- and self.allow_specialized_router_gemm
159
+ and (
160
+ self.allow_specialized_router_gemm
161
+ or (current_platform.is_rocm() and self._router_gemm_no_bias)
162
+ )
163
and out_dtype == torch.float32
164
):
165
self.allow_cublas_router_gemm = self.weight.dtype == torch.bfloat16