Skip to content

vllm.distributed.pp_transport

Alternatives to NCCL for sending a PP stage's intermediate tensors to the next stage, selected by VLLM_PP_TRANSPORT.

Modules:

  • copy_engine –

    PP transport on CUDA copy engines, synchronized on the GPU (one node).

Classes:

  • PPTransport –

    One rank's side of its PP stage boundaries.

Functions:

  • create_pp_transport –

    Collective over the PP group. Returns None (send over NCCL) without PP,

PPTransport

Bases: ABC

One rank's side of its PP stage boundaries.

Methods:

  • recv –

    Receive this step's inputs from the previous stage.

  • send –

    Send this step's outputs to the next stage. The caller may

Source code in vllm/distributed/pp_transport/__init__.py 
20
21
22
23
24
25
26
27
28
29
30
class PPTransport(ABC):
    """One rank's side of its PP stage boundaries."""

    @abstractmethod
    def send(self, tensors: dict, pp_group: "GroupCoordinator") -> None:
        """Send this step's outputs to the next stage. The caller may
        overwrite `tensors` once this returns."""

    @abstractmethod
    def recv(self, pp_group: "GroupCoordinator") -> "IntermediateTensors":
        """Receive this step's inputs from the previous stage."""

recv(pp_group) abstractmethod

Receive this step's inputs from the previous stage.

Source code in vllm/distributed/pp_transport/__init__.py 
28
29
30
@abstractmethod
def recv(self, pp_group: "GroupCoordinator") -> "IntermediateTensors":
    """Receive this step's inputs from the previous stage."""

send(tensors, pp_group) abstractmethod

Send this step's outputs to the next stage. The caller may overwrite tensors once this returns.

Source code in vllm/distributed/pp_transport/__init__.py 
23
24
25
26
@abstractmethod
def send(self, tensors: dict, pp_group: "GroupCoordinator") -> None:
    """Send this step's outputs to the next stage. The caller may
    overwrite `tensors` once this returns."""

create_pp_transport(vllm_config, persistent)

Collective over the PP group. Returns None (send over NCCL) without PP, or under "auto" when copy engines are not usable on some rank; every rank makes the same choice.

Parameters:

  • vllm_config

    (VllmConfig) –

    The engine config.

  • persistent

    (IntermediateTensors | None) –

    The model runner's persistent intermediate tensors (non-first ranks), which received tensors are copied into.

Source code in vllm/distributed/pp_transport/__init__.py 
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
def create_pp_transport(
    vllm_config: "VllmConfig", persistent: "IntermediateTensors | None"
) -> PPTransport | None:
    """Collective over the PP group. Returns None (send over NCCL) without PP,
    or under "auto" when copy engines are not usable on some rank; every rank
    makes the same choice.

    Args:
        vllm_config: The engine config.
        persistent: The model runner's persistent intermediate tensors (non-first
            ranks), which received tensors are copied into.

    """
    from vllm.distributed.parallel_state import get_pp_group

    pp_group = get_pp_group()
    if pp_group.world_size == 1:
        return None
    mode = envs.VLLM_PP_TRANSPORT
    assert mode in ("auto", "copy_engine")
    from vllm.distributed.pp_transport.copy_engine import (
        CopyEngineTransport,
        CopyEngineUnavailable,
    )

    try:
        transport = CopyEngineTransport(
            vllm_config, pp_group, None if persistent is None else persistent.tensors
        )
    except CopyEngineUnavailable as e:
        if mode == "copy_engine":
            raise
        logger.info_once("PP transport: NCCL (%s)", e)
        return None
    logger.info_once("PP transport: copy engines")
    return transport