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 4, 22:30 UTC
[WideEP] Remove naive all2all. Use allgather_reducescatter instead (#33728)
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
Tyler Michael Smith committed
Apr 20, 2026 at 13:53 UTC
81d954f454d45425a0ad0a0a742de2695e4f043a
4 files changed
+10
-130
vllm/distributed/device_communicators/all2all.py
-109
@@ -38,115 +38,6 @@ if has_flashinfer_nvlink_one_sided():
38
logger = init_logger(__name__)
39
40
41
-class NaiveAll2AllManager(All2AllManagerBase):
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, tcp_store_group=None):
50
- super().__init__(cpu_group, tcp_store_group)
51
-
52
- def naive_multicast(
53
- self,
54
- x: torch.Tensor,
55
- cu_tokens_across_sp_cpu: torch.Tensor,
56
- is_sequence_parallel: bool,
57
- ) -> torch.Tensor:
58
- assert len(x.shape) == 2
59
- buffer = torch.empty(
60
- (cu_tokens_across_sp_cpu[-1], x.size(1)), device=x.device, dtype=x.dtype
61
- )
62
-
63
- rank = self.rank if is_sequence_parallel else self.dp_rank
64
- world_size = self.world_size if is_sequence_parallel else self.dp_world_size
65
-
66
- start = 0 if rank == 0 else cu_tokens_across_sp_cpu[rank - 1]
67
- end = cu_tokens_across_sp_cpu[rank]
68
- buffer[start:end, :].copy_(x)
69
- for idx in range(world_size):
70
- start = 0 if idx == 0 else cu_tokens_across_sp_cpu[idx - 1]
71
- end = cu_tokens_across_sp_cpu[idx]
72
- get_ep_group().broadcast(buffer[start:end, :], idx)
73
-
74
- return buffer
75
-
76
- def dispatch_router_logits(
77
- self,
78
- hidden_states: torch.Tensor,
79
- router_logits: torch.Tensor,
80
- is_sequence_parallel: bool = False,
81
- extra_tensors: list[torch.Tensor] | None = None,
82
- ) -> tuple[torch.Tensor, torch.Tensor]:
83
- if extra_tensors is not None:
84
- raise NotImplementedError(
85
- "extra_tensors is not supported for NaiveAll2AllManager"
86
- )
87
- sp_size = self.tp_group.world_size if is_sequence_parallel else 1
88
- dp_metadata = get_forward_context().dp_metadata
89
- assert dp_metadata is not None
90
- cu_tokens_across_sp_cpu = dp_metadata.cu_tokens_across_sp(sp_size)
91
-
92
- hidden_states = self.naive_multicast(
93
- hidden_states, cu_tokens_across_sp_cpu, is_sequence_parallel
94
- )
95
- router_logits = self.naive_multicast(
96
- router_logits, cu_tokens_across_sp_cpu, is_sequence_parallel
97
- )
98
-
99
- return hidden_states, router_logits
100
-
101
- def dispatch(
102
- self,
103
- hidden_states: torch.Tensor,
104
- topk_weights: torch.Tensor,
105
- topk_ids: torch.Tensor,
106
- is_sequence_parallel: bool = False,
107
- extra_tensors: list[torch.Tensor] | None = None,
108
- ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
109
- if extra_tensors is not None:
110
- raise NotImplementedError(
111
- "extra_tensors is not supported for NaiveAll2AllManager"
112
- )
113
- sp_size = self.tp_group.world_size if is_sequence_parallel else 1
114
- dp_metadata = get_forward_context().dp_metadata
115
- assert dp_metadata is not None
116
- cu_tokens_across_sp_cpu = dp_metadata.cu_tokens_across_sp(sp_size)
117
-
118
- hidden_states = self.naive_multicast(
119
- hidden_states, cu_tokens_across_sp_cpu, is_sequence_parallel
120
- )
121
- topk_weights = self.naive_multicast(
122
- topk_weights, cu_tokens_across_sp_cpu, is_sequence_parallel
123
- )
124
- topk_ids = self.naive_multicast(
125
- topk_ids, cu_tokens_across_sp_cpu, is_sequence_parallel
126
- )
127
- return hidden_states, topk_weights, topk_ids
128
-
129
- def combine(
130
- self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False
131
- ) -> torch.Tensor:
132
- ep_rank = self.rank if is_sequence_parallel else self.dp_rank
133
-
134
- dp_metadata = get_forward_context().dp_metadata
135
- assert dp_metadata is not None
136
- sp_size = self.tp_group.world_size if is_sequence_parallel else 1
137
- cu_tokens_across_sp_cpu = dp_metadata.cu_tokens_across_sp(sp_size)
138
-
139
- start = 0 if ep_rank == 0 else cu_tokens_across_sp_cpu[ep_rank - 1]
140
- end = cu_tokens_across_sp_cpu[ep_rank]
141
-
142
- all_hidden_states = get_ep_group().all_reduce(hidden_states)
143
- hidden_states = all_hidden_states[start:end, :]
144
- return hidden_states
145
-
146
- def destroy(self):
147
- pass
148
-
149
-
41
class AgRsAll2AllManager(All2AllManagerBase):
42
"""
43
An implementation of all2all communication based on
vllm/distributed/device_communicators/cpu_communicator.py
+8
-7
@@ -49,18 +49,19 @@ class CpuCommunicator(DeviceCommunicatorBase):
49
self.supports_tensor_dict = isinstance(self.dist_module, _CPUSHMDistributed)
50
51
if self.use_all2all:
52
- if self.all2all_backend != "naive": # type: ignore[has-type]
52
+ if self.all2all_backend not in (
53
+ "naive",
54
+ "allgather_reducescatter",
55
+ ): # type: ignore[has-type]
56
logger.warning(
57
"`%s` all2all manager is not supported on CPU. "
55
- "Falling back to `naive` all2all manager for CPU.",
58
+ "Falling back to `allgather_reducescatter` manager.",
59
self.all2all_backend, # type: ignore[has-type]
60
)
58
- self.all2all_backend = "naive"
59
- if self.all2all_backend == "naive":
60
- from .all2all import NaiveAll2AllManager
61
+ from .all2all import AgRsAll2AllManager
62
62
- self.all2all_manager = NaiveAll2AllManager(self.cpu_group)
63
- logger.info("Using naive all2all manager.")
63
+ self.all2all_manager = AgRsAll2AllManager(self.cpu_group)
64
+ logger.info("Using allgather_reducescatter all2all manager.")
65
66
def _all_group_ranks_share_shm_group_name(self) -> bool:
67
"""
vllm/distributed/device_communicators/cuda_communicator.py
+1
-7
@@ -115,13 +115,7 @@ class CudaCommunicator(DeviceCommunicatorBase):
115
self.qr_comm = QuickAllReduce(group=self.cpu_group, device=self.device)
116
117
if self.use_all2all:
118
- if self.all2all_backend == "naive":
119
- from .all2all import NaiveAll2AllManager
120
-
121
- self.all2all_manager = NaiveAll2AllManager(
122
- self.cpu_group, tcp_store_group
123
- )
124
- elif self.all2all_backend == "allgather_reducescatter":
118
+ if self.all2all_backend in ("naive", "allgather_reducescatter"):
119
from .all2all import AgRsAll2AllManager
120
121
self.all2all_manager = AgRsAll2AllManager(
vllm/distributed/device_communicators/xpu_communicator.py
+1
-7
@@ -23,13 +23,7 @@ class XpuCommunicator(DeviceCommunicatorBase):
23
):
24
super().__init__(cpu_group, device, device_group, unique_name)
25
if self.use_all2all:
26
- if self.all2all_backend == "naive":
27
- from .all2all import NaiveAll2AllManager
28
-
29
- self.all2all_manager = NaiveAll2AllManager(self.cpu_group)
30
- logger.info("Using naive all2all manager.")
31
-
32
- elif self.all2all_backend == "allgather_reducescatter":
26
+ if self.all2all_backend in ("naive", "allgather_reducescatter"):
27
from .all2all import AgRsAll2AllManager
28
29
self.all2all_manager = AgRsAll2AllManager(self.cpu_group)