[ROCm][Perf] Use AITER tuned GEMM for the MoE router gate (backport #50535)

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)

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,
+)