From 47fd76e4cee3293dfb9a99852182de9963cbe403 Mon Sep 17 00:00:00 2001 From: Mikko Tukiainen Date: Thu, 15 Oct 2026 20:53:18 +0300 Subject: [PATCH] [ROCm][Perf] Route softmax MoE top-k, including fused shared-expert sigmoid, through AITER topk_gating (#56226) Signed-off-by: Mikko Tukiainen Co-authored-by: Cursor --- benchmarks/kernels/benchmark_fused_topk.py | 30 +++- tests/kernels/moe/test_rocm_aiter_topk.py | 166 ++++++++++++++++++ .../model_executor/test_enabled_custom_ops.py | 116 ++++++++++++ vllm/_aiter_ops.py | 83 ++++++++- .../aiter_shared_routed_fused_moe_router.py | 14 +- .../fused_moe/router/fused_topk_router.py | 68 ++++++- 6 files changed, 469 insertions(+), 8 deletions(-) diff --git a/benchmarks/kernels/benchmark_fused_topk.py b/benchmarks/kernels/benchmark_fused_topk.py index 72bf2d97cc..ab4d954c87 100644 --- a/benchmarks/kernels/benchmark_fused_topk.py +++ b/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( diff --git a/tests/kernels/moe/test_rocm_aiter_topk.py b/tests/kernels/moe/test_rocm_aiter_topk.py index 11f9cf28ae..827ebe7c04 100644 --- a/tests/kernels/moe/test_rocm_aiter_topk.py +++ b/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) diff --git a/tests/model_executor/test_enabled_custom_ops.py b/tests/model_executor/test_enabled_custom_ops.py index 490284f439..27985bba65 100644 --- a/tests/model_executor/test_enabled_custom_ops.py +++ b/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] ) diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 9e01735441..ba25cc8b5a 100644 --- a/vllm/_aiter_ops.py +++ b/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, diff --git a/vllm/model_executor/layers/fused_moe/router/aiter_shared_routed_fused_moe_router.py b/vllm/model_executor/layers/fused_moe/router/aiter_shared_routed_fused_moe_router.py index 42d4b5d4ad..19e631659a 100644 --- a/vllm/model_executor/layers/fused_moe/router/aiter_shared_routed_fused_moe_router.py +++ b/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, diff --git a/vllm/model_executor/layers/fused_moe/router/fused_topk_router.py b/vllm/model_executor/layers/fused_moe/router/fused_topk_router.py index 28f67b5c69..766940f271 100644 --- a/vllm/model_executor/layers/fused_moe/router/fused_topk_router.py +++ b/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