all2all.py 15 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from typing import Any
4

5
import torch
6
import torch.distributed as dist
7

8
import vllm.envs as envs
9
from vllm.distributed import get_dp_group, get_ep_group
10
from vllm.forward_context import get_forward_context
11
from vllm.logger import init_logger
12
from vllm.utils import has_deep_ep, has_pplx
13
from vllm.utils.flashinfer import has_flashinfer_all2all
14

15
from .base_device_communicator import All2AllManagerBase, Cache
16

17
18
19
20
21
if has_flashinfer_all2all():
    from flashinfer.comm import Mapping
    from flashinfer.comm.mnnvl import MnnvlConfig
    from flashinfer.comm.trtllm_alltoall import MnnvlMoe

22
logger = init_logger(__name__)
23
24


25
class NaiveAll2AllManager(All2AllManagerBase):
26
27
28
29
30
31
32
    """
    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.
    """

33
34
    def __init__(self, cpu_group):
        super().__init__(cpu_group)
35

36
37
38
39
40
41
42
43
44
45
    def naive_multicast(
        self,
        x: torch.Tensor,
        cu_tokens_across_sp_cpu: torch.Tensor,
        is_sequence_parallel: bool,
    ) -> torch.Tensor:
        assert len(x.shape) == 2
        buffer = torch.empty(
            (cu_tokens_across_sp_cpu[-1], x.size(1)), device=x.device, dtype=x.dtype
        )
46

47
        rank = self.rank if is_sequence_parallel else self.dp_rank
48
        world_size = self.world_size if is_sequence_parallel else self.dp_world_size
49
50
51

        start = 0 if rank == 0 else cu_tokens_across_sp_cpu[rank - 1]
        end = cu_tokens_across_sp_cpu[rank]
52
        buffer[start:end, :].copy_(x)
53
54
55
56
        for idx in range(world_size):
            start = 0 if idx == 0 else cu_tokens_across_sp_cpu[idx - 1]
            end = cu_tokens_across_sp_cpu[idx]
            get_ep_group().broadcast(buffer[start:end, :], idx)
57
58
59

        return buffer

60
61
62
63
    def dispatch(
        self,
        hidden_states: torch.Tensor,
        router_logits: torch.Tensor,
64
        is_sequence_parallel: bool = False,
65
66
67
68
69
    ) -> tuple[torch.Tensor, torch.Tensor]:
        sp_size = self.tp_group.world_size if is_sequence_parallel else 1
        dp_metadata = get_forward_context().dp_metadata
        cu_tokens_across_sp_cpu = dp_metadata.cu_tokens_across_sp(sp_size)

70
71
72
73
74
75
        hidden_states = self.naive_multicast(
            hidden_states, cu_tokens_across_sp_cpu, is_sequence_parallel
        )
        router_logits = self.naive_multicast(
            router_logits, cu_tokens_across_sp_cpu, is_sequence_parallel
        )
76
77
        return hidden_states, router_logits

78
79
80
    def combine(
        self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False
    ) -> torch.Tensor:
81
82
83
84
85
86
87
88
89
90
91
        ep_rank = self.rank if is_sequence_parallel else self.dp_rank

        dp_metadata = get_forward_context().dp_metadata
        sp_size = self.tp_group.world_size if is_sequence_parallel else 1
        cu_tokens_across_sp_cpu = dp_metadata.cu_tokens_across_sp(sp_size)

        start = 0 if ep_rank == 0 else cu_tokens_across_sp_cpu[ep_rank - 1]
        end = cu_tokens_across_sp_cpu[ep_rank]

        all_hidden_states = get_ep_group().all_reduce(hidden_states)
        hidden_states = all_hidden_states[start:end, :]
92
93
94
95
        return hidden_states

    def destroy(self):
        pass
96
97


98
99
100
101
102
103
104
105
106
class AgRsAll2AllManager(All2AllManagerBase):
    """
    An implementation of all2all communication based on
    all-gather (dispatch) and reduce-scatter (combine).
    """

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

107
108
109
110
    def dispatch(
        self,
        hidden_states: torch.Tensor,
        router_logits: torch.Tensor,
111
        is_sequence_parallel: bool = False,
112
    ) -> tuple[torch.Tensor, torch.Tensor]:
113
114
115
        """
        Gather hidden_states and router_logits from all dp ranks.
        """
116
        sizes = get_forward_context().dp_metadata.get_chunk_sizes_across_dp_rank()
117
118
119
120

        dist_group = get_ep_group() if is_sequence_parallel else get_dp_group()
        assert sizes[dist_group.rank_in_group] == hidden_states.shape[0]
        hidden_states, router_logits = dist_group.all_gatherv(
121
122
123
124
125
126
            [hidden_states, router_logits],
            dim=0,
            sizes=sizes,
        )
        return hidden_states, router_logits

127
128
129
    def combine(
        self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False
    ) -> torch.Tensor:
130
131
132
        """
        Reduce-scatter hidden_states across all dp ranks.
        """
133
        sizes = get_forward_context().dp_metadata.get_chunk_sizes_across_dp_rank()
134
135

        dist_group = get_ep_group() if is_sequence_parallel else get_dp_group()
136
        hidden_states = dist_group.reduce_scatterv(hidden_states, dim=0, sizes=sizes)
137
138
139
140
141
142
        return hidden_states

    def destroy(self):
        pass


143
144
145
146
147
148
class PPLXAll2AllManager(All2AllManagerBase):
    """
    All2All communication based on PPLX kernels.
    """

    def __init__(self, cpu_group):
149
        assert has_pplx(), (
150
151
152
            "pplx_kernels not found. Please follow https://github.com/vllm-project/vllm/blob/main/tools/ep_kernels/README.md"
            " to install pplx_kernels."
        )
153
154
155
156
157
        super().__init__(cpu_group)

        if self.internode:
            # inter-node communication needs nvshmem,
            # intra-node communication uses p2p mapping directly
158
159
160
161
162
163
            from pplx_kernels.nvshmem import (
                nvshmem_alloc_empty_unique_id,
                nvshmem_get_unique_id,
                nvshmem_init,
            )

164
            logger.debug(
165
166
167
168
169
170
171
172
173
174
175
176
177
178
                "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,
            )
179
180
181
182
183
184
185
            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
186

187
        return self.handle_cache.get_or_create(
188
189
190
            kwargs,
            pplx.AllToAll.internode if self.internode else pplx.AllToAll.intranode,
        )
191

192
193
194
195
    def dispatch(
        self,
        hidden_states: torch.Tensor,
        router_logits: torch.Tensor,
196
        is_sequence_parallel: bool = False,
197
    ) -> tuple[torch.Tensor, torch.Tensor]:
198
199
        raise NotImplementedError

200
201
202
    def combine(
        self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False
    ) -> torch.Tensor:
203
204
205
206
207
208
209
210
211
        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
212

213
214
            logger.debug("PPLX NVSHMEM finalize")
            nvshmem_finalize()
215
216
217
218
219
220
221
222


class DeepEPAll2AllManagerBase(All2AllManagerBase):
    """
    All2All communication based on DeepEP High-Throughput kernels.
    """

    def __init__(self, cpu_group):
223
        assert has_deep_ep(), (
224
225
            "DeepEP kernels not found. Please follow https://github.com/vllm-project/vllm/blob/main/tools/ep_kernels/README.md"
            " to install DeepEP kernels."
226
        )  # noqa
227
228
229
230
231
232
233
234
235
236
        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

237
238
239
240
    def dispatch(
        self,
        hidden_states: torch.Tensor,
        router_logits: torch.Tensor,
241
        is_sequence_parallel: bool = False,
242
    ) -> tuple[torch.Tensor, torch.Tensor]:
243
244
        raise NotImplementedError

245
246
247
    def combine(
        self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False
    ) -> torch.Tensor:
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
        raise NotImplementedError

    def destroy(self):
        pass


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.
264
        num_nvl_bytes = envs.VLLM_DEEPEP_BUFFER_SIZE_MB * 1024 * 1024
265
266
267
268
        num_rdma_bytes = None
        num_qps_per_rank = None

        if self.internode:
269
            num_rdma_bytes = envs.VLLM_DEEPEP_BUFFER_SIZE_MB * 1024 * 1024
270
271
272
273
274
275
276
            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
277
278
279
280
281
282
283
        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,
        )
284
285
286
287

    def get_handle(self, kwargs):
        assert len(kwargs) == 0, (
            "DeepEPHTAll2AllManager expects no arguments. All the required "
288
289
            "args are computed in the Manager itself."
        )
290
291

        import deep_ep
292

293
294
295
        buffer_kwargs = self._make_all2all_kwargs()
        logger.debug("DeepEP all2all args %s", buffer_kwargs)
        handle: deep_ep.Buffer = self.handle_cache.get_or_create(
296
297
            buffer_kwargs, deep_ep.Buffer
        )
298
299
        return handle

300
301
302
303
304
305
306
307
308
309
    def set_num_sms(self, num_sms: int):
        import deep_ep

        # Right now the buffers are sized for only what the kernels were
        # created with. So we can only reduce the number of SMS used
        # but not increase it.
        if num_sms > self.num_sms:
            num_sms = self.num_sms
        deep_ep.Buffer.set_num_sms(num_sms)

310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337

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.
338
        num_nvl_bytes = envs.VLLM_DEEPEP_BUFFER_SIZE_MB * 1024 * 1024
339
        num_qps_per_rank = num_local_experts
340
341
342
343
        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,
344
345
            num_experts=num_global_experts,
        )
346
347

        assert num_rdma_bytes is not None
348
349
350
351
352
353
354
        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,
        )
355
356
357
358
359
360
361

    def get_handle(self, kwargs):
        """
        The kwargs for DeepEPLLAll2AllManager is dictated by
        _make_all2all_kwargs.
        """
        import deep_ep
362

363
364
365
        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(
366
367
            buffer_kwargs, deep_ep.Buffer
        )
368
        return handle
369
370

    # DeepEP LL uses RDMA so no SMs are used for communication
371
    def max_sms_used(self) -> int | None:
372
373
374
375
376
377
378
379
        return 0


class FlashInferAllToAllManager(All2AllManagerBase):
    """
    All2All communication based on flashinfer kernels.
    """

380
381
382
383
384
    # This type lint could be removed after all of the work in
    # https://github.com/vllm-project/vllm/issues/26533 done.
    rank: int
    world_size: int

385
    def __init__(self, cpu_group):
386
387
388
        assert has_flashinfer_all2all(), (
            "flashinfer all2all module not found. Please install/check flashinfer"
        )  # noqa
389
390
        super().__init__(cpu_group)
        logger.debug(
391
392
393
394
            "Initialize for flashinfer All2All rank=%d, world size=%d",
            self.rank,
            self.world_size,
        )
395
396
397
398
399
400
401
402
403
404
405
406
407
408
        self.initialized = False
        self.alltoall_info = None

    def initialize(
        self,
        world_size: int,
        rank: int,
        gpus_per_node: int,
    ):
        """Initialize workspace"""
        if self.initialized:
            return

        self.cleanup()
409
        logger.debug("making map: rank=%d, world size=%d", rank, world_size)
410
411
412
413
414
415
416
417
        self.mapping = Mapping(
            world_size,
            rank,
            gpus_per_node,
            tp_size=world_size,
        )

        from vllm.distributed.device_communicators.mnnvl_compat import (
418
419
420
            CustomCommunicator,
        )

421
422
423
        dp_config = MnnvlConfig(
            comm_backend=CustomCommunicator(get_dp_group().cpu_group),
            fabric_page_size=1 << 29,  # 512MB
424
            allocation_granularity=0,  # Auto-detect
425
426
        )

427
        self.workspace_tensor = MnnvlMoe.get_moe_workspaces(self.mapping, dp_config)
428
        self.prepare_workspace_tensor = MnnvlMoe.get_moe_prepare_workspace(
429
430
            self.mapping, dp_config
        )
431
432
433
434
435
436

        self.world_size = world_size
        self.rank = rank
        self.gpus_per_node = gpus_per_node
        self.initialized = True

437
438
439
        logger.info(
            "FlashInfer All2All initialized for rank %s, size %s", rank, world_size
        )
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461

    def ensure_alltoall_workspace_initialized(self):
        """Ensure workspace is initialized"""
        if not has_flashinfer_all2all():
            return False

        if self.world_size <= 1:
            return False

        if not self.initialized:
            self.initialize(
                world_size=self.world_size,
                rank=self.rank,
                gpus_per_node=torch.cuda.device_count,
            )
        return self.initialized

    def get_handle(self, kwargs):
        return self

    def cleanup(self):
        """Clean up workspace"""
462
463
464
465
466
        if (
            self.initialized
            and self.workspace_tensor is not None
            and self.prepare_workspace_tensor is not None
        ):
467
468
469
470
471
472
473
474
475
            try:
                del self.workspace_tensor
                del self.prepare_workspace_tensor
            except Exception as e:
                logger.warning("Failed to cleanup FlashInfer workspace: %s", e)
            finally:
                self.workspace_tensor = None
                self.prepare_workspace_tensor = None
                self.mapping = None
476
                self.initialized = False