A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Feat][EPLB] A novel static EPLB placement strategy for MoE models. (#23745)
Signed-off-by: bruceszchen <bruceszchen@tencent.com> Signed-off-by: Chen Bruce <bruceszchen@tencent.com> Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Signed-off-by: Chen Bruce <cszwwdz@vip.qq.com> Co-authored-by: lemon412 <lemon412@foxmail.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Chen Bruce committed
Sep 16, 2025 at 18:55 UTC
7ea5c73ad794290b1be12e41227b89941c05752c
4 files changed
+265
-12
tests/distributed/test_expert_placement.py
new
+194
@@ -0,0 +1,194 @@
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.layer import determine_expert_map
7
+
8
+
9
+def verify_round_robin_pattern(expert_map, ep_rank, ep_size,
10
+ global_num_experts):
11
+ """Verify that the expert map follows the round_robin pattern."""
12
+ # Calculate expected local experts (supporting non-divisible cases)
13
+ base_experts = global_num_experts // ep_size
14
+ remainder = global_num_experts % ep_size
15
+
16
+ if ep_rank < remainder:
17
+ local_num_experts = base_experts + 1
18
+ else:
19
+ local_num_experts = base_experts
20
+
21
+ # Expected expert IDs for this rank in round_robin pattern
22
+ # For non-divisible cases, ranks with extra experts start earlier
23
+ expected_expert_ids = []
24
+ for expert_idx in range(local_num_experts):
25
+ global_expert_id = ep_rank + expert_idx * ep_size
26
+ expected_expert_ids.append(global_expert_id)
27
+
28
+ # Check that only expected experts are mapped to this rank
29
+ for global_expert_id in range(global_num_experts):
30
+ if global_expert_id in expected_expert_ids:
31
+ local_expert_id = expert_map[global_expert_id]
32
+ expected_local_id = expected_expert_ids.index(global_expert_id)
33
+ assert (
34
+ local_expert_id == expected_local_id
35
+ ), f"Global expert {global_expert_id} should map to local expert " \
36
+ f"{expected_local_id}, got {local_expert_id}"
37
+ else:
38
+ assert (
39
+ expert_map[global_expert_id] == -1
40
+ ), f"Global expert {global_expert_id} should not be mapped to " \
41
+ f"this rank"
42
+
43
+ # Verify that all local expert IDs are consecutive starting from 0
44
+ local_expert_ids = [
45
+ expert_map[global_id] for global_id in expected_expert_ids
46
+ ]
47
+ expected_local_ids = list(range(local_num_experts))
48
+ assert (
49
+ local_expert_ids == expected_local_ids
50
+ ), f"Expected local expert IDs {expected_local_ids}, got {local_expert_ids}"
51
+
52
+
53
+@pytest.mark.parametrize("expert_placement_strategy", ["round_robin"])
54
+@pytest.mark.parametrize("world_size", [2, 4])
55
+def test_expert_placement_various_sizes(expert_placement_strategy, world_size):
56
+ """Test round_robin expert placement with various expert counts."""
57
+
58
+ # Test with different global_num_experts values
59
+ # Include both divisible and non-divisible cases
60
+ if world_size == 2:
61
+ test_cases = [
62
+ (4, 2), # 4 experts (divisible)
63
+ (8, 2), # 8 experts (divisible)
64
+ (9, 2), # 9 experts (non-divisible)
65
+ (16, 2), # 16 experts (divisible)
66
+ (17, 2), # 17 experts (non-divisible)
67
+ ]
68
+ elif world_size == 4:
69
+ test_cases = [
70
+ (8, 4), # 8 experts (divisible)
71
+ (16, 4), # 16 experts (divisible)
72
+ (18, 4), # 18 experts (non-divisible)
73
+ (32, 4), # 32 experts (divisible)
74
+ (33, 4), # 33 experts (non-divisible)
75
+ ]
76
+ else:
77
+ test_cases = []
78
+
79
+ for test_global_experts, test_ep_size in test_cases:
80
+ # Ensure ep_size matches world_size
81
+ assert (test_ep_size == world_size
82
+ ), f"ep_size {test_ep_size} must equal world_size {world_size}"
83
+
84
+ # Test each rank
85
+ for ep_rank in range(world_size):
86
+ # Calculate expected local experts
87
+ base_experts = test_global_experts // test_ep_size
88
+ remainder = test_global_experts % test_ep_size
89
+ if ep_rank < remainder:
90
+ expected_test_local = base_experts + 1
91
+ else:
92
+ expected_test_local = base_experts
93
+
94
+ test_local_experts, test_expert_map = determine_expert_map(
95
+ ep_size=test_ep_size,
96
+ ep_rank=ep_rank,
97
+ global_num_experts=test_global_experts,
98
+ expert_placement_strategy=expert_placement_strategy,
99
+ )
100
+
101
+ assert (
102
+ test_local_experts == expected_test_local
103
+ ), f"For {test_global_experts} experts on {test_ep_size} ranks, " \
104
+ f"rank {ep_rank}: expected {expected_test_local} local" \
105
+ f"experts, got {test_local_experts}"
106
+
107
+ if test_expert_map is not None:
108
+ assert test_expert_map.shape == (
109
+ test_global_experts,
110
+ ), f"Expected expert map shape ({test_global_experts},), " \
111
+ f"got {test_expert_map.shape}"
112
+
113
+ # Verify round_robin pattern for this test case
114
+ verify_round_robin_pattern(test_expert_map, ep_rank,
115
+ test_ep_size, test_global_experts)
116
+
117
+
118
+@pytest.mark.parametrize("expert_placement_strategy", ["round_robin"])
119
+@pytest.mark.parametrize("world_size", [2, 4])
120
+def test_expert_placement_edge_cases(expert_placement_strategy, world_size):
121
+ """Test edge cases for round_robin expert placement."""
122
+
123
+ # Test case 1: ep_size = 1 (should return None for expert_map)
124
+ local_num_experts, expert_map = determine_expert_map(
125
+ ep_size=1,
126
+ ep_rank=0,
127
+ global_num_experts=8,
128
+ expert_placement_strategy=expert_placement_strategy,
129
+ )
130
+ assert local_num_experts == 8, "For ep_size=1, should get all experts"
131
+ assert expert_map is None, "For ep_size=1, expert_map should be None"
132
+
133
+ # Test case 2: ep_size = 0 (should raise assertion)
134
+ with pytest.raises(AssertionError):
135
+ determine_expert_map(
136
+ ep_size=0,
137
+ ep_rank=0,
138
+ global_num_experts=8,
139
+ expert_placement_strategy=expert_placement_strategy,
140
+ )
141
+
142
+
143
+def test_determine_expert_map_comprehensive():
144
+ """Test of determine_expert_map function with various configurations."""
145
+
146
+ # Test cases: (ep_size, ep_rank, global_num_experts,
147
+ # expert_placement_strategy, expected_local, expected_map_pattern)
148
+ test_cases = [
149
+ # Round robin placement tests
150
+ (2, 0, 8, "round_robin", 4, [0, -1, 1, -1, 2, -1, 3,
151
+ -1]), # rank 0 gets even experts
152
+ (2, 1, 8, "round_robin", 4, [-1, 0, -1, 1, -1, 2, -1,
153
+ 3]), # rank 1 gets odd experts
154
+ (2, 0, 9, "round_robin", 5, [0, -1, 1, -1, 2, -1, 3, -1, 4
155
+ ]), # rank 0 gets 5 experts (even + last)
156
+ (2, 1, 9, "round_robin", 4, [-1, 0, -1, 1, -1, 2, -1, 3,
157
+ -1]), # rank 1 gets 4 experts (odd)
158
+
159
+ # 4-rank tests
160
+ (4, 0, 8, "round_robin", 2, [0, -1, -1, -1, 1, -1, -1,
161
+ -1]), # rank 0 gets experts 0, 4
162
+ (4, 1, 8, "round_robin", 2, [-1, 0, -1, -1, -1, 1, -1,
163
+ -1]), # rank 1 gets experts 1, 5
164
+ (4, 2, 8, "round_robin", 2, [-1, -1, 0, -1, -1, -1, 1,
165
+ -1]), # rank 2 gets experts 2, 6
166
+ (4, 3, 8, "round_robin", 2, [-1, -1, -1, 0, -1, -1, -1,
167
+ 1]), # rank 3 gets experts 3, 7
168
+ ]
169
+
170
+ for ep_size, ep_rank, global_num_experts, expert_placement_strategy, \
171
+ expected_local, expected_map_pattern in test_cases:
172
+ local_num_experts, expert_map = determine_expert_map(
173
+ ep_size=ep_size,
174
+ ep_rank=ep_rank,
175
+ global_num_experts=global_num_experts,
176
+ expert_placement_strategy=expert_placement_strategy,
177
+ )
178
+
179
+ assert local_num_experts == expected_local, \
180
+ f"ep_size={ep_size}, ep_rank={ep_rank}, " \
181
+ f"global_num_experts={global_num_experts}, " \
182
+ f"expert_placement_strategy={expert_placement_strategy}: " \
183
+ f"expected {expected_local} local experts, got {local_num_experts}"
184
+
185
+ if expected_map_pattern is None:
186
+ assert expert_map is None, "Expected expert_map to be None"
187
+ else:
188
+ assert expert_map is not None, "Expected expert_map to not be None"
189
+ actual_map = expert_map.tolist()
190
+ assert actual_map == expected_map_pattern, \
191
+ f"ep_size={ep_size}, ep_rank={ep_rank}, " \
192
+ f"global_num_experts={global_num_experts}, " \
193
+ f"expert_placement_strategy={expert_placement_strategy}: " \
194
+ f"expected map {expected_map_pattern}, got {actual_map}"
vllm/config/parallel.py
+10
@@ -29,6 +29,7 @@ else:
29
30
logger = init_logger(__name__)
31
32
+ExpertPlacementStrategy = Literal["linear", "round_robin"]
33
DistributedExecutorBackend = Literal["ray", "mp", "uni", "external_launcher"]
34
35
@@ -102,6 +103,15 @@ class ParallelConfig:
103
"""Enable expert parallelism load balancing for MoE layers."""
104
eplb_config: EPLBConfig = field(default_factory=EPLBConfig)
105
"""Expert parallelism configuration."""
106
+ expert_placement_strategy: ExpertPlacementStrategy = "linear"
107
+ """The expert placement strategy for MoE layers:\n
108
+ - "linear": Experts are placed in a contiguous manner. For example, with 4
109
+ experts and 2 ranks, rank 0 will have experts [0, 1] and rank 1 will have
110
+ experts [2, 3].\n
111
+ - "round_robin": Experts are placed in a round-robin manner. For example,
112
+ with 4 experts and 2 ranks, rank 0 will have experts [0, 2] and rank 1
113
+ will have experts [1, 3]. This strategy can help improve load balancing
114
+ for grouped expert models with no redundant experts."""
115
num_redundant_experts: Optional[int] = None
116
"""`num_redundant_experts` is deprecated and has been replaced with
117
`eplb_config.num_redundant_experts`. This will be removed in v0.12.0.
vllm/engine/arg_utils.py
+7
@@ -34,6 +34,7 @@ from vllm.config import (BlockSize, CacheConfig, CacheDType, CompilationConfig,
34
SpeculativeConfig, TaskOption, TokenizerMode,
35
VllmConfig, get_attr_docs)
36
from vllm.config.multimodal import MMCacheType, MultiModalConfig
37
+from vllm.config.parallel import ExpertPlacementStrategy
38
from vllm.config.utils import get_field
39
from vllm.logger import init_logger
40
from vllm.platforms import CpuArchEnum, current_platform
@@ -328,6 +329,8 @@ class EngineArgs:
329
enable_expert_parallel: bool = ParallelConfig.enable_expert_parallel
330
eplb_config: EPLBConfig = get_field(ParallelConfig, "eplb_config")
331
enable_eplb: bool = ParallelConfig.enable_eplb
332
+ expert_placement_strategy: ExpertPlacementStrategy = \
333
+ ParallelConfig.expert_placement_strategy
334
num_redundant_experts: int = EPLBConfig.num_redundant_experts
335
eplb_window_size: int = EPLBConfig.window_size
336
eplb_step_interval: int = EPLBConfig.step_interval
@@ -696,6 +699,9 @@ class EngineArgs:
699
**parallel_kwargs["enable_eplb"])
700
parallel_group.add_argument("--eplb-config",
701
**parallel_kwargs["eplb_config"])
702
+ parallel_group.add_argument(
703
+ "--expert-placement-strategy",
704
+ **parallel_kwargs["expert_placement_strategy"])
705
parallel_group.add_argument(
706
"--num-redundant-experts",
707
type=int,
@@ -1335,6 +1341,7 @@ class EngineArgs:
1341
enable_expert_parallel=self.enable_expert_parallel,
1342
enable_eplb=self.enable_eplb,
1343
eplb_config=self.eplb_config,
1344
+ expert_placement_strategy=self.expert_placement_strategy,
1345
max_parallel_loading_workers=self.max_parallel_loading_workers,
1346
disable_custom_all_reduce=self.disable_custom_all_reduce,
1347
ray_workers_use_nsight=self.ray_workers_use_nsight,
vllm/model_executor/layers/fused_moe/layer.py
+54
-12
@@ -4,7 +4,7 @@
4
from abc import abstractmethod
5
from collections.abc import Iterable
6
from enum import Enum
7
-from typing import Callable, Literal, Optional, Union, overload
7
+from typing import Callable, Literal, Optional, Union, get_args, overload
8
9
import torch
10
import torch.nn.functional as F
@@ -12,6 +12,7 @@ from torch.nn.parameter import UninitializedParameter
12
13
import vllm.envs as envs
14
from vllm.config import get_current_vllm_config
15
+from vllm.config.parallel import ExpertPlacementStrategy
16
from vllm.distributed import (get_dp_group, get_ep_group,
17
get_tensor_model_parallel_world_size,
18
tensor_model_parallel_all_reduce)
@@ -675,8 +676,11 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
676
677
678
def determine_expert_map(
678
- ep_size: int, ep_rank: int,
679
- global_num_experts: int) -> tuple[int, Optional[torch.Tensor]]:
679
+ ep_size: int,
680
+ ep_rank: int,
681
+ global_num_experts: int,
682
+ expert_placement_strategy: ExpertPlacementStrategy = "linear",
683
+) -> tuple[int, Optional[torch.Tensor]]:
684
"""
685
Calculates how many experts should be assigned to each rank for EP and
686
creates a mapping from global to local expert index. Experts are
@@ -684,8 +688,11 @@ def determine_expert_map(
688
last rank.
689
690
Args:
687
- ep_size (int): The size of the expert parallel group
688
- global_num_experts (int): The total number of experts in the model.
691
+ ep_size: The size of the expert parallel group
692
+ ep_rank: The rank of the current process in the expert parallel
693
+ group
694
+ global_num_experts: The total number of experts in the model.
695
+ expert_placement_strategy: The expert placement strategy.
696
697
Returns:
698
tuple[int, Optional[torch.Tensor]]: A tuple containing:
@@ -711,9 +718,23 @@ def determine_expert_map(
718
# Create a tensor of size num_experts filled with -1
719
expert_map = torch.full((global_num_experts, ), -1, dtype=torch.int32)
720
# Create an expert map for the local experts
714
- start_idx = ep_rank * base_experts + min(ep_rank, remainder)
715
- expert_map[start_idx:start_idx + local_num_experts] = torch.arange(
716
- 0, local_num_experts, dtype=torch.int32)
721
+ if expert_placement_strategy == "linear":
722
+ start_idx = ep_rank * base_experts + min(ep_rank, remainder)
723
+ expert_map[start_idx:start_idx + local_num_experts] = torch.arange(
724
+ 0, local_num_experts, dtype=torch.int32)
725
+ elif expert_placement_strategy == "round_robin":
726
+ local_log_experts = torch.arange(ep_rank,
727
+ global_num_experts,
728
+ ep_size,
729
+ dtype=torch.int32)
730
+
731
+ expert_map[local_log_experts] = torch.arange(0,
732
+ local_num_experts,
733
+ dtype=torch.int32)
734
+ else:
735
+ raise ValueError("Unsupported expert placement strategy "
736
+ f"'{expert_placement_strategy}', expected one of "
737
+ f"{get_args(ExpertPlacementStrategy)}")
738
return (local_num_experts, expert_map)
739
740
@@ -846,15 +867,36 @@ class FusedMoE(CustomOp):
867
else:
868
assert num_redundant_experts == 0, \
869
"Redundant experts are only supported with EPLB."
870
+
871
+ expert_placement_strategy = (
872
+ vllm_config.parallel_config.expert_placement_strategy)
873
+ if expert_placement_strategy == "round_robin":
874
+ # TODO(Bruce): will support round robin expert placement with
875
+ # EPLB enabled in the future.
876
+ round_robin_supported = ((num_expert_group is not None
877
+ and num_expert_group > 1)
878
+ and num_redundant_experts == 0
879
+ and not self.enable_eplb)
880
+
881
+ if not round_robin_supported:
882
+ logger.warning(
883
+ "Round-robin expert placement is only supported for "
884
+ "models with multiple expert groups and no redundant "
885
+ "experts. Falling back to linear expert placement.")
886
+ expert_placement_strategy = "linear"
887
+
888
self.local_num_experts, self.expert_map = determine_expert_map(
889
ep_size=self.ep_size,
890
ep_rank=self.ep_rank,
852
- global_num_experts=self.global_num_experts)
891
+ global_num_experts=self.global_num_experts,
892
+ expert_placement_strategy=expert_placement_strategy,
893
+ )
894
logger.info_once(
854
- "[EP Rank %s/%s] Expert parallelism is enabled. Local/global"
895
+ "[EP Rank %s/%s] Expert parallelism is enabled. Expert "
896
+ "placement strategy: %s. Local/global"
897
" number of experts: %s/%s. Experts local to global index map:"
856
- " %s.", self.ep_rank, self.ep_size, self.local_num_experts,
857
- self.global_num_experts,
898
+ " %s.", self.ep_rank, self.ep_size, expert_placement_strategy,
899
+ self.local_num_experts, self.global_num_experts,
900
get_compressed_expert_map(self.expert_map))
901
else:
902
self.local_num_experts, self.expert_map = (self.global_num_experts,