A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 11:00 UTC
[core][distributed] add ep group and all2all interface (#18077)
Signed-off-by: youkaichao <youkaichao@gmail.com>
youkaichao committed
May 14, 2025 at 10:46 UTC
6266c57bae0f52e830221007241920e44f9f1ccf
8 files changed
+234
-41
vllm/distributed/device_communicators/all2all.py
new
+93
@@ -0,0 +1,93 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+import torch
3
+
4
+from vllm.forward_context import get_forward_context
5
+
6
+
7
+class All2AllBase:
8
+
9
+ def __init__(self, cpu_group, model):
10
+ self.cpu_group = cpu_group
11
+
12
+ # compute some common properties
13
+ from vllm.distributed.parallel_state import (get_dp_group,
14
+ get_ep_group,
15
+ get_tp_group,
16
+ in_the_same_node_as)
17
+
18
+ # all2all lives in ep group, which is merged from dp and tp group
19
+ self.dp_group = get_dp_group()
20
+ self.tp_group = get_tp_group()
21
+ self.ep_group = get_ep_group()
22
+ self.dp_rank = self.dp_group.rank_in_group
23
+ self.dp_world_size = self.dp_group.world_size
24
+
25
+ # all2all communication often has separate implementations for
26
+ # intra-node and inter-node communication
27
+ self.intranode = in_the_same_node_as(cpu_group, source_rank=0)
28
+ self.internode = not self.intranode
29
+
30
+ def dispatch(self, hidden_states: torch.Tensor,
31
+ router_logits: torch.Tensor):
32
+ raise NotImplementedError
33
+
34
+ def combine(self, hidden_states: torch.Tensor) -> torch.Tensor:
35
+ raise NotImplementedError
36
+
37
+ def destroy(self):
38
+ pass
39
+
40
+
41
+class NaiveAll2All(All2AllBase):
42
+ """
43
+ A naive implementation of all2all communication.
44
+ It uses all-reduce under the hood, which is not
45
+ efficient at all. The main purpose is for testing and
46
+ debugging.
47
+ """
48
+
49
+ def __init__(self, cpu_group, model):
50
+ super().__init__(cpu_group, model)
51
+
52
+ def naive_multicast(self, x: torch.Tensor,
53
+ cu_tokens_across_dp_cpu: torch.Tensor):
54
+ assert (len(x.shape) == 2)
55
+ buffer = torch.empty((cu_tokens_across_dp_cpu[-1], x.size(1)),
56
+ device=x.device,
57
+ dtype=x.dtype)
58
+
59
+ start = 0 if self.dp_rank == 0 else cu_tokens_across_dp_cpu[
60
+ self.dp_rank - 1]
61
+ end = cu_tokens_across_dp_cpu[self.dp_rank]
62
+ buffer[start:end, :].copy_(x)
63
+ for idx in range(self.dp_world_size):
64
+ start = 0 if idx == 0 else cu_tokens_across_dp_cpu[idx - 1]
65
+ end = cu_tokens_across_dp_cpu[idx]
66
+ self.dp_group.broadcast(buffer[start:end, :], idx)
67
+
68
+ return buffer
69
+
70
+ def dispatch(self, hidden_states: torch.Tensor,
71
+ router_logits: torch.Tensor):
72
+ cu_tokens_across_dp_cpu = get_forward_context(
73
+ ).dp_metadata.cu_tokens_across_dp_cpu
74
+
75
+ hidden_states = self.naive_multicast(hidden_states,
76
+ cu_tokens_across_dp_cpu)
77
+ router_logits = self.naive_multicast(router_logits,
78
+ cu_tokens_across_dp_cpu)
79
+ return hidden_states, router_logits
80
+
81
+ def combine(self, hidden_states: torch.Tensor) -> torch.Tensor:
82
+ cu_tokens_across_dp_cpu = get_forward_context(
83
+ ).dp_metadata.cu_tokens_across_dp_cpu
84
+ start = 0 if self.dp_rank == 0 else cu_tokens_across_dp_cpu[
85
+ self.dp_rank - 1]
86
+ end = cu_tokens_across_dp_cpu[self.dp_rank]
87
+
88
+ all_hidden_states = self.dp_group.all_reduce(hidden_states)
89
+ hidden_states = all_hidden_states[start:end, :]
90
+ return hidden_states
91
+
92
+ def destroy(self):
93
+ pass
vllm/distributed/device_communicators/base_device_communicator.py
+25
-1
@@ -1,5 +1,5 @@
1
# SPDX-License-Identifier: Apache-2.0
2
-from typing import Optional
2
+from typing import Optional, Tuple
3
4
import torch
5
import torch.distributed as dist
@@ -149,3 +149,27 @@ class DeviceCommunicatorBase:
149
150
def destroy(self):
151
pass
152
+
153
+ def prepare_communication_buffer_for_model(self,
154
+ model: torch.nn.Module) -> None:
155
+ """
156
+ Prepare the communication buffer for the model.
157
+ This is a no-op in the base class.
158
+ """
159
+ pass
160
+
161
+ def dispatch(
162
+ self, hidden_states: torch.Tensor,
163
+ router_logits: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
164
+ """
165
+ Dispatch the hidden states and router logits to the appropriate device.
166
+ This is a no-op in the base class.
167
+ """
168
+ return hidden_states, router_logits
169
+
170
+ def combine(self, hidden_states: torch.Tensor) -> torch.Tensor:
171
+ """
172
+ Combine the hidden states and router logits from the appropriate device.
173
+ This is a no-op in the base class.
174
+ """
175
+ return hidden_states
vllm/distributed/device_communicators/cuda_communicator.py
+37
-2
@@ -1,10 +1,13 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
-from typing import Optional
3
+from typing import Optional, Tuple
4
5
import torch
6
from torch.distributed import ProcessGroup
7
8
+import vllm.envs as envs
9
+
10
+from .all2all import All2AllBase
11
from .base_device_communicator import DeviceCommunicatorBase
12
13
@@ -23,9 +26,13 @@ class CudaCommunicator(DeviceCommunicatorBase):
26
from vllm.distributed.parallel_state import (
27
_ENABLE_CUSTOM_ALL_REDUCE)
28
use_custom_allreduce = _ENABLE_CUSTOM_ALL_REDUCE
26
- use_pynccl = True
29
+
30
+ # ep does not use pynccl
31
+ use_pynccl = "ep" not in unique_name
32
33
self.use_pynccl = use_pynccl
34
+ self.use_all2all = "ep" in unique_name
35
+ self.all2all_impl: Optional[All2AllBase] = None
36
self.use_custom_allreduce = use_custom_allreduce
37
38
# lazy import to avoid documentation build error
@@ -129,3 +136,31 @@ class CudaCommunicator(DeviceCommunicatorBase):
136
self.pynccl_comm = None
137
if self.ca_comm is not None:
138
self.ca_comm = None
139
+ if self.all2all_impl is not None:
140
+ self.all2all_impl.destroy()
141
+ self.all2all_impl = None
142
+
143
+ def prepare_communication_buffer_for_model(self,
144
+ model: torch.nn.Module) -> None:
145
+ """
146
+ Prepare the communication buffer for the model.
147
+ """
148
+ if not self.use_all2all:
149
+ return
150
+ all2all_backend = envs.VLLM_ALL2ALL_BACKEND
151
+ if all2all_backend == "naive":
152
+ from .all2all import NaiveAll2All
153
+ self.all2all_impl = NaiveAll2All(self.cpu_group, model)
154
+
155
+ def dispatch(
156
+ self, hidden_states: torch.Tensor,
157
+ router_logits: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
158
+ assert self.all2all_impl is not None
159
+ hidden_states, router_logits = self.all2all_impl.dispatch(
160
+ hidden_states, router_logits)
161
+ return hidden_states, router_logits
162
+
163
+ def combine(self, hidden_states: torch.Tensor) -> torch.Tensor:
164
+ assert self.all2all_impl is not None
165
+ hidden_states = self.all2all_impl.combine(hidden_states)
166
+ return hidden_states
vllm/distributed/parallel_state.py
+59
-2
@@ -757,6 +757,22 @@ class GroupCoordinator:
757
if self.mq_broadcaster is not None:
758
self.mq_broadcaster = None
759
760
+ def prepare_communication_buffer_for_model(self, model: torch.nn.Module):
761
+ if self.device_communicator is not None:
762
+ self.device_communicator.prepare_communication_buffer_for_model(
763
+ model)
764
+
765
+ def dispatch(
766
+ self, hidden_states: torch.Tensor,
767
+ router_logits: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
768
+ if self.device_communicator is not None:
769
+ return self.device_communicator.dispatch(hidden_states,
770
+ router_logits)
771
+
772
+ def combine(self, hidden_states) -> torch.Tensor:
773
+ if self.device_communicator is not None:
774
+ return self.device_communicator.combine(hidden_states)
775
+
776
777
_WORLD: Optional[GroupCoordinator] = None
778
@@ -816,6 +832,14 @@ def get_dp_group() -> GroupCoordinator:
832
return _DP
833
834
835
+_EP: Optional[GroupCoordinator] = None
836
+
837
+
838
+def get_ep_group() -> GroupCoordinator:
839
+ assert _EP is not None, ("expert parallel group is not initialized")
840
+ return _EP
841
+
842
+
843
def get_pp_group() -> GroupCoordinator:
844
assert _PP is not None, (
845
"pipeline model parallel group is not initialized")
@@ -1001,10 +1025,21 @@ def initialize_model_parallel(
1025
backend,
1026
group_name="dp")
1027
1028
+ global _EP
1029
+ assert _EP is None, ("expert parallel group is already initialized")
1030
+ group_ranks = all_ranks.transpose(1, 2).reshape(
1031
+ -1, data_parallel_size * tensor_model_parallel_size).unbind(0)
1032
+ group_ranks = [x.tolist() for x in group_ranks]
1033
+ _EP = init_model_parallel_group(group_ranks,
1034
+ get_world_group().local_rank,
1035
+ backend,
1036
+ group_name="ep")
1037
+
1038
logger.info(
1039
"rank %s in world size %s is assigned as "
1006
- "DP rank %s, PP rank %s, TP rank %s", rank, world_size,
1007
- _DP.rank_in_group, _PP.rank_in_group, _TP.rank_in_group)
1040
+ "DP rank %s, PP rank %s, TP rank %s, EP rank %s", rank, world_size,
1041
+ _DP.rank_in_group, _PP.rank_in_group, _TP.rank_in_group,
1042
+ _EP.rank_in_group)
1043
1044
1045
def ensure_model_parallel_initialized(
@@ -1035,6 +1070,23 @@ def ensure_model_parallel_initialized(
1070
f"{pipeline_model_parallel_size=}")
1071
1072
1073
+def prepare_communication_buffer_for_model(model: torch.nn.Module):
1074
+ """Prepare the communication buffer for the model.
1075
+ Traditional communication libraries like NCCL are almost
1076
+ model agnostic. However, emerging new communication libraries like
1077
+ MoE all2all (DeepEP) usually allocate the communication buffer
1078
+ based on the model shape for optimal performance.
1079
+ """
1080
+ if _TP is not None:
1081
+ _TP.prepare_communication_buffer_for_model(model)
1082
+ if _PP is not None:
1083
+ _PP.prepare_communication_buffer_for_model(model)
1084
+ if _DP is not None:
1085
+ _DP.prepare_communication_buffer_for_model(model)
1086
+ if _EP is not None:
1087
+ _EP.prepare_communication_buffer_for_model(model)
1088
+
1089
+
1090
def model_parallel_is_initialized():
1091
"""Check if tensor and pipeline parallel groups are initialized."""
1092
return (_TP is not None and _PP is not None)
@@ -1095,6 +1147,11 @@ def destroy_model_parallel():
1147
_DP.destroy()
1148
_DP = None
1149
1150
+ global _EP
1151
+ if _EP:
1152
+ _EP.destroy()
1153
+ _EP = None
1154
+
1155
1156
def destroy_distributed_environment():
1157
global _WORLD
vllm/distributed/utils.py
+7
-2
@@ -362,6 +362,11 @@ def stateless_destroy_torch_distributed_process_group(
362
stateless_init_torch_distributed_process_group().
363
"""
364
# Lazy import for non-CUDA backends.
365
- from torch.distributed.distributed_c10d import _shutdown_backend
366
- _shutdown_backend(pg)
365
+ try:
366
+ # pytorch <= 2.6
367
+ from torch.distributed.distributed_c10d import _shutdown_backend
368
+ _shutdown_backend(pg)
369
+ except ImportError:
370
+ # pytorch >= 2.7
371
+ pg.shutdown()
372
_unregister_process_group(pg.group_name)
vllm/envs.py
+5
@@ -115,6 +115,7 @@ if TYPE_CHECKING:
115
VLLM_ALLOW_INSECURE_SERIALIZATION: bool = False
116
VLLM_NIXL_SIDE_CHANNEL_HOST: str = "localhost"
117
VLLM_NIXL_SIDE_CHANNEL_PORT: int = 5557
118
+ VLLM_ALL2ALL_BACKEND: str = "naive"
119
120
121
def get_default_cache_root():
@@ -764,6 +765,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
765
# Port used for NIXL handshake between remote agents.
766
"VLLM_NIXL_SIDE_CHANNEL_PORT":
767
lambda: int(os.getenv("VLLM_NIXL_SIDE_CHANNEL_PORT", "5557")),
768
+
769
+ # all2all backend for vllm's expert parallel communication
770
+ "VLLM_ALL2ALL_BACKEND":
771
+ lambda: os.getenv("VLLM_ALL2ALL_BACKEND", "naive"),
772
}
773
774
# end-env-vars-definition
vllm/model_executor/layers/fused_moe/layer.py
+5
-33
@@ -10,7 +10,8 @@ from torch.nn.parameter import UninitializedParameter
10
11
import vllm.envs as envs
12
from vllm.config import get_current_vllm_config
13
-from vllm.distributed import (get_dp_group, get_tensor_model_parallel_rank,
13
+from vllm.distributed import (get_dp_group, get_ep_group,
14
+ get_tensor_model_parallel_rank,
15
get_tensor_model_parallel_world_size,
16
tensor_model_parallel_all_reduce)
17
from vllm.forward_context import ForwardContext, get_forward_context
@@ -832,24 +833,6 @@ class FusedMoE(torch.nn.Module):
833
834
return topk_weights, topk_ids
835
835
- def naive_multicast(self, x: torch.Tensor,
836
- cu_tokens_across_dp_cpu: torch.Tensor):
837
- assert (len(x.shape) == 2)
838
- buffer = torch.empty((cu_tokens_across_dp_cpu[-1], x.size(1)),
839
- device=x.device,
840
- dtype=x.dtype)
841
-
842
- start = 0 if self.dp_rank == 0 else cu_tokens_across_dp_cpu[
843
- self.dp_rank - 1]
844
- end = cu_tokens_across_dp_cpu[self.dp_rank]
845
- buffer[start:end, :].copy_(x)
846
- for idx in range(get_dp_group().world_size):
847
- start = 0 if idx == 0 else cu_tokens_across_dp_cpu[idx - 1]
848
- end = cu_tokens_across_dp_cpu[idx]
849
- get_dp_group().broadcast(buffer[start:end, :], idx)
850
-
851
- return buffer
852
-
836
def forward(self, hidden_states: torch.Tensor,
837
router_logits: torch.Tensor):
838
if self.use_direct_call:
@@ -863,14 +846,8 @@ class FusedMoE(torch.nn.Module):
846
assert self.quant_method is not None
847
848
if self.dp_size > 1:
866
- cu_tokens_across_dp_cpu = get_forward_context(
867
- ).dp_metadata.cu_tokens_across_dp_cpu
868
-
869
- hidden_states = self.naive_multicast(hidden_states,
870
- cu_tokens_across_dp_cpu)
871
- router_logits = self.naive_multicast(router_logits,
872
- cu_tokens_across_dp_cpu)
873
-
849
+ hidden_states, router_logits = get_ep_group().dispatch(
850
+ hidden_states, router_logits)
851
# Matrix multiply.
852
final_hidden_states = self.quant_method.apply(
853
layer=self,
@@ -891,12 +868,7 @@ class FusedMoE(torch.nn.Module):
868
)
869
870
if self.dp_size > 1:
894
- start = 0 if self.dp_rank == 0 else cu_tokens_across_dp_cpu[
895
- self.dp_rank - 1]
896
- end = cu_tokens_across_dp_cpu[self.dp_rank]
897
-
898
- all_hidden_states = get_dp_group().all_reduce(final_hidden_states)
899
- final_hidden_states = all_hidden_states[start:end, :]
871
+ final_hidden_states = get_ep_group().combine(final_hidden_states)
872
873
if self.reduce_results and (self.tp_size > 1 or self.ep_size > 1):
874
# Default set to False. (May have to add shared expert outputs.)
vllm/v1/worker/gpu_model_runner.py
+3
-1
@@ -19,7 +19,8 @@ from vllm.config import (CompilationLevel, VllmConfig,
19
from vllm.distributed.kv_transfer import (get_kv_transfer_group,
20
has_kv_transfer_group)
21
from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorBase_V1
22
-from vllm.distributed.parallel_state import get_pp_group, graph_capture
22
+from vllm.distributed.parallel_state import (
23
+ get_pp_group, graph_capture, prepare_communication_buffer_for_model)
24
from vllm.forward_context import get_forward_context, set_forward_context
25
from vllm.logger import init_logger
26
from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding
@@ -1457,6 +1458,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1458
logger.info("Model loading took %.4f GiB and %.6f seconds",
1459
self.model_memory_usage / GiB_bytes,
1460
time_after_load - time_before_load)
1461
+ prepare_communication_buffer_for_model(self.model)
1462
1463
def _get_prompt_logprobs_dict(
1464
self,