vllm.distributed.device_communicators.mnnvl_compat ¶

class CustomCommunicator(CommBackend):
    def __init__(self, group):
        self._group = group

    def Get_rank(self) -> int:
        return self._group.rank()

    def Get_size(self) -> int:
        return self._group.size()

    def allgather(self, data: int):
        gathered = [None] * self.Get_size()
        dist.all_gather_object(gathered, data, group=self._group)
        return gathered

    # NOTE(rob): CommBackend is an abstract class, and bcast/barrier
    # are unimplemented on vLLM side. If we need to utilize these
    # methods in the future, can create a concrete implementation.
    def bcast(self, data: Any, root: int) -> Any:
        raise NotImplementedError

    def barrier(self) -> None:
        raise NotImplementedError

    def Split(self, color: int, key: int) -> "CustomCommunicator":
        return self