| 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.distributed as dist |
| 6 | from flashinfer.comm.mnnvl import CommBackend as CommBackend |
| 7 | |
| 8 | from vllm.utils.flashinfer import has_flashinfer_nvlink_two_sided |
| 9 | |
| 10 | assert has_flashinfer_nvlink_two_sided(), "Flashinfer alltoallv module cannot be found" |
| 11 | |
| 12 | |
| 13 | class CustomCommunicator(CommBackend): |
| 14 | def __init__(self, group): |
| 15 | self._group = group |
| 16 | |
| 17 | def Get_rank(self) -> int: |
| 18 | return self._group.rank() |
| 19 | |
| 20 | def Get_size(self) -> int: |
| 21 | return self._group.size() |
| 22 | |
| 23 | def allgather(self, data: int): |
| 24 | gathered = [None] * self.Get_size() |
| 25 | dist.all_gather_object(gathered, data, group=self._group) |
| 26 | return gathered |
| 27 | |
| 28 | def bcast(self, data: Any, root: int) -> Any: |
| 29 | obj_list = [data] |
| 30 | # broadcast_object_list mutates obj_list in-place |
| 31 | dist.broadcast_object_list(obj_list, src=root, group=self._group) |
| 32 | return obj_list[0] |
| 33 | |
| 34 | def barrier(self) -> None: |
| 35 | dist.barrier(group=self._group) |
| 36 | |
| 37 | def Split(self, color: int, key: int) -> "CustomCommunicator": |
| 38 | return self |