From 8bf39632b86df21fa9bfbb470ce8305bf67a5838 Mon Sep 17 00:00:00 2001 From: Fangzhou Ai <31551580+Fangzhou-Ai@users.noreply.github.com> Date: Thu, 3 Sep 2026 13:42:25 -0400 Subject: [PATCH] [ROCm][Perf] Add low-M FP32 router GEMM for gfx950 (#54845) --- .buildkite/test-amd.yaml | 5 +- .../kernels/test_gate_linear_rocm_dispatch.py | 92 ++++++- tests/kernels/test_rocm_fp32_router_gemm.py | 228 ++++++++++++++++++ .../layers/fused_moe/router/gate_linear.py | 44 +++- .../fused_moe/router/rocm_fp32_router_gemm.py | 148 ++++++++++++ 5 files changed, 499 insertions(+), 18 deletions(-) create mode 100644 tests/kernels/test_rocm_fp32_router_gemm.py create mode 100644 vllm/model_executor/layers/fused_moe/router/rocm_fp32_router_gemm.py diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 8c58298914..c1f429aaf1 100644 --- a/.buildkite/test-amd.yaml +++ b/.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 diff --git a/tests/kernels/test_gate_linear_rocm_dispatch.py b/tests/kernels/test_gate_linear_rocm_dispatch.py index 865b58e18e..afa2e57f6a 100644 --- a/tests/kernels/test_gate_linear_rocm_dispatch.py +++ b/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 diff --git a/tests/kernels/test_rocm_fp32_router_gemm.py b/tests/kernels/test_rocm_fp32_router_gemm.py new file mode 100644 index 0000000000..899045d39c --- /dev/null +++ b/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) 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 0e4a402599..0f45aaf6fb 100644 --- a/vllm/model_executor/layers/fused_moe/router/gate_linear.py +++ b/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: diff --git a/vllm/model_executor/layers/fused_moe/router/rocm_fp32_router_gemm.py b/vllm/model_executor/layers/fused_moe/router/rocm_fp32_router_gemm.py new file mode 100644 index 0000000000..785bd8a72c --- /dev/null +++ b/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