Repository: vllm-project/vllm
Pull request #18077: [core][distributed] add ep group and all2all interface
Author: login: youkaichao
Labels: name: v1
Created: 2025-05-13T13:01:38Z | Merged: 2025-05-14T02:46:49Z

the all2all interface should be refactored with a follow-up from https://github.com/vllm-project/vllm/pull/15956 .

this PR mainly adds ep (expert parallel) group, and all2all base class, so that backend-dependent code like pplx's `pplx_init` will not invade into common code path (`vllm/distributed/parallel_state.py`) .


Diff:
diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py
new file mode 100644
index 000000000000..b69647b00586
--- /dev/null
+++ b/vllm/distributed/device_communicators/all2all.py
@@ -0,0 +1,93 @@
+# SPDX-License-Identifier: Apache-2.0
+import torch
+
+from vllm.forward_context import get_forward_context
+
+
+class All2AllBase:
+
+    def __init__(self, cpu_group, model):
+        self.cpu_group = cpu_group
+
+        # compute some common properties
+        from vllm.distributed.parallel_state import (get_dp_group,
+                                                     get_ep_group,
+                                                     get_tp_group,
+                                                     in_the_same_node_as)
+
+        # all2all lives in ep group, which is merged from dp and tp group
+        self.dp_group = get_dp_group()
+        self.tp_group = get_tp_group()
+        self.ep_group = get_ep_group()
+        self.dp_rank = self.dp_group.rank_in_group
+        self.dp_world_size = self.dp_group.world_size
+
+        # all2all communication often has separate implementations for
+        # intra-node and inter-node communication
+        self.intranode = in_the_same_node_as(cpu_group, source_rank=0)
+        self.internode = not self.intranode
+
+    def dispatch(self, hidden_states: torch.Tensor,
+                 router_logits: torch.Tensor):
+        raise NotImplementedError
+
+    def combine(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        raise NotImplementedError
+
+    def destroy(self):
+        pass
+
+
+class NaiveAll2All(All2AllBase):
+    """
+    A naive implementation of all2all communication.
+    It uses all-reduce under the hood, which is not
+    efficient at all. The main purpose is for testing and
+    debugging.
+    """
+
+    def __init__(self, cpu_group, model):
+        super().__init__(cpu_group, model)
+
+    def naive_multicast(self, x: torch.Tensor,
+                        cu_tokens_across_dp_cpu: torch.Tensor):
+        assert (len(x.shape) == 2)
+        buffer = torch.empty((cu_tokens_across_dp_cpu[-1], x.size(1)),
+                             device=x.device,
+                             dtype=x.dtype)
+
+        start = 0 if self.dp_rank == 0 else cu_tokens_across_dp_cpu[
+            self.dp_rank - 1]
+        end = cu_tokens_across_dp_cpu[self.dp_rank]
+        buffer[start:end, :].copy_(x)
+        for idx in range(self.dp_world_size):
+            start = 0 if idx == 0 else cu_tokens_across_dp_cpu[idx - 1]
+            end = cu_tokens_across_dp_cpu[idx]
+            self.dp_group.broadcast(buffer[start:end, :], idx)
+
+        return buffer
+
+    def dispatch(self, hidden_states: torch.Tensor,
+                 router_logits: torch.Tensor):
+        cu_tokens_across_dp_cpu = get_forward_context(
+        ).dp_metadata.cu_tokens_across_dp_cpu
+
+        hidden_states = self.naive_multicast(hidden_states,
+                                             cu_tokens_across_dp_cpu)
+        router_logits = self.naive_multicast(router_logits,
+                                             cu_tokens_across_dp_cpu)
+        return hidden_states, router_logits
+
+    def combine(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        cu_tokens_across_dp_cpu = get_forward_context(
+        ).dp_metadata.cu_tokens_across_dp_cpu
+        start = 0 if self.dp_rank == 0 else cu_tokens_across_dp_cpu[
+            self.dp_rank - 1]
+        end = cu_tokens_across_dp_cpu[self.dp_rank]
+
+        all_hidden_states = self.dp_group.all_reduce(hidden_states)
+        hidden_states = all_hidden_states[start:end, :]
+        return hidden_states
+
+    def destroy(self):
+        pass
diff --git a/vllm/distributed/device_communicators/base_device_communicator.py b/vllm/distributed/device_communicators/base_device_communicator.py
index 240313b98c88..c313b66ed8b2 100644
--- a/vllm/distributed/device_communicators/base_device_communicator.py
+++ b/vllm/distributed/device_communicators/base_device_communicator.py
@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-from typing import Optional
+from typing import Optional, Tuple
 
 import torch
 import torch.distributed as dist
@@ -149,3 +149,27 @@ def recv(self,
 
     def destroy(self):
         pass
+
+    def prepare_communication_buffer_for_model(self,
+                                               model: torch.nn.Module) -> None:
+        """
+        Prepare the communication buffer for the model.
+        This is a no-op in the base class.
+        """
+        pass
+
+    def dispatch(
+            self, hidden_states: torch.Tensor,
+            router_logits: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Dispatch the hidden states and router logits to the appropriate device.
+        This is a no-op in the base class.
+        """
+        return hidden_states, router_logits
+
+    def combine(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """
+        Combine the hidden states and router logits from the appropriate device.
+        This is a no-op in the base class.
+        """
+        return hidden_states
diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py
index 8bca278f3888..7a90d63973de 100644
--- a/vllm/distributed/device_communicators/cuda_communicator.py
+++ b/vllm/distributed/device_communicators/cuda_communicator.py
@@ -1,10 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Optional
+from typing import Optional, Tuple
 
 import torch
 from torch.distributed import ProcessGroup
 
+import vllm.envs as envs
+
+from .all2all import All2AllBase
 from .base_device_communicator import DeviceCommunicatorBase
 
 
@@ -23,9 +26,13 @@ def __init__(self,
             from vllm.distributed.parallel_state import (
                 _ENABLE_CUSTOM_ALL_REDUCE)
             use_custom_allreduce = _ENABLE_CUSTOM_ALL_REDUCE
-        use_pynccl = True
+
+        # ep does not use pynccl
+        use_pynccl = "ep" not in unique_name
 
         self.use_pynccl = use_pynccl
+        self.use_all2all = "ep" in unique_name
+        self.all2all_impl: Optional[All2AllBase] = None
         self.use_custom_allreduce = use_custom_allreduce
 
         # lazy import to avoid documentation build error
@@ -129,3 +136,31 @@ def destroy(self):
             self.pynccl_comm = None
         if self.ca_comm is not None:
             self.ca_comm = None
+        if self.all2all_impl is not None:
+            self.all2all_impl.destroy()
+            self.all2all_impl = None
+
+    def prepare_communication_buffer_for_model(self,
+                                               model: torch.nn.Module) -> None:
+        """
+        Prepare the communication buffer for the model.
+        """
+        if not self.use_all2all:
+            return
+        all2all_backend = envs.VLLM_ALL2ALL_BACKEND
+        if all2all_backend == "naive":
+            from .all2all import NaiveAll2All
+            self.all2all_impl = NaiveAll2All(self.cpu_group, model)
+
+    def dispatch(
+            self, hidden_states: torch.Tensor,
+            router_logits: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        assert self.all2all_impl is not None
+        hidden_states, router_logits = self.all2all_impl.dispatch(
+            hidden_states, router_logits)
+        return hidden_states, router_logits
+
+    def combine(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        assert self.all2all_impl is not None
+        hidden_states = self.all2all_impl.combine(hidden_states)
+        return hidden_states
diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py
index cb9658ce1004..4a2a95d94b54 100644
--- a/vllm/distributed/parallel_state.py
+++ b/vllm/distributed/parallel_state.py
@@ -757,6 +757,22 @@ def destroy(self):
         if self.mq_broadcaster is not None:
             self.mq_broadcaster = None
 
+    def prepare_communication_buffer_for_model(self, model: torch.nn.Module):
+        if self.device_communicator is not None:
+            self.device_communicator.prepare_communication_buffer_for_model(
+                model)
+
+    def dispatch(
+            self, hidden_states: torch.Tensor,
+            router_logits: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        if self.device_communicator is not None:
+            return self.device_communicator.dispatch(hidden_states,
+                                                     router_logits)
+
+    def combine(self, hidden_states) -> torch.Tensor:
+        if self.device_communicator is not None:
+            return self.device_communicator.combine(hidden_states)
+
 
 _WORLD: Optional[GroupCoordinator] = None
 
@@ -816,6 +832,14 @@ def get_dp_group() -> GroupCoordinator:
     return _DP
 
 
+_EP: Optional[GroupCoordinator] = None
+
+
+def get_ep_group() -> GroupCoordinator:
+    assert _EP is not None, ("expert parallel group is not initialized")
+    return _EP
+
+
 def get_pp_group() -> GroupCoordinator:
     assert _PP is not None, (
         "pipeline model parallel group is not initialized")
@@ -1001,10 +1025,21 @@ def initialize_model_parallel(
                                     backend,
                                     group_name="dp")
 
+    global _EP
+    assert _EP is None, ("expert parallel group is already initialized")
+    group_ranks = all_ranks.transpose(1, 2).reshape(
+        -1, data_parallel_size * tensor_model_parallel_size).unbind(0)
+    group_ranks = [x.tolist() for x in group_ranks]
+    _EP = init_model_parallel_group(group_ranks,
+                                    get_world_group().local_rank,
+                                    backend,
+                                    group_name="ep")
+
     logger.info(
         "rank %s in world size %s is assigned as "
-        "DP rank %s, PP rank %s, TP rank %s", rank, world_size,
-        _DP.rank_in_group, _PP.rank_in_group, _TP.rank_in_group)
+        "DP rank %s, PP rank %s, TP rank %s, EP rank %s", rank, world_size,
+        _DP.rank_in_group, _PP.rank_in_group, _TP.rank_in_group,
+        _EP.rank_in_group)
 
 
 def ensure_model_parallel_initialized(
@@ -1035,6 +1070,23 @@ def ensure_model_parallel_initialized(
         f"{pipeline_model_parallel_size=}")
 
 
+def prepare_communication_buffer_for_model(model: torch.nn.Module):
+    """Prepare the communication buffer for the model.
+    Traditional communication libraries like NCCL are almost
+    model agnostic. However, emerging new communication libraries like
+    MoE all2all (DeepEP) usually allocate the communication buffer
+    based on the model shape for optimal performance.
+    """
+    if _TP is not None:
+        _TP.prepare_communication_buffer_for_model(model)
+    if _PP is not None:
+        _PP.prepare_communication_buffer_for_model(model)
+    if _DP is not None:
+        _DP.prepare_communication_buffer_for_model(model)
+    if _EP is not None:
+        _EP.prepare_communication_buffer_for_model(model)
+
+
 def model_parallel_is_initialized():
     """Check if tensor and pipeline parallel groups are initialized."""
     return (_TP is not None and _PP is not None)
@@ -1095,6 +1147,11 @@ def destroy_model_parallel():
         _DP.destroy()
     _DP = None
 
+    global _EP
+    if _EP:
+        _EP.destroy()
+    _EP = None
+
 
 def destroy_distributed_environment():
     global _WORLD
diff --git a/vllm/distributed/utils.py b/vllm/distributed/utils.py
index e4d4008cd0a6..1fbf5b284790 100644
--- a/vllm/distributed/utils.py
+++ b/vllm/distributed/utils.py
@@ -361,6 +361,11 @@ def stateless_destroy_torch_distributed_process_group(
         stateless_init_torch_distributed_process_group().
     """
     # Lazy import for non-CUDA backends.
-    from torch.distributed.distributed_c10d import _shutdown_backend
-    _shutdown_backend(pg)
+    try:
+        # pytorch <= 2.6
+        from torch.distributed.distributed_c10d import _shutdown_backend
+        _shutdown_backend(pg)
+    except ImportError:
+        # pytorch >= 2.7
+        pg.shutdown()
     _unregister_process_group(pg.group_name)
diff --git a/vllm/envs.py b/vllm/envs.py
index 0c742bf05623..9d585bf3578e 100644
--- a/vllm/envs.py
+++ b/vllm/envs.py
@@ -115,6 +115,7 @@
     VLLM_ALLOW_INSECURE_SERIALIZATION: bool = False
     VLLM_NIXL_SIDE_CHANNEL_HOST: str = "localhost"
     VLLM_NIXL_SIDE_CHANNEL_PORT: int = 5557
+    VLLM_ALL2ALL_BACKEND: str = "naive"
 
 
 def get_default_cache_root():
@@ -764,6 +765,10 @@ def maybe_convert_int(value: Optional[str]) -> Optional[int]:
     # Port used for NIXL handshake between remote agents.
     "VLLM_NIXL_SIDE_CHANNEL_PORT":
     lambda: int(os.getenv("VLLM_NIXL_SIDE_CHANNEL_PORT", "5557")),
+
+    # all2all backend for vllm's expert parallel communication
+    "VLLM_ALL2ALL_BACKEND":
+    lambda: os.getenv("VLLM_ALL2ALL_BACKEND", "naive"),
 }
 
 # end-env-vars-definition
diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py
index eae029b33e80..d745a15e3d22 100644
--- a/vllm/model_executor/layers/fused_moe/layer.py
+++ b/vllm/model_executor/layers/fused_moe/layer.py
@@ -10,7 +10,8 @@
 
 import vllm.envs as envs
 from vllm.config import get_current_vllm_config
-from vllm.distributed import (get_dp_group, get_tensor_model_parallel_rank,
+from vllm.distributed import (get_dp_group, get_ep_group,
+                              get_tensor_model_parallel_rank,
                               get_tensor_model_parallel_world_size,
                               tensor_model_parallel_all_reduce)
 from vllm.forward_context import ForwardContext, get_forward_context
@@ -832,24 +833,6 @@ def select_experts(hidden_states: torch.Tensor,
 
         return topk_weights, topk_ids
 
-    def naive_multicast(self, x: torch.Tensor,
-                        cu_tokens_across_dp_cpu: torch.Tensor):
-        assert (len(x.shape) == 2)
-        buffer = torch.empty((cu_tokens_across_dp_cpu[-1], x.size(1)),
-                             device=x.device,
-                             dtype=x.dtype)
-
-        start = 0 if self.dp_rank == 0 else cu_tokens_across_dp_cpu[
-            self.dp_rank - 1]
-        end = cu_tokens_across_dp_cpu[self.dp_rank]
-        buffer[start:end, :].copy_(x)
-        for idx in range(get_dp_group().world_size):
-            start = 0 if idx == 0 else cu_tokens_across_dp_cpu[idx - 1]
-            end = cu_tokens_across_dp_cpu[idx]
-            get_dp_group().broadcast(buffer[start:end, :], idx)
-
-        return buffer
-
     def forward(self, hidden_states: torch.Tensor,
                 router_logits: torch.Tensor):
         if self.use_direct_call:
@@ -863,14 +846,8 @@ def forward_impl(self, hidden_states: torch.Tensor,
         assert self.quant_method is not None
 
         if self.dp_size > 1:
-            cu_tokens_across_dp_cpu = get_forward_context(
-            ).dp_metadata.cu_tokens_across_dp_cpu
-
-            hidden_states = self.naive_multicast(hidden_states,
-                                                 cu_tokens_across_dp_cpu)
-            router_logits = self.naive_multicast(router_logits,
-                                                 cu_tokens_across_dp_cpu)
-
+            hidden_states, router_logits = get_ep_group().dispatch(
+                hidden_states, router_logits)
         # Matrix multiply.
         final_hidden_states = self.quant_method.apply(
             layer=self,
@@ -891,12 +868,7 @@ def forward_impl(self, hidden_states: torch.Tensor,
         )
 
         if self.dp_size > 1:
-            start = 0 if self.dp_rank == 0 else cu_tokens_across_dp_cpu[
-                self.dp_rank - 1]
-            end = cu_tokens_across_dp_cpu[self.dp_rank]
-
-            all_hidden_states = get_dp_group().all_reduce(final_hidden_states)
-            final_hidden_states = all_hidden_states[start:end, :]
+            final_hidden_states = get_ep_group().combine(final_hidden_states)
 
         if self.reduce_results and (self.tp_size > 1 or self.ep_size > 1):
             # Default set to False. (May have to add shared expert outputs.)
diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py
index 31895cc0832a..0e8f9eee449f 100644
--- a/vllm/v1/worker/gpu_model_runner.py
+++ b/vllm/v1/worker/gpu_model_runner.py
@@ -19,7 +19,8 @@
 from vllm.distributed.kv_transfer import (get_kv_transfer_group,
                                           has_kv_transfer_group)
 from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorBase_V1
-from vllm.distributed.parallel_state import get_pp_group, graph_capture
+from vllm.distributed.parallel_state import (
+    get_pp_group, graph_capture, prepare_communication_buffer_for_model)
 from vllm.forward_context import get_forward_context, set_forward_context
 from vllm.logger import init_logger
 from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding
@@ -1458,6 +1459,7 @@ def load_model(self) -> None:
         logger.info("Model loading took %.4f GiB and %.6f seconds",
                     self.model_memory_usage / GiB_bytes,
                     time_after_load - time_before_load)
+        prepare_communication_buffer_for_model(self.model)
 
     def _get_prompt_logprobs_dict(
         self,


Reviews:
user: login: youkaichao
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-13T13:15:19Z
user: login: youkaichao
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-13T13:20:58Z
user: login: tlrmchlsmth
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-13T16:49:52Z
user: login: tlrmchlsmth
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-13T16:58:03Z
user: login: youkaichao
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-13T17:06:16Z
user: login: tlrmchlsmth
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-13T17:08:42Z
user: login: youkaichao
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-13T17:14:10Z
user: login: bnellnm
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-13T18:01:20Z
user: login: bnellnm
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-13T19:29:38Z
user: login: tlrmchlsmth
state: APPROVED
author_association: MEMBER
submitted_at: 2025-05-14T00:09:34Z
user: login: youkaichao
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-14T02:46:20Z
user: login: youkaichao
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-14T02:46:38Z

Review comments:
diff_hunk: @@ -0,0 +1,80 @@
+# SPDX-License-Identifier: Apache-2.0
+import torch
+
+from vllm.forward_context import get_forward_context
+
+
+class All2AllBase:
+
+    def __init__(self, cpu_group, model):
path: vllm/distributed/device_communicators/all2all.py
user: login: youkaichao
body: for how to use `cpu_group` to initialize pplx, see https://github.com/ppl-ai/pplx-kernels/issues/18 as an example.

here we have access to the model instance, so we can also get moe configs here, and we can assume all moe layers have the same config, without `_all_to_all_cache` .
created_at: 2025-05-13T13:15:19Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,80 @@
+# SPDX-License-Identifier: Apache-2.0
+import torch
+
+from vllm.forward_context import get_forward_context
+
+
+class All2AllBase:
+
+    def __init__(self, cpu_group, model):
path: vllm/distributed/device_communicators/all2all.py
user: login: youkaichao
body: from @nandor , 

```python
uid = nvshmem_get_unique_id() if rank == 0 else nvshmem_alloc_empty_unique_id()
torch.distributed.broadcast(uid, src=0)
nvshmem_init(uid, rank, world_size)
```

is not necessary for intranode code.

we can tell if we are in single node with https://github.com/vllm-project/vllm/blob/79a1d25bbd09420ad7d8631671a70cbcaf496cfb/vllm/distributed/parallel_state.py#L1127 
created_at: 2025-05-13T13:20:58Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,93 @@
+# SPDX-License-Identifier: Apache-2.0
+import torch
+
+from vllm.forward_context import get_forward_context
+
+
+class All2AllBase:
+
+    def __init__(self, cpu_group, model):
+        self.cpu_group = cpu_group
+
+        # compute some common properties
+        from vllm.distributed.parallel_state import (get_dp_group,
+                                                     get_ep_group,
+                                                     get_tp_group,
+                                                     in_the_same_node_as)
+
+        # all2all lives in ep group, which is merged from dp and tp group
+        self.dp_group = get_dp_group()
+        self.tp_group = get_tp_group()
+        self.ep_group = get_ep_group()
+        self.dp_rank = self.dp_group.rank_in_group
+        self.dp_world_size = self.dp_group.world_size
+
+        # all2all communication often has separate implementations for
+        # intra-node and inter-node communication
+        self.intranode = in_the_same_node_as(cpu_group, source_rank=0)
+        self.internode = not self.intranode
+
+    def dispatch(self, hidden_states: torch.Tensor,
+                 router_logits: torch.Tensor):
+        raise NotImplementedError
+
+    def combine(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        raise NotImplementedError
path: vllm/distributed/device_communicators/all2all.py
user: login: tlrmchlsmth
body: Is there a function signature for dispatch and combine that will work across all kernels?
* The naive multicast implementation
* pplx-kernels
* DeepEP

I think it would to look like:
```
def dispatch(self, 
             hidden_states: torch.Tensor,
             hs_scales: torch.Tensor, # Quantized scales for hidden_states
             router_logits: torch.Tensor,  
             topk_weights: torch.Tensor,  
             topk_ids: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor] # out_hidden_states and out_scales
```
(There are other parameters to the pplx-kernels that are needed but I think they can be encapsulated in a PPLXAll2All wrapper class)

@varun-sundar-rabindranath @bnellnm WDYT?
created_at: 2025-05-13T16:49:52Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,93 @@
+# SPDX-License-Identifier: Apache-2.0
+import torch
+
+from vllm.forward_context import get_forward_context
+
+
+class All2AllBase:
+
+    def __init__(self, cpu_group, model):
+        self.cpu_group = cpu_group
+
+        # compute some common properties
+        from vllm.distributed.parallel_state import (get_dp_group,
+                                                     get_ep_group,
+                                                     get_tp_group,
+                                                     in_the_same_node_as)
+
+        # all2all lives in ep group, which is merged from dp and tp group
+        self.dp_group = get_dp_group()
+        self.tp_group = get_tp_group()
+        self.ep_group = get_ep_group()
+        self.dp_rank = self.dp_group.rank_in_group
+        self.dp_world_size = self.dp_group.world_size
+
+        # all2all communication often has separate implementations for
+        # intra-node and inter-node communication
+        self.intranode = in_the_same_node_as(cpu_group, source_rank=0)
+        self.internode = not self.intranode
+
+    def dispatch(self, hidden_states: torch.Tensor,
+                 router_logits: torch.Tensor):
+        raise NotImplementedError
+
+    def combine(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        raise NotImplementedError
path: vllm/distributed/device_communicators/all2all.py
user: login: tlrmchlsmth
body: Another consideration here is that the output format of the different A2A implementations will be different.

* For naive All2All, the output shape would be `[total_tokens_across_dp, hidden_size]`
* For pplx All2All, it will be `[num_experts_per_rank, max_tokens_per_expert, hidden_size]`, with padding along axis 1
* For DeepEP I am not sure what the output format will look like
created_at: 2025-05-13T16:58:03Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,93 @@
+# SPDX-License-Identifier: Apache-2.0
+import torch
+
+from vllm.forward_context import get_forward_context
+
+
+class All2AllBase:
+
+    def __init__(self, cpu_group, model):
+        self.cpu_group = cpu_group
+
+        # compute some common properties
+        from vllm.distributed.parallel_state import (get_dp_group,
+                                                     get_ep_group,
+                                                     get_tp_group,
+                                                     in_the_same_node_as)
+
+        # all2all lives in ep group, which is merged from dp and tp group
+        self.dp_group = get_dp_group()
+        self.tp_group = get_tp_group()
+        self.ep_group = get_ep_group()
+        self.dp_rank = self.dp_group.rank_in_group
+        self.dp_world_size = self.dp_group.world_size
+
+        # all2all communication often has separate implementations for
+        # intra-node and inter-node communication
+        self.intranode = in_the_same_node_as(cpu_group, source_rank=0)
+        self.internode = not self.intranode
+
+    def dispatch(self, hidden_states: torch.Tensor,
+                 router_logits: torch.Tensor):
+        raise NotImplementedError
+
+    def combine(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        raise NotImplementedError
path: vllm/distributed/device_communicators/all2all.py
user: login: youkaichao
body: for a generic dispatch/combine interface for different backends (pplx, naive, deepep), I think we can have:
```python
def dispatch(self, tensors: List[torch.Tensor]) -> List[torch.Tensor]:
def combine(self, tensors: List[torch.Tensor]) -> List[torch.Tensor]:
```
the backend can decide what does each tensor mean, and the moe layer can use assertion to make sure the list contains data it needs.
created_at: 2025-05-13T17:06:16Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,93 @@
+# SPDX-License-Identifier: Apache-2.0
+import torch
+
+from vllm.forward_context import get_forward_context
+
+
+class All2AllBase:
+
+    def __init__(self, cpu_group, model):
+        self.cpu_group = cpu_group
+
+        # compute some common properties
+        from vllm.distributed.parallel_state import (get_dp_group,
+                                                     get_ep_group,
+                                                     get_tp_group,
+                                                     in_the_same_node_as)
+
+        # all2all lives in ep group, which is merged from dp and tp group
+        self.dp_group = get_dp_group()
+        self.tp_group = get_tp_group()
+        self.ep_group = get_ep_group()
+        self.dp_rank = self.dp_group.rank_in_group
+        self.dp_world_size = self.dp_group.world_size
+
+        # all2all communication often has separate implementations for
+        # intra-node and inter-node communication
+        self.intranode = in_the_same_node_as(cpu_group, source_rank=0)
+        self.internode = not self.intranode
+
+    def dispatch(self, hidden_states: torch.Tensor,
+                 router_logits: torch.Tensor):
+        raise NotImplementedError
+
+    def combine(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        raise NotImplementedError
path: vllm/distributed/device_communicators/all2all.py
user: login: tlrmchlsmth
body: What's the reason for having the base class in this case? It seems that the caller will have to know exactly the implementation it's using
created_at: 2025-05-13T17:08:42Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,93 @@
+# SPDX-License-Identifier: Apache-2.0
+import torch
+
+from vllm.forward_context import get_forward_context
+
+
+class All2AllBase:
+
+    def __init__(self, cpu_group, model):
+        self.cpu_group = cpu_group
+
+        # compute some common properties
+        from vllm.distributed.parallel_state import (get_dp_group,
+                                                     get_ep_group,
+                                                     get_tp_group,
+                                                     in_the_same_node_as)
+
+        # all2all lives in ep group, which is merged from dp and tp group
+        self.dp_group = get_dp_group()
+        self.tp_group = get_tp_group()
+        self.ep_group = get_ep_group()
+        self.dp_rank = self.dp_group.rank_in_group
+        self.dp_world_size = self.dp_group.world_size
+
+        # all2all communication often has separate implementations for
+        # intra-node and inter-node communication
+        self.intranode = in_the_same_node_as(cpu_group, source_rank=0)
+        self.internode = not self.intranode
+
+    def dispatch(self, hidden_states: torch.Tensor,
+                 router_logits: torch.Tensor):
+        raise NotImplementedError
+
+    def combine(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        raise NotImplementedError
path: vllm/distributed/device_communicators/all2all.py
user: login: youkaichao
body: It's mainly to manage the lifecycle of the communication library, and separate the implementations.

it is true that the dispatch / combine apis are quite different, but what matters in this pr is to unify the construction, backend selection, and deconstruction. as a byproduct, we need to have a generic dispatch / combine function.
created_at: 2025-05-13T17:14:10Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,93 @@
+# SPDX-License-Identifier: Apache-2.0
+import torch
+
+from vllm.forward_context import get_forward_context
+
+
+class All2AllBase:
+
+    def __init__(self, cpu_group, model):
+        self.cpu_group = cpu_group
+
+        # compute some common properties
+        from vllm.distributed.parallel_state import (get_dp_group,
+                                                     get_ep_group,
+                                                     get_tp_group,
+                                                     in_the_same_node_as)
+
+        # all2all lives in ep group, which is merged from dp and tp group
+        self.dp_group = get_dp_group()
+        self.tp_group = get_tp_group()
+        self.ep_group = get_ep_group()
+        self.dp_rank = self.dp_group.rank_in_group
+        self.dp_world_size = self.dp_group.world_size
+
+        # all2all communication often has separate implementations for
+        # intra-node and inter-node communication
+        self.intranode = in_the_same_node_as(cpu_group, source_rank=0)
+        self.internode = not self.intranode
+
+    def dispatch(self, hidden_states: torch.Tensor,
+                 router_logits: torch.Tensor):
+        raise NotImplementedError
+
+    def combine(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        raise NotImplementedError
path: vllm/distributed/device_communicators/all2all.py
user: login: bnellnm
body: I agree with @tlrmchlsmth that it would be good if these interfaces matched the pplx AllToAll interfaces.
created_at: 2025-05-13T18:01:19Z
author_association: COLLABORATOR
diff_hunk: @@ -361,6 +361,11 @@ def stateless_destroy_torch_distributed_process_group(
         stateless_init_torch_distributed_process_group().
     """
     # Lazy import for non-CUDA backends.
-    from torch.distributed.distributed_c10d import _shutdown_backend
-    _shutdown_backend(pg)
+    try:
+        # pytorch <= 2.6
+        from torch.distributed.distributed_c10d import _shutdown_backend
+        _shutdown_backend(pg)
+    except ImportError:
+        # pytorch >= 2.7
+        pg.shutdown()
path: vllm/distributed/utils.py
user: login: bnellnm
body: I used this in the PPLX PR.  I think this is a bit nicer than try/catch.
```
    if is_torch_equal_or_newer("2.7"):
        pg.shutdown()
    else:
        # Lazy import for non-CUDA backends.                                                                                                                            
        from torch.distributed.distributed_c10d import _shutdown_backend
        _shutdown_backend(pg)
```
created_at: 2025-05-13T19:29:38Z
author_association: COLLABORATOR
diff_hunk: @@ -361,6 +361,11 @@ def stateless_destroy_torch_distributed_process_group(
         stateless_init_torch_distributed_process_group().
     """
     # Lazy import for non-CUDA backends.
-    from torch.distributed.distributed_c10d import _shutdown_backend
-    _shutdown_backend(pg)
+    try:
+        # pytorch <= 2.6
+        from torch.distributed.distributed_c10d import _shutdown_backend
+        _shutdown_backend(pg)
+    except ImportError:
+        # pytorch >= 2.7
+        pg.shutdown()
path: vllm/distributed/utils.py
user: login: youkaichao
body: feel free to change it in the moe modularization pr!
created_at: 2025-05-14T02:46:20Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,93 @@
+# SPDX-License-Identifier: Apache-2.0
+import torch
+
+from vllm.forward_context import get_forward_context
+
+
+class All2AllBase:
+
+    def __init__(self, cpu_group, model):
+        self.cpu_group = cpu_group
+
+        # compute some common properties
+        from vllm.distributed.parallel_state import (get_dp_group,
+                                                     get_ep_group,
+                                                     get_tp_group,
+                                                     in_the_same_node_as)
+
+        # all2all lives in ep group, which is merged from dp and tp group
+        self.dp_group = get_dp_group()
+        self.tp_group = get_tp_group()
+        self.ep_group = get_ep_group()
+        self.dp_rank = self.dp_group.rank_in_group
+        self.dp_world_size = self.dp_group.world_size
+
+        # all2all communication often has separate implementations for
+        # intra-node and inter-node communication
+        self.intranode = in_the_same_node_as(cpu_group, source_rank=0)
+        self.internode = not self.intranode
+
+    def dispatch(self, hidden_states: torch.Tensor,
+                 router_logits: torch.Tensor):
+        raise NotImplementedError
+
+    def combine(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        raise NotImplementedError
path: vllm/distributed/device_communicators/all2all.py
user: login: youkaichao
body: feel free to change the interface in the moe modularization pr!
created_at: 2025-05-14T02:46:38Z
author_association: MEMBER

PR comments:
user: login: varun-sundar-rabindranath
created_at: 2025-05-13T13:35:38Z
body: Thanks @youkaichao  - The abstractions generally look good . But the dispatch-combine calls for pplx input a very different set of information than the naive all2all implementation. I think we can cross that bridge when we add the pplx All2All implementation to the abstraction - But just adding it as a note. 
author_association: CONTRIBUTOR
user: login: youkaichao
created_at: 2025-05-13T13:40:19Z
body: > But the dispatch-combine calls for pplx input a very different set of information than the naive all2all implementation.

agree, right now the interface is just porting code from the naive implementation. we should definitely change it when we add the pplx All2All implementation to the abstraction.
author_association: MEMBER