[Perf][MoE] Optimize deepep_v2 receiver CPU Overhead (#51114)
Signed-off-by: Lucas Wilkinson <wilkinson.lucas@gmail.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
Lucas Wilkinson committed
Aug 17, 2026 at 14:21 UTC
75dde08d3fd0cf4217c90e8499881d60fe426aa0
1 file changed
+25
-16
vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py
+25
-16
@@ -80,9 +80,23 @@ class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
80
# DBO microbatching: one handle slot per micro-batch.
81
self.handles: list[deep_ep.EPHandle | None] = [None, None]
82
83
+ # arange(num_local_experts) + rank_expert_offset. Rank-constant, so it
84
+ # is built once per device instead of once per layer per step.
85
+ self._global_expert_ids_cache: torch.Tensor | None = None
86
+
87
def num_dispatchers(self) -> int:
88
return self.num_dispatchers_
89
90
+ def _global_expert_ids(self, num_local: int, device: torch.device) -> torch.Tensor:
91
+ ids = self._global_expert_ids_cache
92
+ if ids is None or ids.numel() != num_local or ids.device != device:
93
+ ids = (
94
+ torch.arange(num_local, dtype=torch.int64, device=device)
95
+ + self.rank_expert_offset
96
+ )
97
+ self._global_expert_ids_cache = ids
98
+ return ids
99
+
100
def output_is_reduced(self) -> bool:
101
return True
102
@@ -196,23 +210,23 @@ class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
210
else:
211
expert_x, expert_x_scale = recv_x, None
212
213
+ expert_tokens_meta = mk.ExpertTokensMetadata.make_from_list(
214
+ recv_expert_num_tokens,
215
+ device=expert_x.device,
216
+ )
217
+
218
if recv_topk_idx is None:
219
# do_expand=True (prefill mode): build topk_ids from
220
# per-expert token counts.
221
total_tokens = sum(recv_expert_num_tokens)
222
if total_tokens > 0:
204
- recv_topk_idx = torch.empty(
205
- total_tokens,
206
- dtype=torch.int64,
207
- device=expert_x.device,
223
+ recv_topk_idx = torch.repeat_interleave(
224
+ self._global_expert_ids(
225
+ len(recv_expert_num_tokens), expert_x.device
226
+ ),
227
+ expert_tokens_meta.expert_num_tokens,
228
+ output_size=total_tokens,
229
)
209
- offset = 0
210
- for i, count in enumerate(recv_expert_num_tokens):
211
- if count > 0:
212
- recv_topk_idx[offset : offset + count].fill_(
213
- i + self.rank_expert_offset
214
- )
215
- offset += count
230
else:
231
recv_topk_idx = torch.empty(
232
0,
@@ -243,11 +257,6 @@ class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
257
if recv_topk_weights is not None and recv_topk_weights.ndim == 1:
258
recv_topk_weights = recv_topk_weights.unsqueeze(1)
259
246
- expert_tokens_meta = mk.ExpertTokensMetadata.make_from_list(
247
- recv_expert_num_tokens,
248
- device=expert_x.device,
249
- )
250
-
260
if not quant_config.is_block_quantized and not defer_input_quant:
261
expert_x_scale = None
262
if expert_x.numel() != 0: