| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for the AITER biased_grouped_topk expert-group count. |
| 4 | |
| 5 | The AITER kernel is only instantiated for ``NUM_GRP`` in {1, 2, 4, 8}, so the |
| 6 | group count derived from ``num_experts`` has to be rounded to a supported |
| 7 | value that still divides ``num_experts``. Grouping is a no-op on this path |
| 8 | (``topk_group == num_expert_group``), so any such value routes identically; |
| 9 | the constraint is purely about which kernel exists. |
| 10 | """ |
| 11 | |
| 12 | import pytest |
| 13 | import torch |
| 14 | |
| 15 | from vllm._aiter_ops import rocm_aiter_ops |
| 16 | from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import ( |
| 17 | AITER_MAX_EXPERTS_PER_GROUP as MAX_EXPERTS_PER_GROUP, |
| 18 | ) |
| 19 | from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import ( |
| 20 | AITER_SUPPORTED_NUM_GRP as SUPPORTED_NUM_GRP, |
| 21 | ) |
| 22 | from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import ( |
| 23 | _aiter_can_use_biased_grouped_topk, |
| 24 | _aiter_get_num_expert_group, |
| 25 | ) |
| 26 | from vllm.model_executor.layers.fused_moe.router.grouped_topk_router import ( |
| 27 | grouped_topk, |
| 28 | ) |
| 29 | from vllm.platforms import current_platform |
| 30 | from vllm.platforms.rocm import on_cdna |
| 31 | |
| 32 | pytestmark = pytest.mark.skipif( |
| 33 | not current_platform.is_rocm(), reason="ROCm-specific tests" |
| 34 | ) |
| 35 | |
| 36 | |
| 37 | @pytest.mark.parametrize( |
| 38 | ("num_experts", "expected"), |
| 39 | [ |
| 40 | (8, 1), # fits in a single group |
| 41 | (32, 1), # exactly at the per-group limit |
| 42 | (64, 2), |
| 43 | (128, 4), |
| 44 | (256, 8), |
| 45 | (72, 8), # naive ceil gives 3 -> unsupported, rounds up to 8 |
| 46 | (96, 8), # naive ceil gives 3 -> unsupported, rounds up to 8 |
| 47 | (160, 8), # naive ceil gives 5 -> unsupported, rounds up to 8 |
| 48 | (192, 8), # naive ceil gives 6 -> unsupported, rounds up to 8 |
| 49 | # No supported NUM_GRP divides these within the group-size limit, so |
| 50 | # the naive value is kept and the call site declines the kernel. |
| 51 | (33, 3), |
| 52 | (129, 43), |
| 53 | (257, 257), |
| 54 | (320, 10), |
| 55 | (384, 12), |
| 56 | ], |
| 57 | ) |
| 58 | def test_group_count_and_guard(num_experts, expected): |
| 59 | g = _aiter_get_num_expert_group(num_experts) |
| 60 | assert g == expected |
| 61 | # The router asserts both unconditionally; rounding must never break them. |
| 62 | assert num_experts % g == 0 |
| 63 | assert num_experts // g <= MAX_EXPERTS_PER_GROUP |
| 64 | |
| 65 | supported = g in SUPPORTED_NUM_GRP |
| 66 | # Ample topk, so only the membership check can decline. |
| 67 | assert _aiter_can_use_biased_grouped_topk(num_experts, topk=64) == supported |
| 68 | if supported: |
| 69 | assert _aiter_can_use_biased_grouped_topk(num_experts, topk=g) |
| 70 | assert not _aiter_can_use_biased_grouped_topk(num_experts, topk=g - 1) |
| 71 | |
| 72 | |
| 73 | @pytest.mark.skipif(not on_cdna(), reason="CDNA ROCm only") |
| 74 | @pytest.mark.parametrize( |
| 75 | "num_experts", |
| 76 | [ |
| 77 | 64, # g == 2, not rounded -- covers the no-op claim itself |
| 78 | 128, # g == 4, not rounded |
| 79 | 256, # g == 8, not rounded |
| 80 | 96, # naive 3 -> rounded to 8 |
| 81 | 192, # naive 6 -> rounded to 8 |
| 82 | ], |
| 83 | ) |
| 84 | def test_rounded_group_count_routes_like_the_reference(num_experts): |
| 85 | """The rounded group count must route identically to ``grouped_topk``. |
| 86 | |
| 87 | Gating is continuous random, so exact score ties -- the one place the two |
| 88 | may legitimately differ -- have measure zero. |
| 89 | """ |
| 90 | torch.manual_seed(num_experts) |
| 91 | device = "cuda" |
| 92 | topk = 8 |
| 93 | num_tokens = 83 # not a multiple of any warp/tile size |
| 94 | g = _aiter_get_num_expert_group(num_experts) |
| 95 | assert _aiter_can_use_biased_grouped_topk(num_experts, topk) |
| 96 | |
| 97 | gating = torch.randn(num_tokens, num_experts, dtype=torch.float32, device=device) |
| 98 | bias = torch.randn(num_experts, dtype=torch.float32, device=device) |
| 99 | |
| 100 | ref_weights, ref_ids = grouped_topk( |
| 101 | hidden_states=torch.empty(num_tokens, 1, device=device), |
| 102 | gating_output=gating, |
| 103 | topk=topk, |
| 104 | renormalize=True, |
| 105 | num_expert_group=g, |
| 106 | topk_group=g, |
| 107 | scoring_func="sigmoid", |
| 108 | e_score_correction_bias=bias, |
| 109 | ) |
| 110 | |
| 111 | weights = torch.empty(num_tokens, topk, dtype=torch.float32, device=device) |
| 112 | ids = torch.empty(num_tokens, topk, dtype=torch.int32, device=device) |
| 113 | rocm_aiter_ops.biased_grouped_topk( |
| 114 | gating, |
| 115 | bias, |
| 116 | weights, |
| 117 | ids, |
| 118 | num_expert_group=g, |
| 119 | topk_group=g, |
| 120 | need_renorm=True, |
| 121 | ) |
| 122 | |
| 123 | # Neither side promises an order within a token's top-k, so sort by expert |
| 124 | # id and carry the weights through the same permutation -- sorting the two |
| 125 | # independently would not catch a weight attached to the wrong expert. |
| 126 | order = ids.argsort(dim=-1) |
| 127 | ref_order = ref_ids.argsort(dim=-1) |
| 128 | torch.testing.assert_close( |
| 129 | ids.gather(1, order), ref_ids.to(torch.int32).gather(1, ref_order) |
| 130 | ) |
| 131 | torch.testing.assert_close( |
| 132 | weights.gather(1, order), |
| 133 | ref_weights.gather(1, ref_order), |
| 134 | atol=1e-4, |
| 135 | rtol=1e-4, |
| 136 | ) |