Skip to content

vllm.distributed.communication_op

Functions:

tensor_model_parallel_all_gather(input_, dim=-1)

All-gather the input tensor across model parallel group.

Source code in vllm/distributed/communication_op.py 
17
18
19
20
21
def tensor_model_parallel_all_gather(
    input_: torch.Tensor, dim: int = -1
) -> torch.Tensor:
    """All-gather the input tensor across model parallel group."""
    return get_tp_group().all_gather(input_, dim)

tensor_model_parallel_all_reduce(input_)

All-reduce the input tensor across model parallel group.

Source code in vllm/distributed/communication_op.py 
12
13
14
def tensor_model_parallel_all_reduce(input_: torch.Tensor) -> torch.Tensor:
    """All-reduce the input tensor across model parallel group."""
    return get_tp_group().all_reduce(input_)

tensor_model_parallel_gather(input_, dst=0, dim=-1)

Gather the input tensor across model parallel group.

Source code in vllm/distributed/communication_op.py 
31
32
33
34
35
def tensor_model_parallel_gather(
    input_: torch.Tensor, dst: int = 0, dim: int = -1
) -> torch.Tensor | None:
    """Gather the input tensor across model parallel group."""
    return get_tp_group().gather(input_, dst, dim)

tensor_model_parallel_reduce_scatter(input_, dim=-1)

Reduce-Scatter the input tensor across model parallel group.

Source code in vllm/distributed/communication_op.py 
24
25
26
27
28
def tensor_model_parallel_reduce_scatter(
    input_: torch.Tensor, dim: int = -1
) -> torch.Tensor:
    """Reduce-Scatter the input tensor across model parallel group."""
    return get_tp_group().reduce_scatter(input_, dim)