Skip to content

vllm.distributed.device_communicators.mnnvl_compat

CustomCommunicator

Bases: CommBackend

Source code in vllm/distributed/device_communicators/mnnvl_compat.py 
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
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

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

_group instance-attribute 

_group = group

Get_rank 

Get_rank() -> int
Source code in vllm/distributed/device_communicators/mnnvl_compat.py 
16
17
def Get_rank(self) -> int:
    return self._group.rank()

Get_size 

Get_size() -> int
Source code in vllm/distributed/device_communicators/mnnvl_compat.py 
19
20
def Get_size(self) -> int:
    return self._group.size()

Split 

Split(color: int, key: int) -> CustomCommunicator
Source code in vllm/distributed/device_communicators/mnnvl_compat.py 
27
28
def Split(self, color: int, key: int) -> 'CustomCommunicator':
    return self

__init__ 

__init__(group)
Source code in vllm/distributed/device_communicators/mnnvl_compat.py 
13
14
def __init__(self, group):
    self._group = group

allgather 

allgather(data: int)
Source code in vllm/distributed/device_communicators/mnnvl_compat.py 
22
23
24
25
def allgather(self, data: int):
    gathered = [None] * self.Get_size()
    dist.all_gather_object(gathered, data, group=self._group)
    return gathered