Skip to content

vllm.distributed.device_communicators.all2all

logger module-attribute 

logger = init_logger(__name__)

DeepEPAll2AllManagerBase

Bases: All2AllManagerBase

All2All communication based on DeepEP High-Throughput kernels.

Source code in vllm/distributed/device_communicators/all2all.py 
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
class DeepEPAll2AllManagerBase(All2AllManagerBase):
    """
    All2All communication based on DeepEP High-Throughput kernels.
    """

    def __init__(self, cpu_group):
        assert has_deep_ep(
        ), "DeepEP kernels not found. Please follow https://github.com/vllm-project/vllm/blob/main/tools/ep_kernels/README.md to install DeepEP kernels."  # noqa
        super().__init__(cpu_group)
        self.handle_cache = Cache()

        # This is the DeepEP default. Stick to it till we can establish
        # reasonable defaults based on profiling.
        self.num_sms = 20

    def get_handle(self, kwargs):
        raise NotImplementedError

    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

handle_cache instance-attribute 

handle_cache = Cache()

num_sms instance-attribute 

num_sms = 20

__init__ 

__init__(cpu_group)
Source code in vllm/distributed/device_communicators/all2all.py 
135
136
137
138
139
140
141
142
143
def __init__(self, cpu_group):
    assert has_deep_ep(
    ), "DeepEP kernels not found. Please follow https://github.com/vllm-project/vllm/blob/main/tools/ep_kernels/README.md to install DeepEP kernels."  # noqa
    super().__init__(cpu_group)
    self.handle_cache = Cache()

    # This is the DeepEP default. Stick to it till we can establish
    # reasonable defaults based on profiling.
    self.num_sms = 20

combine 

combine(hidden_states: Tensor) -> Tensor
Source code in vllm/distributed/device_communicators/all2all.py 
152
153
def combine(self, hidden_states: torch.Tensor) -> torch.Tensor:
    raise NotImplementedError

destroy 

destroy()
Source code in vllm/distributed/device_communicators/all2all.py 
155
156
def destroy(self):
    pass

dispatch 

dispatch(hidden_states: Tensor, router_logits: Tensor)
Source code in vllm/distributed/device_communicators/all2all.py 
148
149
150
def dispatch(self, hidden_states: torch.Tensor,
             router_logits: torch.Tensor):
    raise NotImplementedError

get_handle 

get_handle(kwargs)
Source code in vllm/distributed/device_communicators/all2all.py 
145
146
def get_handle(self, kwargs):
    raise NotImplementedError

DeepEPHTAll2AllManager

Bases: DeepEPAll2AllManagerBase

All2All communication based on DeepEP High-Throughput kernels.

Source code in vllm/distributed/device_communicators/all2all.py 
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
class DeepEPHTAll2AllManager(DeepEPAll2AllManagerBase):
    """
    All2All communication based on DeepEP High-Throughput kernels.
    """

    def __init__(self, cpu_group):
        super().__init__(cpu_group)

    def _make_all2all_kwargs(self) -> dict[Any, Any]:
        # Defaults for internode and intranode are taken from DeepEP tests.
        num_nvl_bytes = 1024 * 1024 * 1024
        num_rdma_bytes = None
        num_qps_per_rank = None

        if self.internode:
            num_rdma_bytes = 1024 * 1024 * 1024
            num_qps_per_rank = self.num_sms // 2
        else:
            num_rdma_bytes = 0
            num_qps_per_rank = 1

        assert num_rdma_bytes is not None
        assert num_qps_per_rank is not None
        return dict(group=self.cpu_group,
                    num_nvl_bytes=num_nvl_bytes,
                    num_rdma_bytes=num_rdma_bytes,
                    low_latency_mode=False,
                    num_qps_per_rank=num_qps_per_rank)

    def get_handle(self, kwargs):

        assert len(kwargs) == 0, (
            "DeepEPHTAll2AllManager expects no arguments. All the required "
            "args are computed in the Manager itself.")

        import deep_ep
        buffer_kwargs = self._make_all2all_kwargs()
        logger.debug("DeepEP all2all args %s", buffer_kwargs)
        handle: deep_ep.Buffer = self.handle_cache.get_or_create(
            buffer_kwargs, deep_ep.Buffer)
        # It is dangerous to set num sms outside this function. num_sms is not
        # a part of the hash-key that identifies this object. If we are in a
        # situation where we make objects with different num_sms, the hash key
        # in get_or_create must be updated.
        handle.set_num_sms(self.num_sms)
        return handle

__init__ 

__init__(cpu_group)
Source code in vllm/distributed/device_communicators/all2all.py 
164
165
def __init__(self, cpu_group):
    super().__init__(cpu_group)

_make_all2all_kwargs 

_make_all2all_kwargs() -> dict[Any, Any]
Source code in vllm/distributed/device_communicators/all2all.py 
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
def _make_all2all_kwargs(self) -> dict[Any, Any]:
    # Defaults for internode and intranode are taken from DeepEP tests.
    num_nvl_bytes = 1024 * 1024 * 1024
    num_rdma_bytes = None
    num_qps_per_rank = None

    if self.internode:
        num_rdma_bytes = 1024 * 1024 * 1024
        num_qps_per_rank = self.num_sms // 2
    else:
        num_rdma_bytes = 0
        num_qps_per_rank = 1

    assert num_rdma_bytes is not None
    assert num_qps_per_rank is not None
    return dict(group=self.cpu_group,
                num_nvl_bytes=num_nvl_bytes,
                num_rdma_bytes=num_rdma_bytes,
                low_latency_mode=False,
                num_qps_per_rank=num_qps_per_rank)

get_handle 

get_handle(kwargs)
Source code in vllm/distributed/device_communicators/all2all.py 
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
def get_handle(self, kwargs):

    assert len(kwargs) == 0, (
        "DeepEPHTAll2AllManager expects no arguments. All the required "
        "args are computed in the Manager itself.")

    import deep_ep
    buffer_kwargs = self._make_all2all_kwargs()
    logger.debug("DeepEP all2all args %s", buffer_kwargs)
    handle: deep_ep.Buffer = self.handle_cache.get_or_create(
        buffer_kwargs, deep_ep.Buffer)
    # It is dangerous to set num sms outside this function. num_sms is not
    # a part of the hash-key that identifies this object. If we are in a
    # situation where we make objects with different num_sms, the hash key
    # in get_or_create must be updated.
    handle.set_num_sms(self.num_sms)
    return handle

DeepEPLLAll2AllManager

Bases: DeepEPAll2AllManagerBase

All2All communication based on DeepEP Low-Latency kernels.

Source code in vllm/distributed/device_communicators/all2all.py 
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
class DeepEPLLAll2AllManager(DeepEPAll2AllManagerBase):
    """
    All2All communication based on DeepEP Low-Latency kernels.
    """

    def __init__(self, cpu_group):
        super().__init__(cpu_group)

    def _make_all2all_kwargs(
        self,
        max_num_tokens_per_dp_rank: int,
        token_hidden_size: int,
        num_ep_ranks: int,
        num_global_experts: int,
        num_local_experts: int,
    ) -> dict[Any, Any]:
        """
        max_num_tokens_per_dp_rank : the maximum number of tokens a DP rank
          can dispatch all the ranks must hold the same value.
        token_hidden_size: the hidden dimension of each token.
        num_ep_ranks: the number of EP group ranks.
        num_global_experts: Number of experts in the model.
        num_local_experts: Number of experts in an EP rank.
        """
        import deep_ep

        # Defaults for internode and intranode are taken from DeepEP tests.
        num_nvl_bytes = 1024 * 1024 * 1024
        num_qps_per_rank = num_local_experts
        num_rdma_bytes = deep_ep.Buffer.get_low_latency_rdma_size_hint(
            num_max_dispatch_tokens_per_rank=max_num_tokens_per_dp_rank,
            hidden=token_hidden_size,
            num_ranks=num_ep_ranks,
            num_experts=num_global_experts)

        assert num_rdma_bytes is not None
        return dict(group=self.cpu_group,
                    num_nvl_bytes=num_nvl_bytes,
                    num_rdma_bytes=num_rdma_bytes,
                    low_latency_mode=True,
                    num_qps_per_rank=num_qps_per_rank)

    def get_handle(self, kwargs):
        """
        The kwargs for DeepEPLLAll2AllManager is dictated by
        _make_all2all_kwargs.
        """
        import deep_ep
        buffer_kwargs = self._make_all2all_kwargs(**kwargs)
        logger.debug("DeepEP all2all args %s", buffer_kwargs)
        handle: deep_ep.Buffer = self.handle_cache.get_or_create(
            buffer_kwargs, deep_ep.Buffer)
        # It is dangerous to set num sms outside this function. num_sms is not
        # a part of the hash-key that identifies this object. If we are in a
        # situation where we make objects with different num_sms, the hash key
        # in get_or_create must be updated.
        handle.set_num_sms(self.num_sms)
        return handle

__init__ 

__init__(cpu_group)
Source code in vllm/distributed/device_communicators/all2all.py 
212
213
def __init__(self, cpu_group):
    super().__init__(cpu_group)

_make_all2all_kwargs 

_make_all2all_kwargs(
    max_num_tokens_per_dp_rank: int,
    token_hidden_size: int,
    num_ep_ranks: int,
    num_global_experts: int,
    num_local_experts: int,
) -> dict[Any, Any]
the maximum number of tokens a DP rank

can dispatch all the ranks must hold the same value.

token_hidden_size: the hidden dimension of each token. num_ep_ranks: the number of EP group ranks. num_global_experts: Number of experts in the model. num_local_experts: Number of experts in an EP rank.

Source code in vllm/distributed/device_communicators/all2all.py 
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
def _make_all2all_kwargs(
    self,
    max_num_tokens_per_dp_rank: int,
    token_hidden_size: int,
    num_ep_ranks: int,
    num_global_experts: int,
    num_local_experts: int,
) -> dict[Any, Any]:
    """
    max_num_tokens_per_dp_rank : the maximum number of tokens a DP rank
      can dispatch all the ranks must hold the same value.
    token_hidden_size: the hidden dimension of each token.
    num_ep_ranks: the number of EP group ranks.
    num_global_experts: Number of experts in the model.
    num_local_experts: Number of experts in an EP rank.
    """
    import deep_ep

    # Defaults for internode and intranode are taken from DeepEP tests.
    num_nvl_bytes = 1024 * 1024 * 1024
    num_qps_per_rank = num_local_experts
    num_rdma_bytes = deep_ep.Buffer.get_low_latency_rdma_size_hint(
        num_max_dispatch_tokens_per_rank=max_num_tokens_per_dp_rank,
        hidden=token_hidden_size,
        num_ranks=num_ep_ranks,
        num_experts=num_global_experts)

    assert num_rdma_bytes is not None
    return dict(group=self.cpu_group,
                num_nvl_bytes=num_nvl_bytes,
                num_rdma_bytes=num_rdma_bytes,
                low_latency_mode=True,
                num_qps_per_rank=num_qps_per_rank)

get_handle 

get_handle(kwargs)

The kwargs for DeepEPLLAll2AllManager is dictated by _make_all2all_kwargs.

Source code in vllm/distributed/device_communicators/all2all.py 
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
def get_handle(self, kwargs):
    """
    The kwargs for DeepEPLLAll2AllManager is dictated by
    _make_all2all_kwargs.
    """
    import deep_ep
    buffer_kwargs = self._make_all2all_kwargs(**kwargs)
    logger.debug("DeepEP all2all args %s", buffer_kwargs)
    handle: deep_ep.Buffer = self.handle_cache.get_or_create(
        buffer_kwargs, deep_ep.Buffer)
    # It is dangerous to set num sms outside this function. num_sms is not
    # a part of the hash-key that identifies this object. If we are in a
    # situation where we make objects with different num_sms, the hash key
    # in get_or_create must be updated.
    handle.set_num_sms(self.num_sms)
    return handle

NaiveAll2AllManager

Bases: All2AllManagerBase

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.

Source code in vllm/distributed/device_communicators/all2all.py 
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
class NaiveAll2AllManager(All2AllManagerBase):
    """
    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):
        super().__init__(cpu_group)

    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

__init__ 

__init__(cpu_group)
Source code in vllm/distributed/device_communicators/all2all.py 
30
31
def __init__(self, cpu_group):
    super().__init__(cpu_group)

combine 

combine(hidden_states: Tensor) -> Tensor
Source code in vllm/distributed/device_communicators/all2all.py 
62
63
64
65
66
67
68
69
70
71
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

destroy 

destroy()
Source code in vllm/distributed/device_communicators/all2all.py 
73
74
def destroy(self):
    pass

dispatch 

dispatch(hidden_states: Tensor, router_logits: Tensor)
Source code in vllm/distributed/device_communicators/all2all.py 
51
52
53
54
55
56
57
58
59
60
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

naive_multicast 

naive_multicast(x: Tensor, cu_tokens_across_dp_cpu: Tensor)
Source code in vllm/distributed/device_communicators/all2all.py 
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
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

PPLXAll2AllManager

Bases: All2AllManagerBase

All2All communication based on PPLX kernels.

Source code in vllm/distributed/device_communicators/all2all.py 
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
class PPLXAll2AllManager(All2AllManagerBase):
    """
    All2All communication based on PPLX kernels.
    """

    def __init__(self, cpu_group):
        assert has_pplx(
        ), "pplx_kernels not found. Please follow https://github.com/vllm-project/vllm/blob/main/tools/ep_kernels/README.md to install pplx_kernels."  # noqa
        super().__init__(cpu_group)

        if self.internode:
            # inter-node communication needs nvshmem,
            # intra-node communication uses p2p mapping directly
            from pplx_kernels.nvshmem import (nvshmem_alloc_empty_unique_id,
                                              nvshmem_get_unique_id,
                                              nvshmem_init)
            logger.debug(
                "Initialize NVSHMEM for pplx_kernels: "
                "rank=%d, world size=%d", self.rank, self.world_size)
            uid = nvshmem_get_unique_id(
            ) if self.rank == 0 else nvshmem_alloc_empty_unique_id()
            dist.broadcast(uid,
                           src=dist.get_process_group_ranks(self.cpu_group)[0],
                           group=self.cpu_group)
            logger.debug("PPLX NVSHMEM UID = %s", uid)
            nvshmem_init(uid, self.rank, self.world_size)

        self.handle_cache = Cache()

    def get_handle(self, kwargs):
        import pplx_kernels as pplx
        return self.handle_cache.get_or_create(
            kwargs, pplx.AllToAll.internode
            if self.internode else pplx.AllToAll.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):
        with self.handle_cache._lock:
            for _, handle in self.handle_cache._cache.items():
                handle.destroy()

        if self.internode:
            from pplx_kernels.nvshmem import nvshmem_finalize
            logger.debug("PPLX NVSHMEM finalize")
            nvshmem_finalize()

handle_cache instance-attribute 

handle_cache = Cache()

__init__ 

__init__(cpu_group)
Source code in vllm/distributed/device_communicators/all2all.py 
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
def __init__(self, cpu_group):
    assert has_pplx(
    ), "pplx_kernels not found. Please follow https://github.com/vllm-project/vllm/blob/main/tools/ep_kernels/README.md to install pplx_kernels."  # noqa
    super().__init__(cpu_group)

    if self.internode:
        # inter-node communication needs nvshmem,
        # intra-node communication uses p2p mapping directly
        from pplx_kernels.nvshmem import (nvshmem_alloc_empty_unique_id,
                                          nvshmem_get_unique_id,
                                          nvshmem_init)
        logger.debug(
            "Initialize NVSHMEM for pplx_kernels: "
            "rank=%d, world size=%d", self.rank, self.world_size)
        uid = nvshmem_get_unique_id(
        ) if self.rank == 0 else nvshmem_alloc_empty_unique_id()
        dist.broadcast(uid,
                       src=dist.get_process_group_ranks(self.cpu_group)[0],
                       group=self.cpu_group)
        logger.debug("PPLX NVSHMEM UID = %s", uid)
        nvshmem_init(uid, self.rank, self.world_size)

    self.handle_cache = Cache()

combine 

combine(hidden_states: Tensor) -> Tensor
Source code in vllm/distributed/device_communicators/all2all.py 
116
117
def combine(self, hidden_states: torch.Tensor) -> torch.Tensor:
    raise NotImplementedError

destroy 

destroy()
Source code in vllm/distributed/device_communicators/all2all.py 
119
120
121
122
123
124
125
126
127
def destroy(self):
    with self.handle_cache._lock:
        for _, handle in self.handle_cache._cache.items():
            handle.destroy()

    if self.internode:
        from pplx_kernels.nvshmem import nvshmem_finalize
        logger.debug("PPLX NVSHMEM finalize")
        nvshmem_finalize()

dispatch 

dispatch(hidden_states: Tensor, router_logits: Tensor)
Source code in vllm/distributed/device_communicators/all2all.py 
112
113
114
def dispatch(self, hidden_states: torch.Tensor,
             router_logits: torch.Tensor):
    raise NotImplementedError

get_handle 

get_handle(kwargs)
Source code in vllm/distributed/device_communicators/all2all.py 
106
107
108
109
110
def get_handle(self, kwargs):
    import pplx_kernels as pplx
    return self.handle_cache.get_or_create(
        kwargs, pplx.AllToAll.internode
        if self.internode else pplx.AllToAll.intranode)