[A2A] Add DeepEPv2 to SP Supported List (#57210)
Signed-off-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Robert Shaw <robshaw@redhat.com>
Robert Shaw committed
Sep 16, 2026 at 21:29 UTC
ed769b4a4520d9ca46657bd9e1ac414a3bbb9cdd
3 files changed
+8
-2
vllm/config/parallel.py
+1
@@ -717,6 +717,7 @@ class ParallelConfig:
717
"allgather_reducescatter",
718
"deepep_high_throughput",
719
"deepep_low_latency",
720
+ "deepep_v2",
721
"flashinfer_nvlink_one_sided",
722
"mori_high_throughput",
723
"mori_low_latency",
vllm/model_executor/layers/fused_moe/all2all_utils.py
+1
@@ -276,6 +276,7 @@ def maybe_make_prepare_finalize(
276
num_topk=moe.experts_per_token,
277
use_fp8_dispatch=use_fp8_dispatch,
278
use_cudagraph=use_cudagraph,
279
+ sp_size=moe.moe_parallel_config.sp_size,
280
)
281
282
elif moe.use_mori_kernels:
vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py
+6
-2
@@ -18,7 +18,7 @@ from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
18
swizzle_mxfp8_scale,
19
)
20
from vllm.triton_utils import tl, triton
21
-from vllm.utils.math_utils import round_up
21
+from vllm.utils.math_utils import cdiv, round_up
22
from vllm.v1.worker.ubatching import (
23
dbo_current_ubatch_id,
24
dbo_enabled,
@@ -114,6 +114,7 @@ class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
114
num_topk: int,
115
use_fp8_dispatch: bool = False,
116
use_cudagraph: bool = False,
117
+ sp_size: int = 1,
118
):
119
super().__init__()
120
self.buffer = buffer
@@ -124,6 +125,7 @@ class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
125
self.num_topk = num_topk
126
self.use_fp8_dispatch = use_fp8_dispatch
127
self.use_cudagraph = use_cudagraph
128
+ self.sp_size = sp_size
129
130
# DBO microbatching: one handle slot per micro-batch.
131
self.handles: list[deep_ep.EPHandle | None] = [None, None]
@@ -192,12 +194,14 @@ class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
194
# high concurrency). Round up to a power of 2 instead: this bounds the
195
# set to ~log2(max_num_batched_tokens) values (compiled once, then
196
# cached) while staying small for decode (e.g. 1 token -> 1) and capped
195
- # at the buffer's init capacity for prefill.
197
+ # at the buffer's init capacity for prefill. With sequence parallelism,
198
+ # each EP rank holds a ceil(n / sp_size) shard of its DP rank's batch.
199
num_max_tokens_per_rank = None
200
if not do_expand:
201
dp_meta = get_forward_context().dp_metadata
202
if dp_meta is not None:
203
n = int(dp_meta.num_tokens_across_dp_cpu.max())
204
+ n = cdiv(n, self.sp_size)
205
else:
206
n = tokens.shape[0]
207
num_max_tokens_per_rank = 1 << max(n - 1, 0).bit_length()