[DeepEP V2] Bound num_max_tokens_per_rank in do_expand=False (#46404)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: Roy Wang <jasonailu87@gmail.com> Co-authored-by: gnovack <novackgm@gmail.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Woosuk Kwon committed
Jun 22, 2026 at 18:14 UTC
e48592066ee4a435c9ac5316edbecb887596de02
1 file changed
+24
vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py
+24
@@ -6,6 +6,7 @@ import deep_ep
6
import torch
7
8
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
9
+from vllm.forward_context import get_forward_context
10
from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
11
from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
12
TopKWeightAndReduceContiguous,
@@ -116,6 +117,28 @@ class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
117
do_expand = not self.use_cudagraph
118
do_cpu_sync = not self.use_cudagraph
119
120
+ # In do_expand=False mode, the recv buffer is the worst case
121
+ # R * num_max_tokens_per_rank. Defaulting to the buffer's init value
122
+ # (= max_num_batched_tokens) makes the experts process ~R*8192 rows even
123
+ # for a handful of decode tokens. Bound it to the actual DP-padded batch
124
+ # size (uniform across ranks): max(num_tokens_across_dp).
125
+ #
126
+ # DeepEP JIT-compiles a separate dispatch kernel per distinct
127
+ # num_max_tokens_per_rank, so feeding it the raw per-step size would make
128
+ # it recompile for every batch size (a cicc storm that starves the GPU at
129
+ # high concurrency). Round up to a power of 2 instead: this bounds the
130
+ # set to ~log2(max_num_batched_tokens) values (compiled once, then
131
+ # cached) while staying small for decode (e.g. 1 token -> 1) and capped
132
+ # at the buffer's init capacity for prefill.
133
+ num_max_tokens_per_rank = None
134
+ if not do_expand:
135
+ dp_meta = get_forward_context().dp_metadata
136
+ if dp_meta is not None:
137
+ n = int(dp_meta.num_tokens_across_dp_cpu.max())
138
+ else:
139
+ n = tokens.shape[0]
140
+ num_max_tokens_per_rank = 1 << max(n - 1, 0).bit_length()
141
+
142
(
143
recv_x,
144
recv_topk_idx,
@@ -127,6 +150,7 @@ class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
150
topk_idx=rank_topk_ids,
151
topk_weights=rank_topk_weights,
152
num_experts=num_experts,
153
+ num_max_tokens_per_rank=num_max_tokens_per_rank,
154
do_expand=do_expand,
155
do_cpu_sync=do_cpu_sync,
156
async_with_compute_stream=False,