[ROCm][Perf] Add low-M FP32 router GEMM for gfx950 (#54845)

.buildkite/test-amd.yaml

@@ -3467,11 +3467,14 @@ steps:
   - tests/kernels/core
   - tests/kernels/test_top_k_per_row.py
   - tests/kernels/test_concat_mla_q.py
+  - tests/kernels/test_rocm_fp32_router_gemm.py
+  - vllm/model_executor/layers/fused_moe/router/gate_linear.py
+  - vllm/model_executor/layers/fused_moe/router/rocm_fp32_router_gemm.py
   - vllm/model_executor/layers/rotary_embedding/
   - vllm/_aiter_ops.py
   - vllm/platforms/rocm.py
   commands:
-  - 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
+  - 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
 
 - label: ":amd: (MI355) MiniMax Reduce RMS Kernels" # TBD
   timeout_in_minutes: 45

tests/kernels/test_gate_linear_rocm_dispatch.py

@@ -1,16 +1,17 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Platform-agnostic eligibility tests for GateLinear's fused router GEMM.
+"""Platform-agnostic eligibility tests for GateLinear's ROCm router GEMMs.
 
-These assert the ``allow_cublas_router_gemm`` dispatch flag directly, so they
-run device-free by mocking the platform predicates. The flag decides whether
-the bf16xbf16->fp32 router GEMM uses ``torch.mm``'s fused out_dtype epilogue
-(one kernel) or falls back to a bf16 matmul plus a standalone bf16->fp32 copy.
+These assert the dispatch flags directly, so they run device-free by mocking
+the platform predicates. ``allow_cublas_router_gemm`` selects the
+bf16xbf16->fp32 ``torch.mm`` epilogue, while ``allow_fp32_router_gemm`` selects
+the gfx950 low-M kernel with fp32 weights and output.
 
 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 pytest
 import torch
 
 import vllm.model_executor.layers.fused_moe.router.gate_linear as gate_linear_mod
@@ -25,6 +26,10 @@ def _make_gate(
     bias: bool = False,
     params_dtype: torch.dtype = torch.bfloat16,
     out_dtype: torch.dtype | None = torch.float32,
+    input_size: int = 2048,
+    output_size: int = 64,
+    on_gfx950: bool = False,
+    parallel_world_size: int = 1,
 ) -> GateLinear:
     """Build a GateLinear with platform predicates mocked, no GPU needed."""
     for target in (
@@ -37,20 +42,22 @@ def _make_gate(
         )
         monkeypatch.setattr(
             f"{target}.get_tensor_model_parallel_world_size",
-            lambda: 1,
+            lambda: parallel_world_size,
         )
 
     platform = gate_linear_mod.current_platform
     monkeypatch.setattr(platform, "is_cuda", lambda: is_cuda)
     monkeypatch.setattr(platform, "is_rocm", lambda: is_rocm)
-    # Force the CUDA specialized-kernel gate off so ROCm eligibility is the
-    # only thing under test (these are the SM90/SM100 capability checks).
     monkeypatch.setattr(platform, "is_device_capability", lambda *a, **k: False)
     monkeypatch.setattr(platform, "is_device_capability_family", lambda *a, **k: False)
+    if is_rocm:
+        import vllm.platforms.rocm as rocm_platform
+
+        monkeypatch.setattr(rocm_platform, "on_gfx950", lambda: on_gfx950)
 
     return GateLinear(
-        input_size=2048,
-        output_size=64,
+        input_size=input_size,
+        output_size=output_size,
         bias=bias,
         out_dtype=out_dtype,
         params_dtype=params_dtype,
@@ -96,3 +103,68 @@ def test_rocm_set_out_dtype_respects_bias_guard(monkeypatch):
     gate = _make_gate(monkeypatch, is_rocm=True, bias=True, out_dtype=None)
     gate.set_out_dtype(torch.float32)
     assert not gate.allow_cublas_router_gemm
+
+
+@pytest.mark.parametrize(
+    ("input_size", "output_size"),
+    [(3072, 256), (4096, 8), (4096, 192), (6144, 128), (6144, 256)],
+)
+def test_rocm_gfx950_enables_fp32_router_gemm(
+    monkeypatch, input_size: int, output_size: int
+) -> None:
+    gate = _make_gate(
+        monkeypatch,
+        is_rocm=True,
+        params_dtype=torch.float32,
+        input_size=input_size,
+        output_size=output_size,
+        on_gfx950=True,
+    )
+    assert gate.allow_fp32_router_gemm
+    assert not gate.allow_cublas_router_gemm
+
+
+@pytest.mark.parametrize("ep_size", [2, 4, 8])
+def test_rocm_fp32_router_gemm_is_replicated_across_ep_sizes(
+    monkeypatch, ep_size: int
+) -> None:
+    gate = _make_gate(
+        monkeypatch,
+        is_rocm=True,
+        params_dtype=torch.float32,
+        input_size=6144,
+        output_size=128,
+        on_gfx950=True,
+        parallel_world_size=ep_size,
+    )
+    assert gate.weight.shape == (128, 6144)
+    assert gate.allow_fp32_router_gemm
+
+
+@pytest.mark.parametrize(
+    ("input_size", "output_size", "params_dtype", "bias", "on_gfx950"),
+    [
+        pytest.param(6144, 128, torch.float32, False, False, id="gfx942"),
+        pytest.param(2048, 64, torch.float32, False, True, id="shape"),
+        pytest.param(6144, 128, torch.bfloat16, False, True, id="bf16-weight"),
+        pytest.param(6144, 128, torch.float32, True, True, id="bias"),
+    ],
+)
+def test_rocm_fp32_router_gemm_rejects_unsupported_configs(
+    monkeypatch,
+    input_size: int,
+    output_size: int,
+    params_dtype: torch.dtype,
+    bias: bool,
+    on_gfx950: bool,
+) -> None:
+    gate = _make_gate(
+        monkeypatch,
+        is_rocm=True,
+        params_dtype=params_dtype,
+        bias=bias,
+        input_size=input_size,
+        output_size=output_size,
+        on_gfx950=on_gfx950,
+    )
+    assert not gate.allow_fp32_router_gemm

tests/kernels/test_rocm_fp32_router_gemm.py

@@ -0,0 +1,228 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Correctness tests for the gfx950 low-M fp32 router GEMM."""
+
+import pytest
+import torch
+
+from vllm.model_executor.layers.fused_moe.router.gate_linear import (
+    fp32_router_gemm_dispatch_impl,
+)
+from vllm.model_executor.layers.fused_moe.router.rocm_fp32_router_gemm import (
+    rocm_fp32_router_gemm,
+)
+from vllm.platforms import current_platform
+
+SHAPES = [
+    (3072, 256),
+    (4096, 8),
+    (4096, 192),
+    (6144, 128),
+    (6144, 256),
+]
+MAX_TOKENS = 32
+ATOL = 5e-4
+RTOL = 0.0
+
+
+def _on_gfx950() -> bool:
+    if not current_platform.is_rocm():
+        return False
+    try:
+        from vllm.platforms.rocm import on_gfx950
+
+        return on_gfx950()
+    except Exception:
+        return False
+
+
+pytestmark = pytest.mark.skipif(
+    not _on_gfx950(), reason="rocm_fp32_router_gemm requires ROCm gfx950"
+)
+
+
+def _reference(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+    return torch.nn.functional.linear(x.float(), weight)
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
+@pytest.mark.parametrize(("hidden_size", "num_experts"), SHAPES)
+@pytest.mark.parametrize("num_tokens", range(MAX_TOKENS + 1))
+@torch.inference_mode()
+def test_rocm_fp32_router_gemm_matches_reference(
+    num_tokens: int,
+    hidden_size: int,
+    num_experts: int,
+    dtype: torch.dtype,
+) -> None:
+    torch.manual_seed(41 + num_tokens + hidden_size + num_experts)
+    device = torch.device("cuda")
+    x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
+    weight = torch.randn(num_experts, hidden_size, dtype=torch.float32, device=device)
+
+    output = rocm_fp32_router_gemm(x, weight)
+    expected = _reference(x, weight)
+
+    assert output.shape == (num_tokens, num_experts)
+    assert output.dtype == torch.float32
+    assert output.device == x.device
+    torch.testing.assert_close(output, expected, atol=ATOL, rtol=RTOL)
+
+
+@pytest.mark.parametrize("num_tokens", [1, 4, 16, 24, 32])
+@pytest.mark.parametrize(("hidden_size", "num_experts"), SHAPES)
+@torch.inference_mode()
+def test_rocm_fp32_router_gemm_preserves_topk(
+    num_tokens: int, hidden_size: int, num_experts: int
+) -> None:
+    torch.manual_seed(1000 + num_tokens + hidden_size + num_experts)
+    device = torch.device("cuda")
+    x = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device=device)
+    weight = torch.randn(num_experts, hidden_size, dtype=torch.float32, device=device)
+
+    top_k = 2 if (hidden_size, num_experts) == (4096, 8) else 8
+    actual = rocm_fp32_router_gemm(x, weight).topk(top_k, dim=-1).indices
+    reference = _reference(x, weight)
+    reference_values, expected = reference.topk(top_k, dim=-1)
+    for token in range(num_tokens):
+        actual_set = set(actual[token].tolist())
+        expected_set = set(expected[token].tolist())
+        if actual_set == expected_set:
+            continue
+        kth_value = reference_values[token, -1].item()
+        for expert in actual_set.symmetric_difference(expected_set):
+            gap = abs(reference[token, expert].item() - kth_value)
+            assert gap < 1e-3
+
+
+@torch.inference_mode()
+def test_rocm_fp32_router_gemm_rejects_invalid_inputs() -> None:
+    device = torch.device("cuda")
+    hidden_size, num_experts = 6144, 128
+    x = torch.randn(4, hidden_size, dtype=torch.bfloat16, device=device)
+    weight = torch.randn(num_experts, hidden_size, dtype=torch.float32, device=device)
+
+    with pytest.raises(ValueError, match="num_tokens"):
+        rocm_fp32_router_gemm(
+            torch.randn(
+                MAX_TOKENS + 1,
+                hidden_size,
+                dtype=torch.bfloat16,
+                device=device,
+            ),
+            weight,
+        )
+
+    with pytest.raises(ValueError, match="shape"):
+        rocm_fp32_router_gemm(x[:, :2048].contiguous(), weight[:, :2048].contiguous())
+
+    with pytest.raises(ValueError, match="contiguous"):
+        wide_x = torch.randn(4, hidden_size * 2, dtype=torch.bfloat16, device=device)
+        rocm_fp32_router_gemm(wide_x[:, ::2], weight)
+
+    with pytest.raises(ValueError, match="float32"):
+        rocm_fp32_router_gemm(x, weight.to(torch.bfloat16))
+
+    with pytest.raises(ValueError, match="bfloat16 or float32"):
+        rocm_fp32_router_gemm(x.to(torch.float16), weight)
+
+
+@torch.inference_mode()
+def test_rocm_fp32_router_gemm_cuda_graph_observes_input_mutation() -> None:
+    torch.manual_seed(2026)
+    device = torch.device("cuda")
+    hidden_size, num_experts = 6144, 128
+    static_x = torch.randn(16, hidden_size, dtype=torch.bfloat16, device=device)
+    weight = torch.randn(num_experts, hidden_size, dtype=torch.float32, device=device)
+
+    rocm_fp32_router_gemm(static_x, weight)
+    torch.accelerator.synchronize()
+
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured_output = rocm_fp32_router_gemm(static_x, weight)
+    torch.accelerator.synchronize()
+
+    first_input = torch.randn_like(static_x)
+    static_x.copy_(first_input)
+    graph.replay()
+    torch.accelerator.synchronize()
+    first_output = captured_output.clone()
+
+    second_input = torch.randn_like(static_x)
+    static_x.copy_(second_input)
+    graph.replay()
+    torch.accelerator.synchronize()
+    second_output = captured_output.clone()
+
+    torch.testing.assert_close(
+        first_output, _reference(first_input, weight), atol=ATOL, rtol=RTOL
+    )
+    torch.testing.assert_close(
+        second_output, _reference(second_input, weight), atol=ATOL, rtol=RTOL
+    )
+    assert not torch.equal(first_output, second_output)
+
+
+@pytest.mark.parametrize("num_tokens", [4, 32, 33])
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
+@torch.inference_mode()
+def test_rocm_fp32_router_gemm_custom_op_dispatch(
+    num_tokens: int, dtype: torch.dtype
+) -> None:
+    torch.manual_seed(3100 + num_tokens)
+    device = torch.device("cuda")
+    hidden_size, num_experts = 6144, 128
+    x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
+    weight = torch.randn(num_experts, hidden_size, dtype=torch.float32, device=device)
+
+    output = torch.ops.vllm.fp32_router_gemm_dispatch(x, weight, False)
+    direct_output = fp32_router_gemm_dispatch_impl(x, weight, False)
+
+    torch.testing.assert_close(output, direct_output, atol=ATOL, rtol=RTOL)
+    torch.testing.assert_close(output, _reference(x, weight), atol=ATOL, rtol=RTOL)
+
+
+@torch.inference_mode()
+def test_rocm_fp32_router_gemm_dispatch_accepts_noncontiguous_input(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    torch.manual_seed(4100)
+    device = torch.device("cuda")
+    hidden_size, num_experts = 6144, 128
+    wide_x = torch.randn(4, hidden_size * 2, dtype=torch.bfloat16, device=device)
+    x = wide_x[:, ::2]
+    weight = torch.randn(num_experts, hidden_size, dtype=torch.float32, device=device)
+    assert not x.is_contiguous()
+
+    expected = _reference(x, weight)
+
+    def fail_fallback(*args, **kwargs):
+        raise AssertionError("noncontiguous low-M input used the linear fallback")
+
+    monkeypatch.setattr(torch.nn.functional, "linear", fail_fallback)
+    output = torch.ops.vllm.fp32_router_gemm_dispatch(x, weight, False)
+
+    torch.testing.assert_close(output, expected, atol=ATOL, rtol=RTOL)
+
+
+@torch.inference_mode()
+def test_rocm_fp32_router_gemm_dynamic_compile_dispatch() -> None:
+    torch.manual_seed(5100)
+    device = torch.device("cuda")
+    hidden_size, num_experts = 6144, 128
+    weight = torch.randn(num_experts, hidden_size, dtype=torch.float32, device=device)
+
+    def dispatch(x: torch.Tensor) -> torch.Tensor:
+        return torch.ops.vllm.fp32_router_gemm_dispatch(x, weight, False)
+
+    compiled_dispatch = torch.compile(dispatch, dynamic=True, fullgraph=True)
+    for num_tokens in (4, 16, 33, 5, 32):
+        x = torch.randn(
+            num_tokens,
+            hidden_size,
+            dtype=torch.bfloat16,
+            device=device,
+        )
+        output = compiled_dispatch(x)
+        torch.testing.assert_close(output, _reference(x, weight), atol=ATOL, rtol=RTOL)

vllm/model_executor/layers/fused_moe/router/gate_linear.py

@@ -20,8 +20,8 @@ class GateLinear(ReplicatedLinear):
 
     1. cuteDSL ll_bf16_gemm (SM90+, M<=16, bf16 in, fp32 out,
        K divisible by 8)
-    2. fp32 specialized kernel  (SM90+, bf16/fp32 in, fp32 out, M<=32,
-       (H, E) in {(3072, 256), (6144, 128)})
+    2. fp32 specialized kernel (SM90+ or gfx950, bf16/fp32 in, fp32 out,
+       M<=32, model-specific shapes)
     3. experimental bf16x3 CuteDSL kernel (opt-in, SM100, bf16 in, fp32 weight)
     4. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 weight + fp32 out_dtype)
     5. F.linear via ReplicatedLinear (ultimate fallback)
@@ -48,6 +48,21 @@ class GateLinear(ReplicatedLinear):
     ):
         is_hopper = current_platform.is_device_capability((9, 0))
         is_blackwell = current_platform.is_device_capability_family(100)
+        is_gfx950 = False
+        if current_platform.is_rocm():
+            from vllm.platforms.rocm import on_gfx950
+
+            is_gfx950 = on_gfx950()
+        is_rocm_fp32_shape = False
+        if is_gfx950:
+            from vllm.model_executor.layers.fused_moe.router.rocm_fp32_router_gemm import (  # noqa: E501
+                ROCM_FP32_ROUTER_GEMM_SUPPORTED_SHAPES,
+            )
+
+            is_rocm_fp32_shape = (
+                input_size,
+                output_size,
+            ) in ROCM_FP32_ROUTER_GEMM_SUPPORTED_SHAPES
         can_use_specialized_kernels = (
             current_platform.is_cuda() and (is_hopper or is_blackwell) and not bias
         )
@@ -69,7 +84,7 @@ class GateLinear(ReplicatedLinear):
 
         self.allow_specialized_router_gemm = can_use_specialized_kernels
 
-        # fp32 specialized kernel eligibility (SM90+, exact dims, fp32 weight)
+        # fp32 specialized kernel eligibility (exact dims, fp32 weight)
         vllm_config = get_current_vllm_config_or_none()
         enable_bf16x3_router_gemm = (
             vllm_config is not None
@@ -78,9 +93,14 @@ class GateLinear(ReplicatedLinear):
         self.allow_fp32_router_gemm = (
             not bias
             and self.weight.dtype == torch.float32
-            and current_platform.is_cuda()
-            and (is_hopper or is_blackwell)
-            and (input_size, output_size) in self.FP32_SUPPORTED_SHAPES
+            and (
+                (
+                    current_platform.is_cuda()
+                    and (is_hopper or is_blackwell)
+                    and (input_size, output_size) in self.FP32_SUPPORTED_SHAPES
+                )
+                or (is_gfx950 and is_rocm_fp32_shape)
+            )
         )
         self.allow_bf16x3_router_gemm = (
             not bias
@@ -166,7 +186,7 @@ class GateLinear(ReplicatedLinear):
             output = ll_bf16_gemm(x, self.weight)
             return output, None
 
-        # Tier 2: fp32 specialized kernel (H=3072, E=256, M<=32)
+        # 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.
         if self.allow_fp32_router_gemm and x.dtype in (
@@ -217,6 +237,16 @@ def fp32_router_gemm_dispatch_impl(
     does not support runtime dispatching on num_tokens.
     """
     if x.shape[0] <= _FP32_ROUTER_GEMM_MAX_TOKENS:
+        if current_platform.is_rocm():
+            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)
         return ops.fp32_router_gemm(x, weight)
 
     if allow_bf16x3_router_gemm and x.dtype == torch.bfloat16:

vllm/model_executor/layers/fused_moe/router/rocm_fp32_router_gemm.py

@@ -0,0 +1,148 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Low-token FP32 router GEMM for gfx950."""
+
+import torch
+
+from vllm.platforms import current_platform
+from vllm.platforms.rocm import on_gfx950
+from vllm.triton_utils import tl, triton
+
+_MAX_TOKENS = 32
+ROCM_FP32_ROUTER_GEMM_SUPPORTED_SHAPES = frozenset(
+    {
+        (3072, 256),
+        (4096, 8),
+        (4096, 192),
+        (6144, 128),
+        (6144, 256),
+    }
+)
+
+
+@triton.jit
+def _rocm_fp32_router_gemm_kernel(
+    hidden_states_ptr,
+    router_weight_ptr,
+    output_ptr,
+    M,
+    K: tl.constexpr,
+    N: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    pid_m = pid // N
+    pid_n = pid - pid_m * N
+    offsets_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offsets_k = tl.arange(0, BLOCK_K)
+    partials = tl.zeros((BLOCK_M, BLOCK_K), dtype=tl.float32)
+
+    for k_start in range(0, K, BLOCK_K):
+        offsets_k_block = k_start + offsets_k
+        hidden_states = tl.load(
+            hidden_states_ptr + offsets_m[:, None] * K + offsets_k_block[None, :],
+            mask=offsets_m[:, None] < M,
+            other=0.0,
+        ).to(tl.float32)
+        router_weight = tl.load(router_weight_ptr + pid_n * K + offsets_k_block)
+        partials += hidden_states * router_weight[None, :]
+
+    accumulator = tl.sum(partials, axis=1)
+    tl.store(
+        output_ptr + offsets_m * N + pid_n,
+        accumulator,
+        mask=offsets_m < M,
+    )
+
+
+def _launch_config(
+    hidden_size: int,
+    num_experts: int,
+    num_tokens: int,
+) -> tuple[int, int, int]:
+    if (hidden_size, num_experts) == (4096, 8):
+        return 1, 1024, 4
+
+    if num_tokens <= 4:
+        return 1, 1024, 8
+    if num_tokens <= 16:
+        block_m = 8 if (hidden_size, num_experts) == (6144, 256) else 4
+        return block_m, 1024, 8
+    if (hidden_size, num_experts) == (3072, 256):
+        return 8, 1024, 8
+    return 4, 2048, 8
+
+
+def can_use_rocm_fp32_router_gemm(
+    hidden_states: torch.Tensor,
+    router_weight: torch.Tensor,
+) -> bool:
+    """Return whether the tensors match the tuned gfx950 fast path."""
+    try:
+        _validate_inputs(hidden_states, router_weight)
+    except (RuntimeError, ValueError):
+        return False
+    return True
+
+
+def _validate_inputs(
+    hidden_states: torch.Tensor,
+    router_weight: torch.Tensor,
+) -> None:
+    if not current_platform.is_rocm() or not on_gfx950():
+        raise RuntimeError("rocm_fp32_router_gemm requires ROCm gfx950")
+    if hidden_states.dim() != 2 or router_weight.dim() != 2:
+        raise ValueError("hidden_states and router_weight must be 2D tensors")
+    if hidden_states.dtype not in (torch.bfloat16, torch.float32):
+        raise ValueError("hidden_states must have dtype bfloat16 or float32")
+    if router_weight.dtype != torch.float32:
+        raise ValueError("router_weight must have dtype float32")
+    if hidden_states.device.type != "cuda" or router_weight.device.type != "cuda":
+        raise ValueError("hidden_states and router_weight must be GPU tensors")
+    if hidden_states.device != router_weight.device:
+        raise ValueError("hidden_states and router_weight must be on the same device")
+    if not hidden_states.is_contiguous() or not router_weight.is_contiguous():
+        raise ValueError("hidden_states and router_weight must be contiguous")
+    shape = (hidden_states.shape[1], router_weight.shape[0])
+    if (
+        shape not in ROCM_FP32_ROUTER_GEMM_SUPPORTED_SHAPES
+        or router_weight.shape[1] != shape[0]
+    ):
+        raise ValueError(
+            "supported (hidden_size, num_experts) shape pairs are "
+            "(3072, 256), (4096, 8), (4096, 192), (6144, 128), "
+            "and (6144, 256)"
+        )
+    if not 0 <= hidden_states.shape[0] <= _MAX_TOKENS:
+        raise ValueError(f"num_tokens must be in [0, {_MAX_TOKENS}]")
+
+
+def rocm_fp32_router_gemm(
+    hidden_states: torch.Tensor,
+    router_weight: torch.Tensor,
+) -> torch.Tensor:
+    """Compute ``hidden_states @ router_weight.T`` with FP32 accumulation."""
+    _validate_inputs(hidden_states, router_weight)
+    num_tokens = hidden_states.shape[0]
+    hidden_size = hidden_states.shape[1]
+    num_experts = router_weight.shape[0]
+    output = hidden_states.new_empty((num_tokens, num_experts), dtype=torch.float32)
+    if num_tokens == 0:
+        return output
+
+    block_m, block_k, num_warps = _launch_config(hidden_size, num_experts, num_tokens)
+    grid = (triton.cdiv(num_tokens, block_m) * num_experts,)
+    _rocm_fp32_router_gemm_kernel[grid](
+        hidden_states,
+        router_weight,
+        output,
+        M=num_tokens,
+        K=hidden_size,
+        N=num_experts,
+        BLOCK_M=block_m,
+        BLOCK_K=block_k,
+        num_warps=num_warps,
+        num_stages=1,
+    )
+    return output