iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[ROCm][Perf] Any-shape low-latency MoE router GEMM on MI355 (#61895)

Browse Source 
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Co-authored-by: Claude <noreply@anthropic.com>
This commit is contained in:
Andreas Karatzas
2026-10-15 23:22:31 -05:00
committed by GitHub
co-authored by Claude
parent 642c18d64c
commit b466fbacd2
6 changed files with 296 additions and 75 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.buildkite/test-amd.yaml
+18 -2
View File
@@ -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]
.buildkite/test_areas/kernels.yaml
+18 -3
View File
@@ -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
tests/kernels/test_gate_linear_rocm_dispatch.py
+1 -2
View File
@@ -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"),
],
tests/kernels/test_rocm_fp32_router_gemm.py
+171 -5
View File
@@ -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
)
vllm/model_executor/layers/fused_moe/router/gate_linear.py
+20 -22
View File
@@ -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:
vllm/model_executor/layers/fused_moe/router/rocm_fp32_router_gemm.py
+68 -41
View File
@@ -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,
)