[ROCm][Perf] Route softmax MoE top-k, including fused shared-expert sigmoid, through AITER topk_gating (#56226)

Signed-off-by: Mikko Tukiainen <Mikko.Tukiainen@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

benchmarks/kernels/benchmark_fused_topk.py

@@ -11,7 +11,7 @@ from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 num_tokens_range = [2**i for i in range(0, 8, 2)]
 num_experts_range = [16, 32, 64, 128, 256, 512]
-topk_range = [3, 4]
+topk_range = [3, 4, 8, 10]
 configs = list(itertools.product(num_tokens_range, num_experts_range, topk_range))
 
 
@@ -39,9 +39,9 @@ def get_benchmark(scoring_func):
             x_names=["num_tokens", "num_experts", "topk"],
             x_vals=[list(_) for _ in configs],
             line_arg="provider",
-            line_vals=["torch", "vllm"],
-            line_names=["Torch", "vLLM"],
-            styles=[("blue", "-"), ("red", "-")],
+            line_vals=["torch", "vllm", "aiter_gating"],
+            line_names=["Torch", "vLLM", "AITER topk_gating"],
+            styles=[("blue", "-"), ("red", "-"), ("green", "-")],
             ylabel="us",
             plot_name=f"fused-topk-perf-{scoring_func}",
             args={},
@@ -70,6 +70,28 @@ def get_benchmark(scoring_func):
                 ),
                 quantiles=quantiles,
             )
+        elif provider == "aiter_gating":
+            from vllm._aiter_ops import rocm_aiter_ops
+
+            if scoring_func != "softmax" or not rocm_aiter_ops.is_topk_gating_enabled():
+                return float("nan"), float("nan"), float("nan")
+            topk_weights = torch.empty(
+                (num_tokens, topk), dtype=torch.float32, device="cuda"
+            )
+            topk_ids = torch.empty((num_tokens, topk), dtype=torch.int32, device="cuda")
+            token_expert_indices = torch.empty(
+                (num_tokens, topk), dtype=torch.int32, device="cuda"
+            )
+            ms, min_ms, max_ms = triton.testing.do_bench(
+                lambda: rocm_aiter_ops.topk_gating(
+                    topk_weights,
+                    topk_ids,
+                    token_expert_indices,
+                    gating_output,
+                    renormalize,
+                ),
+                quantiles=quantiles,
+            )
         else:
             ms, min_ms, max_ms = triton.testing.do_bench(
                 lambda: fused_topk(

tests/kernels/moe/test_rocm_aiter_topk.py

@@ -219,3 +219,169 @@ def test_rocm_aiter_grouped_topk_torch_compile_compatibility():
         topk_weights_original, topk_weights_compiled, rtol=1e-2, atol=1e-2
     )
     assert torch.allclose(topk_ids_original, topk_ids_compiled)
+
+
+def test_rocm_aiter_topk_gating_custom_op_registration():
+    assert callable(torch.ops.vllm.rocm_aiter_topk_gating)
+
+
+def _sort_routing(weights: torch.Tensor, ids: torch.Tensor):
+    ids_sorted, order = torch.sort(ids, dim=-1)
+    weights_sorted = torch.gather(weights, 1, order)
+    return weights_sorted, ids_sorted
+
+
+@pytest.mark.parametrize("num_tokens", [1, 4, 16, 64])
+@pytest.mark.parametrize("num_experts", [256, 512])
+@pytest.mark.parametrize("topk", [8, 10])
+@pytest.mark.parametrize("renormalize", [True, False])
+def test_rocm_aiter_topk_gating_matches_softmax_reference(
+    num_tokens: int,
+    num_experts: int,
+    topk: int,
+    renormalize: bool,
+):
+    """Numerical check vs softmax + topk, including Qwen3.8 (E=512, k=10)."""
+    torch.manual_seed(0)
+    gating_output = torch.randn(
+        (num_tokens, num_experts), dtype=torch.bfloat16, device="cuda"
+    )
+    scores = torch.softmax(gating_output.float(), dim=-1)
+    cutoff = torch.topk(scores, k=topk, dim=-1).values[:, -1:]
+
+    topk_weights = torch.empty((num_tokens, topk), dtype=torch.float32, device="cuda")
+    topk_ids = torch.empty((num_tokens, topk), dtype=torch.int32, device="cuda")
+    token_expert_indices = torch.empty(
+        (num_tokens, topk), dtype=torch.int32, device="cuda"
+    )
+    torch.ops.vllm.rocm_aiter_topk_gating(
+        topk_weights,
+        topk_ids,
+        token_expert_indices,
+        gating_output,
+        renormalize,
+        0,
+        "",
+    )
+
+    # BF16 logits can tie at the top-k boundary, so a different tied expert is
+    # valid. Check selection against the cutoff and weights against the scores
+    # of the experts the kernel actually selected.
+    assert ((topk_ids >= 0) & (topk_ids < num_experts)).all()
+    sorted_ids = torch.sort(topk_ids, dim=-1).values
+    assert (sorted_ids[:, 1:] != sorted_ids[:, :-1]).all()
+    selected_scores = torch.gather(scores, 1, topk_ids.long())
+    assert (selected_scores >= cutoff).all()
+    expected_weights = selected_scores
+    if renormalize:
+        expected_weights = expected_weights / expected_weights.sum(dim=-1, keepdim=True)
+    torch.testing.assert_close(topk_weights, expected_weights, atol=2e-2, rtol=2e-2)
+
+
+def test_rocm_aiter_topk_gating_shared_expert_suffix():
+    """Fused shared expert: softmax top-k on the prefix, sigmoid appended.
+
+    Matches the Qwen3.8 FSE launch: 512 routed experts, top-10, one shared
+    column, weights wide enough for the suffix, and the shared id prefilled.
+    """
+    torch.manual_seed(0)
+    num_tokens, num_experts, topk, num_shared = 4, 512, 10, 1
+    routed = torch.randn((num_tokens, num_experts), dtype=torch.bfloat16, device="cuda")
+    shared = torch.randn((num_tokens, num_shared), dtype=torch.bfloat16, device="cuda")
+    shared[0, 0] = 40
+    gating_output = torch.cat([routed, shared], dim=-1).contiguous()
+    width = topk + num_shared
+    shared_ids = torch.arange(
+        num_experts, num_experts + num_shared, device="cuda", dtype=torch.int32
+    )
+    topk_ids = torch.full((num_tokens, width), -1, dtype=torch.int32, device="cuda")
+    topk_ids[:, topk:] = shared_ids
+    topk_weights = torch.empty((num_tokens, width), dtype=torch.float32, device="cuda")
+    token_expert_indices = torch.empty(
+        (num_tokens, topk), dtype=torch.int32, device="cuda"
+    )
+
+    torch.ops.vllm.rocm_aiter_topk_gating(
+        topk_weights,
+        topk_ids[:, :topk],
+        token_expert_indices,
+        gating_output,
+        True,
+        num_shared,
+        "sigmoid",
+    )
+
+    scores = torch.softmax(routed.float(), dim=-1)
+    cutoff = torch.topk(scores, k=topk, dim=-1).values[:, -1:]
+    routed_ids = topk_ids[:, :topk]
+    assert ((routed_ids >= 0) & (routed_ids < num_experts)).all()
+    sorted_ids = torch.sort(routed_ids, dim=-1).values
+    assert (sorted_ids[:, 1:] != sorted_ids[:, :-1]).all()
+    selected = torch.gather(scores, 1, routed_ids.long())
+    assert (selected >= cutoff).all()
+    expected = selected / selected.sum(dim=-1, keepdim=True)
+    torch.testing.assert_close(topk_weights[:, :topk], expected, atol=2e-2, rtol=2e-2)
+    torch.testing.assert_close(
+        topk_weights[:, topk:],
+        torch.sigmoid(shared.float()),
+        atol=2e-2,
+        rtol=2e-2,
+    )
+    assert torch.equal(topk_ids[:, topk:], shared_ids.expand_as(topk_ids[:, topk:]))
+
+
+def test_rocm_aiter_topk_gating_torch_compile_compatibility():
+    token, expert, topk, renormalize = 4, 512, 10, True
+    gating_output = torch.randn((token, expert), dtype=torch.bfloat16, device="cuda")
+    device = gating_output.device
+
+    def gating_fn(gating_output, topk_weights, topk_ids, token_expert_indices):
+        return torch.ops.vllm.rocm_aiter_topk_gating(
+            topk_weights,
+            topk_ids,
+            token_expert_indices,
+            gating_output,
+            renormalize,
+            0,
+            "",
+        )
+
+    torch.library.opcheck(
+        torch.ops.vllm.rocm_aiter_topk_gating,
+        (
+            torch.empty((token, topk), dtype=torch.float32, device=device),
+            torch.empty((token, topk), dtype=torch.int32, device=device),
+            torch.empty((token, topk), dtype=torch.int32, device=device),
+            gating_output,
+        ),
+        kwargs={
+            "renormalize": renormalize,
+            "num_shared_experts": 0,
+            "shared_expert_scoring_func": "",
+        },
+        test_utils=("test_faketensor"),
+    )
+
+    compiled_fn = torch.compile(
+        gating_fn,
+        fullgraph=True,
+        backend="inductor",
+        mode="reduce-overhead",
+        dynamic=False,
+    )
+
+    def alloc():
+        return (
+            torch.empty((token, topk), dtype=torch.float32, device=device),
+            torch.empty((token, topk), dtype=torch.int32, device=device),
+            torch.empty((token, topk), dtype=torch.int32, device=device),
+        )
+
+    w0, i0, t0 = alloc()
+    w1, i1, t1 = alloc()
+    gating_fn(gating_output, w0, i0, t0)
+    compiled_fn(gating_output, w1, i1, t1)
+    w0, i0 = _sort_routing(w0, i0)
+    w1, i1 = _sort_routing(w1, i1)
+    assert torch.allclose(w0, w1, rtol=1e-2, atol=1e-2)
+    assert torch.equal(i0, i1)

tests/model_executor/test_enabled_custom_ops.py

@@ -139,6 +139,122 @@ def test_topk_softmax_dispatch(use_rocm_aiter: bool):
         assert topk_func == vllm_topk_softmax
 
 
+def _topk_gating_launch(
+    num_tokens: int = 4,
+    num_experts: int = 512,
+    topk: int = 10,
+    num_shared_experts: int = 0,
+    weights_width: int | None = None,
+    ids_width: int | None = None,
+    contiguous: bool = True,
+) -> dict:
+    """Tensors for one softmax top-k launch, as dispatch_topk_softmax_func sees them."""
+    width = topk + num_shared_experts
+    weights = torch.empty(num_tokens, weights_width or width)
+    ids = torch.empty(num_tokens, ids_width or width, dtype=torch.int32)
+    ids = ids[:, :topk]
+    if contiguous:
+        gating = torch.empty(num_tokens, num_experts + num_shared_experts)
+    else:
+        gating = torch.empty(num_experts + num_shared_experts + 3, num_tokens).t()
+    return dict(topk_weights=weights, topk_indices=ids, gating_output=gating)
+
+
+@pytest.mark.parametrize(
+    "launch, num_shared_experts, scoring_func, expect_gating",
+    [
+        pytest.param({}, 0, "", True, id="plain-softmax"),
+        pytest.param({"num_tokens": 4096}, 0, "", True, id="max-tokens"),
+        pytest.param({"num_tokens": 4097}, 0, "", False, id="too-many-tokens"),
+        pytest.param({"contiguous": False}, 0, "", False, id="non-contiguous"),
+        pytest.param({}, 0, "sigmoid", False, id="scoring-without-shared"),
+        pytest.param({"num_shared_experts": 1}, 1, "sigmoid", True, id="one-shared"),
+        pytest.param({"num_shared_experts": 8}, 8, "sigmoid", True, id="eight-shared"),
+        pytest.param(
+            {"num_shared_experts": 3}, 3, "sigmoid", False, id="unsupported-count"
+        ),
+        pytest.param({"num_shared_experts": 1}, 1, "", False, id="no-shared-scoring"),
+        pytest.param(
+            {"num_shared_experts": 1}, 1, "softmax", False, id="wrong-shared-scoring"
+        ),
+        pytest.param(
+            {"num_shared_experts": 1, "weights_width": 10},
+            1,
+            "sigmoid",
+            False,
+            id="weights-too-narrow",
+        ),
+        pytest.param(
+            {"num_shared_experts": 1, "ids_width": 12},
+            1,
+            "sigmoid",
+            False,
+            id="row-stride-mismatch",
+        ),
+        pytest.param(
+            {"num_tokens": 4097, "num_shared_experts": 1},
+            1,
+            "sigmoid",
+            False,
+            id="shared-too-many-tokens",
+        ),
+    ],
+)
+@pytest.mark.parametrize("gating_enabled", [True, False])
+def test_topk_softmax_dispatch_aiter_topk_gating(
+    monkeypatch: pytest.MonkeyPatch,
+    launch: dict,
+    num_shared_experts: int,
+    scoring_func: str,
+    expect_gating: bool,
+    gating_enabled: bool,
+):
+    """AITER launches use topk_gating only where it is enabled, and only when it
+    supports them.
+
+    Every other AITER launch must reach the legacy topk_softmax, so this is the
+    one place the choice is made.
+    """
+    monkeypatch.setattr(
+        rocm_aiter_ops, "is_topk_gating_enabled", lambda: gating_enabled
+    )
+    topk_func = dispatch_topk_softmax_func(
+        True,
+        **_topk_gating_launch(**launch),
+        num_shared_experts=num_shared_experts,
+        shared_expert_scoring_func=scoring_func,
+    )
+    if expect_gating and gating_enabled:
+        assert topk_func == rocm_aiter_ops.topk_gating
+    else:
+        assert topk_func == rocm_aiter_ops.topk_softmax
+
+
+@pytest.mark.parametrize(
+    "cdna, fused_moe_enabled, expected",
+    [
+        pytest.param(3, True, True, id="cdna3"),
+        pytest.param(4, True, True, id="cdna4"),
+        pytest.param(2, True, False, id="cdna2"),
+        pytest.param(3, False, False, id="cdna3-aiter-moe-off"),
+        pytest.param(4, False, False, id="cdna4-aiter-moe-off"),
+    ],
+)
+def test_topk_gating_enabled_by_cdna(
+    monkeypatch: pytest.MonkeyPatch,
+    cdna: int,
+    fused_moe_enabled: bool,
+    expected: bool,
+):
+    """topk_gating is on for CDNA3 and CDNA4 with AITER MoE enabled, nowhere else."""
+    monkeypatch.setattr("vllm._aiter_ops.is_aiter_found_and_supported", lambda: True)
+    monkeypatch.setattr(
+        rocm_aiter_ops, "is_fused_moe_enabled", lambda: fused_moe_enabled
+    )
+    monkeypatch.setattr("vllm.platforms.rocm.get_cdna_version", lambda: cdna)
+    assert bool(rocm_aiter_ops.is_topk_gating_enabled()) is expected
+
+
 @pytest.mark.parametrize(
     "use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False]
 )

vllm/_aiter_ops.py

@@ -526,6 +526,46 @@ def _rocm_aiter_topk_softmax_impl(
     )
 
 
+def _rocm_aiter_topk_gating_impl(
+    topk_weights: torch.Tensor,
+    topk_indices: torch.Tensor,
+    token_expert_indices: torch.Tensor,
+    gating_output: torch.Tensor,
+    renormalize: bool,
+    num_shared_experts: int = 0,
+    shared_expert_scoring_func: str = "",
+) -> None:
+    # Callers pick this op only for launches AITER's topk_gating supports (see
+    # dispatch_topk_softmax_func); everything else goes to topk_softmax.
+    from aiter.ops.topk import topk_gating
+
+    # AITER renormalizes the selected softmax mass inside this launch,
+    # preserving vLLM's `renormalize=True` semantics without a second
+    # pointwise kernel, and sigmoids the trailing shared-expert columns.
+    topk_gating(
+        topk_weights,
+        topk_indices,
+        gating_output,
+        correction_bias=None,
+        need_renorm=renormalize,
+        routed_scaling_factor=1.0,
+        score_func="softmax",
+        num_shared_experts=num_shared_experts,
+    )
+
+
+def _rocm_aiter_topk_gating_fake(
+    topk_weights: torch.Tensor,
+    topk_indices: torch.Tensor,
+    token_expert_indices: torch.Tensor,
+    gating_output: torch.Tensor,
+    renormalize: bool,
+    num_shared_experts: int = 0,
+    shared_expert_scoring_func: str = "",
+) -> None:
+    pass
+
+
 def _rocm_aiter_topk_sigmoid_impl(
     topk_weights: torch.Tensor,
     topk_indices: torch.Tensor,
@@ -2020,7 +2060,7 @@ class rocm_aiter_ops:
     Operations:
         - GEMM operations: gemm_a8w8, gemm_a8w8_blockscale
         - Fused MoE: fused_moe, asm_moe_tkw1
-        - Routing: topk_softmax, biased_grouped_topk, grouped_topk
+        - Routing: topk_softmax, topk_gating, biased_grouped_topk, grouped_topk
         - MLA decode: mla_decode_fwd
         - Quantization: per_tensor_quant, per_token_quant, group_fp8_quant
         - Triton ops: triton_rotary_embed, triton_fp8_bmm, triton_gemm_a8w8_blockscale
@@ -2198,6 +2238,18 @@ class rocm_aiter_ops:
     def is_fused_moe_enabled(cls) -> bool:
         return cls._AITER_ENABLED and cls._FMOE_ENABLED
 
+    @classmethod
+    @if_aiter_supported
+    def is_topk_gating_enabled(cls) -> bool:
+        """Use AITER ``topk_gating`` for softmax routing.
+
+        CDNA3 and CDNA4 only (gfx942, gfx950): those are the targets it was
+        measured on. Other AITER targets keep the legacy ``topk_softmax`` launcher.
+        """
+        from vllm.platforms.rocm import get_cdna_version
+
+        return cls.is_fused_moe_enabled() and get_cdna_version() in (3, 4)
+
     @classmethod
     @if_aiter_supported
     def is_fusion_moe_shared_experts_enabled(cls) -> bool:
@@ -2851,6 +2903,14 @@ class rocm_aiter_ops:
                 dispatch_key=current_platform.dispatch_key,
             )
 
+            direct_register_custom_op(
+                op_name="rocm_aiter_topk_gating",
+                op_func=_rocm_aiter_topk_gating_impl,
+                mutates_args=["topk_weights", "topk_indices"],
+                fake_impl=_rocm_aiter_topk_gating_fake,
+                dispatch_key=current_platform.dispatch_key,
+            )
+
             direct_register_custom_op(
                 op_name="rocm_aiter_topk_sigmoid",
                 op_func=_rocm_aiter_topk_sigmoid_impl,
@@ -3382,6 +3442,27 @@ class rocm_aiter_ops:
         )
         return topk_weights, topk_indices
 
+    @staticmethod
+    def topk_gating(
+        topk_weights: torch.Tensor,
+        topk_indices: torch.Tensor,
+        token_expert_indices: torch.Tensor,
+        gating_output: torch.Tensor,
+        renormalize: bool,
+        num_shared_experts: int = 0,
+        shared_expert_scoring_func: str = "",
+    ) -> tuple[torch.Tensor, ...]:
+        torch.ops.vllm.rocm_aiter_topk_gating(
+            topk_weights,
+            topk_indices,
+            token_expert_indices,
+            gating_output,
+            renormalize,
+            num_shared_experts,
+            shared_expert_scoring_func,
+        )
+        return topk_weights, topk_indices
+
     @staticmethod
     def topk_sigmoid(
         topk_weights: torch.Tensor,

vllm/model_executor/layers/fused_moe/router/aiter_shared_routed_fused_moe_router.py

@@ -84,7 +84,14 @@ class AiterSharedRoutedFusedMoERouter(BaseRouter):
             total_topk_weights_slice = total_topk_weights[:M]
             topk_ids_slice = total_topk_ids[:M, :topk]
 
-            topk_func = dispatch_topk_softmax_func(use_rocm_aiter=True)
+            topk_func = dispatch_topk_softmax_func(
+                use_rocm_aiter=True,
+                topk_weights=total_topk_weights_slice,
+                topk_indices=topk_ids_slice,
+                gating_output=router_logits,
+                num_shared_experts=num_fse,
+                shared_expert_scoring_func="sigmoid",
+            )
             topk_func(
                 total_topk_weights_slice,
                 topk_ids_slice,
@@ -110,7 +117,10 @@ class AiterSharedRoutedFusedMoERouter(BaseRouter):
         )
 
         topk_func = dispatch_topk_softmax_func(
-            use_rocm_aiter=rocm_aiter_ops.is_fused_moe_enabled()
+            use_rocm_aiter=rocm_aiter_ops.is_fused_moe_enabled(),
+            topk_weights=topk_weights,
+            topk_indices=topk_ids,
+            gating_output=routing_logits,
         )
         topk_weights, topk_ids = topk_func(
             topk_weights,

vllm/model_executor/layers/fused_moe/router/fused_topk_router.py

@@ -61,10 +61,73 @@ def vllm_topk_sigmoid(
     return topk_weights, topk_indices
 
 
+# At E=512 / top-10, AITER's topk_gating is faster than the legacy launcher
+# through this many tokens; its one-row generic prefill path regresses above it.
+_AITER_TOPK_GATING_MAX_TOKENS = 4096
+
+
+def _aiter_topk_gating_supported(
+    topk_weights: torch.Tensor,
+    topk_indices: torch.Tensor,
+    gating_output: torch.Tensor,
+    num_shared_experts: int,
+    shared_expert_scoring_func: str,
+) -> bool:
+    """Whether AITER's ``topk_gating`` can serve this softmax top-k launch.
+
+    It needs contiguous gating rows and at most ``_AITER_TOPK_GATING_MAX_TOKENS``
+    tokens. Fused shared experts are scored with sigmoid after the routed
+    top-k: 1, 2, 4 or 8 of them, with ``topk_weights`` and ``topk_indices``
+    sharing a row stride and ``topk_weights`` wide enough for the shared
+    columns. Shared expert ids stay as the caller prefilled them.
+    """
+    if (
+        not gating_output.is_contiguous()
+        or gating_output.shape[0] > _AITER_TOPK_GATING_MAX_TOKENS
+    ):
+        return False
+    if num_shared_experts == 0:
+        return not shared_expert_scoring_func
+    return (
+        shared_expert_scoring_func == "sigmoid"
+        and num_shared_experts in (1, 2, 4, 8)
+        and topk_weights.stride(0) == topk_indices.stride(0)
+        and topk_weights.shape[-1] >= topk_indices.shape[-1] + num_shared_experts
+    )
+
+
 def dispatch_topk_softmax_func(
     use_rocm_aiter: bool = False,
+    *,
+    topk_weights: torch.Tensor | None = None,
+    topk_indices: torch.Tensor | None = None,
+    gating_output: torch.Tensor | None = None,
+    num_shared_experts: int = 0,
+    shared_expert_scoring_func: str = "",
 ) -> Callable[..., tuple[torch.Tensor, ...]]:
+    """Pick the softmax top-k implementation for one launch.
+
+    On ROCm with AITER, ``topk_gating`` is chosen on gfx942 and gfx950 when the
+    launch described by the tensors and shared-expert arguments is one it supports;
+    every other AITER launch uses the legacy ``topk_softmax``. The tensors are
+    optional: without them the launch cannot be judged and the legacy launcher
+    is used.
+    """
     if use_rocm_aiter:
+        if (
+            rocm_aiter_ops.is_topk_gating_enabled()
+            and topk_weights is not None
+            and topk_indices is not None
+            and gating_output is not None
+            and _aiter_topk_gating_supported(
+                topk_weights,
+                topk_indices,
+                gating_output,
+                num_shared_experts,
+                shared_expert_scoring_func,
+            )
+        ):
+            return rocm_aiter_ops.topk_gating
         return rocm_aiter_ops.topk_softmax
     return vllm_topk_softmax
 
@@ -104,7 +167,10 @@ def fused_topk(
 
     if scoring_func == "softmax":
         topk_func = dispatch_topk_softmax_func(
-            use_rocm_aiter=rocm_aiter_ops.is_fused_moe_enabled()
+            use_rocm_aiter=rocm_aiter_ops.is_fused_moe_enabled(),
+            topk_weights=topk_weights,
+            topk_indices=topk_ids,
+            gating_output=gating_output,
         )
         topk_weights, topk_ids = topk_func(
             topk_weights, topk_ids, token_expert_indices, gating_output, renormalize