From ed769b4a4520d9ca46657bd9e1ac414a3bbb9cdd Mon Sep 17 00:00:00 2001 From: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Date: Wed, 16 Sep 2026 21:29:10 -0400 Subject: [PATCH] [A2A] Add DeepEPv2 to SP Supported List (#57210) Signed-off-by: Robert Shaw Co-authored-by: Robert Shaw --- vllm/config/parallel.py | 1 + vllm/model_executor/layers/fused_moe/all2all_utils.py | 1 + .../layers/fused_moe/prepare_finalize/deepep_v2.py | 8 ++++++-- 3 files changed, 8 insertions(+), 2 deletions(-) diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index f66af2f30b..f0a2b7e457 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -717,6 +717,7 @@ class ParallelConfig: "allgather_reducescatter", "deepep_high_throughput", "deepep_low_latency", + "deepep_v2", "flashinfer_nvlink_one_sided", "mori_high_throughput", "mori_low_latency", diff --git a/vllm/model_executor/layers/fused_moe/all2all_utils.py b/vllm/model_executor/layers/fused_moe/all2all_utils.py index 6984e34e56..775f595bfa 100644 --- a/vllm/model_executor/layers/fused_moe/all2all_utils.py +++ b/vllm/model_executor/layers/fused_moe/all2all_utils.py @@ -276,6 +276,7 @@ def maybe_make_prepare_finalize( num_topk=moe.experts_per_token, use_fp8_dispatch=use_fp8_dispatch, use_cudagraph=use_cudagraph, + sp_size=moe.moe_parallel_config.sp_size, ) elif moe.use_mori_kernels: diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py index 8bd215423d..a584ab5f52 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py @@ -18,7 +18,7 @@ from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( swizzle_mxfp8_scale, ) from vllm.triton_utils import tl, triton -from vllm.utils.math_utils import round_up +from vllm.utils.math_utils import cdiv, round_up from vllm.v1.worker.ubatching import ( dbo_current_ubatch_id, dbo_enabled, @@ -114,6 +114,7 @@ class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): num_topk: int, use_fp8_dispatch: bool = False, use_cudagraph: bool = False, + sp_size: int = 1, ): super().__init__() self.buffer = buffer @@ -124,6 +125,7 @@ class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): self.num_topk = num_topk self.use_fp8_dispatch = use_fp8_dispatch self.use_cudagraph = use_cudagraph + self.sp_size = sp_size # DBO microbatching: one handle slot per micro-batch. self.handles: list[deep_ep.EPHandle | None] = [None, None] @@ -192,12 +194,14 @@ class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): # high concurrency). Round up to a power of 2 instead: this bounds the # set to ~log2(max_num_batched_tokens) values (compiled once, then # cached) while staying small for decode (e.g. 1 token -> 1) and capped - # at the buffer's init capacity for prefill. + # at the buffer's init capacity for prefill. With sequence parallelism, + # each EP rank holds a ceil(n / sp_size) shard of its DP rank's batch. num_max_tokens_per_rank = None if not do_expand: dp_meta = get_forward_context().dp_metadata if dp_meta is not None: n = int(dp_meta.num_tokens_across_dp_cpu.max()) + n = cdiv(n, self.sp_size) else: n = tokens.shape[0] num_max_tokens_per_rank = 1 << max(n - 1, 0).bit_length()