buffer.py 42.3 KB
Newer Older
Chenggang Zhao's avatar
Chenggang Zhao committed
1
import os
lijian6's avatar
lijian6 committed
2
3
from typing import Callable, List, Optional, Tuple, Union

Chenggang Zhao's avatar
Chenggang Zhao committed
4
5
6
import torch
import torch.distributed as dist

lijian6's avatar
lijian6 committed
7
8
9
from . import deep_ep_cpp
from .deep_ep_cpp import Config, EventHandle

10
from .utils import EventOverlap, check_nvlink_connections
Chenggang Zhao's avatar
Chenggang Zhao committed
11
12
13
14
15
16


class Buffer:
    """
    The core expert-parallel (EP) communication buffers for Mixture of Experts (MoE) model, which supports:
        - high-throughput intranode all-to-all (dispatch and combine, using NVLink)
17
18
        - high-throughput internode all-to-all (dispatch and combine, using RDMA and NVLink)
        - low-latency all-to-all (dispatch and combine, using RDMA)
Chenggang Zhao's avatar
Chenggang Zhao committed
19
20
21
22
23
24
25
26
27
28
29
30
31

    Attributes:
        num_sms: the SMs used in high-throughput kernels.
        rank: the local rank number.
        group_size: the number of ranks in the group.
        group: the communication group.
        num_nvl_bytes: the buffer size for intranode NVLink communication.
        num_rdma_bytes: the buffer size for internode (also for intranode with low-latency mode) RDMA communication.
        runtime: the C++ runtime.
    """

    num_sms: int = 20

lijian6's avatar
lijian6 committed
32
33
34
35
36
37
38
39
40
41
42
43
    def __init__(
        self,
        group: dist.ProcessGroup,
        num_nvl_bytes: int = 0,
        num_rdma_bytes: int = 0,
        low_latency_mode: bool = False,
        num_qps_per_rank: int = 24,
        allow_nvlink_for_low_latency_mode: bool = True,
        allow_mnnvl: bool = False,
        explicitly_destroy: bool = False,
        use_default_stream_as_comm_stream: bool = True,
    ) -> None:
Chenggang Zhao's avatar
Chenggang Zhao committed
44
45
46
47
48
49
50
51
52
53
        """
        Initialize the communication buffer.

        Arguments:
            group: the communication group.
            num_nvl_bytes: the buffer size for intranode NVLink communication.
            num_rdma_bytes: the buffer size for internode (also for intranode with low-latency mode) RDMA communication.
            low_latency_mode: whether to enable low-latency mode.
            num_qps_per_rank: the number of QPs for RDMA, the low-latency mode requires that this number equals
                to the number of local experts.
Chenggang Zhao's avatar
Chenggang Zhao committed
54
55
56
57
            allow_nvlink_for_low_latency_mode: whether allow NVLink traffic for low-latency mode, you should notice
                this is somehow incompatible with the hook-based overlapping.
                Warning: PCIe connections may lead to errors due to memory ordering issues,
                please make sure all connections are via NVLink.
fzyzcjy's avatar
more  
fzyzcjy committed
58
            allow_mnnvl: whether to allow MNNVL
59
60
61
            explicitly_destroy: If this flag is set to True, you need to explicitly call `destroy()` to release resources;
                otherwise, the resources will be released by the destructor.
                Note: Releasing resources in the destructor may cause Python's exception handling process to hang.
Chenggang Zhao's avatar
Chenggang Zhao committed
62
        """
63
        check_nvlink_connections(group)
Chenggang Zhao's avatar
Chenggang Zhao committed
64
65

        # Initialize the CPP runtime
lijian6's avatar
lijian6 committed
66
67
68
        self.rank = group.rank()
        self.group_size = group.size()
        self.group = group
Chenggang Zhao's avatar
Chenggang Zhao committed
69
70
71
        self.num_nvl_bytes = num_nvl_bytes
        self.num_rdma_bytes = num_rdma_bytes
        self.low_latency_mode = low_latency_mode
72
        self.explicitly_destroy = explicitly_destroy
lijian6's avatar
lijian6 committed
73
74
75
76
77
78
79
80
81
        self.runtime = deep_ep_cpp.Buffer(
            self.rank,
            self.group_size,
            num_nvl_bytes,
            num_rdma_bytes,
            low_latency_mode,
            explicitly_destroy,
            use_default_stream_as_comm_stream,
        )
Chenggang Zhao's avatar
Chenggang Zhao committed
82
83

        # Synchronize device IDs
lijian6's avatar
lijian6 committed
84
85
86
        device_ids = [
            None,
        ] * self.group_size
Chenggang Zhao's avatar
Chenggang Zhao committed
87
        local_device_id = self.runtime.get_local_device_id()
lijian6's avatar
lijian6 committed
88
        dist.all_gather_object(device_ids, local_device_id, group)
Chenggang Zhao's avatar
Chenggang Zhao committed
89
90

        # Synchronize IPC handles
lijian6's avatar
lijian6 committed
91
92
93
        ipc_handles = [
            None,
        ] * self.group_size
Chenggang Zhao's avatar
Chenggang Zhao committed
94
        local_ipc_handle = self.runtime.get_local_ipc_handle()
lijian6's avatar
lijian6 committed
95
        dist.all_gather_object(ipc_handles, local_ipc_handle, group)
Chenggang Zhao's avatar
Chenggang Zhao committed
96
97
98
99

        # Synchronize NVSHMEM unique IDs
        root_unique_id = None
        if self.runtime.get_num_rdma_ranks() > 1 or low_latency_mode:
sky's avatar
sky committed
100
            # Enable IBGDA
101
            assert num_qps_per_rank > 0
lijian6's avatar
lijian6 committed
102
103
104
            os.environ["NVSHMEM_DISABLE_P2P"] = "0" if allow_nvlink_for_low_latency_mode else "1"
            os.environ["NVSHMEM_IB_ENABLE_IBGDA"] = "1"
            os.environ["NVSHMEM_IBGDA_NUM_RC_PER_PE"] = f"{num_qps_per_rank}"
105
            # Make sure QP depth is always larger than the number of on-flight WRs, so that we can skip WQ slot check
lijian6's avatar
lijian6 committed
106
            os.environ["NVSHMEM_QP_DEPTH"] = os.environ.get("NVSHMEM_QP_DEPTH", "1024")
107
108
109

            # Reduce gpu memory usage
            # 6 default teams + 1 extra team
lijian6's avatar
lijian6 committed
110
            os.environ["NVSHMEM_MAX_TEAMS"] = "7"
111
            # Disable NVLink SHArP
lijian6's avatar
lijian6 committed
112
            os.environ["NVSHMEM_DISABLE_NVLS"] = "1"
113
            # NOTES: NVSHMEM initialization requires at least 256 MiB
lijian6's avatar
lijian6 committed
114
            os.environ["NVSHMEM_CUMEM_GRANULARITY"] = f"{2 ** 29}"
Chenggang Zhao's avatar
Chenggang Zhao committed
115

fzyzcjy's avatar
more  
fzyzcjy committed
116
117
            if not allow_mnnvl:
                # Disable multi-node NVLink detection
lijian6's avatar
lijian6 committed
118
                os.environ["NVSHMEM_DISABLE_MNNVL"] = "1"
119

Chenggang Zhao's avatar
Chenggang Zhao committed
120
            # Synchronize using the root ID
lijian6's avatar
lijian6 committed
121
122
123
124
125
126
            nvshmem_unique_ids = [
                None,
            ] * self.group_size
            if (low_latency_mode and self.rank == 0) or (
                not low_latency_mode and self.runtime.get_rdma_rank() == 0
            ):
Chenggang Zhao's avatar
Chenggang Zhao committed
127
                root_unique_id = self.runtime.get_local_nvshmem_unique_id()
lijian6's avatar
lijian6 committed
128
129
130
131
            dist.all_gather_object(nvshmem_unique_ids, root_unique_id, group)
            root_unique_id = nvshmem_unique_ids[
                0 if low_latency_mode else self.runtime.get_root_rdma_rank(True)
            ]
Chenggang Zhao's avatar
Chenggang Zhao committed
132
133
134
135
136

        # Make CPP runtime available
        self.runtime.sync(device_ids, ipc_handles, root_unique_id)
        assert self.runtime.is_available()

137
138
139
    def destroy(self):
        """
        Destroy the cpp runtime and release resources.
sky's avatar
sky committed
140

141
142
        """

lijian6's avatar
lijian6 committed
143
        assert self.explicitly_destroy, "`explicitly_destroy` flag must be set"
144
145
146
147

        self.runtime.destroy()
        self.runtime = None

lijian6's avatar
lijian6 committed
148
149
150
    # @staticmethod
    # def is_sm90_compiled():
    #     return deep_ep_cpp.is_sm90_compiled()
151

Chenggang Zhao's avatar
Chenggang Zhao committed
152
153
154
155
156
157
158
159
160
    @staticmethod
    def set_num_sms(new_num_sms: int) -> None:
        """
        Set the number of SMs to use in high-throughput kernels.

        Arguments:
            new_num_sms: the new number to be set.
        """

lijian6's avatar
lijian6 committed
161
        assert new_num_sms % 2 == 0, "The SM count must be even"
Chenggang Zhao's avatar
Chenggang Zhao committed
162
163
164
165
166
167
168
169
170
171
172
173
174
        Buffer.num_sms = new_num_sms

    @staticmethod
    def capture() -> EventOverlap:
        """
        Capture a CUDA event on the current stream, i.e. `torch.cuda.current_stream()`.

        Returns:
            event: the captured event.
        """
        return EventOverlap(EventHandle())

    @staticmethod
lijian6's avatar
lijian6 committed
175
176
177
    def get_low_latency_rdma_size_hint(
        num_max_dispatch_tokens_per_rank: int, hidden: int, num_ranks: int, num_experts: int
    ) -> int:
Chenggang Zhao's avatar
Chenggang Zhao committed
178
179
180
181
182
183
184
185
186
187
188
189
        """
        Get a minimum size requirement for the RDMA buffer. The size calculation will be done with BF16.

        Arguments:
            num_max_dispatch_tokens_per_rank: the maximum number of tokens to dispatch, all the ranks must hold the same value.
            hidden: the hidden dimension of each token.
            num_ranks: the number of EP group ranks.
            num_experts: the number of all experts.

        Returns:
            size: the RDMA buffer size recommended.
        """
lijian6's avatar
lijian6 committed
190
191
192
        return deep_ep_cpp.get_low_latency_rdma_size_hint(
            num_max_dispatch_tokens_per_rank, hidden, num_ranks, num_experts
        )
sky's avatar
sky committed
193

Shangyan Zhou's avatar
Shangyan Zhou committed
194
195
196
197
198
    def get_comm_stream(self) -> torch.Stream:
        """
        Get the communication stream.

        Returns:
sky's avatar
sky committed
199
            stream: the communication stream.
Shangyan Zhou's avatar
Shangyan Zhou committed
200
201
        """
        ts: torch.Stream = self.runtime.get_comm_stream()
lijian6's avatar
lijian6 committed
202
203
204
205
206
207
208
209
210
211
212
        return torch.cuda.Stream(
            stream_id=ts.stream_id, device_index=ts.device_index, device_type=ts.device_type
        )

    def get_local_buffer_tensor(
        self,
        dtype: torch.dtype,
        size: Optional[torch.Size] = None,
        offset: int = 0,
        use_rdma_buffer: bool = False,
    ) -> torch.Tensor:
Chenggang Zhao's avatar
Chenggang Zhao committed
213
214
215
216
217
218
219
220
221
222
223
224
225
226
        """
        Get the raw buffer (slice supported) as a PyTorch tensor.

        Argument:
            dtype: the data type (PyTorch `dtype`) for the tensor.
            size: the slice size (by elements) to get from the buffer.
            offset: the offset of the beginning element.
            use_rdma_buffer: whether to return the RDMA buffer.
        """
        tensor = self.runtime.get_local_buffer_tensor(dtype, offset, use_rdma_buffer)
        if size is None:
            return tensor

        assert tensor.numel() >= size.numel()
lijian6's avatar
lijian6 committed
227
        return tensor[: size.numel()].view(size)
Chenggang Zhao's avatar
Chenggang Zhao committed
228

Shangyan Zhou's avatar
Shangyan Zhou committed
229
230
231
232
233
234
235
236
237
238
    @staticmethod
    def _unpack_bias(bias: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]):
        bias_0, bias_1 = None, None
        if isinstance(bias, torch.Tensor):
            bias_0 = bias
        elif isinstance(bias, tuple):
            assert len(bias) == 2
            bias_0, bias_1 = bias
        return bias_0, bias_1

Chenggang Zhao's avatar
Chenggang Zhao committed
239
240
241
242
243
244
245
246
247
248
249
250
    @staticmethod
    def get_dispatch_config(num_ranks: int) -> Config:
        """
        Get a recommended dispatch config.

        Argument:
            num_ranks: the number of ranks.

        Returns:
            config: the recommended config.
        """

Shifang Xu's avatar
Shifang Xu committed
251
        # TODO: automatically tune
Chenggang Zhao's avatar
Chenggang Zhao committed
252
        config_map = {
253
254
            2: Config(Buffer.num_sms, 24, 256, 6, 128),
            4: Config(Buffer.num_sms, 6, 256, 6, 128),
Chenggang Zhao's avatar
Chenggang Zhao committed
255
            8: Config(Buffer.num_sms, 6, 256, 6, 128),
Shangyan Zhou's avatar
Shangyan Zhou committed
256
            16: Config(Buffer.num_sms, 36, 288, 20, 128),
Chenggang Zhao's avatar
Chenggang Zhao committed
257
            24: Config(Buffer.num_sms, 8, 288, 32, 128),
lijian6's avatar
lijian6 committed
258
259
260
            32: Config(Buffer.num_sms, 32, 288, 32, 128),
            64: Config(Buffer.num_sms, 20, 288, 28, 128),
            128: Config(Buffer.num_sms, 20, 560, 32, 128),
Chenggang Zhao's avatar
Chenggang Zhao committed
261
262
263
            144: Config(Buffer.num_sms, 32, 720, 12, 128),
            160: Config(Buffer.num_sms, 28, 720, 12, 128),
        }
lijian6's avatar
lijian6 committed
264
        assert num_ranks in config_map, f"Unsupported number of EP ranks: {num_ranks}"
Chenggang Zhao's avatar
Chenggang Zhao committed
265
266
267
268
269
270
271
272
273
274
275
276
277
278
        return config_map[num_ranks]

    @staticmethod
    def get_combine_config(num_ranks: int) -> Config:
        """
        Get a recommended combine config.

        Argument:
            num_ranks: the number of ranks.

        Returns:
            config: the recommended config.
        """

Shifang Xu's avatar
Shifang Xu committed
279
        # TODO: automatically tune
Chenggang Zhao's avatar
Chenggang Zhao committed
280
        config_map = {
281
282
283
            2: Config(Buffer.num_sms, 10, 256, 6, 128),
            4: Config(Buffer.num_sms, 9, 256, 6, 128),
            8: Config(Buffer.num_sms, 4, 256, 6, 128),
Shangyan Zhou's avatar
Shangyan Zhou committed
284
            16: Config(Buffer.num_sms, 4, 288, 12, 128),
285
286
            24: Config(Buffer.num_sms, 1, 288, 8, 128),
            32: Config(Buffer.num_sms, 1, 288, 8, 128),
lijian6's avatar
lijian6 committed
287
288
            64: Config(Buffer.num_sms, 1, 288, 20, 128),
            128: Config(Buffer.num_sms, 1, 560, 12, 128),
Chenggang Zhao's avatar
Chenggang Zhao committed
289
290
291
            144: Config(Buffer.num_sms, 2, 720, 8, 128),
            160: Config(Buffer.num_sms, 2, 720, 8, 128),
        }
lijian6's avatar
lijian6 committed
292
        assert num_ranks in config_map, f"Unsupported number of EP ranks: {num_ranks}"
Chenggang Zhao's avatar
Chenggang Zhao committed
293
294
295
        return config_map[num_ranks]

    # noinspection PyTypeChecker
lijian6's avatar
lijian6 committed
296
297
298
299
300
301
302
303
    def get_dispatch_layout(
        self,
        topk_idx: torch.Tensor,
        num_experts: int,
        previous_event: Optional[EventOverlap] = None,
        async_finish: bool = False,
        allocate_on_comm_stream: bool = False,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor, torch.Tensor, EventOverlap]:
Chenggang Zhao's avatar
Chenggang Zhao committed
304
305
306
307
        """
        Calculate the layout required for later communication.

        Arguments:
lijian6's avatar
lijian6 committed
308
309
            topk_idx: `[num_tokens, num_topk]`, dtype must be `torch.int64`, the expert indices selected by each token,
                `-1` means no selections.
Chenggang Zhao's avatar
Chenggang Zhao committed
310
311
312
313
314
315
316
317
318
319
320
321
322
            num_experts: the number of experts.
            previous_event: the event to wait before actually executing the kernel.
            async_finish: the current stream will not wait for the communication kernels to be finished if set.
            allocate_on_comm_stream: control whether all the allocated tensors' ownership to be on the communication stream.

        Returns:
            num_tokens_per_rank: `[num_ranks]` with `torch.int`, the number of tokens to be sent to each rank.
            num_tokens_per_rdma_rank: `[num_rdma_ranks]` with `torch.int`, the number of tokens to be sent to each RDMA
                rank (with the same GPU index), return `None` for intranode settings.
            num_tokens_per_expert: `[num_experts]` with `torch.int`, the number of tokens to be sent to each expert.
            is_token_in_rank: `[num_tokens, num_ranks]` with `torch.bool`, whether a token be sent to a rank.
            event: the event after executing the kernel (valid only if `async_finish` is set).
        """
lijian6's avatar
lijian6 committed
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
        num_tokens_per_rank, num_tokens_per_rdma_rank, num_tokens_per_expert, is_token_in_rank, event = (
            self.runtime.get_dispatch_layout(
                topk_idx,
                num_experts,
                getattr(previous_event, "event", None),
                async_finish,
                allocate_on_comm_stream,
            )
        )
        return (
            num_tokens_per_rank,
            num_tokens_per_rdma_rank,
            num_tokens_per_expert,
            is_token_in_rank,
            EventOverlap(event),
        )
Chenggang Zhao's avatar
Chenggang Zhao committed
339
340

    # noinspection PyTypeChecker
lijian6's avatar
lijian6 committed
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
    def dispatch(
        self,
        x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
        handle: Optional[Tuple] = None,
        num_tokens_per_rank: Optional[torch.Tensor] = None,
        num_tokens_per_rdma_rank: Optional[torch.Tensor] = None,
        is_token_in_rank: Optional[torch.Tensor] = None,
        num_tokens_per_expert: Optional[torch.Tensor] = None,
        topk_idx: Optional[torch.Tensor] = None,
        topk_weights: Optional[torch.Tensor] = None,
        expert_alignment: int = 1,
        num_worst_tokens: int = 0,
        config: Optional[Config] = None,
        previous_event: Optional[EventOverlap] = None,
        async_finish: bool = False,
        allocate_on_comm_stream: bool = False,
        num_recv_tokens_per_expert_as_cuda: bool = False,
    ) -> Tuple[
        Union[Tuple[torch.Tensor, torch.Tensor], torch.Tensor],
        Optional[torch.Tensor],
        Optional[torch.Tensor],
        List[int],
        torch.Tensor,
        Tuple,
        EventOverlap,
    ]:
Chenggang Zhao's avatar
Chenggang Zhao committed
367
368
369
370
        """
        Dispatch tokens to different ranks, both intranode and internode settings are supported.
        Intranode kernels require all the ranks should be visible via NVLink.
        Internode kernels require the ranks in a node should be visible via NVLink, while the ranks with the same GPU
371
            index should be visible via RDMA.
Chenggang Zhao's avatar
Chenggang Zhao committed
372
373
374
375
376
377
378
379
380
381
382
383

        Arguments:
            x: `torch.Tensor` or tuple of `torch.Tensor`, for the first type, the shape must be `[num_tokens, hidden]`,
                and type must be `torch.bfloat16`; for the second type, the first element of the tuple must be shaped as
                `[num_tokens, hidden]` with type `torch.float8_e4m3fn`, the second must be `[num_tokens, hidden // 128]`
                 (requiring divisible) with type `torch.float`.
            handle: an optional communication handle, if set, the CPU will reuse the layout information to save some time.
            num_tokens_per_rank: `[num_ranks]` with `torch.int`, the number of tokens to be sent to each rank.
            num_tokens_per_rdma_rank: `[num_rdma_ranks]` with `torch.int`, the number of tokens to be sent to each RDMA
                rank (with the same GPU index), return `None` for intranode settings.
            is_token_in_rank: `[num_tokens, num_ranks]` with `torch.bool`, whether a token be sent to a rank.
            num_tokens_per_expert: `[num_experts]` with `torch.int`, the number of tokens to be sent to each expert.
lijian6's avatar
lijian6 committed
384
385
            topk_idx: `[num_tokens, num_topk]` with `torch.int64`, the expert indices selected by each token,
                `-1` means no selections.
Chenggang Zhao's avatar
Chenggang Zhao committed
386
387
            topk_weights: `[num_tokens, num_topk]` with `torch.float`, the expert weights of each token to dispatch.
            expert_alignment: align the number of tokens received by each local expert to this variable.
388
389
            num_worst_tokens: the worst number of tokens to receive, if specified, there will be no CPU sync, and it
                will be CUDA-graph compatible. Please also notice that this flag is for intranode only.
Chenggang Zhao's avatar
Chenggang Zhao committed
390
391
392
393
            config: the performance tuning config.
            previous_event: the event to wait before actually executing the kernel.
            async_finish: the current stream will not wait for the communication kernels to be finished if set.
            allocate_on_comm_stream: control whether all the allocated tensors' ownership to be on the communication stream.
lijian6's avatar
lijian6 committed
394
            num_recv_tokens_per_expert_as_cuda: control return num_recv_tokens_per_expert as cuda tensor or python list.
Chenggang Zhao's avatar
Chenggang Zhao committed
395
396
397
398
399
        Returns:
            recv_x: received tokens, the same type and tuple as the input `x`, but the number of tokens equals to the
                received token count.
            recv_topk_idx: received expert indices.
            recv_topk_weights: received expert weights.
lijian6's avatar
lijian6 committed
400
            num_recv_tokens_per_expert: Python list or cuda tensor shaped `[num_local_experts]`, the received token count by
Chenggang Zhao's avatar
Chenggang Zhao committed
401
402
                each local expert, aligned to the input `expert_alignment`. If `num_worst_tokens` is specified, the list
                will be empty.
Chenggang Zhao's avatar
Chenggang Zhao committed
403
404
405
406
407
408
409
410
            handle: the returned communication handle.
            event: the event after executing the kernel (valid only if `async_finish` is set).
        """
        # Default config
        config = self.get_dispatch_config(self.group_size) if config is None else config

        # Internode
        if self.runtime.get_num_rdma_ranks() > 1:
lijian6's avatar
lijian6 committed
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
            assert num_worst_tokens == 0, "Internode dispatch does not support `num_worst_tokens > 0`"
            return self.internode_dispatch(
                x,
                handle,
                num_tokens_per_rank,
                num_tokens_per_rdma_rank,
                is_token_in_rank,
                num_tokens_per_expert,
                topk_idx,
                topk_weights,
                expert_alignment,
                config,
                previous_event,
                async_finish,
                allocate_on_comm_stream,
            )
Chenggang Zhao's avatar
Chenggang Zhao committed
427
428
429
430
431

        # Launch the kernel with cached or non-cached mode
        x, x_scales = x if isinstance(x, tuple) else (x, None)
        if handle is not None:
            assert topk_idx is None and topk_weights is None
lijian6's avatar
lijian6 committed
432
433
434
435
436
437
438
439
            (
                rank_prefix_matrix,
                channel_prefix_matrix,
                recv_channel_prefix_matrix,
                recv_src_idx,
                is_token_in_rank,
                send_head,
            ) = handle
Chenggang Zhao's avatar
Chenggang Zhao committed
440
            num_recv_tokens = recv_src_idx.size(0)
lijian6's avatar
lijian6 committed
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
            recv_x, recv_x_scales, _, _, _, _, _, _, _, _, _, event = self.runtime.intranode_dispatch(
                x,
                x_scales,
                None,
                None,
                None,
                is_token_in_rank,
                None,
                num_recv_tokens,
                rank_prefix_matrix,
                channel_prefix_matrix,
                expert_alignment,
                num_worst_tokens,
                config,
                getattr(previous_event, "event", None),
                async_finish,
                allocate_on_comm_stream,
            )
            return (
                (recv_x, recv_x_scales) if x_scales is not None else recv_x,
                None,
                None,
                None,
                None,
                EventOverlap(event),
            )
Chenggang Zhao's avatar
Chenggang Zhao committed
467
        else:
lijian6's avatar
lijian6 committed
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
            assert (
                num_tokens_per_rank is not None
                and is_token_in_rank is not None
                and num_tokens_per_expert is not None
            )
            (
                recv_x,
                recv_x_scales,
                recv_topk_idx,
                recv_topk_weights,
                num_recv_tokens_per_expert_list,
                num_recv_tokens_per_expert_cuda,
                rank_prefix_matrix,
                channel_prefix_matrix,
                recv_channel_prefix_matrix,
                recv_src_idx,
                send_head,
                event,
            ) = self.runtime.intranode_dispatch(
                x,
                x_scales,
                topk_idx,
                topk_weights,
                num_tokens_per_rank,
                is_token_in_rank,
                num_tokens_per_expert,
                0,
                None,
                None,
                expert_alignment,
                num_worst_tokens,
                config,
                getattr(previous_event, "event", None),
                async_finish,
                allocate_on_comm_stream,
            )
            handle = (
                rank_prefix_matrix,
                channel_prefix_matrix,
                recv_channel_prefix_matrix,
                recv_src_idx,
                is_token_in_rank,
                send_head,
            )
            return (
                (recv_x, recv_x_scales) if x_scales is not None else recv_x,
                recv_topk_idx,
                recv_topk_weights,
                (
                    num_recv_tokens_per_expert_cuda
                    if num_recv_tokens_per_expert_as_cuda
                    else num_recv_tokens_per_expert_list
                ),
                handle,
                EventOverlap(event),
            )
Chenggang Zhao's avatar
Chenggang Zhao committed
524
525

    # noinspection PyTypeChecker
lijian6's avatar
lijian6 committed
526
527
528
529
530
531
532
533
534
535
536
    def combine(
        self,
        x: torch.Tensor,
        handle: Tuple,
        topk_weights: Optional[torch.Tensor] = None,
        bias: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]] = None,
        config: Optional[Config] = None,
        previous_event: Optional[EventOverlap] = None,
        async_finish: bool = False,
        allocate_on_comm_stream: bool = False,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], EventOverlap]:
Chenggang Zhao's avatar
Chenggang Zhao committed
537
538
539
540
541
        """
        Combine (reduce) tokens (addition **without** weights) from different ranks, both intranode and internode
            settings are supported.
        Intranode kernels require all the ranks should be visible via NVLink.
        Internode kernels require the ranks in a node should be visible via NVLink, while the ranks with the same GPU
542
            index should be visible via RDMA.
Chenggang Zhao's avatar
Chenggang Zhao committed
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562

        Arguments:
            x: `[num_tokens, hidden]` with `torch.bfloat16`, the tokens to send for reducing to its original ranks.
            handle: a must-set communication handle, you can obtain this from the dispatch function.
            topk_weights: `[num_tokens, num_topk]` with `torch.float`, the tokens' top-k weights for reducing to its original ranks.
            config: the performance tuning config.
            previous_event: the event to wait before actually executing the kernel.
            async_finish: the current stream will not wait for the communication kernels to be finished if set.
            allocate_on_comm_stream: control whether all the allocated tensors' ownership to be on the communication stream.

        Returns:
            recv_x: the reduced token from its dispatched ranks.
            recv_topk_weights: the reduced top-k weights from its dispatch ranks.
            event: the event after executing the kernel (valid only if `async_finish` is set).
        """
        # Default config
        config = self.get_combine_config(self.group_size) if config is None else config

        # Internode
        if self.runtime.get_num_rdma_ranks() > 1:
lijian6's avatar
lijian6 committed
563
564
565
            return self.internode_combine(
                x, handle, topk_weights, bias, config, previous_event, async_finish, allocate_on_comm_stream
            )
Chenggang Zhao's avatar
Chenggang Zhao committed
566
567
568

        # NOTES: the second `_` is for the sending side, so we should use the third one
        rank_prefix_matrix, _, channel_prefix_matrix, src_idx, is_recv_token_in_rank, send_head = handle
Shangyan Zhou's avatar
Shangyan Zhou committed
569
        bias_0, bias_1 = Buffer._unpack_bias(bias)
Chenggang Zhao's avatar
Chenggang Zhao committed
570
571
572

        # Launch the kernel
        recv_x, recv_topk_weights, event = self.runtime.intranode_combine(
lijian6's avatar
lijian6 committed
573
574
575
576
577
578
579
580
581
582
583
584
585
            x,
            topk_weights,
            bias_0,
            bias_1,
            src_idx,
            rank_prefix_matrix,
            channel_prefix_matrix,
            send_head,
            config,
            getattr(previous_event, "event", None),
            async_finish,
            allocate_on_comm_stream,
        )
Chenggang Zhao's avatar
Chenggang Zhao committed
586
587
588
        return recv_x, recv_topk_weights, EventOverlap(event)

    # noinspection PyTypeChecker
lijian6's avatar
lijian6 committed
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
    def internode_dispatch(
        self,
        x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
        handle: Optional[Tuple] = None,
        num_tokens_per_rank: Optional[torch.Tensor] = None,
        num_tokens_per_rdma_rank: Optional[torch.Tensor] = None,
        is_token_in_rank: Optional[torch.Tensor] = None,
        num_tokens_per_expert: Optional[torch.Tensor] = None,
        topk_idx: Optional[torch.Tensor] = None,
        topk_weights: Optional[torch.Tensor] = None,
        expert_alignment: int = 1,
        config: Optional[Config] = None,
        previous_event: Optional[EventOverlap] = None,
        async_finish: bool = False,
        allocate_on_comm_stream: bool = False,
    ) -> Tuple[
        Union[Tuple[torch.Tensor, torch.Tensor], torch.Tensor],
        Optional[torch.Tensor],
        Optional[torch.Tensor],
        List[int],
        Tuple,
        EventOverlap,
    ]:
Chenggang Zhao's avatar
Chenggang Zhao committed
612
613
614
615
616
617
618
619
620
621
        """
        Internode dispatch implementation, for more details, please refer to the `dispatch` docs.
        Normally, you should not directly call this function.
        """
        assert config is not None

        # Launch the kernel with cached or non-cached mode
        x, x_scales = x if isinstance(x, tuple) else (x, None)
        if handle is not None:
            assert topk_idx is None and topk_weights is None
lijian6's avatar
lijian6 committed
622
623
624
625
626
627
628
629
630
631
632
633
            (
                is_token_in_rank,
                rdma_channel_prefix_matrix,
                gbl_channel_prefix_matrix,
                recv_rdma_channel_prefix_matrix,
                recv_rdma_rank_prefix_sum,
                recv_gbl_channel_prefix_matrix,
                recv_gbl_rank_prefix_sum,
                recv_src_meta,
                send_rdma_head,
                send_nvl_head,
            ) = handle
Chenggang Zhao's avatar
Chenggang Zhao committed
634
635
            num_recv_tokens = recv_src_meta.size(0)
            num_rdma_recv_tokens = send_nvl_head.size(0)
lijian6's avatar
lijian6 committed
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
            recv_x, recv_x_scales, _, _, _, _, _, _, _, _, _, _, _, _, event = (
                self.runtime.internode_dispatch(
                    x,
                    x_scales,
                    topk_idx,
                    topk_weights,
                    None,
                    None,
                    is_token_in_rank,
                    None,
                    num_recv_tokens,
                    num_rdma_recv_tokens,
                    rdma_channel_prefix_matrix,
                    recv_rdma_rank_prefix_sum,
                    gbl_channel_prefix_matrix,
                    recv_gbl_rank_prefix_sum,
                    expert_alignment,
                    config,
                    getattr(previous_event, "event", None),
                    async_finish,
                    allocate_on_comm_stream,
                )
            )
            return (
                (recv_x, recv_x_scales) if x_scales is not None else recv_x,
                None,
                None,
                None,
                None,
                EventOverlap(event),
            )
Chenggang Zhao's avatar
Chenggang Zhao committed
667
        else:
lijian6's avatar
lijian6 committed
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
            assert (
                num_tokens_per_rank is not None
                and is_token_in_rank is not None
                and num_tokens_per_expert is not None
            )
            (
                recv_x,
                recv_x_scales,
                recv_topk_idx,
                recv_topk_weights,
                num_recv_tokens_per_expert_list,
                rdma_channel_prefix_matrix,
                gbl_channel_prefix_matrix,
                recv_rdma_channel_prefix_matrix,
                recv_rdma_rank_prefix_sum,
                recv_gbl_channel_prefix_matrix,
                recv_gbl_rank_prefix_sum,
                recv_src_meta,
                send_rdma_head,
                send_nvl_head,
                event,
            ) = self.runtime.internode_dispatch(
                x,
                x_scales,
                topk_idx,
                topk_weights,
                num_tokens_per_rank,
                num_tokens_per_rdma_rank,
                is_token_in_rank,
                num_tokens_per_expert,
                0,
                0,
                None,
                None,
                None,
                None,
                expert_alignment,
                config,
                getattr(previous_event, "event", None),
                async_finish,
                allocate_on_comm_stream,
            )
            handle = (
                is_token_in_rank,
                rdma_channel_prefix_matrix,
                gbl_channel_prefix_matrix,
                recv_rdma_channel_prefix_matrix,
                recv_rdma_rank_prefix_sum,
                recv_gbl_channel_prefix_matrix,
                recv_gbl_rank_prefix_sum,
                recv_src_meta,
                send_rdma_head,
                send_nvl_head,
            )
            return (
                (recv_x, recv_x_scales) if x_scales is not None else recv_x,
                recv_topk_idx,
                recv_topk_weights,
                num_recv_tokens_per_expert_list,
                handle,
                EventOverlap(event),
            )
Chenggang Zhao's avatar
Chenggang Zhao committed
730
731

    # noinspection PyTypeChecker
lijian6's avatar
lijian6 committed
732
733
734
735
736
737
738
739
740
741
742
    def internode_combine(
        self,
        x: torch.Tensor,
        handle: Union[tuple, list],
        topk_weights: Optional[torch.Tensor] = None,
        bias: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]] = None,
        config: Optional[Config] = None,
        previous_event: Optional[EventOverlap] = None,
        async_finish: bool = False,
        allocate_on_comm_stream: bool = False,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], EventOverlap]:
Chenggang Zhao's avatar
Chenggang Zhao committed
743
744
745
746
747
748
        """
        Internode combine implementation, for more details, please refer to the `combine` docs.
        Normally, you should not directly call this function.
        """
        assert config is not None

Shangyan Zhou's avatar
Shangyan Zhou committed
749
        # Unpack handle and bias
lijian6's avatar
lijian6 committed
750
751
752
753
754
755
756
757
758
759
760
761
        (
            is_combined_token_in_rank,
            _,
            _,
            rdma_channel_prefix_matrix,
            rdma_rank_prefix_sum,
            gbl_channel_prefix_matrix,
            gbl_rank_prefix_sum,
            src_meta,
            send_rdma_head,
            send_nvl_head,
        ) = handle
Shangyan Zhou's avatar
Shangyan Zhou committed
762
        bias_0, bias_1 = Buffer._unpack_bias(bias)
Chenggang Zhao's avatar
Chenggang Zhao committed
763
764
765

        # Launch the kernel
        combined_x, combined_topk_weights, event = self.runtime.internode_combine(
lijian6's avatar
lijian6 committed
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
            x,
            topk_weights,
            bias_0,
            bias_1,
            src_meta,
            is_combined_token_in_rank,
            rdma_channel_prefix_matrix,
            rdma_rank_prefix_sum,
            gbl_channel_prefix_matrix,
            send_rdma_head,
            send_nvl_head,
            config,
            getattr(previous_event, "event", None),
            async_finish,
            allocate_on_comm_stream,
        )
Chenggang Zhao's avatar
Chenggang Zhao committed
782
783
        return combined_x, combined_topk_weights, EventOverlap(event)

lijian6's avatar
lijian6 committed
784
785
786
    def clean_low_latency_buffer(
        self, num_max_dispatch_tokens_per_rank: int, hidden: int, num_experts: int
    ) -> None:
Chenggang Zhao's avatar
Chenggang Zhao committed
787
788
789
790
791
792
793
794
795
796
797
798
799
800
        """
        As low-latency kernels require part of the buffer to be zero-initialized, so it is vital to clean the buffer
            if the buffer is dirty at some time.
        For example, after running the normal dispatch/combine, you must run this function before executing any
            low-latency kernel.

        Arguments:
            num_max_dispatch_tokens_per_rank: the maximum number of tokens to dispatch, all the ranks must hold the same value.
            hidden: the hidden dimension of each token.
            num_experts: the number of all experts.
        """
        self.runtime.clean_low_latency_buffer(num_max_dispatch_tokens_per_rank, hidden, num_experts)

    # noinspection PyTypeChecker
lijian6's avatar
lijian6 committed
801
802
803
804
805
806
807
808
809
810
811
812
813
814
    def low_latency_dispatch(
        self,
        x: torch.Tensor,
        topk_idx: torch.Tensor,
        num_max_dispatch_tokens_per_rank: int,
        num_experts: int,
        cumulative_local_expert_recv_stats: Optional[torch.Tensor] = None,
        dispatch_wait_recv_cost_stats: Optional[torch.Tensor] = None,
        use_fp8: bool = True,
        round_scale: bool = False,
        use_ue8m0: bool = False,
        async_finish: bool = False,
        return_recv_hook: bool = False,
    ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], torch.Tensor, Tuple, EventOverlap, Callable]:
Chenggang Zhao's avatar
Chenggang Zhao committed
815
        """
816
        A low-latency implementation for dispatching with IBGDA.
Chenggang Zhao's avatar
Chenggang Zhao committed
817
818
        This kernel requires all the ranks (no matter intranode or internode) should be visible via RDMA
            (specifically, IBGDA must be enabled).
Shifang Xu's avatar
Shifang Xu committed
819
        Warning: as there are only two buffers, and the returned tensors reuse the buffer, you cannot hold more than 2
820
            low-latency kernels' result tensors at a single moment.
Chenggang Zhao's avatar
Chenggang Zhao committed
821
822
823
824

        Arguments:
            x: `torch.Tensor` with `torch.bfloat16`, shaped as `[num_tokens, hidden]`, only several hidden shapes are
                supported. The number of tokens to be dispatched must be less than `num_max_dispatch_tokens_per_rank`.
lijian6's avatar
lijian6 committed
825
826
            topk_idx: `torch.Tensor` with `torch.int64`, shaped as `[num_tokens, num_topk]`, only several top-k shapes
                are supported. `-1` indices (not selecting any expert) are supported.
Chenggang Zhao's avatar
Chenggang Zhao committed
827
828
            num_max_dispatch_tokens_per_rank: the maximum number of tokens to dispatch, all the ranks must hold the same value.
            num_experts: the number of all experts.
829
830
831
            cumulative_local_expert_recv_stats: a cumulative expert count tensor for statistics, which should have shape
                `[num_local_experts]` and be typed as `torch.int`. This is useful for online service EP load balance
                monitoring.
832
833
            dispatch_wait_recv_cost_stats: a cumulative time spent waiting to receive each token tensor for statistics,
                which should have shape `[num_ranks, num_ranks]` and be typed as `torch.int64`.
lijian6's avatar
lijian6 committed
834
                This is useful for detecting and pre-cisely localizing slow anomalies.
835
            use_fp8: whether to enable FP8 casting, with this, the received data will be a tuple of FP8 tensor and scaling factors.
Shifang Xu's avatar
Shifang Xu committed
836
837
            round_scale: whether round the scaling factors into power of 2.
            use_ue8m0: whether use UE8M0 as scaling factor format (available only with `round_scale=True`).
Chenggang Zhao's avatar
Chenggang Zhao committed
838
839
840
            async_finish: the current stream will not wait for the communication kernels to be finished if set.
            return_recv_hook: return a receiving hook if set. If set, the kernel will just do the RDMA request issues,
                but **without actually receiving the data**. You must call the received hook to make sure the data's arrival.
Shifang Xu's avatar
Shifang Xu committed
841
                If you do not set this flag, the kernel will ensure the data's arrival.
Chenggang Zhao's avatar
Chenggang Zhao committed
842
843

        Returns:
844
845
            recv_x: a tensor or tuple with received tokens for each expert.
                With `use_fp8=True`: the first element is a `torch.Tensor` shaped as
Chenggang Zhao's avatar
Chenggang Zhao committed
846
847
                `[num_local_experts, num_max_dispatch_tokens_per_rank * num_ranks, hidden]` with `torch.float8_e4m3fn`.
                The second tensor is the corresponding scales for the first element with shape
Shifang Xu's avatar
Shifang Xu committed
848
849
850
                `[num_local_experts, num_max_dispatch_tokens_per_rank * num_ranks, hidden // 128]` with `torch.float`,
                if `use_ue8m0=False`. With `use_ue8m0=True`, the second one is packed and shaped as
                `[num_local_experts, num_max_dispatch_tokens_per_rank * num_ranks, hidden // 512]` with type `torch.int`.
Chenggang Zhao's avatar
Chenggang Zhao committed
851
                Notice that, the last-two-dimension of the scaling tensors are in column-major for TMA compatibility.
852
853
                With `use_fp8=False`, the result would be a tensor shaped as
                `[num_local_experts, num_max_dispatch_tokens_per_rank * num_ranks, hidden]` with `torch.bfloat16`.
Chenggang Zhao's avatar
Chenggang Zhao committed
854
                Moreover, not all tokens are valid, only some of the `num_max_dispatch_tokens_per_rank * num_ranks` are,
855
                as we do not synchronize CPU received count with GPU (also not incompatible with CUDA graph if synced).
Chenggang Zhao's avatar
Chenggang Zhao committed
856
            recv_count: a tensor shaped `[num_local_experts]` with type `torch.int`, indicating how many tokens each
857
                expert receives. As mentioned before, not all tokens are valid in `recv_x`.
Chenggang Zhao's avatar
Chenggang Zhao committed
858
859
860
861
            handle: the communication handle to be used in the `low_latency_combine` function.
            event: the event after executing the kernel (valid only if `async_finish` is set).
            hook: the receiving hook function (valid only if `return_recv_hook` is set).
        """
lijian6's avatar
lijian6 committed
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
        (
            packed_recv_x,
            packed_recv_x_scales,
            packed_recv_count,
            packed_recv_src_info,
            packed_recv_layout_range,
            event,
            hook,
        ) = self.runtime.low_latency_dispatch(
            x,
            topk_idx,
            cumulative_local_expert_recv_stats,
            dispatch_wait_recv_cost_stats,
            num_max_dispatch_tokens_per_rank,
            num_experts,
            use_fp8,
            round_scale,
            use_ue8m0,
            async_finish,
            return_recv_hook,
        )
        handle = (
            packed_recv_src_info,
            packed_recv_layout_range,
            num_max_dispatch_tokens_per_rank,
            x.size(1),
            num_experts,
        )
        tensors_to_record = (
            x,
            topk_idx,
            packed_recv_x,
            packed_recv_x_scales,
            packed_recv_count,
            packed_recv_src_info,
            packed_recv_layout_range,
            cumulative_local_expert_recv_stats,
        )
        return (
            (packed_recv_x, packed_recv_x_scales) if use_fp8 else packed_recv_x,
            packed_recv_count,
            handle,
            EventOverlap(event, tensors_to_record if async_finish else None),
            hook,
        )
Chenggang Zhao's avatar
Chenggang Zhao committed
907
908

    # noinspection PyTypeChecker
lijian6's avatar
lijian6 committed
909
910
911
912
913
914
915
916
917
918
919
920
921
    def low_latency_combine(
        self,
        x: torch.Tensor,
        topk_idx: torch.Tensor,
        topk_weights: torch.Tensor,
        handle: tuple,
        use_logfmt: bool = False,
        zero_copy: bool = False,
        async_finish: bool = False,
        return_recv_hook: bool = False,
        out: Optional[torch.Tensor] = None,
        combine_wait_recv_cost_stats: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, EventOverlap, Callable]:
Chenggang Zhao's avatar
Chenggang Zhao committed
922
923
924
925
        """
        A low-latency implementation for combining tokens (reduce **with weights**) with IBGDA.
        This kernel requires all the ranks (no matter intranode or internode) should be visible via RDMA
            (specifically, IBGDA must be enabled).
Shifang Xu's avatar
Shifang Xu committed
926
927
        Warning: as there are only two buffers, and the returned tensors reuse the buffer, you cannot hold more than 2
            low-latency kernels' result tensors at a single moment.
Chenggang Zhao's avatar
Chenggang Zhao committed
928
929
930
931

        Arguments:
            x: `[num_local_experts, num_max_dispatch_tokens_per_rank * num_ranks, hidden]` with `torch.bfloat16`,
                the local calculated tokens to be sent to this original rank and reduced.
lijian6's avatar
lijian6 committed
932
933
934
            topk_idx: `[num_combined_tokens, num_topk]` with `torch.int64`, the expert indices selected by the dispatched
                tokens. `-1` indices (not selecting any expert) are supported. Note that, `num_combined_tokens` equals
                to the number of dispatched tokens.
Chenggang Zhao's avatar
Chenggang Zhao committed
935
936
937
            topk_weights: `[num_combined_tokens, num_topk]` with `torch.float`, the expert weights selected by the dispatched
                tokens. The received tokens will be reduced with the weights in this tensor.
            handle: the communication handle given by the `dispatch` function.
938
            use_logfmt: whether to use an internal "LogFMT with dynamic per-64-channel cast" format (10 bits).
939
940
            zero_copy: whether the tensor is already copied into the RDMA buffer, should be cooperative
                with `get_next_low_latency_combine_buffer`.
Chenggang Zhao's avatar
Chenggang Zhao committed
941
942
943
            async_finish: the current stream will not wait for the communication kernels to be finished if set.
            return_recv_hook: return a receiving hook if set. If set, the kernel will just do the RDMA request issues,
                but **without actually receiving the data**. You must call the received hook to make sure the data's arrival.
Shifang Xu's avatar
Shifang Xu committed
944
                If you do not set this flag, the kernel will ensure the data's arrival.
945
            out: the in-place output tensor, if set, the kernel will write the result to this tensor and return it directly.
946
947
948
            combine_wait_recv_cost_stats: a cumulative time spent waiting to receive each token tensor for statistics,
                which should have shape `[num_ranks, num_ranks]` and be typed as `torch.int64`.
                This is useful for detecting and pre-cisely localizing slow anomalies.
Chenggang Zhao's avatar
Chenggang Zhao committed
949
950

        Returns:
fzyzcjy's avatar
fzyzcjy committed
951
            combined_x: the reduced token tensor, with shape `[num_combined_tokens, hidden]` and type `torch.bfloat16`.
Chenggang Zhao's avatar
Chenggang Zhao committed
952
953
954
            event: the event after executing the kernel (valid only if `async_finish` is set).
            hook: the receiving hook function (valid only if `return_recv_hook` is set).
        """
955
        src_info, layout_range, num_max_dispatch_tokens_per_rank, hidden, num_experts = handle
lijian6's avatar
lijian6 committed
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
        combined_x, event, hook = self.runtime.low_latency_combine(
            x,
            topk_idx,
            topk_weights,
            src_info,
            layout_range,
            combine_wait_recv_cost_stats,
            num_max_dispatch_tokens_per_rank,
            num_experts,
            use_logfmt,
            zero_copy,
            async_finish,
            return_recv_hook,
            out,
        )
Chenggang Zhao's avatar
Chenggang Zhao committed
971
972
        tensors_to_record = (x, topk_idx, topk_weights, src_info, layout_range, combined_x)
        return combined_x, EventOverlap(event, tensors_to_record if async_finish else None), hook
973
974
975
976
977
978
979
980
981
982
983
984
985
986

    def get_next_low_latency_combine_buffer(self, handle: object):
        """
        Get the raw registered RDMA buffer tensor for next low-latency combine, so that the next combine kernel can skip the copying.

        Arguments:
            handle: the communication handle given by the `dispatch` function.

        Returns:
            buffer: the raw RDMA low-latency buffer as a BF16 PyTorch tensor with shape
                `[num_local_experts, num_ranks * num_max_dispatch_tokens_per_rank, hidden]`, you should fill this buffer
                by yourself.
        """
        src_info, layout_range, num_max_dispatch_tokens_per_rank, hidden, num_experts = handle
lijian6's avatar
lijian6 committed
987
988
989
        return self.runtime.get_next_low_latency_combine_buffer(
            num_max_dispatch_tokens_per_rank, hidden, num_experts
        )