[ROCm][Perf] Add low-M FP32 router GEMM for gfx950 (#54845)
Fangzhou Ai committed
Sep 3, 2026 at 13:42 UTC
8bf39632b86df21fa9bfbb470ce8305bf67a5838
5 files changed
+499
-18
.buildkite/test-amd.yaml
+4
-1
@@ -3467,11 +3467,14 @@ steps:
3467
- tests/kernels/core
3468
- tests/kernels/test_top_k_per_row.py
3469
- tests/kernels/test_concat_mla_q.py
3470
+ - tests/kernels/test_rocm_fp32_router_gemm.py
3471
+ - vllm/model_executor/layers/fused_moe/router/gate_linear.py
3472
+ - vllm/model_executor/layers/fused_moe/router/rocm_fp32_router_gemm.py
3473
- vllm/model_executor/layers/rotary_embedding/
3474
- vllm/_aiter_ops.py
3475
- vllm/platforms/rocm.py
3476
commands:
3474
- - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
3477
+ - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py kernels/test_rocm_fp32_router_gemm.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
3478
3479
- label: ":amd: (MI355) MiniMax Reduce RMS Kernels" # TBD
3480
timeout_in_minutes: 45
tests/kernels/test_gate_linear_rocm_dispatch.py
+82
-10
@@ -1,16 +1,17 @@
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.
3
+"""Platform-agnostic eligibility tests for GateLinear's ROCm router GEMMs.
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.
5
+These assert the dispatch flags directly, so they run device-free by mocking
6
+the platform predicates. ``allow_cublas_router_gemm`` selects the
7
+bf16xbf16->fp32 ``torch.mm`` epilogue, while ``allow_fp32_router_gemm`` selects
8
+the gfx950 low-M kernel with fp32 weights and output.
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 pytest
15
import torch
16
17
import vllm.model_executor.layers.fused_moe.router.gate_linear as gate_linear_mod
@@ -25,6 +26,10 @@ def _make_gate(
26
bias: bool = False,
27
params_dtype: torch.dtype = torch.bfloat16,
28
out_dtype: torch.dtype | None = torch.float32,
29
+ input_size: int = 2048,
30
+ output_size: int = 64,
31
+ on_gfx950: bool = False,
32
+ parallel_world_size: int = 1,
33
) -> GateLinear:
34
"""Build a GateLinear with platform predicates mocked, no GPU needed."""
35
for target in (
@@ -37,20 +42,22 @@ def _make_gate(
42
)
43
monkeypatch.setattr(
44
f"{target}.get_tensor_model_parallel_world_size",
40
- lambda: 1,
45
+ lambda: parallel_world_size,
46
)
47
48
platform = gate_linear_mod.current_platform
49
monkeypatch.setattr(platform, "is_cuda", lambda: is_cuda)
50
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).
51
monkeypatch.setattr(platform, "is_device_capability", lambda *a, **k: False)
52
monkeypatch.setattr(platform, "is_device_capability_family", lambda *a, **k: False)
53
+ if is_rocm:
54
+ import vllm.platforms.rocm as rocm_platform
55
+
56
+ monkeypatch.setattr(rocm_platform, "on_gfx950", lambda: on_gfx950)
57
58
return GateLinear(
52
- input_size=2048,
53
- output_size=64,
59
+ input_size=input_size,
60
+ output_size=output_size,
61
bias=bias,
62
out_dtype=out_dtype,
63
params_dtype=params_dtype,
@@ -96,3 +103,68 @@ def test_rocm_set_out_dtype_respects_bias_guard(monkeypatch):
103
gate = _make_gate(monkeypatch, is_rocm=True, bias=True, out_dtype=None)
104
gate.set_out_dtype(torch.float32)
105
assert not gate.allow_cublas_router_gemm
106
+
107
+
108
+@pytest.mark.parametrize(
109
+ ("input_size", "output_size"),
110
+ [(3072, 256), (4096, 8), (4096, 192), (6144, 128), (6144, 256)],
111
+)
112
+def test_rocm_gfx950_enables_fp32_router_gemm(
113
+ monkeypatch, input_size: int, output_size: int
114
+) -> None:
115
+ gate = _make_gate(
116
+ monkeypatch,
117
+ is_rocm=True,
118
+ params_dtype=torch.float32,
119
+ input_size=input_size,
120
+ output_size=output_size,
121
+ on_gfx950=True,
122
+ )
123
+ assert gate.allow_fp32_router_gemm
124
+ assert not gate.allow_cublas_router_gemm
125
+
126
+
127
+@pytest.mark.parametrize("ep_size", [2, 4, 8])
128
+def test_rocm_fp32_router_gemm_is_replicated_across_ep_sizes(
129
+ monkeypatch, ep_size: int
130
+) -> None:
131
+ gate = _make_gate(
132
+ monkeypatch,
133
+ is_rocm=True,
134
+ params_dtype=torch.float32,
135
+ input_size=6144,
136
+ output_size=128,
137
+ on_gfx950=True,
138
+ parallel_world_size=ep_size,
139
+ )
140
+ assert gate.weight.shape == (128, 6144)
141
+ assert gate.allow_fp32_router_gemm
142
+
143
+
144
+@pytest.mark.parametrize(
145
+ ("input_size", "output_size", "params_dtype", "bias", "on_gfx950"),
146
+ [
147
+ pytest.param(6144, 128, torch.float32, False, False, id="gfx942"),
148
+ pytest.param(2048, 64, torch.float32, False, True, id="shape"),
149
+ pytest.param(6144, 128, torch.bfloat16, False, True, id="bf16-weight"),
150
+ pytest.param(6144, 128, torch.float32, True, True, id="bias"),
151
+ ],
152
+)
153
+def test_rocm_fp32_router_gemm_rejects_unsupported_configs(
154
+ monkeypatch,
155
+ input_size: int,
156
+ output_size: int,
157
+ params_dtype: torch.dtype,
158
+ bias: bool,
159
+ on_gfx950: bool,
160
+) -> None:
161
+ gate = _make_gate(
162
+ monkeypatch,
163
+ is_rocm=True,
164
+ params_dtype=params_dtype,
165
+ bias=bias,
166
+ input_size=input_size,
167
+ output_size=output_size,
168
+ on_gfx950=on_gfx950,
169
+ )
170
+ assert not gate.allow_fp32_router_gemm
tests/kernels/test_rocm_fp32_router_gemm.py
new
+228
@@ -0,0 +1,228 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Correctness tests for the gfx950 low-M fp32 router GEMM."""
4
+
5
+import pytest
6
+import torch
7
+
8
+from vllm.model_executor.layers.fused_moe.router.gate_linear import (
9
+ fp32_router_gemm_dispatch_impl,
10
+)
11
+from vllm.model_executor.layers.fused_moe.router.rocm_fp32_router_gemm import (
12
+ rocm_fp32_router_gemm,
13
+)
14
+from vllm.platforms import current_platform
15
+
16
+SHAPES = [
17
+ (3072, 256),
18
+ (4096, 8),
19
+ (4096, 192),
20
+ (6144, 128),
21
+ (6144, 256),
22
+]
23
+MAX_TOKENS = 32
24
+ATOL = 5e-4
25
+RTOL = 0.0
26
+
27
+
28
+def _on_gfx950() -> bool:
29
+ if not current_platform.is_rocm():
30
+ return False
31
+ try:
32
+ from vllm.platforms.rocm import on_gfx950
33
+
34
+ return on_gfx950()
35
+ except Exception:
36
+ return False
37
+
38
+
39
+pytestmark = pytest.mark.skipif(
40
+ not _on_gfx950(), reason="rocm_fp32_router_gemm requires ROCm gfx950"
41
+)
42
+
43
+
44
+def _reference(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
45
+ return torch.nn.functional.linear(x.float(), weight)
46
+
47
+
48
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
49
+@pytest.mark.parametrize(("hidden_size", "num_experts"), SHAPES)
50
+@pytest.mark.parametrize("num_tokens", range(MAX_TOKENS + 1))
51
+@torch.inference_mode()
52
+def test_rocm_fp32_router_gemm_matches_reference(
53
+ num_tokens: int,
54
+ hidden_size: int,
55
+ num_experts: int,
56
+ dtype: torch.dtype,
57
+) -> None:
58
+ torch.manual_seed(41 + num_tokens + hidden_size + num_experts)
59
+ device = torch.device("cuda")
60
+ x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
61
+ weight = torch.randn(num_experts, hidden_size, dtype=torch.float32, device=device)
62
+
63
+ output = rocm_fp32_router_gemm(x, weight)
64
+ expected = _reference(x, weight)
65
+
66
+ assert output.shape == (num_tokens, num_experts)
67
+ assert output.dtype == torch.float32
68
+ assert output.device == x.device
69
+ torch.testing.assert_close(output, expected, atol=ATOL, rtol=RTOL)
70
+
71
+
72
+@pytest.mark.parametrize("num_tokens", [1, 4, 16, 24, 32])
73
+@pytest.mark.parametrize(("hidden_size", "num_experts"), SHAPES)
74
+@torch.inference_mode()
75
+def test_rocm_fp32_router_gemm_preserves_topk(
76
+ num_tokens: int, hidden_size: int, num_experts: int
77
+) -> None:
78
+ torch.manual_seed(1000 + num_tokens + hidden_size + num_experts)
79
+ device = torch.device("cuda")
80
+ x = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device=device)
81
+ weight = torch.randn(num_experts, hidden_size, dtype=torch.float32, device=device)
82
+
83
+ top_k = 2 if (hidden_size, num_experts) == (4096, 8) else 8
84
+ actual = rocm_fp32_router_gemm(x, weight).topk(top_k, dim=-1).indices
85
+ reference = _reference(x, weight)
86
+ reference_values, expected = reference.topk(top_k, dim=-1)
87
+ for token in range(num_tokens):
88
+ actual_set = set(actual[token].tolist())
89
+ expected_set = set(expected[token].tolist())
90
+ if actual_set == expected_set:
91
+ continue
92
+ kth_value = reference_values[token, -1].item()
93
+ for expert in actual_set.symmetric_difference(expected_set):
94
+ gap = abs(reference[token, expert].item() - kth_value)
95
+ assert gap < 1e-3
96
+
97
+
98
+@torch.inference_mode()
99
+def test_rocm_fp32_router_gemm_rejects_invalid_inputs() -> None:
100
+ device = torch.device("cuda")
101
+ hidden_size, num_experts = 6144, 128
102
+ x = torch.randn(4, hidden_size, dtype=torch.bfloat16, device=device)
103
+ weight = torch.randn(num_experts, hidden_size, dtype=torch.float32, device=device)
104
+
105
+ with pytest.raises(ValueError, match="num_tokens"):
106
+ rocm_fp32_router_gemm(
107
+ torch.randn(
108
+ MAX_TOKENS + 1,
109
+ hidden_size,
110
+ dtype=torch.bfloat16,
111
+ device=device,
112
+ ),
113
+ weight,
114
+ )
115
+
116
+ with pytest.raises(ValueError, match="shape"):
117
+ rocm_fp32_router_gemm(x[:, :2048].contiguous(), weight[:, :2048].contiguous())
118
+
119
+ with pytest.raises(ValueError, match="contiguous"):
120
+ wide_x = torch.randn(4, hidden_size * 2, dtype=torch.bfloat16, device=device)
121
+ rocm_fp32_router_gemm(wide_x[:, ::2], weight)
122
+
123
+ with pytest.raises(ValueError, match="float32"):
124
+ rocm_fp32_router_gemm(x, weight.to(torch.bfloat16))
125
+
126
+ with pytest.raises(ValueError, match="bfloat16 or float32"):
127
+ rocm_fp32_router_gemm(x.to(torch.float16), weight)
128
+
129
+
130
+@torch.inference_mode()
131
+def test_rocm_fp32_router_gemm_cuda_graph_observes_input_mutation() -> None:
132
+ torch.manual_seed(2026)
133
+ device = torch.device("cuda")
134
+ hidden_size, num_experts = 6144, 128
135
+ static_x = torch.randn(16, hidden_size, dtype=torch.bfloat16, device=device)
136
+ weight = torch.randn(num_experts, hidden_size, dtype=torch.float32, device=device)
137
+
138
+ rocm_fp32_router_gemm(static_x, weight)
139
+ torch.accelerator.synchronize()
140
+
141
+ graph = torch.cuda.CUDAGraph()
142
+ with torch.cuda.graph(graph):
143
+ captured_output = rocm_fp32_router_gemm(static_x, weight)
144
+ torch.accelerator.synchronize()
145
+
146
+ first_input = torch.randn_like(static_x)
147
+ static_x.copy_(first_input)
148
+ graph.replay()
149
+ torch.accelerator.synchronize()
150
+ first_output = captured_output.clone()
151
+
152
+ second_input = torch.randn_like(static_x)
153
+ static_x.copy_(second_input)
154
+ graph.replay()
155
+ torch.accelerator.synchronize()
156
+ second_output = captured_output.clone()
157
+
158
+ torch.testing.assert_close(
159
+ first_output, _reference(first_input, weight), atol=ATOL, rtol=RTOL
160
+ )
161
+ torch.testing.assert_close(
162
+ second_output, _reference(second_input, weight), atol=ATOL, rtol=RTOL
163
+ )
164
+ assert not torch.equal(first_output, second_output)
165
+
166
+
167
+@pytest.mark.parametrize("num_tokens", [4, 32, 33])
168
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
169
+@torch.inference_mode()
170
+def test_rocm_fp32_router_gemm_custom_op_dispatch(
171
+ num_tokens: int, dtype: torch.dtype
172
+) -> None:
173
+ torch.manual_seed(3100 + num_tokens)
174
+ device = torch.device("cuda")
175
+ hidden_size, num_experts = 6144, 128
176
+ x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
177
+ weight = torch.randn(num_experts, hidden_size, dtype=torch.float32, device=device)
178
+
179
+ output = torch.ops.vllm.fp32_router_gemm_dispatch(x, weight, False)
180
+ direct_output = fp32_router_gemm_dispatch_impl(x, weight, False)
181
+
182
+ torch.testing.assert_close(output, direct_output, atol=ATOL, rtol=RTOL)
183
+ torch.testing.assert_close(output, _reference(x, weight), atol=ATOL, rtol=RTOL)
184
+
185
+
186
+@torch.inference_mode()
187
+def test_rocm_fp32_router_gemm_dispatch_accepts_noncontiguous_input(
188
+ monkeypatch: pytest.MonkeyPatch,
189
+) -> None:
190
+ torch.manual_seed(4100)
191
+ device = torch.device("cuda")
192
+ hidden_size, num_experts = 6144, 128
193
+ wide_x = torch.randn(4, hidden_size * 2, dtype=torch.bfloat16, device=device)
194
+ x = wide_x[:, ::2]
195
+ weight = torch.randn(num_experts, hidden_size, dtype=torch.float32, device=device)
196
+ assert not x.is_contiguous()
197
+
198
+ expected = _reference(x, weight)
199
+
200
+ def fail_fallback(*args, **kwargs):
201
+ raise AssertionError("noncontiguous low-M input used the linear fallback")
202
+
203
+ monkeypatch.setattr(torch.nn.functional, "linear", fail_fallback)
204
+ output = torch.ops.vllm.fp32_router_gemm_dispatch(x, weight, False)
205
+
206
+ torch.testing.assert_close(output, expected, atol=ATOL, rtol=RTOL)
207
+
208
+
209
+@torch.inference_mode()
210
+def test_rocm_fp32_router_gemm_dynamic_compile_dispatch() -> None:
211
+ torch.manual_seed(5100)
212
+ device = torch.device("cuda")
213
+ hidden_size, num_experts = 6144, 128
214
+ weight = torch.randn(num_experts, hidden_size, dtype=torch.float32, device=device)
215
+
216
+ def dispatch(x: torch.Tensor) -> torch.Tensor:
217
+ return torch.ops.vllm.fp32_router_gemm_dispatch(x, weight, False)
218
+
219
+ compiled_dispatch = torch.compile(dispatch, dynamic=True, fullgraph=True)
220
+ for num_tokens in (4, 16, 33, 5, 32):
221
+ x = torch.randn(
222
+ num_tokens,
223
+ hidden_size,
224
+ dtype=torch.bfloat16,
225
+ device=device,
226
+ )
227
+ output = compiled_dispatch(x)
228
+ torch.testing.assert_close(output, _reference(x, weight), atol=ATOL, rtol=RTOL)
vllm/model_executor/layers/fused_moe/router/gate_linear.py
+37
-7
@@ -20,8 +20,8 @@ class GateLinear(ReplicatedLinear):
20
21
1. cuteDSL ll_bf16_gemm (SM90+, M<=16, bf16 in, fp32 out,
22
K divisible by 8)
23
- 2. fp32 specialized kernel (SM90+, bf16/fp32 in, fp32 out, M<=32,
24
- (H, E) in {(3072, 256), (6144, 128)})
23
+ 2. fp32 specialized kernel (SM90+ or gfx950, bf16/fp32 in, fp32 out,
24
+ M<=32, model-specific shapes)
25
3. experimental bf16x3 CuteDSL kernel (opt-in, SM100, bf16 in, fp32 weight)
26
4. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 weight + fp32 out_dtype)
27
5. F.linear via ReplicatedLinear (ultimate fallback)
@@ -48,6 +48,21 @@ class GateLinear(ReplicatedLinear):
48
):
49
is_hopper = current_platform.is_device_capability((9, 0))
50
is_blackwell = current_platform.is_device_capability_family(100)
51
+ is_gfx950 = False
52
+ if current_platform.is_rocm():
53
+ from vllm.platforms.rocm import on_gfx950
54
+
55
+ is_gfx950 = on_gfx950()
56
+ is_rocm_fp32_shape = False
57
+ if is_gfx950:
58
+ from vllm.model_executor.layers.fused_moe.router.rocm_fp32_router_gemm import ( # noqa: E501
59
+ ROCM_FP32_ROUTER_GEMM_SUPPORTED_SHAPES,
60
+ )
61
+
62
+ is_rocm_fp32_shape = (
63
+ input_size,
64
+ output_size,
65
+ ) in ROCM_FP32_ROUTER_GEMM_SUPPORTED_SHAPES
66
can_use_specialized_kernels = (
67
current_platform.is_cuda() and (is_hopper or is_blackwell) and not bias
68
)
@@ -69,7 +84,7 @@ class GateLinear(ReplicatedLinear):
84
85
self.allow_specialized_router_gemm = can_use_specialized_kernels
86
72
- # fp32 specialized kernel eligibility (SM90+, exact dims, fp32 weight)
87
+ # fp32 specialized kernel eligibility (exact dims, fp32 weight)
88
vllm_config = get_current_vllm_config_or_none()
89
enable_bf16x3_router_gemm = (
90
vllm_config is not None
@@ -78,9 +93,14 @@ class GateLinear(ReplicatedLinear):
93
self.allow_fp32_router_gemm = (
94
not bias
95
and self.weight.dtype == torch.float32
81
- and current_platform.is_cuda()
82
- and (is_hopper or is_blackwell)
83
- and (input_size, output_size) in self.FP32_SUPPORTED_SHAPES
96
+ and (
97
+ (
98
+ current_platform.is_cuda()
99
+ and (is_hopper or is_blackwell)
100
+ and (input_size, output_size) in self.FP32_SUPPORTED_SHAPES
101
+ )
102
+ or (is_gfx950 and is_rocm_fp32_shape)
103
+ )
104
)
105
self.allow_bf16x3_router_gemm = (
106
not bias
@@ -166,7 +186,7 @@ class GateLinear(ReplicatedLinear):
186
output = ll_bf16_gemm(x, self.weight)
187
return output, None
188
169
- # Tier 2: fp32 specialized kernel (H=3072, E=256, M<=32)
189
+ # Tier 2: fp32 specialized kernel (model-specific shapes, M<=32)
190
# Dispatch is wrapped in a custom op so that torch.compile/CUDA-graph
191
# capture does not freeze the runtime num_tokens branch.
192
if self.allow_fp32_router_gemm and x.dtype in (
@@ -217,6 +237,16 @@ def fp32_router_gemm_dispatch_impl(
237
does not support runtime dispatching on num_tokens.
238
"""
239
if x.shape[0] <= _FP32_ROUTER_GEMM_MAX_TOKENS:
240
+ if current_platform.is_rocm():
241
+ from vllm.model_executor.layers.fused_moe.router.rocm_fp32_router_gemm import ( # noqa: E501
242
+ can_use_rocm_fp32_router_gemm,
243
+ rocm_fp32_router_gemm,
244
+ )
245
+
246
+ x = x.contiguous()
247
+ if can_use_rocm_fp32_router_gemm(x, weight):
248
+ return rocm_fp32_router_gemm(x, weight)
249
+ return torch.nn.functional.linear(x.float(), weight)
250
return ops.fp32_router_gemm(x, weight)
251
252
if allow_bf16x3_router_gemm and x.dtype == torch.bfloat16:
vllm/model_executor/layers/fused_moe/router/rocm_fp32_router_gemm.py
new
+148
@@ -0,0 +1,148 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Low-token FP32 router GEMM for gfx950."""
4
+
5
+import torch
6
+
7
+from vllm.platforms import current_platform
8
+from vllm.platforms.rocm import on_gfx950
9
+from vllm.triton_utils import tl, triton
10
+
11
+_MAX_TOKENS = 32
12
+ROCM_FP32_ROUTER_GEMM_SUPPORTED_SHAPES = frozenset(
13
+ {
14
+ (3072, 256),
15
+ (4096, 8),
16
+ (4096, 192),
17
+ (6144, 128),
18
+ (6144, 256),
19
+ }
20
+)
21
+
22
+
23
+@triton.jit
24
+def _rocm_fp32_router_gemm_kernel(
25
+ hidden_states_ptr,
26
+ router_weight_ptr,
27
+ output_ptr,
28
+ M,
29
+ K: tl.constexpr,
30
+ N: tl.constexpr,
31
+ BLOCK_M: tl.constexpr,
32
+ BLOCK_K: tl.constexpr,
33
+):
34
+ pid = tl.program_id(0)
35
+ pid_m = pid // N
36
+ pid_n = pid - pid_m * N
37
+ offsets_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
38
+ offsets_k = tl.arange(0, BLOCK_K)
39
+ partials = tl.zeros((BLOCK_M, BLOCK_K), dtype=tl.float32)
40
+
41
+ for k_start in range(0, K, BLOCK_K):
42
+ offsets_k_block = k_start + offsets_k
43
+ hidden_states = tl.load(
44
+ hidden_states_ptr + offsets_m[:, None] * K + offsets_k_block[None, :],
45
+ mask=offsets_m[:, None] < M,
46
+ other=0.0,
47
+ ).to(tl.float32)
48
+ router_weight = tl.load(router_weight_ptr + pid_n * K + offsets_k_block)
49
+ partials += hidden_states * router_weight[None, :]
50
+
51
+ accumulator = tl.sum(partials, axis=1)
52
+ tl.store(
53
+ output_ptr + offsets_m * N + pid_n,
54
+ accumulator,
55
+ mask=offsets_m < M,
56
+ )
57
+
58
+
59
+def _launch_config(
60
+ hidden_size: int,
61
+ num_experts: int,
62
+ num_tokens: int,
63
+) -> tuple[int, int, int]:
64
+ if (hidden_size, num_experts) == (4096, 8):
65
+ return 1, 1024, 4
66
+
67
+ if num_tokens <= 4:
68
+ return 1, 1024, 8
69
+ if num_tokens <= 16:
70
+ block_m = 8 if (hidden_size, num_experts) == (6144, 256) else 4
71
+ return block_m, 1024, 8
72
+ if (hidden_size, num_experts) == (3072, 256):
73
+ return 8, 1024, 8
74
+ return 4, 2048, 8
75
+
76
+
77
+def can_use_rocm_fp32_router_gemm(
78
+ hidden_states: torch.Tensor,
79
+ router_weight: torch.Tensor,
80
+) -> bool:
81
+ """Return whether the tensors match the tuned gfx950 fast path."""
82
+ try:
83
+ _validate_inputs(hidden_states, router_weight)
84
+ except (RuntimeError, ValueError):
85
+ return False
86
+ return True
87
+
88
+
89
+def _validate_inputs(
90
+ hidden_states: torch.Tensor,
91
+ router_weight: torch.Tensor,
92
+) -> None:
93
+ if not current_platform.is_rocm() or not on_gfx950():
94
+ raise RuntimeError("rocm_fp32_router_gemm requires ROCm gfx950")
95
+ if hidden_states.dim() != 2 or router_weight.dim() != 2:
96
+ raise ValueError("hidden_states and router_weight must be 2D tensors")
97
+ if hidden_states.dtype not in (torch.bfloat16, torch.float32):
98
+ raise ValueError("hidden_states must have dtype bfloat16 or float32")
99
+ if router_weight.dtype != torch.float32:
100
+ raise ValueError("router_weight must have dtype float32")
101
+ if hidden_states.device.type != "cuda" or router_weight.device.type != "cuda":
102
+ raise ValueError("hidden_states and router_weight must be GPU tensors")
103
+ if hidden_states.device != router_weight.device:
104
+ raise ValueError("hidden_states and router_weight must be on the same device")
105
+ if not hidden_states.is_contiguous() or not router_weight.is_contiguous():
106
+ raise ValueError("hidden_states and router_weight must be contiguous")
107
+ shape = (hidden_states.shape[1], router_weight.shape[0])
108
+ if (
109
+ shape not in ROCM_FP32_ROUTER_GEMM_SUPPORTED_SHAPES
110
+ or router_weight.shape[1] != shape[0]
111
+ ):
112
+ raise ValueError(
113
+ "supported (hidden_size, num_experts) shape pairs are "
114
+ "(3072, 256), (4096, 8), (4096, 192), (6144, 128), "
115
+ "and (6144, 256)"
116
+ )
117
+ if not 0 <= hidden_states.shape[0] <= _MAX_TOKENS:
118
+ raise ValueError(f"num_tokens must be in [0, {_MAX_TOKENS}]")
119
+
120
+
121
+def rocm_fp32_router_gemm(
122
+ hidden_states: torch.Tensor,
123
+ router_weight: torch.Tensor,
124
+) -> torch.Tensor:
125
+ """Compute ``hidden_states @ router_weight.T`` with FP32 accumulation."""
126
+ _validate_inputs(hidden_states, router_weight)
127
+ num_tokens = hidden_states.shape[0]
128
+ hidden_size = hidden_states.shape[1]
129
+ num_experts = router_weight.shape[0]
130
+ output = hidden_states.new_empty((num_tokens, num_experts), dtype=torch.float32)
131
+ if num_tokens == 0:
132
+ return output
133
+
134
+ block_m, block_k, num_warps = _launch_config(hidden_size, num_experts, num_tokens)
135
+ grid = (triton.cdiv(num_tokens, block_m) * num_experts,)
136
+ _rocm_fp32_router_gemm_kernel[grid](
137
+ hidden_states,
138
+ router_weight,
139
+ output,
140
+ M=num_tokens,
141
+ K=hidden_size,
142
+ N=num_experts,
143
+ BLOCK_M=block_m,
144
+ BLOCK_K=block_k,
145
+ num_warps=num_warps,
146
+ num_stages=1,
147
+ )
148
+ return output