[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>
This commit is contained in:
Mikko Tukiainen
co-authored by
Cursor
Cursor
parent
4f49b55c1f
commit
47fd76e4ce
@@ -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)
|
||||
|
||||
@@ -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]
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user