diff --git a/tests/kernels/test_rocm_aiter_router_gemm.py b/tests/kernels/test_rocm_aiter_router_gemm.py new file mode 100644 index 000000000000..26a9833a8995 --- /dev/null +++ b/tests/kernels/test_rocm_aiter_router_gemm.py @@ -0,0 +1,143 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the ROCm AITER router-gate GEMM: bf16 x bf16 -> out_dtype. + +This is the ``GateLinear`` tier used on gfx950, where the MoE router gate is a +skinny GEMM (M=num_tokens, N=num_experts, K=hidden_size). Correctness baseline +is a float64 matmul; what ultimately matters is that numeric error never flips +the top-k expert selection by more than the kernel's own error. + +The tier runs AITER's tuned bf16 kernel and casts, so the result carries bf16 +output precision even when ``out_dtype`` is fp32, and the tuned kernels for +these shapes reduce over split-K, which contributes more error than that +rounding does. Either way the absolute error scales with the output magnitude, +which itself grows as sqrt(hidden_size) for unit-variance inputs. Tolerances +below are expressed as a fraction of the largest reference logit so they hold +across all three shapes rather than being tuned per shape. +""" + +import pytest +import torch + +from vllm._aiter_ops import rocm_aiter_ops +from vllm.platforms import current_platform + +# Register torch.ops.vllm.rocm_aiter_router_gemm. +import vllm.model_executor.layers.fused_moe.router.gate_linear # noqa: F401 isort: skip + +# (hidden_size, num_experts): GLM-5/5.2, DeepSeek-V3 and Kimi-K2 routers. +SHAPES = [(6144, 256), (7168, 256), (7168, 384)] +NUM_TOKENS = [1, 2, 4, 8, 16, 32, 64, 128] + +# Error as a fraction of the largest reference logit. Measured worst case for +# the tuned bf16 kernel across every shape/token count here is 8.8e-3; keep +# ~2x headroom for kernel-selection changes between AITER versions. +REL_TO_PEAK = 2e-2 +# A top-k membership change needs error on both the promoted and demoted logit, +# so the tie window is twice the single-value tolerance. +TIE_REL_TO_PEAK = 2 * REL_TO_PEAK + + +def _requires_aiter_tgemm(): + if not current_platform.is_rocm(): + pytest.skip("AITER router GEMM requires ROCm") + if not rocm_aiter_ops.is_tgemm_enabled(): + pytest.skip("AITER tuned GEMM not enabled (needs AITER linear + gfx950)") + + +def _run(x: torch.Tensor, weight: torch.Tensor, out_dtype: torch.dtype): + return torch.ops.vllm.rocm_aiter_router_gemm(x, weight, out_dtype) + + +def _assert_close_to_peak(out: torch.Tensor, ref: torch.Tensor, tol: float): + """Compare after normalizing by the peak reference magnitude. + + ``assert_close``'s rtol is per-element, which is meaningless here: a logit + that lands near zero is a cancelling sum of ~K large products, so its own + magnitude says nothing about the error it accumulated. + """ + scale = ref.abs().max().clamp(min=torch.finfo(torch.float32).tiny) + torch.testing.assert_close( + out.float() / scale, ref.float() / scale, atol=tol, rtol=0 + ) + + +@pytest.mark.parametrize("hidden_dim,num_experts", SHAPES) +@pytest.mark.parametrize("num_tokens", NUM_TOKENS) +@pytest.mark.parametrize("out_dtype", [torch.float32, torch.bfloat16]) +def test_matches_reference( + num_tokens: int, hidden_dim: int, num_experts: int, out_dtype: torch.dtype +): + """bf16 activation x bf16 weight should track a float64 reference.""" + _requires_aiter_tgemm() + torch.manual_seed(42) + device = torch.device("cuda") + x = torch.randn(num_tokens, hidden_dim, dtype=torch.bfloat16, device=device) + weight = torch.randn(num_experts, hidden_dim, dtype=torch.bfloat16, device=device) + + out = _run(x, weight, out_dtype) + ref = x.double() @ weight.double().t() + + assert out.shape == (num_tokens, num_experts) + assert out.dtype == out_dtype + _assert_close_to_peak(out, ref, REL_TO_PEAK) + + +@pytest.mark.parametrize("hidden_dim,num_experts", SHAPES) +@pytest.mark.parametrize("num_tokens", NUM_TOKENS) +def test_topk_routing_consistency(num_tokens: int, hidden_dim: int, num_experts: int): + """The gate feeds top-k expert selection, so numeric error only matters if + it changes the selected experts. Experts whose reference logit sits within + the kernel's error of the k-th value are genuinely tied, so swapping them + is acceptable; anything outside that window is a real routing bug.""" + _requires_aiter_tgemm() + top_k = 8 + device = torch.device("cuda") + for seed in range(5): + torch.manual_seed(1000 + seed) + x = torch.randn(num_tokens, hidden_dim, dtype=torch.bfloat16, device=device) + weight = torch.randn( + num_experts, hidden_dim, dtype=torch.bfloat16, device=device + ) + + out = _run(x, weight, torch.float32) + ref = x.double() @ weight.double().t() + tie_window = TIE_REL_TO_PEAK * ref.abs().max().item() + kernel_idx = out.topk(top_k, dim=-1).indices + ref_vals, ref_idx = ref.topk(top_k, dim=-1) + for t in range(num_tokens): + got = set(kernel_idx[t].tolist()) + want = set(ref_idx[t].tolist()) + if got == want: + continue + kth = ref_vals[t, -1].item() + for e in got.symmetric_difference(want): + gap = abs(ref[t, e].item() - kth) + assert gap < tie_window, ( + f"top-{top_k} mismatch beyond tie tolerance: token {t}, " + f"expert {e}, gap {gap:.3e} > {tie_window:.3e}" + ) + + +@pytest.mark.parametrize("hidden_dim,num_experts", SHAPES) +@pytest.mark.parametrize("num_tokens", [1, 16, 128]) +def test_matches_fp32_fallback(hidden_dim: int, num_experts: int, num_tokens: int): + """The AITER tier must stay close to the fp32 fallback it replaces. + + When ``force_fp32_compute`` is set and no specialized kernel is available, + ``GateLinear`` keeps fp32 weights and upcasts the activation, so the gate + runs as an fp32 GEMM. Both operands still hold bf16-representable values, so + the differences are in the reduction and the output: the tuned kernel + accumulates over split-K rather than in one fp32 pass, which dominates, and + the result is rounded to bf16 before the cast. + """ + _requires_aiter_tgemm() + torch.manual_seed(7) + device = torch.device("cuda") + x = torch.randn(num_tokens, hidden_dim, dtype=torch.bfloat16, device=device) + weight = torch.randn(num_experts, hidden_dim, dtype=torch.bfloat16, device=device) + + out = _run(x, weight, torch.float32) + fp32_fallback = torch.nn.functional.linear(x.float(), weight.float()) + + _assert_close_to_peak(out, fp32_fallback, REL_TO_PEAK) 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 ba46c1de65d0..52ccb431c20f 100644 --- a/vllm/model_executor/layers/fused_moe/router/gate_linear.py +++ b/vllm/model_executor/layers/fused_moe/router/gate_linear.py @@ -4,6 +4,7 @@ from torch.nn.parameter import Parameter import vllm._custom_ops as ops +from vllm._aiter_ops import rocm_aiter_ops from vllm.config import get_current_vllm_config_or_none from vllm.logger import init_logger from vllm.model_executor.custom_op import PluggableLayer @@ -24,7 +25,8 @@ class GateLinear(ReplicatedLinear): 3. fp32 specialized kernel (SM90+, bf16/fp32 in, fp32 out, M<=32, (H, E) in {(3072, 256), (6144, 128), (6144, 256)}) 4. experimental bf16x3 CuteDSL kernel (opt-in, SM100, bf16 in, fp32 weight) - 5. cuBLAS bf16×bf16→fp32 (SM90+ + bf16 weight + fp32 out_dtype) + 5. bf16×bf16→fp32 GEMM: cuBLAS (SM90+ + bf16 weight + fp32 out_dtype) + or AITER tuned GEMM (ROCm gfx950, bf16 weight) 6. F.linear via ReplicatedLinear (ultimate fallback) The ``out_dtype`` attribute is mutable and can be set after init @@ -61,11 +63,18 @@ def __init__( can_use_specialized_kernels = ( current_platform.is_cuda() and (is_hopper or is_blackwell) and not bias ) + can_use_aiter_tuned_gemm = ( + not bias + and current_platform.is_rocm() + and bool(rocm_aiter_ops.is_tgemm_enabled()) + ) # If fp32 compute is required and no specialized kernel is available, # store weights in fp32 so the fallback linear path computes in fp32. + # The AITER tier is the exception: it needs bf16 weights and casts the + # logits instead, so fp32 weights here would silently disable it. if force_fp32_compute and not can_use_specialized_kernels: - params_dtype = torch.float32 + params_dtype = torch.bfloat16 if can_use_aiter_tuned_gemm else torch.float32 super().__init__( input_size, @@ -144,6 +153,13 @@ def __init__( and is_available() ) + # AITER tuned GEMM eligibility + self.allow_aiter_router_gemm = ( + can_use_aiter_tuned_gemm and self.weight.dtype == torch.bfloat16 + ) + if self.allow_aiter_router_gemm: + logger.info_once("Enabled the AITER tuned GEMM router gate.") + def set_out_dtype(self, out_dtype: torch.dtype) -> None: """Set output dtype for the router logits after init. @@ -218,10 +234,20 @@ def forward( output = bf16x3_router_gemm(x, self.weight) return output, None - # Tier 5: 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) - return output, None + # Tier 5: bf16→fp32 GEMM, cuBLAS on CUDA or the AITER tuned GEMM on ROCm. + # At most one is ever eligible: cuBLAS needs is_cuda(), AITER is_rocm(). + if x.dtype == torch.bfloat16: + if self.allow_cublas_router_gemm: + output = torch.mm(x, self.weight.T, out_dtype=torch.float32) + return output, None + + if self.allow_aiter_router_gemm: + output = torch.ops.vllm.rocm_aiter_router_gemm( + x, + self.weight, + self.out_dtype if self.out_dtype is not None else self.weight.dtype, + ) + return output, None # Tier 6: F.linear (ReplicatedLinear) if self.out_dtype is not None and x.dtype != self.weight.dtype: @@ -273,3 +299,40 @@ def fp32_router_gemm_dispatch_fake( op_func=fp32_router_gemm_dispatch_impl, fake_impl=fp32_router_gemm_dispatch_fake, ) + + +def rocm_aiter_router_gemm_impl( + x: torch.Tensor, + weight: torch.Tensor, + out_dtype: torch.dtype, +) -> torch.Tensor: + """AITER tuned GEMM for the MoE router gate on ROCm. + + Wrapped in a custom op so torch.compile treats the AITER tuned-config + lookup as opaque instead of specializing on it. + + AITER keys its tuned configs on the output dtype as well as the shape, and + only ships bf16-output entries for the gate's shapes. Asking for out_dtype + directly would miss the table whenever the router runs in fp32 and fall back + to an untuned solution, so take the tuned bf16 kernel and cast the tiny + num_tokens x num_experts output instead. + """ + from aiter.tuned_gemm import tgemm + + out = tgemm.mm(x, weight, None) + return out if out.dtype == out_dtype else out.to(out_dtype) + + +def rocm_aiter_router_gemm_fake( + x: torch.Tensor, + weight: torch.Tensor, + out_dtype: torch.dtype, +) -> torch.Tensor: + return x.new_empty((x.shape[0], weight.shape[0]), dtype=out_dtype) + + +direct_register_custom_op( + op_name="rocm_aiter_router_gemm", + op_func=rocm_aiter_router_gemm_impl, + fake_impl=rocm_aiter_router_gemm_fake, +)