| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import pytest |
| 5 | |
| 6 | from vllm.model_executor.layers.fused_moe.expert_map_manager import ( |
| 7 | determine_expert_map, |
| 8 | ) |
| 9 | |
| 10 | |
| 11 | def verify_round_robin_pattern(expert_map, ep_rank, ep_size, global_num_experts): |
| 12 | """Verify that the expert map follows the round_robin pattern.""" |
| 13 | # Calculate expected local experts (supporting non-divisible cases) |
| 14 | base_experts = global_num_experts // ep_size |
| 15 | remainder = global_num_experts % ep_size |
| 16 | |
| 17 | local_num_experts = base_experts + 1 if ep_rank < remainder else base_experts |
| 18 | |
| 19 | # Expected expert IDs for this rank in round_robin pattern |
| 20 | # For non-divisible cases, ranks with extra experts start earlier |
| 21 | expected_expert_ids = [] |
| 22 | for expert_idx in range(local_num_experts): |
| 23 | global_expert_id = ep_rank + expert_idx * ep_size |
| 24 | expected_expert_ids.append(global_expert_id) |
| 25 | |
| 26 | # Check that only expected experts are mapped to this rank |
| 27 | for global_expert_id in range(global_num_experts): |
| 28 | if global_expert_id in expected_expert_ids: |
| 29 | local_expert_id = expert_map[global_expert_id] |
| 30 | expected_local_id = expected_expert_ids.index(global_expert_id) |
| 31 | assert local_expert_id == expected_local_id, ( |
| 32 | f"Global expert {global_expert_id} should map to local expert " |
| 33 | f"{expected_local_id}, got {local_expert_id}" |
| 34 | ) |
| 35 | else: |
| 36 | assert expert_map[global_expert_id] == -1, ( |
| 37 | f"Global expert {global_expert_id} should not be mapped to this rank" |
| 38 | ) |
| 39 | |
| 40 | # Verify that all local expert IDs are consecutive starting from 0 |
| 41 | local_expert_ids = [expert_map[global_id] for global_id in expected_expert_ids] |
| 42 | expected_local_ids = list(range(local_num_experts)) |
| 43 | assert local_expert_ids == expected_local_ids, ( |
| 44 | f"Expected local expert IDs {expected_local_ids}, got {local_expert_ids}" |
| 45 | ) |
| 46 | |
| 47 | |
| 48 | @pytest.mark.parametrize("expert_placement_strategy", ["round_robin"]) |
| 49 | @pytest.mark.parametrize("world_size", [2, 4]) |
| 50 | def test_expert_placement_various_sizes(expert_placement_strategy, world_size): |
| 51 | """Test round_robin expert placement with various expert counts.""" |
| 52 | # Test with different global_num_experts values |
| 53 | # Include both divisible and non-divisible cases |
| 54 | if world_size == 2: |
| 55 | test_cases = [ |
| 56 | (4, 2), # 4 experts (divisible) |
| 57 | (8, 2), # 8 experts (divisible) |
| 58 | (9, 2), # 9 experts (non-divisible) |
| 59 | (16, 2), # 16 experts (divisible) |
| 60 | (17, 2), # 17 experts (non-divisible) |
| 61 | ] |
| 62 | elif world_size == 4: |
| 63 | test_cases = [ |
| 64 | (8, 4), # 8 experts (divisible) |
| 65 | (16, 4), # 16 experts (divisible) |
| 66 | (18, 4), # 18 experts (non-divisible) |
| 67 | (32, 4), # 32 experts (divisible) |
| 68 | (33, 4), # 33 experts (non-divisible) |
| 69 | ] |
| 70 | else: |
| 71 | test_cases = [] |
| 72 | |
| 73 | for test_global_experts, test_ep_size in test_cases: |
| 74 | # Ensure ep_size matches world_size |
| 75 | assert test_ep_size == world_size, ( |
| 76 | f"ep_size {test_ep_size} must equal world_size {world_size}" |
| 77 | ) |
| 78 | |
| 79 | # Test each rank |
| 80 | for ep_rank in range(world_size): |
| 81 | # Calculate expected local experts |
| 82 | base_experts = test_global_experts // test_ep_size |
| 83 | remainder = test_global_experts % test_ep_size |
| 84 | if ep_rank < remainder: |
| 85 | expected_test_local = base_experts + 1 |
| 86 | else: |
| 87 | expected_test_local = base_experts |
| 88 | |
| 89 | test_local_experts, test_expert_map, _ = determine_expert_map( |
| 90 | ep_size=test_ep_size, |
| 91 | ep_rank=ep_rank, |
| 92 | global_num_experts=test_global_experts, |
| 93 | expert_placement_strategy=expert_placement_strategy, |
| 94 | ) |
| 95 | |
| 96 | assert test_local_experts == expected_test_local, ( |
| 97 | f"For {test_global_experts} experts on {test_ep_size} ranks, " |
| 98 | f"rank {ep_rank}: expected {expected_test_local} local" |
| 99 | f"experts, got {test_local_experts}" |
| 100 | ) |
| 101 | |
| 102 | if test_expert_map is not None: |
| 103 | assert test_expert_map.shape == (test_global_experts,), ( |
| 104 | f"Expected expert map shape ({test_global_experts},), " |
| 105 | f"got {test_expert_map.shape}" |
| 106 | ) |
| 107 | |
| 108 | # Verify round_robin pattern for this test case |
| 109 | verify_round_robin_pattern( |
| 110 | test_expert_map, ep_rank, test_ep_size, test_global_experts |
| 111 | ) |
| 112 | |
| 113 | |
| 114 | @pytest.mark.parametrize("expert_placement_strategy", ["round_robin"]) |
| 115 | @pytest.mark.parametrize("world_size", [2, 4]) |
| 116 | def test_expert_placement_edge_cases(expert_placement_strategy, world_size): |
| 117 | """Test edge cases for round_robin expert placement.""" |
| 118 | # Test case 1: ep_size = 1 (should return None for expert_map) |
| 119 | local_num_experts, expert_map, _ = determine_expert_map( |
| 120 | ep_size=1, |
| 121 | ep_rank=0, |
| 122 | global_num_experts=8, |
| 123 | expert_placement_strategy=expert_placement_strategy, |
| 124 | ) |
| 125 | assert local_num_experts == 8, "For ep_size=1, should get all experts" |
| 126 | assert expert_map is None, "For ep_size=1, expert_map should be None" |
| 127 | |
| 128 | # Test case 2: ep_size = 0 (should raise assertion) |
| 129 | with pytest.raises(AssertionError): |
| 130 | determine_expert_map( |
| 131 | ep_size=0, |
| 132 | ep_rank=0, |
| 133 | global_num_experts=8, |
| 134 | expert_placement_strategy=expert_placement_strategy, |
| 135 | ) |
| 136 | |
| 137 | |
| 138 | def test_determine_expert_map_comprehensive(): |
| 139 | """Test of determine_expert_map function with various configurations.""" |
| 140 | # Test cases: (ep_size, ep_rank, global_num_experts, |
| 141 | # expert_placement_strategy, expected_local, expected_map_pattern) |
| 142 | test_cases = [ |
| 143 | # Round robin placement tests |
| 144 | ( |
| 145 | 2, |
| 146 | 0, |
| 147 | 8, |
| 148 | "round_robin", |
| 149 | 4, |
| 150 | [0, -1, 1, -1, 2, -1, 3, -1], |
| 151 | ), # rank 0 gets even experts |
| 152 | ( |
| 153 | 2, |
| 154 | 1, |
| 155 | 8, |
| 156 | "round_robin", |
| 157 | 4, |
| 158 | [-1, 0, -1, 1, -1, 2, -1, 3], |
| 159 | ), # rank 1 gets odd experts |
| 160 | ( |
| 161 | 2, |
| 162 | 0, |
| 163 | 9, |
| 164 | "round_robin", |
| 165 | 5, |
| 166 | [0, -1, 1, -1, 2, -1, 3, -1, 4], |
| 167 | ), # rank 0 gets 5 experts (even + last) |
| 168 | ( |
| 169 | 2, |
| 170 | 1, |
| 171 | 9, |
| 172 | "round_robin", |
| 173 | 4, |
| 174 | [-1, 0, -1, 1, -1, 2, -1, 3, -1], |
| 175 | ), # rank 1 gets 4 experts (odd) |
| 176 | # 4-rank tests |
| 177 | ( |
| 178 | 4, |
| 179 | 0, |
| 180 | 8, |
| 181 | "round_robin", |
| 182 | 2, |
| 183 | [0, -1, -1, -1, 1, -1, -1, -1], |
| 184 | ), # rank 0 gets experts 0, 4 |
| 185 | ( |
| 186 | 4, |
| 187 | 1, |
| 188 | 8, |
| 189 | "round_robin", |
| 190 | 2, |
| 191 | [-1, 0, -1, -1, -1, 1, -1, -1], |
| 192 | ), # rank 1 gets experts 1, 5 |
| 193 | ( |
| 194 | 4, |
| 195 | 2, |
| 196 | 8, |
| 197 | "round_robin", |
| 198 | 2, |
| 199 | [-1, -1, 0, -1, -1, -1, 1, -1], |
| 200 | ), # rank 2 gets experts 2, 6 |
| 201 | ( |
| 202 | 4, |
| 203 | 3, |
| 204 | 8, |
| 205 | "round_robin", |
| 206 | 2, |
| 207 | [-1, -1, -1, 0, -1, -1, -1, 1], |
| 208 | ), # rank 3 gets experts 3, 7 |
| 209 | ] |
| 210 | |
| 211 | for ( |
| 212 | ep_size, |
| 213 | ep_rank, |
| 214 | global_num_experts, |
| 215 | expert_placement_strategy, |
| 216 | expected_local, |
| 217 | expected_map_pattern, |
| 218 | ) in test_cases: |
| 219 | local_num_experts, expert_map, _ = determine_expert_map( |
| 220 | ep_size=ep_size, |
| 221 | ep_rank=ep_rank, |
| 222 | global_num_experts=global_num_experts, |
| 223 | expert_placement_strategy=expert_placement_strategy, |
| 224 | ) |
| 225 | |
| 226 | assert local_num_experts == expected_local, ( |
| 227 | f"ep_size={ep_size}, ep_rank={ep_rank}, " |
| 228 | f"global_num_experts={global_num_experts}, " |
| 229 | f"expert_placement_strategy={expert_placement_strategy}: " |
| 230 | f"expected {expected_local} local experts, got {local_num_experts}" |
| 231 | ) |
| 232 | |
| 233 | if expected_map_pattern is None: |
| 234 | assert expert_map is None, "Expected expert_map to be None" |
| 235 | else: |
| 236 | assert expert_map is not None, "Expected expert_map to not be None" |
| 237 | actual_map = expert_map.tolist() |
| 238 | assert actual_map == expected_map_pattern, ( |
| 239 | f"ep_size={ep_size}, ep_rank={ep_rank}, " |
| 240 | f"global_num_experts={global_num_experts}, " |
| 241 | f"expert_placement_strategy={expert_placement_strategy}: " |
| 242 | f"expected map {expected_map_pattern}, got {actual_map}" |
| 243 | ) |