From b466fbacd2ff68f0ed46a47f2551455ccbe76cef Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Thu, 15 Oct 2026 23:22:31 -0500 Subject: [PATCH] [ROCm][Perf] Any-shape low-latency MoE router GEMM on MI355 (#61895) Signed-off-by: Andreas Karatzas Co-authored-by: Claude --- .buildkite/test-amd.yaml | 20 +- .buildkite/test_areas/kernels.yaml | 21 ++- .../kernels/test_gate_linear_rocm_dispatch.py | 3 +- tests/kernels/test_rocm_fp32_router_gemm.py | 176 +++++++++++++++++- .../layers/fused_moe/router/gate_linear.py | 42 ++--- .../fused_moe/router/rocm_fp32_router_gemm.py | 109 +++++++---- 6 files changed, 296 insertions(+), 75 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 56cba27ceb..7a9876032d 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -4704,16 +4704,32 @@ steps: kernels/test_fused_inv_rope_fp8_quant.py kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py kernels/test_fused_quant_activation.py - kernels/test_gate_linear_rocm_dispatch.py kernels/test_gemm_rs_ar.py kernels/test_kpool_decode_update_batched.py kernels/test_minimax_m3_amd_ops.py kernels/test_minimax_m3_sparse_attn_fp8_scale.py kernels/test_onednn.py kernels/test_relu2_fp8_quant.py - kernels/test_rocm_fp32_router_gemm.py --deselect=tests/kernels/test_bf16_skinny_gemm.py::test_every_dsv3_routed_shape_is_instantiated +- label: ":amd: (MI355) Gate Linear Router Kernels" + timeout_in_minutes: 20 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_1 + optional: true + working_dir: "/vllm-workspace/" + source_file_dependencies: + - 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/fused_moe/router/bf16x3_router_gemm_rocm.py + - vllm/platforms/rocm.py + - tests/kernels/test_rocm_fp32_router_gemm.py + - tests/kernels/test_gate_linear_rocm_dispatch.py + - tests/kernels/test_bf16x3_router_gemm_rocm.py + commands: + - pytest -v -s tests/kernels/test_rocm_fp32_router_gemm.py tests/kernels/test_gate_linear_rocm_dispatch.py tests/kernels/test_bf16x3_router_gemm_rocm.py + - label: ":amd: (MI355 DPX) Attention DiffKV Kernels" timeout_in_minutes: 35 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 1c1633c06e..12b8516926 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -290,7 +290,6 @@ 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 @@ -305,14 +304,12 @@ steps: kernels/test_fused_inv_rope_fp8_quant.py kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py kernels/test_fused_quant_activation.py - kernels/test_gate_linear_rocm_dispatch.py kernels/test_gemm_rs_ar.py kernels/test_kpool_decode_update_batched.py kernels/test_minimax_m3_amd_ops.py kernels/test_minimax_m3_sparse_attn_fp8_scale.py kernels/test_onednn.py kernels/test_relu2_fp8_quant.py - kernels/test_rocm_fp32_router_gemm.py --deselect=tests/kernels/test_bf16_skinny_gemm.py::test_every_dsv3_routed_shape_is_instantiated - label: ":nvidia: (L4) Attention Kernels Shard %N" @@ -725,6 +722,24 @@ steps: - nvidia-smi - pytest -v -s tests/kernels/test_ll_bf16_gemm.py - pytest -v -s tests/kernels/test_ll_fp32w_gemm.py + mirror: + amd: + label: ":amd: (MI355) Gate Linear Router Kernels" + dind: false + device: mi355_1 + timeout_in_minutes: 20 + depends_on: + - image-build-amd + source_file_dependencies: + - 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/fused_moe/router/bf16x3_router_gemm_rocm.py + - vllm/platforms/rocm.py + - tests/kernels/test_rocm_fp32_router_gemm.py + - tests/kernels/test_gate_linear_rocm_dispatch.py + - tests/kernels/test_bf16x3_router_gemm_rocm.py + commands: + - pytest -v -s tests/kernels/test_rocm_fp32_router_gemm.py tests/kernels/test_gate_linear_rocm_dispatch.py tests/kernels/test_bf16x3_router_gemm_rocm.py - label: ":nvidia: (DGX) Spark b12x Linear Kernels Nightly" key: b12x-linear-kernels-dgx-spark-nightly diff --git a/tests/kernels/test_gate_linear_rocm_dispatch.py b/tests/kernels/test_gate_linear_rocm_dispatch.py index 06d69fc985..22fc355415 100644 --- a/tests/kernels/test_gate_linear_rocm_dispatch.py +++ b/tests/kernels/test_gate_linear_rocm_dispatch.py @@ -109,7 +109,7 @@ def test_rocm_set_out_dtype_respects_bias_guard(monkeypatch): @pytest.mark.parametrize( ("input_size", "output_size"), - [(3072, 256), (4096, 8), (4096, 192), (6144, 128), (6144, 256)], + [(3072, 256), (4096, 8), (6144, 128), (2048, 64), (2560, 512), (7168, 384)], ) def test_rocm_gfx950_enables_fp32_router_gemm( monkeypatch, input_size: int, output_size: int @@ -147,7 +147,6 @@ def test_rocm_fp32_router_gemm_is_replicated_across_ep_sizes( ("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"), ], diff --git a/tests/kernels/test_rocm_fp32_router_gemm.py b/tests/kernels/test_rocm_fp32_router_gemm.py index 899045d39c..efa6e6aeb5 100644 --- a/tests/kernels/test_rocm_fp32_router_gemm.py +++ b/tests/kernels/test_rocm_fp32_router_gemm.py @@ -5,6 +5,9 @@ import pytest import torch +from vllm.model_executor.layers.fused_moe.router import ( + rocm_fp32_router_gemm as rocm_router, +) from vllm.model_executor.layers.fused_moe.router.gate_linear import ( fp32_router_gemm_dispatch_impl, ) @@ -19,6 +22,11 @@ SHAPES = [ (4096, 192), (6144, 128), (6144, 256), + (2560, 512), + (2880, 128), + (5120, 16), + (7168, 384), + (1000, 100), ] MAX_TOKENS = 32 ATOL = 5e-4 @@ -80,7 +88,7 @@ def test_rocm_fp32_router_gemm_preserves_topk( 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 + top_k = 2 if num_experts <= 16 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) @@ -113,15 +121,22 @@ def test_rocm_fp32_router_gemm_rejects_invalid_inputs() -> None: weight, ) - with pytest.raises(ValueError, match="shape"): - rocm_fp32_router_gemm(x[:, :2048].contiguous(), weight[:, :2048].contiguous()) + with pytest.raises(ValueError, match="hidden sizes"): + rocm_fp32_router_gemm(x, weight[:, :2048].contiguous()) + + with pytest.raises(ValueError, match="num_experts"): + big_weight = torch.empty(1024, 7168, dtype=torch.float32, device=device) + rocm_fp32_router_gemm( + torch.empty(MAX_TOKENS, 7168, dtype=torch.bfloat16, device=device), + big_weight, + ) 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="float32 or bfloat16"): + rocm_fp32_router_gemm(x, weight.to(torch.float16)) with pytest.raises(ValueError, match="bfloat16 or float32"): rocm_fp32_router_gemm(x.to(torch.float16), weight) @@ -226,3 +241,154 @@ def test_rocm_fp32_router_gemm_dynamic_compile_dispatch() -> None: ) output = compiled_dispatch(x) torch.testing.assert_close(output, _reference(x, weight), atol=ATOL, rtol=RTOL) + + +@pytest.mark.parametrize("hidden_size", [64, 100, 513, 1000, 2049, 3000]) +@pytest.mark.parametrize("num_tokens", [1, 5, 32]) +@torch.inference_mode() +def test_rocm_fp32_router_gemm_any_hidden_size(num_tokens: int, hidden_size: int): + """Hidden sizes that do not divide the K tile exercise the masked tail.""" + torch.manual_seed(6100 + num_tokens + hidden_size) + x = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device="cuda") + weight = torch.randn(64, hidden_size, dtype=torch.float32, device="cuda") + torch.testing.assert_close( + rocm_fp32_router_gemm(x, weight), _reference(x, weight), atol=ATOL, rtol=RTOL + ) + + +@pytest.mark.parametrize("case", ["large", "near_zero", "zero_input", "ones"]) +@torch.inference_mode() +def test_rocm_fp32_router_gemm_special_values(case: str) -> None: + torch.manual_seed(7100) + x = torch.randn(4, 5120, dtype=torch.bfloat16, device="cuda") + weight = torch.randn(128, 5120, dtype=torch.float32, device="cuda") + if case in ("large", "near_zero"): + scale = 100.0 if case == "large" else 1e-4 + x = (x * scale).to(torch.bfloat16) + weight *= scale + elif case == "zero_input": + x.zero_() + else: + x.fill_(1.0) + output = rocm_fp32_router_gemm(x, weight) + expected = _reference(x, weight) + assert torch.isfinite(output).all() + if case == "zero_input": + assert torch.equal(output, torch.zeros_like(output)) + else: + assert ((output - expected).norm() / expected.norm()).item() < 1e-5 + + +@pytest.mark.parametrize("num_tokens", [1, 4, 16, 32]) +@torch.inference_mode() +def test_rocm_fp32_router_gemm_deterministic(num_tokens: int) -> None: + torch.manual_seed(8100 + num_tokens) + x = torch.randn(num_tokens, 7168, dtype=torch.bfloat16, device="cuda") + weight = torch.randn(384, 7168, dtype=torch.float32, device="cuda") + first = rocm_fp32_router_gemm(x, weight) + assert torch.equal(first, rocm_fp32_router_gemm(x, weight)) + + +@pytest.mark.parametrize("num_tokens", [1, 4, 16, 32]) +@pytest.mark.parametrize(("hidden_size", "num_experts"), SHAPES) +@torch.inference_mode() +def test_rocm_fp32_router_gemm_bf16_weight( + num_tokens: int, hidden_size: int, num_experts: int +) -> None: + torch.manual_seed(9100 + num_tokens + hidden_size + num_experts) + x = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device="cuda") + weight = torch.randn(num_experts, hidden_size, dtype=torch.bfloat16, device="cuda") + # Past this size hipBLASLt beats the kernel for BF16 weights on MI355. + eligible = num_tokens * num_experts <= 8192 + assert rocm_router.can_use_rocm_fp32_router_gemm(x, weight) == eligible + if eligible: + torch.testing.assert_close( + rocm_fp32_router_gemm(x, weight), + _reference(x.float(), weight.float()), + atol=ATOL, + rtol=RTOL, + ) + + +def _make_gate_linear( + monkeypatch: pytest.MonkeyPatch, + input_size, + output_size, + params_dtype=torch.float32, +): + for module in ("layers.linear", "parameter"): + monkeypatch.setattr( + f"vllm.model_executor.{module}.get_tensor_model_parallel_rank", lambda: 0 + ) + monkeypatch.setattr( + f"vllm.model_executor.{module}.get_tensor_model_parallel_world_size", + lambda: 1, + ) + from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear + + gate = GateLinear( + input_size=input_size, + output_size=output_size, + bias=False, + out_dtype=torch.float32, + params_dtype=params_dtype, + ).cuda() + with torch.no_grad(): + gate.weight.normal_() + gate.quant_method.process_weights_after_loading(gate) + return gate + + +@pytest.mark.parametrize("x_dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("num_tokens", [1, 32, 33]) +@torch.inference_mode() +def test_gate_linear_uses_rocm_fp32_router_gemm( + monkeypatch: pytest.MonkeyPatch, num_tokens: int, x_dtype: torch.dtype +) -> None: + """A router shape outside the old tuned list takes the fast path for M<=32.""" + gate = _make_gate_linear(monkeypatch, 2560, 512) + assert gate.allow_fp32_router_gemm + calls = [] + kernel = rocm_router.rocm_fp32_router_gemm + + def spy(hidden_states, router_weight): + calls.append(hidden_states.shape[0]) + return kernel(hidden_states, router_weight) + + monkeypatch.setattr(rocm_router, "rocm_fp32_router_gemm", spy) + x = torch.randn(num_tokens, 2560, dtype=x_dtype, device="cuda") + output, bias = gate(x) + + assert bias is None + assert calls == ([num_tokens] if num_tokens <= MAX_TOKENS else []) + torch.testing.assert_close(output, _reference(x, gate.weight), atol=ATOL, rtol=RTOL) + + +@pytest.mark.parametrize( + ("num_tokens", "num_experts", "uses_kernel"), + [(1, 256, True), (32, 256, True), (33, 256, False), (32, 512, False)], +) +@torch.inference_mode() +def test_gate_linear_bf16_weight_uses_rocm_router_gemm( + monkeypatch: pytest.MonkeyPatch, num_tokens, num_experts, uses_kernel +) -> None: + """BF16-weight gates with FP32 logits use the kernel only where it beats + hipBLASLt; everything else keeps the torch.mm path.""" + gate = _make_gate_linear(monkeypatch, 7168, num_experts, torch.bfloat16) + assert gate.allow_cublas_router_gemm + calls = [] + kernel = rocm_router.rocm_fp32_router_gemm + + def spy(hidden_states, router_weight): + calls.append(hidden_states.shape[0]) + return kernel(hidden_states, router_weight) + + monkeypatch.setattr(rocm_router, "rocm_fp32_router_gemm", spy) + x = torch.randn(num_tokens, 7168, dtype=torch.bfloat16, device="cuda") + output, _ = gate(x) + + assert output.dtype == torch.float32 + assert calls == ([num_tokens] if uses_kernel else []) + torch.testing.assert_close( + output, _reference(x.float(), gate.weight.float()), 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 8ecd244b61..c6b1da890b 100644 --- a/vllm/model_executor/layers/fused_moe/router/gate_linear.py +++ b/vllm/model_executor/layers/fused_moe/router/gate_linear.py @@ -24,14 +24,15 @@ class GateLinear(ReplicatedLinear): 1. cuteDSL ll_bf16_gemm (SM90+, M<=16, bf16 in, fp32 out, any shape w/ K divisible by 8) - 2. fp32 specialized kernel (SM90+ or gfx950, bf16/fp32 in, fp32 out, - M<=32, model-specific shapes) + 2. fp32 specialized kernel (bf16/fp32 in, fp32 out, M<=32; SM90+ + model-specific shapes, gfx950 any shape) 3. cuteDSL ll_fp32w_gemm (SM90+, M<=32, bf16/fp32 in, fp32 out, any shape) 4. bf16x3 CuteDSL kernel (SM100, bf16 in, fp32 weight) 5. ROCm bf16x3 router GEMM (gfx950, bf16 in, fp32 weight, fp32 out, M>=2048) - 6. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 weight + fp32 out_dtype) + 6. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 weight + fp32 out_dtype); gfx950 + routes low M through the tier 2 kernel 7. F.linear via ReplicatedLinear (ultimate fallback) The ``out_dtype`` attribute is mutable and can be set after init @@ -73,16 +74,6 @@ class GateLinear(ReplicatedLinear): 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 batch_invariant = envs.VLLM_BATCH_INVARIANT can_use_specialized_kernels = ( current_platform.is_cuda() @@ -107,6 +98,7 @@ class GateLinear(ReplicatedLinear): return_bias=return_bias, ) self.out_dtype = out_dtype + self.is_gfx950 = is_gfx950 # A quantized gate exposes no plain ``weight``, so every specialized # tier below is disabled and forward falls back to ReplicatedLinear. @@ -127,7 +119,7 @@ class GateLinear(ReplicatedLinear): and (is_hopper or is_blackwell) and (input_size, output_size) in self.FP32_SUPPORTED_SHAPES ) - or (is_gfx950 and is_rocm_fp32_shape) + or is_gfx950 ) ) self.allow_bf16x3_router_gemm = ( @@ -349,7 +341,10 @@ class GateLinear(ReplicatedLinear): # Tier 6: 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) + if self.is_gfx950: + output = torch.ops.vllm.fp32_router_gemm_dispatch(x, self.weight, False) + else: + output = torch.mm(x, self.weight.T, out_dtype=torch.float32) return self._return(output) # Tier 7: F.linear (ReplicatedLinear) @@ -465,17 +460,20 @@ def fp32_router_gemm_dispatch_impl( This must be wrapped in a custom op because our torch.compile integration 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, - ) + 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, + ) + if x.shape[0] <= _FP32_ROUTER_GEMM_MAX_TOKENS: 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) + if weight.dtype == torch.bfloat16: + return torch.mm(x, weight.T, out_dtype=torch.float32) + return torch.nn.functional.linear(x.float(), weight) + if x.shape[0] <= _FP32_ROUTER_GEMM_MAX_TOKENS: 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 index 785bd8a72c..f76318ae66 100644 --- 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 @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Low-token FP32 router GEMM for gfx950.""" +"""Low-token router GEMM with FP32 output for gfx950 (FP32 or BF16 weights).""" import torch @@ -9,15 +9,10 @@ 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), - } -) +# From this many multiply-adds on, F.linear wins on MI355. +_MAX_MACS = 1 << 27 +# Above this num_tokens * num_experts, hipBLASLt wins for BF16 weights. +_MAX_BF16_TOKEN_EXPERTS = 8192 @triton.jit @@ -30,6 +25,7 @@ def _rocm_fp32_router_gemm_kernel( N: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_K: tl.constexpr, + EVEN_K: tl.constexpr, ): pid = tl.program_id(0) pid_m = pid // N @@ -40,12 +36,27 @@ def _rocm_fp32_router_gemm_kernel( 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) + if EVEN_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).to( + tl.float32 + ) + else: + k_mask = offsets_k_block < K + hidden_states = tl.load( + hidden_states_ptr + offsets_m[:, None] * K + offsets_k_block[None, :], + mask=(offsets_m[:, None] < M) & k_mask[None, :], + other=0.0, + ).to(tl.float32) + router_weight = tl.load( + router_weight_ptr + pid_n * K + offsets_k_block, + mask=k_mask, + other=0.0, + ).to(tl.float32) partials += hidden_states * router_weight[None, :] accumulator = tl.sum(partials, axis=1) @@ -57,28 +68,36 @@ def _rocm_fp32_router_gemm_kernel( def _launch_config( - hidden_size: int, - num_experts: int, - num_tokens: int, + num_experts: int, num_tokens: int, bf16_weight: bool ) -> tuple[int, int, int]: - if (hidden_size, num_experts) == (4096, 8): - return 1, 1024, 4 + """Return (BLOCK_M, BLOCK_K, num_warps), swept on MI355 for M<=32. + One program computes BLOCK_M tokens for one expert, so narrow routers keep + BLOCK_M=1 to fill the GPU and wide routers grow BLOCK_M to reuse weights. + """ + if bf16_weight: + return (1 if num_tokens <= 4 else 2), 1024, 2 + if num_experts < 64: + return 1, 1024, 4 + wide = num_experts > 256 + if num_tokens <= 2: + return 1, 512, 2 if num_tokens <= 4: - return 1, 1024, 8 + return (2, 1024, 2) if wide else (2, 2048, 4) + if num_tokens <= 8: + return (4, 512, 2) if wide else (4, 1024, 4) 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 + return (8, 512, 2) if wide else (4, 1024, 4) + if num_experts >= 768: + return 16, 256, 2 + return (8, 256, 2) if wide else (4, 512, 2) 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.""" + """Return whether the tensors can use the gfx950 fast path.""" try: _validate_inputs(hidden_states, router_weight) except (RuntimeError, ValueError): @@ -96,26 +115,30 @@ def _validate_inputs( 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 router_weight.dtype not in (torch.float32, torch.bfloat16): + raise ValueError("router_weight must have dtype float32 or bfloat16") 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 router_weight.shape[1] != hidden_states.shape[1]: + raise ValueError("router_weight and hidden_states hidden sizes differ") if not 0 <= hidden_states.shape[0] <= _MAX_TOKENS: raise ValueError(f"num_tokens must be in [0, {_MAX_TOKENS}]") + if hidden_states.shape[0] * router_weight.numel() >= _MAX_MACS: + raise ValueError( + f"num_tokens * num_experts * hidden_size must be < {_MAX_MACS}" + ) + if ( + router_weight.dtype == torch.bfloat16 + and hidden_states.shape[0] * router_weight.shape[0] > _MAX_BF16_TOKEN_EXPERTS + ): + raise ValueError( + f"num_tokens * num_experts must be <= {_MAX_BF16_TOKEN_EXPERTS} " + "for bfloat16 weights" + ) def rocm_fp32_router_gemm( @@ -131,7 +154,10 @@ def rocm_fp32_router_gemm( if num_tokens == 0: return output - block_m, block_k, num_warps = _launch_config(hidden_size, num_experts, num_tokens) + block_m, block_k, num_warps = _launch_config( + num_experts, num_tokens, router_weight.dtype == torch.bfloat16 + ) + block_k = min(block_k, triton.next_power_of_2(hidden_size)) grid = (triton.cdiv(num_tokens, block_m) * num_experts,) _rocm_fp32_router_gemm_kernel[grid]( hidden_states, @@ -142,6 +168,7 @@ def rocm_fp32_router_gemm( N=num_experts, BLOCK_M=block_m, BLOCK_K=block_k, + EVEN_K=hidden_size % block_k == 0, num_warps=num_warps, num_stages=1, )