[A2A] Add DeepEPv2 to SP Supported List (#57210)

Signed-off-by: Robert Shaw <robshaw@redhat.com>
Co-authored-by: Robert Shaw <robshaw@redhat.com>

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",

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:

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()