[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>
2026-10-15 17:53:18 +00:00
parent 4f49b55c1f
commit 47fd76e4ce
6 changed files with 469 additions and 8 deletions
@@ -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]
 )