[Feature] Support sequence parallel without the need for DP, 1.9%~5.0% E2E Throughput Improvement (#47070)

Signed-off-by: yewentao256 <zhyanwentao@126.com>
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
Co-authored-by: Canlin Guo <canlinguosdu@gmail.com>

vllm/config/parallel.py

@@ -632,8 +632,8 @@ class ParallelConfig:
 
     # The all_reduce at the end of attention (during o_proj) means that
     # inputs are replicated across each rank of the tensor parallel group.
-    # If using expert-parallelism with DeepEP All2All ops, replicated
-    # tokens results in useless duplicate computation and communication.
+    # If using expert-parallelism, replicated tokens results in useless
+    # duplicate computation and communication.
     #
     # In this case, ensure the input to the experts is sequence parallel
     # to avoid the excess work.
@@ -652,7 +652,6 @@ class ParallelConfig:
             )
             and self.enable_expert_parallel
             and self.tensor_parallel_size > 1
-            and self.data_parallel_size > 1
         )
 
     @property

vllm/distributed/device_communicators/base_device_communicator.py

@@ -175,10 +175,11 @@ class DeviceCommunicatorBase:
 
         config = get_current_vllm_config_or_none()
         if config is not None:
-            # as long as we use data parallel (coupled data parallel
-            # where all data parallel ranks execute forward together),
-            # we initialize the all2all manager used in expert parallel.
-            use_ep = config.parallel_config.data_parallel_size > 1
+            # initialize the all2all manager for DP or sequence-parallel EP.
+            use_ep = (
+                config.parallel_config.data_parallel_size > 1
+                or config.parallel_config.use_sequence_parallel_moe
+            )
             all2all_backend = config.parallel_config.all2all_backend
 
         self.is_ep_communicator = unique_name.split(":")[0] == "ep"

vllm/forward_context.py

@@ -83,7 +83,10 @@ class DPMetadata:
         num_tokens_across_dp_cpu: torch.Tensor,
     ) -> "DPMetadata":
         assert num_tokens_across_dp_cpu is not None
-        assert parallel_config.data_parallel_size > 1
+        assert (
+            parallel_config.data_parallel_size > 1
+            or parallel_config.use_sequence_parallel_moe
+        )
         assert parallel_config.is_moe_model is not False
         dp_rank = parallel_config.data_parallel_rank
         batchsize = num_tokens
@@ -277,14 +280,20 @@ def set_forward_context(
 
     dp_metadata: DPMetadata | None = None
     if (
-        vllm_config.parallel_config.data_parallel_size > 1
+        (
+            vllm_config.parallel_config.data_parallel_size > 1
+            or vllm_config.parallel_config.use_sequence_parallel_moe
+        )
         and vllm_config.parallel_config.is_moe_model is not False
         and (attn_metadata is not None or num_tokens is not None)
     ):
         # If num_tokens_across_dp hasn't already been initialized, then
         # initialize it here. Both DP padding and Microbatching will be
         # disabled.
-        if num_tokens_across_dp is None:
+        if (
+            num_tokens_across_dp is None
+            and vllm_config.parallel_config.data_parallel_size > 1
+        ):
             assert ubatch_slices is None
             assert num_tokens is not None
             _, num_tokens_across_dp, _ = coordinate_batch_across_dp(
@@ -293,6 +302,9 @@ def set_forward_context(
                 allow_microbatching=False,
             )
             assert num_tokens_across_dp is not None
+        elif num_tokens_across_dp is None:
+            assert num_tokens is not None
+            num_tokens_across_dp = torch.tensor([num_tokens], dtype=torch.int32)
         dp_metadata = DPMetadata.make(
             vllm_config.parallel_config, num_tokens or 0, num_tokens_across_dp
         )

vllm/model_executor/layers/fused_moe/config.py

@@ -1038,7 +1038,7 @@ class FusedMoEParallelConfig:
 
     @property
     def use_all2all_kernels(self):
-        return self.dp_size > 1 and self.use_ep
+        return self.use_ep and (self.dp_size > 1 or self.is_sequence_parallel)
 
     @property
     def use_deepep_ht_kernels(self):

vllm/model_executor/layers/fused_moe/runner/moe_runner.py

@@ -725,8 +725,8 @@ class MoERunner(MoERunnerInterface):
     @property
     def do_naive_dispatch_combine(self) -> bool:
         return (
-            self.moe_config.dp_size > 1 and not self._quant_method.supports_internal_mk
-        )
+            self.moe_config.dp_size > 1 or self.moe_config.is_sequence_parallel
+        ) and not self._quant_method.supports_internal_mk
 
     def _maybe_dispatch(
         self,

vllm/model_executor/models/deepseek_v2.py

@@ -1305,15 +1305,11 @@ class DeepseekV2DecoderLayer(nn.Module):
                 residual *= 1.0 / self.routed_scaling_factor
 
         if self.use_sequence_parallel_moe:
-            sp_remainder = (
-                hidden_states.shape[0] % get_tensor_model_parallel_world_size()
-            )
+            tp_world_size = get_tensor_model_parallel_world_size()
+            # small trick using minus, eg. -17 % 8 = 7
+            sp_pad = (-hidden_states.shape[0]) % tp_world_size
             # pad if not divisible by world size
-            if sp_remainder:
-                sp_pad = get_tensor_model_parallel_world_size() - sp_remainder
-                hidden_states = torch.nn.functional.pad(
-                    hidden_states, (0, 0, 0, sp_pad)
-                )
+            hidden_states = torch.nn.functional.pad(hidden_states, (0, 0, 0, sp_pad))
             hidden_states = tensor_model_parallel_reduce_scatter(hidden_states, 0)
             if not input_is_sequence_parallel:
                 residual = sequence_parallel_chunk(residual)

vllm/model_executor/models/gpt_oss.py

@@ -192,7 +192,10 @@ class MLPBlock(torch.nn.Module):
         quant_config = vllm_config.quant_config
         parallel_config = vllm_config.parallel_config
 
-        self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe
+        self.is_sequence_parallel = (
+            parallel_config.use_sequence_parallel_moe
+            and vllm_config.lora_config is None
+        )
 
         self.layer_idx = layer_idx
         self.num_experts = config.num_local_experts