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 3, 20:15 UTC
[MoE Refactor] Move expert map related code into ExpertMapManager class (#41046)
Signed-off-by: Bill Nell <bnell@redhat.com> Signed-off-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Signed-off-by: Robert Shaw <robertgshaw2@gmail.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Co-authored-by: Robert Shaw <robertgshaw2@gmail.com>
bnellnm committed
May 12, 2026 at 09:18 UTC
206eaed08d25fb26f0ffb4de26fbadcbeec3dc7b
19 files changed
+593
-323
tests/distributed/test_expert_placement.py
+3
-1
@@ -3,7 +3,9 @@
3
4
import pytest
5
6
-from vllm.model_executor.layers.fused_moe.layer import determine_expert_map
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):
tests/kernels/moe/test_moe.py
+1
-1
@@ -1588,7 +1588,7 @@ def test_unquantized_bf16_flashinfer_trtllm_backend(
1588
layer.apply_router_weight_on_input = False
1589
layer.routed_scaling_factor = None
1590
layer.shared_experts = None
1591
- layer._maybe_init_expert_routing_tables = lambda: None
1591
+ layer._expert_routing_tables = lambda: None
1592
1593
quant_method.process_weights_after_loading(layer)
1594
tests/kernels/moe/test_moe_permute_unpermute.py
+3
-1
@@ -10,7 +10,9 @@ import pytest
10
import torch
11
12
from vllm.model_executor.layers.fused_moe import fused_topk
13
-from vllm.model_executor.layers.fused_moe.layer import determine_expert_map
13
+from vllm.model_executor.layers.fused_moe.expert_map_manager import (
14
+ determine_expert_map,
15
+)
16
from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (
17
moe_permute,
18
moe_permute_unpermute_supported,
vllm/model_executor/layers/fused_moe/expert_map_manager.py
new
+516
@@ -0,0 +1,516 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+"""Expert Map Manager for MoE layers.
5
+
6
+This module contains the ExpertMapManager class which manages expert ID
7
+mappings and placement strategies for Expert Parallelism in MoE models.
8
+"""
9
+
10
+import torch
11
+
12
+from vllm.config.parallel import ExpertPlacementStrategy
13
+from vllm.logger import init_logger
14
+from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig
15
+from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
16
+ init_aiter_topK_meta_data,
17
+)
18
+
19
+logger = init_logger(__name__)
20
+
21
+
22
+def determine_expert_map(
23
+ ep_size: int,
24
+ ep_rank: int,
25
+ global_num_experts: int,
26
+ expert_placement_strategy: ExpertPlacementStrategy = "linear",
27
+ num_fused_shared_experts: int = 0,
28
+ return_expert_mask: bool = False,
29
+) -> tuple[int, torch.Tensor | None, torch.Tensor | None]:
30
+ """
31
+ Calculates how many experts should be assigned to each rank for EP and
32
+ creates a mapping from global to local expert index. Experts are
33
+ distributed evenly across ranks. Any remaining are assigned to the
34
+ last rank.
35
+
36
+ Args:
37
+ ep_size: The size of the expert parallel group
38
+ ep_rank: The rank of the current process in the expert parallel
39
+ group
40
+ global_num_experts: The total number of experts in the model.
41
+ expert_placement_strategy: The expert placement strategy.
42
+ num_fused_shared_experts: Number of fused shared experts (for AITER)
43
+ return_expert_mask: Whether to return expert mask for AITER
44
+
45
+ Returns:
46
+ tuple[int, Optional[torch.Tensor], Optional[torch.Tensor]]: A tuple containing:
47
+ - local_num_experts (int): The number of experts assigned
48
+ to the current rank.
49
+ - expert_map (Optional[torch.Tensor]): A tensor of shape
50
+ (global_num_experts,) mapping from global to local index.
51
+ Contains -1 for experts not assigned to the current rank.
52
+ Returns None if ep_size is 1.
53
+ - expert_mask (Optional[torch.Tensor]): A tensor of shape
54
+ (global_num_experts + num_fused_shared_experts + 1,)
55
+ containing 1 for experts assigned to the current rank
56
+ and 0 for sentinel.
57
+ Returns None if ep_size is 1.
58
+ Used only when AITER MOE is enabled.
59
+ """
60
+ from typing import get_args
61
+
62
+ assert ep_size > 0
63
+ if ep_size == 1:
64
+ return (global_num_experts, None, None)
65
+
66
+ # Distribute experts as evenly as possible to each rank.
67
+ base_experts = global_num_experts // ep_size
68
+ remainder = global_num_experts % ep_size
69
+ local_num_experts = base_experts + 1 if ep_rank < remainder else base_experts
70
+
71
+ # Create a tensor of size num_experts filled with -1
72
+ expert_map = torch.full((global_num_experts,), -1, dtype=torch.int32)
73
+
74
+ # Create an expert map for the local experts
75
+ if expert_placement_strategy == "linear":
76
+ start_idx = ep_rank * base_experts + min(ep_rank, remainder)
77
+ expert_map[start_idx : start_idx + local_num_experts] = torch.arange(
78
+ 0, local_num_experts, dtype=torch.int32
79
+ )
80
+ elif expert_placement_strategy == "round_robin":
81
+ local_log_experts = torch.arange(
82
+ ep_rank, global_num_experts, ep_size, dtype=torch.int32
83
+ )
84
+
85
+ expert_map[local_log_experts] = torch.arange(
86
+ 0, local_num_experts, dtype=torch.int32
87
+ )
88
+ else:
89
+ raise ValueError(
90
+ "Unsupported expert placement strategy "
91
+ f"'{expert_placement_strategy}', expected one of "
92
+ f"{get_args(ExpertPlacementStrategy)}"
93
+ )
94
+
95
+ expert_mask = None
96
+ if return_expert_mask:
97
+ expert_mask = torch.ones(
98
+ (global_num_experts + num_fused_shared_experts + 1,), dtype=torch.int32
99
+ )
100
+ expert_mask[-1] = 0
101
+ expert_mask[:global_num_experts] = expert_map > -1
102
+ expert_map = torch.cat(
103
+ (
104
+ expert_map,
105
+ torch.tensor(
106
+ [local_num_experts + i for i in range(num_fused_shared_experts)],
107
+ dtype=torch.int32,
108
+ ),
109
+ ),
110
+ dim=0,
111
+ )
112
+
113
+ return (local_num_experts, expert_map, expert_mask)
114
+
115
+
116
+def determine_expert_placement_strategy(
117
+ expert_placement_strategy: ExpertPlacementStrategy,
118
+ moe_parallel_config: FusedMoEParallelConfig,
119
+ num_expert_group: int | None,
120
+ num_redundant_experts: int,
121
+ enable_eplb: bool,
122
+) -> ExpertPlacementStrategy:
123
+ if expert_placement_strategy == "round_robin":
124
+ round_robin_supported = (
125
+ (num_expert_group is not None and num_expert_group > 1)
126
+ and num_redundant_experts == 0
127
+ and not enable_eplb
128
+ )
129
+
130
+ if not round_robin_supported:
131
+ logger.warning(
132
+ "Round-robin expert placement is only supported for "
133
+ "models with multiple expert groups and no redundant "
134
+ "experts. Falling back to linear expert placement."
135
+ )
136
+ return "linear"
137
+ if (
138
+ moe_parallel_config.use_all2all_kernels
139
+ and not moe_parallel_config.needs_round_robin_routing_tables
140
+ ):
141
+ logger.warning(
142
+ "Round-robin expert placement currently only supports "
143
+ "the DeepEP low-latency or NIXL EP backend, but '%s' was configured. "
144
+ "Falling back to linear expert placement.",
145
+ moe_parallel_config.all2all_backend,
146
+ )
147
+ return "linear"
148
+
149
+ return expert_placement_strategy
150
+
151
+
152
+class ExpertMapManager:
153
+ """
154
+ Manages expert ID mappings and placement for Expert Parallelism.
155
+
156
+ Responsibilities:
157
+ - Calculate local vs global expert counts
158
+ - Map between global, local, and physical expert IDs
159
+ - Manage placement strategies (linear, round_robin)
160
+ - Maintain routing tables for round-robin placement
161
+ - Support dynamic reconfiguration of EP topology
162
+
163
+ When expert_map is required:
164
+ - Expert Parallelism (EP) is enabled, i.e., when ep_size > 1
165
+ - EP disabled (ep_size == 1): expert_map is None
166
+ * All experts are local to the current rank
167
+ * No mapping is needed
168
+ - EP enabled (ep_size > 1): expert_map is created
169
+ * Maps global expert IDs to local expert IDs
170
+ * Shape: (global_num_experts,)
171
+ * Contains the local expert index for experts on this rank, -1 for experts
172
+ on other ranks
173
+ * Used by kernels to handle distributed expert execution
174
+ - Kernel support varies:
175
+ * Supports expert_map: fused_moe, fused_marlin_moe, fused_humming_moe,
176
+ rocm_aiter_fused_moe, deep_gemm_moe, xpu_moe, gpt_oss_triton_kernels_moe
177
+ * Does not support: flashinfer_cutlass_moe, fused_batched_moe, most cutlass_moe
178
+ variants, trtllm_* kernels
179
+ * When kernel doesn't support expert_map: The modular kernel method sets
180
+ expert_map=None even if EP is enabled
181
+ """
182
+
183
+ def __init__(
184
+ self,
185
+ max_num_batched_tokens: int,
186
+ top_k: int,
187
+ global_num_experts: int,
188
+ num_redundant_experts: int,
189
+ num_expert_group: int | None,
190
+ moe_parallel_config: FusedMoEParallelConfig,
191
+ placement_strategy: ExpertPlacementStrategy,
192
+ enable_eplb: bool,
193
+ num_fused_shared_experts: int = 0,
194
+ rocm_aiter_enabled: bool = False,
195
+ ):
196
+ """
197
+ Initialize expert map manager.
198
+
199
+ Args:
200
+ global_num_experts: Total number of experts across all ranks
201
+ moe_parallel_config: MoE parallel configuration (contains ep_size,
202
+ ep_rank, backend flags)
203
+ placement_strategy: Strategy for placing experts ('linear' or 'round_robin')
204
+ num_fused_shared_experts: Number of fused shared experts (for AITER)
205
+ rocm_aiter_enabled: Whether ROCm AITER fusion is enabled
206
+ """
207
+ self.global_num_experts = global_num_experts
208
+ self.moe_parallel_config = moe_parallel_config
209
+ self.num_fused_shared_experts = num_fused_shared_experts
210
+ self.rocm_aiter_enabled = rocm_aiter_enabled
211
+ self.top_k = top_k
212
+ self.max_num_batched_tokens = max_num_batched_tokens
213
+
214
+ if moe_parallel_config.use_ep:
215
+ # Determine expert placement strategy before creating manager
216
+ placement_strategy = determine_expert_placement_strategy(
217
+ expert_placement_strategy=placement_strategy,
218
+ moe_parallel_config=moe_parallel_config,
219
+ num_expert_group=num_expert_group,
220
+ num_redundant_experts=num_redundant_experts,
221
+ enable_eplb=enable_eplb,
222
+ )
223
+
224
+ # Determine effective placement strategy
225
+ self._placement_strategy = self._determine_placement_strategy(
226
+ placement_strategy
227
+ )
228
+
229
+ # Calculate expert mappings
230
+ self._calculate_expert_maps()
231
+
232
+ # Initialize routing tables if needed
233
+ self._routing_tables = self._init_routing_tables()
234
+
235
+ self._init_aiter_shared_experts_topK_buffer()
236
+
237
+ if self.use_ep and self.rocm_aiter_enabled:
238
+ expert_mask = self.expert_mask
239
+ assert expert_mask is None or torch.all(
240
+ (expert_mask == 0) | (expert_mask == 1)
241
+ ), "Aiter Fused MoE kernel only supports expert_map with 0 and 1s."
242
+
243
+ # Log EP configuration
244
+ if self.use_ep:
245
+ logger.info_once(
246
+ "[EP Rank %s/%s] Expert parallelism is enabled. Expert "
247
+ "placement strategy: %s. Local/global"
248
+ " number of experts: %s/%s. Experts local to global index map:"
249
+ " %s.",
250
+ self.ep_rank,
251
+ self.ep_size,
252
+ self.placement_strategy,
253
+ self.local_num_experts,
254
+ self.global_num_experts,
255
+ self.get_compressed_map_string(),
256
+ )
257
+
258
+ def _init_aiter_shared_experts_topK_buffer(self):
259
+ if self.num_fused_shared_experts > 0:
260
+ dp_size = self.moe_parallel_config.dp_size
261
+ init_aiter_topK_meta_data(
262
+ n_routed_experts=self.global_num_experts,
263
+ n_shared_experts=self.num_fused_shared_experts,
264
+ top_k=self.top_k,
265
+ tp_rank=self.ep_rank if self.use_ep else self.tp_rank,
266
+ tp_size=self.ep_size if self.use_ep else self.tp_size,
267
+ shared_experts_score=1.0,
268
+ max_num_tokens=self.max_num_batched_tokens * dp_size,
269
+ is_EP=self.use_ep,
270
+ )
271
+
272
+ @property
273
+ def use_ep(self) -> int:
274
+ return self.moe_parallel_config.use_ep
275
+
276
+ @property
277
+ def ep_size(self) -> int:
278
+ return self.moe_parallel_config.ep_size
279
+
280
+ @property
281
+ def ep_rank(self) -> int:
282
+ return self.moe_parallel_config.ep_rank
283
+
284
+ @property
285
+ def tp_size(self) -> int:
286
+ return self.moe_parallel_config.tp_size
287
+
288
+ @property
289
+ def tp_rank(self) -> int:
290
+ return self.moe_parallel_config.tp_rank
291
+
292
+ @property
293
+ def local_num_experts(self) -> int:
294
+ return self._local_num_experts
295
+
296
+ @property
297
+ def expert_map(self) -> torch.Tensor | None:
298
+ """
299
+ Mapping from global expert ID to local expert ID.
300
+
301
+ Returns tensor of shape (global_num_experts,) where:
302
+ - expert_map[global_id] = local_id if expert is on this rank
303
+ - expert_map[global_id] = -1 if expert is not on this rank
304
+
305
+ Returns None if EP is not enabled (ep_size == 1).
306
+ """
307
+ return self._expert_map
308
+
309
+ @property
310
+ def expert_mask(self) -> torch.Tensor | None:
311
+ """
312
+ Expert mask for AITER fusion (ROCm-specific).
313
+
314
+ Returns tensor of shape (global_num_experts + num_fused_shared + 1,)
315
+ where 1 indicates expert is on this rank, 0 otherwise.
316
+ """
317
+ return self._expert_mask
318
+
319
+ @property
320
+ def placement_strategy(self) -> ExpertPlacementStrategy:
321
+ """Expert placement strategy ('linear' or 'round_robin')."""
322
+ return self._placement_strategy
323
+
324
+ @property
325
+ def routing_tables(
326
+ self,
327
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None:
328
+ """
329
+ Routing tables for round-robin placement.
330
+
331
+ Returns (global_to_physical, physical_to_global, local_to_global)
332
+ or None if not using round-robin or tables not needed.
333
+ """
334
+ return self._routing_tables
335
+
336
+ def map_global_to_local(self, global_id: int) -> int:
337
+ """
338
+ Map global expert ID to local expert ID.
339
+
340
+ Args:
341
+ global_id: Global expert ID (0 to global_num_experts - 1)
342
+
343
+ Returns:
344
+ Local expert ID (0 to local_num_experts - 1)
345
+
346
+ Raises:
347
+ ValueError: If expert is not on this rank
348
+ """
349
+ if self._expert_map is None:
350
+ return global_id
351
+
352
+ return self._expert_map[global_id].item()
353
+
354
+ def is_local_expert(self, global_id: int) -> bool:
355
+ """Check if expert is assigned to this rank."""
356
+ if self._expert_map is None:
357
+ return True
358
+ return self._expert_map[global_id] != -1
359
+
360
+ def get_local_expert_ids(self) -> list[int]:
361
+ """Get list of global IDs for experts on this rank."""
362
+ if self._expert_map is None:
363
+ return list(range(self.global_num_experts))
364
+
365
+ return torch.where(self._expert_map != -1)[0].tolist()
366
+
367
+ def update(
368
+ self,
369
+ moe_parallel_config: FusedMoEParallelConfig,
370
+ global_num_experts: int,
371
+ ) -> None:
372
+ """
373
+ Update expert mappings for new EP configuration.
374
+
375
+ Used during dynamic reconfiguration (e.g., elastic scaling).
376
+
377
+ Args:
378
+ global_num_experts: New total number of experts across all ranks
379
+ moe_parallel_config: New MoE parallel configuration (contains ep_size,
380
+ ep_rank, backend flags)
381
+ """
382
+ self.moe_parallel_config = moe_parallel_config
383
+ self.global_num_experts = global_num_experts
384
+
385
+ if self._expert_map is not None:
386
+ device = self._expert_map.device
387
+ elif self._expert_mask is not None:
388
+ device = self._expert_mask.device
389
+ else:
390
+ raise AssertionError("_expert_map or _expert_mask must be present.")
391
+
392
+ with device:
393
+ self._calculate_expert_maps()
394
+ self._routing_tables = self._init_routing_tables()
395
+
396
+ # Reinitialize AITER buffer if needed and parameters provided
397
+ self._init_aiter_shared_experts_topK_buffer()
398
+
399
+ def get_compressed_map_string(self) -> str:
400
+ """
401
+ Get compressed string representation of expert map for logging.
402
+
403
+ Returns string mapping local to global expert IDs.
404
+ """
405
+ if self._expert_map is None:
406
+ return f"[0..{self.global_num_experts - 1}]"
407
+
408
+ global_indices = torch.where(self._expert_map != -1)[0]
409
+ local_indices = self._expert_map[global_indices]
410
+ return ", ".join(
411
+ f"{local_index.item()}->{global_index.item()}"
412
+ for local_index, global_index in zip(local_indices, global_indices)
413
+ )
414
+
415
+ # Private methods
416
+
417
+ def _determine_placement_strategy(
418
+ self, requested_strategy: ExpertPlacementStrategy
419
+ ) -> ExpertPlacementStrategy:
420
+ """Determine effective placement strategy based on config."""
421
+ if requested_strategy != "round_robin":
422
+ return requested_strategy
423
+
424
+ # Round-robin requires specific conditions
425
+ if self.ep_size == 1:
426
+ return "linear"
427
+
428
+ if (
429
+ self.moe_parallel_config.use_all2all_kernels
430
+ and not self.moe_parallel_config.needs_round_robin_routing_tables
431
+ ):
432
+ logger.warning(
433
+ "Round-robin placement requires DeepEP-ll or NIXL backend. "
434
+ "Falling back to linear."
435
+ )
436
+ return "linear"
437
+
438
+ return "round_robin"
439
+
440
+ def _calculate_expert_maps(self) -> None:
441
+ """Calculate expert mappings based on placement strategy."""
442
+ (
443
+ self._local_num_experts,
444
+ self._expert_map,
445
+ self._expert_mask,
446
+ ) = determine_expert_map(
447
+ ep_size=self.ep_size,
448
+ ep_rank=self.ep_rank,
449
+ global_num_experts=self.global_num_experts,
450
+ expert_placement_strategy=self._placement_strategy,
451
+ num_fused_shared_experts=self.num_fused_shared_experts,
452
+ return_expert_mask=self.rocm_aiter_enabled,
453
+ )
454
+
455
+ self._local_num_experts += self.num_fused_shared_experts
456
+
457
+ def _init_routing_tables(
458
+ self,
459
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None:
460
+ """
461
+ Ensure routing tables are initialized if needed for round-robin.
462
+
463
+ This is a public method that can be called to explicitly initialize
464
+ routing tables. It's safe to call multiple times (idempotent).
465
+ """
466
+ if self._placement_strategy != "round_robin":
467
+ return None
468
+
469
+ if not self.moe_parallel_config.needs_round_robin_routing_tables:
470
+ return None
471
+
472
+ if self._expert_map is None:
473
+ return None
474
+
475
+ return self._init_round_robin_expert_routing_tables()
476
+
477
+ def _init_round_robin_expert_routing_tables(
478
+ self,
479
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
480
+ """Build routing tables for round-robin placement."""
481
+ assert self.num_fused_shared_experts == 0, (
482
+ "Round robin not supported for AITER."
483
+ )
484
+
485
+ global_indices = torch.arange(
486
+ self.global_num_experts,
487
+ dtype=torch.long,
488
+ )
489
+ owner = torch.remainder(global_indices, self.ep_size)
490
+ local_index = torch.div(global_indices, self.ep_size, rounding_mode="floor")
491
+
492
+ base = self.global_num_experts // self.ep_size
493
+ remainder = self.global_num_experts % self.ep_size
494
+ physical_offset = owner * base
495
+
496
+ if remainder > 0:
497
+ remainder_tensor = torch.tensor(
498
+ remainder,
499
+ dtype=torch.long,
500
+ )
501
+ physical_offset = physical_offset + torch.minimum(owner, remainder_tensor)
502
+
503
+ global_to_physical = physical_offset + local_index
504
+ physical_to_global = torch.empty_like(global_to_physical)
505
+ physical_to_global[global_to_physical] = global_indices
506
+
507
+ local_global = torch.arange(
508
+ self.ep_rank,
509
+ self.global_num_experts,
510
+ self.ep_size,
511
+ dtype=torch.long,
512
+ )
513
+ if local_global.numel() != self._local_num_experts:
514
+ local_global = local_global[: self._local_num_experts]
515
+
516
+ return (global_to_physical, physical_to_global, local_global)
vllm/model_executor/layers/fused_moe/layer.py
+54
-304
@@ -3,13 +3,13 @@
3
4
from collections.abc import Callable, Iterable
5
from enum import Enum
6
-from typing import Literal, cast, get_args, overload
6
+from typing import Literal, cast, overload
7
8
import torch
9
from torch.nn.parameter import UninitializedParameter
10
11
from vllm._aiter_ops import rocm_aiter_ops
12
-from vllm.config import VllmConfig, get_current_vllm_config
12
+from vllm.config import get_current_vllm_config
13
from vllm.config.parallel import ExpertPlacementStrategy
14
from vllm.distributed import (
15
get_dp_group,
@@ -26,8 +26,8 @@ from vllm.model_executor.layers.fused_moe.config import (
26
FusedMoEQuantConfig,
27
RoutingMethodType,
28
)
29
-from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
30
- init_aiter_topK_meta_data,
29
+from vllm.model_executor.layers.fused_moe.expert_map_manager import (
30
+ ExpertMapManager,
31
)
32
from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
33
FusedMoEMethodBase,
@@ -68,152 +68,6 @@ class FusedMoeWeightScaleSupported(Enum):
68
BLOCK = "block"
69
70
71
-def determine_expert_map(
72
- ep_size: int,
73
- ep_rank: int,
74
- global_num_experts: int,
75
- expert_placement_strategy: ExpertPlacementStrategy = "linear",
76
- num_fused_shared_experts: int = 0,
77
- return_expert_mask: bool = False,
78
-) -> tuple[int, torch.Tensor | None, torch.Tensor | None]:
79
- """
80
- Calculates how many experts should be assigned to each rank for EP and
81
- creates a mapping from global to local expert index. Experts are
82
- distributed evenly across ranks. Any remaining are assigned to the
83
- last rank.
84
-
85
- Args:
86
- ep_size: The size of the expert parallel group
87
- ep_rank: The rank of the current process in the expert parallel
88
- group
89
- global_num_experts: The total number of experts in the model.
90
- expert_placement_strategy: The expert placement strategy.
91
-
92
- Returns:
93
- tuple[int, Optional[torch.Tensor]]: A tuple containing:
94
- - local_num_experts (int): The number of experts assigned
95
- to the current rank.
96
- - expert_map (Optional[torch.Tensor]): A tensor of shape
97
- (global_num_experts,) mapping from global to local index.
98
- Contains -1 for experts not assigned to the current rank.
99
- Returns None if ep_size is 1.
100
- - expert_mask (Optional[torch.Tensor]): A tensor of shape
101
- (global_num_experts + num_fused_shared_experts + 1,)
102
- containing 1 for experts assigned to the current rank
103
- and 0 for sentinel.
104
- Returns None if ep_size is 1.
105
- Used only when AITER MOE is enabled.
106
- """
107
- assert ep_size > 0
108
- if ep_size == 1:
109
- return (global_num_experts, None, None)
110
-
111
- # Distribute experts as evenly as possible to each rank.
112
- base_experts = global_num_experts // ep_size
113
- remainder = global_num_experts % ep_size
114
- local_num_experts = base_experts + 1 if ep_rank < remainder else base_experts
115
-
116
- # Create a tensor of size num_experts filled with -1
117
- expert_map = torch.full((global_num_experts,), -1, dtype=torch.int32)
118
- # Create an expert map for the local experts
119
- if expert_placement_strategy == "linear":
120
- start_idx = ep_rank * base_experts + min(ep_rank, remainder)
121
- expert_map[start_idx : start_idx + local_num_experts] = torch.arange(
122
- 0, local_num_experts, dtype=torch.int32
123
- )
124
- elif expert_placement_strategy == "round_robin":
125
- local_log_experts = torch.arange(
126
- ep_rank, global_num_experts, ep_size, dtype=torch.int32
127
- )
128
-
129
- expert_map[local_log_experts] = torch.arange(
130
- 0, local_num_experts, dtype=torch.int32
131
- )
132
- else:
133
- raise ValueError(
134
- "Unsupported expert placement strategy "
135
- f"'{expert_placement_strategy}', expected one of "
136
- f"{get_args(ExpertPlacementStrategy)}"
137
- )
138
-
139
- expert_mask = None
140
- if return_expert_mask:
141
- expert_mask = torch.ones(
142
- (global_num_experts + num_fused_shared_experts + 1,), dtype=torch.int32
143
- )
144
- expert_mask[-1] = 0
145
- expert_mask[:global_num_experts] = expert_map > -1
146
- expert_map = torch.cat(
147
- (
148
- expert_map,
149
- torch.tensor(
150
- [local_num_experts + i for i in range(num_fused_shared_experts)],
151
- dtype=torch.int32,
152
- ),
153
- ),
154
- dim=0,
155
- )
156
-
157
- return (local_num_experts, expert_map, expert_mask)
158
-
159
-
160
-def determine_expert_placement_strategy(
161
- expert_placement_strategy: ExpertPlacementStrategy,
162
- moe_parallel_config: FusedMoEParallelConfig,
163
- num_expert_group: int | None,
164
- num_redundant_experts: int,
165
- enable_eplb: bool,
166
-) -> ExpertPlacementStrategy:
167
- if expert_placement_strategy == "round_robin":
168
- round_robin_supported = (
169
- (num_expert_group is not None and num_expert_group > 1)
170
- and num_redundant_experts == 0
171
- and not enable_eplb
172
- )
173
-
174
- if not round_robin_supported:
175
- logger.warning(
176
- "Round-robin expert placement is only supported for "
177
- "models with multiple expert groups and no redundant "
178
- "experts. Falling back to linear expert placement."
179
- )
180
- return "linear"
181
- if (
182
- moe_parallel_config.use_all2all_kernels
183
- and not moe_parallel_config.needs_round_robin_routing_tables
184
- ):
185
- logger.warning(
186
- "Round-robin expert placement currently only supports "
187
- "the DeepEP low-latency or NIXL EP backend, but '%s' was configured. "
188
- "Falling back to linear expert placement.",
189
- moe_parallel_config.all2all_backend,
190
- )
191
- return "linear"
192
-
193
- return expert_placement_strategy
194
-
195
-
196
-def get_compressed_expert_map(expert_map: torch.Tensor) -> str:
197
- """
198
- Compresses the expert map by removing any -1 entries.
199
-
200
- Args:
201
- expert_map (torch.Tensor): A tensor of shape (global_num_experts,)
202
- mapping from global to local index. Contains -1 for experts not
203
- assigned to the current rank.
204
-
205
- Returns:
206
- str: A string mapping from local to global index.
207
- Using str to support hashing for logging once only.
208
- """
209
- global_indices = torch.where(expert_map != -1)[0]
210
- local_indices = expert_map[global_indices]
211
- return ", ".join(
212
- f"{local_index.item()}->{global_index.item()}"
213
- for local_index, global_index in zip(local_indices, global_indices)
214
- )
215
-
216
-
71
# --8<-- [start:fused_moe]
72
@PluggableLayer.register("fused_moe")
73
class FusedMoE(PluggableLayer):
@@ -394,55 +248,27 @@ class FusedMoE(PluggableLayer):
248
"Redundant experts are only supported with EPLB."
249
)
250
397
- self.expert_placement_strategy = determine_expert_placement_strategy(
398
- expert_placement_strategy=self.expert_placement_strategy,
399
- moe_parallel_config=self.moe_parallel_config,
400
- num_expert_group=num_expert_group,
401
- num_redundant_experts=num_redundant_experts,
402
- enable_eplb=self.enable_eplb,
403
- )
251
+ max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens
252
405
- self._expert_map: torch.Tensor | None
406
- local_num_experts, expert_map, expert_mask = determine_expert_map(
407
- ep_size=self.ep_size,
408
- ep_rank=self.ep_rank,
409
- global_num_experts=self.global_num_experts,
410
- expert_placement_strategy=self.expert_placement_strategy,
411
- num_fused_shared_experts=self.num_fused_shared_experts,
412
- return_expert_mask=self.rocm_aiter_fmoe_enabled,
413
- )
414
- self.local_num_experts = local_num_experts
415
- self.register_buffer("_expert_map", expert_map)
416
- self.register_buffer("expert_mask", expert_mask)
417
- self._maybe_init_expert_routing_tables()
418
- logger.info_once(
419
- "[EP Rank %s/%s] Expert parallelism is enabled. Expert "
420
- "placement strategy: %s. Local/global"
421
- " number of experts: %s/%s. Experts local to global index map:"
422
- " %s.",
423
- self.ep_rank,
424
- self.ep_size,
425
- self.expert_placement_strategy,
426
- self.local_num_experts,
427
- self.global_num_experts,
428
- get_compressed_expert_map(self._expert_map),
429
- )
430
- else:
431
- self.local_num_experts, self._expert_map, self.expert_mask = (
432
- self.global_num_experts,
433
- None,
434
- None,
435
- )
253
+ # Create ExpertMapManager to handle expert mapping and placement for EP.
254
+ # See ExpertMapManager for a detailed description of what it does and when
255
+ # it is required.
256
+ self.expert_map_manager = ExpertMapManager(
257
+ max_num_batched_tokens=max_num_batched_tokens,
258
+ top_k=top_k,
259
+ global_num_experts=self.global_num_experts,
260
+ num_redundant_experts=num_redundant_experts,
261
+ num_expert_group=num_expert_group,
262
+ moe_parallel_config=self.moe_parallel_config,
263
+ placement_strategy=self.expert_placement_strategy,
264
+ enable_eplb=self.enable_eplb,
265
+ num_fused_shared_experts=self.num_fused_shared_experts,
266
+ rocm_aiter_enabled=self.rocm_aiter_fmoe_enabled,
267
+ )
268
437
- self.top_k = top_k
269
+ self.update_expert_map_info()
270
439
- self._init_aiter_shared_experts_topK_buffer(
440
- vllm_config=vllm_config, dp_size=dp_size_
441
- )
442
- if self.use_ep and self.rocm_aiter_fmoe_enabled:
443
- assert self.expert_mask is None or torch.all(
444
- (expert_mask == 0) | (expert_mask == 1)
445
- ), "Aiter Fused MoE kernel only supports expert_map with 0 and 1s."
271
+ self.top_k = top_k
272
273
assert intermediate_size % self.tp_size == 0
274
intermediate_size_per_partition = intermediate_size // self.tp_size
@@ -510,7 +336,7 @@ class FusedMoE(PluggableLayer):
336
in_dtype=moe_in_dtype,
337
moe_backend=vllm_config.kernel_config.moe_backend,
338
router_logits_dtype=router_logits_dtype,
513
- max_num_tokens=vllm_config.scheduler_config.max_num_batched_tokens,
339
+ max_num_tokens=max_num_batched_tokens,
340
has_bias=has_bias,
341
is_act_and_mul=is_act_and_mul,
342
is_lora_enabled=vllm_config.lora_config is not None,
@@ -642,11 +468,8 @@ class FusedMoE(PluggableLayer):
468
return None
469
470
self.ensure_moe_quant_config_init()
645
- # routing_tables only needed for round-robin expert placement with
646
- # DeepEP all2all backend.
647
- routing_tables = self._maybe_init_expert_routing_tables()
471
prepare_finalize = self.base_quant_method.maybe_make_prepare_finalize(
649
- routing_tables=routing_tables
472
+ routing_tables=self._expert_routing_tables()
473
)
474
if prepare_finalize is not None:
475
logger.debug(
@@ -698,16 +521,26 @@ class FusedMoE(PluggableLayer):
521
# By default, router/gate is called before FusedMoE forward pass
522
return self.runner.is_internal_router()
523
701
- def _maybe_init_expert_routing_tables(
524
+ def update_expert_map_info(self):
525
+ # Update local attributes from ExpertMapManager
526
+ self.local_num_experts = self.expert_map_manager.local_num_experts
527
+ self.expert_placement_strategy = self.expert_map_manager.placement_strategy
528
+ self.register_buffer("_expert_map", self.expert_map_manager.expert_map)
529
+ self.register_buffer("expert_mask", self.expert_map_manager.expert_mask)
530
+
531
+ # Get routing tables from ExpertMapManager
532
+ routing_tables = self.expert_map_manager.routing_tables
533
+ if routing_tables is not None:
534
+ # Register routing tables as buffers for this layer
535
+ global_to_physical, physical_to_global, local_global = routing_tables
536
+ self.register_buffer("expert_global_to_physical", global_to_physical)
537
+ self.register_buffer("expert_physical_to_global", physical_to_global)
538
+ self.register_buffer("expert_local_to_global", local_global)
539
+
540
+ def _expert_routing_tables(
541
self,
542
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None:
704
- # Currently routing_tables only needed for round-robin expert placement
705
- # with DeepEP-ll or NIXL EP all2all backends.
706
- if self.expert_placement_strategy != "round_robin" or (
707
- not self.moe_parallel_config.needs_round_robin_routing_tables
708
- ):
709
- return None
710
-
543
+ # Return cached routing tables if already registered as buffers
544
if hasattr(self, "expert_global_to_physical"):
545
return cast(
546
tuple[torch.Tensor, torch.Tensor, torch.Tensor],
@@ -717,85 +550,21 @@ class FusedMoE(PluggableLayer):
550
self.expert_local_to_global,
551
),
552
)
553
+ return None
554
721
- if self._expert_map is None:
722
- return None
723
-
724
- routing_tables = self.ensure_round_robin_expert_routing_tables(
555
+ def update_expert_map(self):
556
+ # Update ExpertMapManager with new EP configuration
557
+ # The moe_parallel_config (including ep_size and ep_rank)
558
+ # should already be updated.
559
+ # Note: ExpertMapManager.update() recalculates expert maps and
560
+ # reinitializes routing tables internally.
561
+ self.expert_map_manager.update(
562
+ self.moe_parallel_config,
563
global_num_experts=self.global_num_experts,
726
- ep_size=self.ep_size,
727
- ep_rank=self.ep_rank,
728
- local_num_experts=self.local_num_experts,
729
- device=self._expert_map.device,
730
- )
731
-
732
- global_to_physical, physical_to_global, local_global = routing_tables
733
- self.register_buffer("expert_global_to_physical", global_to_physical)
734
- self.register_buffer("expert_physical_to_global", physical_to_global)
735
- self.register_buffer("expert_local_to_global", local_global)
736
-
737
- return routing_tables
738
-
739
- @staticmethod
740
- def ensure_round_robin_expert_routing_tables(
741
- global_num_experts: int,
742
- ep_size: int,
743
- ep_rank: int,
744
- local_num_experts: int,
745
- device: torch.device | None = None,
746
- ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
747
- device_kwargs = {"device": device} if device is not None else {}
748
- global_indices = torch.arange(
749
- global_num_experts, dtype=torch.long, **device_kwargs
750
- )
751
- owner = torch.remainder(global_indices, ep_size)
752
- local_index = torch.div(global_indices, ep_size, rounding_mode="floor")
753
- base = global_num_experts // ep_size
754
- remainder = global_num_experts % ep_size
755
- physical_offset = owner * base
756
- if remainder > 0:
757
- remainder_tensor = torch.tensor(
758
- remainder, dtype=torch.long, **device_kwargs
759
- )
760
- physical_offset = physical_offset + torch.minimum(owner, remainder_tensor)
761
-
762
- global_to_physical = physical_offset + local_index
763
- physical_to_global = torch.empty_like(global_to_physical)
764
- physical_to_global[global_to_physical] = global_indices
765
-
766
- local_global = torch.arange(
767
- ep_rank,
768
- global_num_experts,
769
- ep_size,
770
- dtype=torch.long,
771
- **device_kwargs,
564
)
773
- if local_global.numel() != local_num_experts:
774
- local_global = local_global[:local_num_experts]
775
-
776
- return (global_to_physical, physical_to_global, local_global)
565
778
- def update_expert_map(self):
779
- # ep_size and ep_rank should already be updated
780
- assert self._expert_map is not None
781
- with self._expert_map.device:
782
- local_num_experts, expert_map, expert_mask = determine_expert_map(
783
- ep_size=self.ep_size,
784
- ep_rank=self.ep_rank,
785
- global_num_experts=self.global_num_experts,
786
- expert_placement_strategy=self.expert_placement_strategy,
787
- num_fused_shared_experts=self.num_fused_shared_experts,
788
- return_expert_mask=self.rocm_aiter_fmoe_enabled,
789
- )
790
- self.local_num_experts = local_num_experts
791
- self.register_buffer("_expert_map", expert_map)
792
- self.register_buffer("expert_mask", expert_mask)
793
- self._maybe_init_expert_routing_tables()
794
- if self.aiter_fmoe_shared_expert_enabled:
795
- self._init_aiter_shared_experts_topK_buffer(
796
- vllm_config=get_current_vllm_config(),
797
- dp_size=get_dp_group().world_size,
798
- )
566
+ # Update local attributes from ExpertMapManager
567
+ self.update_expert_map_info()
568
569
def _load_per_tensor_weight_scale(
570
self,
@@ -1063,26 +832,7 @@ class FusedMoE(PluggableLayer):
832
expert_data.copy_(loaded_weight)
833
834
def _map_global_expert_id_to_local_expert_id(self, expert_id: int) -> int:
1066
- if self._expert_map is None:
1067
- return expert_id
1068
- return self._expert_map[expert_id].item()
1069
-
1070
- def _init_aiter_shared_experts_topK_buffer(
1071
- self, vllm_config: VllmConfig, dp_size: int
1072
- ):
1073
- if self.num_fused_shared_experts > 0:
1074
- init_aiter_topK_meta_data(
1075
- n_routed_experts=self.global_num_experts,
1076
- n_shared_experts=self.num_fused_shared_experts,
1077
- top_k=self.top_k,
1078
- tp_rank=self.ep_rank if self.use_ep else self.tp_rank,
1079
- tp_size=self.ep_size if self.use_ep else self.tp_size,
1080
- shared_experts_score=1.0,
1081
- max_num_tokens=vllm_config.scheduler_config.max_num_batched_tokens
1082
- * dp_size,
1083
- is_EP=self.use_ep,
1084
- )
1085
- self.local_num_experts += self.num_fused_shared_experts
835
+ return self.expert_map_manager.map_global_to_local(expert_id)
836
837
@overload
838
def weight_loader(
vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
+1
-1
@@ -191,7 +191,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
191
moe_config=self.moe,
192
backend=self.unquantized_backend,
193
experts_cls=self.experts_cls,
194
- routing_tables=layer._maybe_init_expert_routing_tables(),
194
+ routing_tables=layer._expert_routing_tables(),
195
shared_experts=layer.shared_experts,
196
)
197
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py
+1
-1
@@ -195,7 +195,7 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):
195
experts_cls=self.experts_cls,
196
mxfp4_backend=self.mxfp4_backend,
197
shared_experts=layer.shared_experts,
198
- routing_tables=layer._maybe_init_expert_routing_tables(),
198
+ routing_tables=layer._expert_routing_tables(),
199
)
200
201
def apply(
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py
+1
-1
@@ -236,7 +236,7 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
236
moe_config=self.moe,
237
experts_cls=self.experts_cls,
238
shared_experts=layer.shared_experts,
239
- routing_tables=layer._maybe_init_expert_routing_tables(),
239
+ routing_tables=layer._expert_routing_tables(),
240
)
241
self.moe_kernel.fused_experts.process_weights_after_loading(layer)
242
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_fp8.py
+1
-1
@@ -336,7 +336,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
336
moe_config=self.moe,
337
fp8_backend=self.fp8_backend,
338
experts_cls=self.experts_cls,
339
- routing_tables=layer._maybe_init_expert_routing_tables(),
339
+ routing_tables=layer._expert_routing_tables(),
340
shared_experts=layer.shared_experts,
341
)
342
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_int8.py
+1
-1
@@ -147,7 +147,7 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
147
moe_quant_config=self.moe_quant_config,
148
moe_config=self.moe,
149
experts_cls=self.experts_cls,
150
- routing_tables=layer._maybe_init_expert_routing_tables(),
150
+ routing_tables=layer._expert_routing_tables(),
151
shared_experts=layer.shared_experts,
152
)
153
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_mxfp8.py
+1
-1
@@ -138,7 +138,7 @@ class CompressedTensorsW8A8Mxfp8MoEMethod(CompressedTensorsMoEMethod):
138
moe_config=self.moe,
139
fp8_backend=self.fp8_backend,
140
experts_cls=self.experts_cls,
141
- routing_tables=layer._maybe_init_expert_routing_tables(),
141
+ routing_tables=layer._expert_routing_tables(),
142
shared_experts=layer.shared_experts,
143
)
144
vllm/model_executor/layers/quantization/fp8.py
+1
-1
@@ -786,7 +786,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
786
moe_config=self.moe,
787
fp8_backend=self.fp8_backend,
788
experts_cls=self.experts_cls,
789
- routing_tables=layer._maybe_init_expert_routing_tables(),
789
+ routing_tables=layer._expert_routing_tables(),
790
shared_experts=layer.shared_experts,
791
)
792
vllm/model_executor/layers/quantization/gptq_marlin.py
+1
-1
@@ -769,7 +769,7 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
769
w2_g_idx=layer.w2_g_idx,
770
w13_g_idx_sort_indices=layer.w13_g_idx_sort_indices,
771
w2_g_idx_sort_indices=layer.w2_g_idx_sort_indices,
772
- routing_tables=layer._maybe_init_expert_routing_tables(),
772
+ routing_tables=layer._expert_routing_tables(),
773
shared_experts=layer.shared_experts,
774
)
775
vllm/model_executor/layers/quantization/modelopt.py
+2
-2
@@ -901,7 +901,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
901
moe_config=self.moe,
902
fp8_backend=self.fp8_backend,
903
experts_cls=self.experts_cls,
904
- routing_tables=layer._maybe_init_expert_routing_tables(),
904
+ routing_tables=layer._expert_routing_tables(),
905
shared_experts=layer.shared_experts,
906
)
907
@@ -1590,7 +1590,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1590
moe_config=self.moe,
1591
experts_cls=self.experts_cls,
1592
shared_experts=layer.shared_experts,
1593
- routing_tables=layer._maybe_init_expert_routing_tables(),
1593
+ routing_tables=layer._expert_routing_tables(),
1594
)
1595
self.moe_kernel.fused_experts.process_weights_after_loading(layer)
1596
vllm/model_executor/layers/quantization/mxfp4.py
+2
-2
@@ -369,7 +369,7 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase):
369
moe_config=self.moe,
370
mxfp4_backend=self.mxfp4_backend,
371
experts_cls=self.experts_cls,
372
- routing_tables=layer._maybe_init_expert_routing_tables(),
372
+ routing_tables=layer._expert_routing_tables(),
373
shared_experts=layer.shared_experts,
374
layer=layer,
375
)
@@ -702,7 +702,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
702
moe_config=self.moe,
703
mxfp4_backend=self.mxfp4_backend,
704
experts_cls=self.experts_cls,
705
- routing_tables=layer._maybe_init_expert_routing_tables(),
705
+ routing_tables=layer._expert_routing_tables(),
706
shared_experts=layer.shared_experts,
707
layer=layer,
708
)
vllm/model_executor/layers/quantization/online/fp8.py
+1
-1
@@ -354,7 +354,7 @@ class _Fp8OnlineMoEBase(OnlineMoEMethodBase):
354
moe_config=self.moe,
355
fp8_backend=self.fp8_backend,
356
experts_cls=self.experts_cls,
357
- routing_tables=layer._maybe_init_expert_routing_tables(),
357
+ routing_tables=layer._expert_routing_tables(),
358
shared_experts=layer.shared_experts,
359
)
360
vllm/model_executor/layers/quantization/online/int8.py
+1
-1
@@ -99,7 +99,7 @@ class Int8OnlineMoEMethod(OnlineMoEMethodBase):
99
moe_quant_config=self.moe_quant_config,
100
moe_config=self.moe,
101
experts_cls=self.experts_cls,
102
- routing_tables=layer._maybe_init_expert_routing_tables(),
102
+ routing_tables=layer._expert_routing_tables(),
103
shared_experts=layer.shared_experts,
104
)
105
vllm/model_executor/layers/quantization/online/mxfp8.py
+1
-1
@@ -199,7 +199,7 @@ class Mxfp8OnlineMoEMethod(OnlineMoEMethodBase):
199
moe_config=self.moe,
200
fp8_backend=self.fp8_backend,
201
experts_cls=self.experts_cls,
202
- routing_tables=layer._maybe_init_expert_routing_tables(),
202
+ routing_tables=layer._expert_routing_tables(),
203
shared_experts=layer.shared_experts,
204
)
205
vllm/model_executor/layers/quantization/quark/quark_moe.py
+1
-1
@@ -1340,7 +1340,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1340
moe_config=self.moe,
1341
mxfp4_backend=self.mxfp4_backend,
1342
experts_cls=self.experts_cls,
1343
- routing_tables=layer._maybe_init_expert_routing_tables(),
1343
+ routing_tables=layer._expert_routing_tables(),
1344
shared_experts=layer.shared_experts,
1345
)
1346