[Bugfix][R3] Size monolithic routing replay buffer for DP (#50874)
Signed-off-by: tbarnatan <tbarnatan@nvidia.com> Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: zjy0516 <riverclouds.zhu@qq.com>
TomerBN-Nvidia committed
Aug 13, 2026 at 10:18 UTC
399f97424e61ea918135bd80176491646cb69d17
2 files changed
+30
-5
vllm/model_executor/layers/fused_moe/modular_kernel.py
+8
-1
@@ -1023,8 +1023,15 @@ class FusedMoEExpertsMonolithic(FusedMoEExperts):
1023
if capture_fn is None:
1024
self._routing_replay_buffer = None
1025
return
1026
+ # Allocate for per-rank batches gathered across the DP or EP group.
1027
+ dispatch_group_size = (
1028
+ self.moe_config.ep_size
1029
+ if self.moe_config.use_ep
1030
+ else self.moe_config.dp_size
1031
+ )
1032
+ max_num_replay_tokens = self.moe_config.max_num_tokens * dispatch_group_size
1033
self._routing_replay_buffer = torch.empty(
1027
- (self.moe_config.max_num_tokens, self.moe_config.experts_per_token),
1034
+ (max_num_replay_tokens, self.moe_config.experts_per_token),
1035
dtype=torch.int16,
1036
device=self.moe_config.device,
1037
)
vllm/model_executor/layers/fused_moe/routed_experts_capturer.py
+22
-4
@@ -108,15 +108,20 @@ class RoutedExpertsCapturer:
108
def capture(self, layer_id: int, topk_ids: torch.Tensor) -> None:
109
"""Capture expert routing decisions for a specific layer.
110
111
- Under data parallelism, ``topk_ids`` may have three different batch
111
+ Under data parallelism, ``topk_ids`` may have four different batch
112
layouts depending on where the DP combine happens and whether
113
- Sequence Parallelism (SP) is active for the MoE layer:
113
+ Expert Parallelism (EP) or Sequence Parallelism (SP) is active for the
114
+ MoE layer:
115
- ``n == total`` (naive dispatch): all DP ranks' tokens are
116
concatenated before routing; we slice out this rank's span
117
using the cumulative per-rank counts.
118
- ``n == token_num_per_dp`` (modular-kernel path): DP combine
119
happens inside ``quant_method.apply``; ``select_experts`` only
120
ever sees this rank's tokens, so we take the whole tensor.
121
+ - ``n == sum(dp_metadata.local_sizes)`` (naive DP+EP dispatch):
122
+ sequence-parallel shards from every DP rank are gathered through
123
+ the flattened EP group. The shard sizes include CUDA-graph / SP
124
+ padding, so we use them to locate this DP rank's unpadded rows.
125
- ``n == ceil(token_num_per_dp / tp_size)`` (SP + modular-kernel
126
path): tokens were split along dim=0 across the TP group by
127
``_sequence_parallel_context``
@@ -143,6 +148,8 @@ class RoutedExpertsCapturer:
148
token_num_per_dp = int(num_tokens_dp[self.dp_rank].item())
149
total = int(num_tokens_dp.sum().item())
150
n = topk_ids.shape[0]
151
+ shard_sizes = getattr(ctx.dp_metadata, "local_sizes", None)
152
+ gathered_size = sum(shard_sizes) if shard_sizes is not None else -1
153
154
if n == total:
155
# Naive dispatch: all DP ranks' tokens concatenated
@@ -157,6 +164,17 @@ class RoutedExpertsCapturer:
164
# rank's tokens, take the whole tensor.
165
start_loc = 0
166
end_loc = token_num_per_dp
167
+ elif shard_sizes is not None and n == gathered_size:
168
+ # Naive DP+EP dispatch gathers the sequence-parallel shards in
169
+ # flattened DP-then-TP order. ``local_sizes`` is the exact
170
+ # all-gatherv layout, including padding. Select this DP rank's
171
+ # contiguous shard group, then trim its trailing padding.
172
+ num_dp_ranks = len(num_tokens_dp)
173
+ assert len(shard_sizes) % num_dp_ranks == 0
174
+ shards_per_dp_rank = len(shard_sizes) // num_dp_ranks
175
+ first_shard = self.dp_rank * shards_per_dp_rank
176
+ start_loc = sum(shard_sizes[:first_shard])
177
+ end_loc = start_loc + token_num_per_dp
178
elif (
179
self.tp_size > 1
180
and n != token_num_per_dp
@@ -189,8 +207,8 @@ class RoutedExpertsCapturer:
207
raise AssertionError(
208
"RoutedExpertsCapturer: unexpected topk_ids batch "
209
f"dim {n} (expected {total}, {token_num_per_dp}, "
192
- f"or {sp_expected} for dp_rank={self.dp_rank}, "
193
- f"tp_size={self.tp_size})"
210
+ f"{gathered_size}, or {sp_expected} for "
211
+ f"dp_rank={self.dp_rank}, tp_size={self.tp_size})"
212
)
213
214
if layer_id >= self.device_buffer.shape[1]: