vllm.v1.worker.gpu.pp_utils

Pipeline Parallelism utils for V2 Model Runner.

Classes:

• PPHandler –

  Runs the PP sampled-token broadcast/recv on a side stream so the

• PendingRecv –

  Per-step slot data for a deferred postprocess on the main stream.

Functions:

• compute_need_sampled_mask –

  Return a bool array of shape [input_batch.num_reqs] marking requests

PPHandler

Runs the PP sampled-token broadcast/recv on a side stream so the default stream isn't gated by the matching peer call. Step T's recv is consumed at step T+pp_size via get_prev_sampled_outputs.

Uses a dedicated NCCL communicator (sibling of the PP device_group) for the broadcast so it does not serialize on the wire with the inter-stage hidden-state p2p send/recv ops.

Methods:

• add_upstream_aux_hidden_states –

  Turn the last stage's model output, which carries only the aux hidden

• broadcast_drafts –

  Broadcast draft proposals so non-last ranks can embed real token ids.

• get_prev_sampled_outputs –

  Consume the entry from pp_size steps ago and wait for its recv event,

• receive –

  Returns True iff sampled tokens need to be gathered from all

• relay_aux_hidden_states –

  Forward earlier stages' aux hidden states next to this stage's.

Source code in vllm/v1/worker/gpu/pp_utils.py

class PPHandler:
    """Runs the PP sampled-token broadcast/recv on a side stream so the
    default stream isn't gated by the matching peer call. Step T's recv is
    consumed at step T+pp_size via `get_prev_sampled_outputs`.

    Uses a dedicated NCCL communicator (sibling of the PP `device_group`)
    for the broadcast so it does not serialize on the wire with the
    inter-stage hidden-state p2p send/recv ops.
    """

    def __init__(
        self, max_num_reqs: int, num_speculative_steps: int, device: torch.device
    ):
        self.is_last_rank = get_pp_group().is_last_rank
        self.last_rank = get_pp_group().last_rank
        self.max_sample_len = num_speculative_steps + 1
        self.num_speculative_steps = num_speculative_steps
        self.device = device
        self.main_stream = torch.cuda.current_stream(device)
        self.broadcast_stream = torch.cuda.Stream(device)

        # On non-last ranks, a FIFO with one entry per in-flight step: the entry
        # pushed by step T's `receive` is consumed pp_size steps later. Pre-seeded
        # with pp_size None placeholders so the first pp_size consumes are no-ops.
        # None means no postprocess is pending for that step (broadcast skipped).
        # Only XPU can disable microbatching via VLLM_XPU_PP_MICROBATCH.
        ring_depth = get_pp_group().world_size
        if current_platform.is_xpu() and not envs.VLLM_XPU_PP_MICROBATCH:
            ring_depth = 1
        self.queue: deque[PendingRecv | None] = (
            deque() if self.is_last_rank else deque([None] * ring_depth)
        )

        # Per req-index generation counter, incremented every time a request
        # index is freed in RequestStats. Used for invalidating freed req data
        # between PP decodes.
        self.req_idx_gen_np = np.zeros(max_num_reqs, dtype=np.int32)

        # Dedicated subgroup for the sampled-token broadcast.
        self.broadcast_group = get_pp_group().make_sibling_device_group(
            group_desc="pp_broadcast"
        )
        # Whether sampled tokens were broadcast this step.
        self.sampled_tokens_broadcast = False
        # Aux hidden states taken by earlier stages, and by this non-last stage.
        self.upstream_aux_keys: tuple[str, ...] = ()
        self.local_aux_keys: tuple[str, ...] = ()
        self.num_aux_hidden_states = 0

    def on_req_idx_freed(self, req_idx: int) -> None:
        self.req_idx_gen_np[req_idx] += 1

    def _sample_len(self, input_batch: InputBatch) -> int:
        # Width of the sampled-token broadcast. Must be the same on all ranks.
        return self.max_sample_len if input_batch.num_draft_tokens > 0 else 1

    def configure_aux_hidden_state_relay(self, model: torch.nn.Module) -> None:
        from vllm.model_executor.models.interfaces import get_eagle_decoder

        decoder = get_eagle_decoder(model)
        assert decoder is not None
        self.upstream_aux_keys, self.local_aux_keys = (
            decoder.aux_hidden_state_slot_keys()
        )
        self.num_aux_hidden_states = len(decoder.aux_hidden_state_layers)

    def relay_aux_hidden_states(
        self,
        intermediate_tensors: IntermediateTensors | None,
        output_intermediate_tensors: IntermediateTensors,
    ) -> IntermediateTensors:
        """Forward earlier stages' aux hidden states next to this stage's."""
        missing = [
            key
            for key in self.local_aux_keys
            if key not in output_intermediate_tensors.tensors
        ]
        if missing:
            raise RuntimeError(
                f"The model did not send aux hidden states {missing} to the next "
                "PP stage; its forward must add pack_local_aux_hidden_states() "
                "to the IntermediateTensors it returns."
            )
        if not self.upstream_aux_keys:
            return output_intermediate_tensors
        assert intermediate_tensors is not None
        return IntermediateTensors(
            output_intermediate_tensors.tensors
            | {key: intermediate_tensors[key] for key in self.upstream_aux_keys}
        )

    def add_upstream_aux_hidden_states(
        self,
        intermediate_tensors: IntermediateTensors | None,
        model_output: torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]],
    ) -> tuple[torch.Tensor, list[torch.Tensor]]:
        """Turn the last stage's model output, which carries only the aux hidden
        states this stage took, into the full list the speculator reads. Call it
        once per forward."""
        if isinstance(model_output, tuple):
            hidden_states, aux_hidden_states = model_output
        else:
            hidden_states, aux_hidden_states = model_output, []
        if self.upstream_aux_keys:
            assert intermediate_tensors is not None
            aux_hidden_states = [
                intermediate_tensors[key] for key in self.upstream_aux_keys
            ] + aux_hidden_states
        if len(aux_hidden_states) != self.num_aux_hidden_states:
            raise RuntimeError(
                f"The last PP stage has {len(aux_hidden_states)} aux hidden states "
                f"for {self.num_aux_hidden_states} requested layers; the model's "
                "forward must take them at global layer indices."
            )
        return hidden_states, aux_hidden_states

    def get_prev_sampled_outputs(self) -> dict[str, torch.Tensor | None] | None:
        """Consume the entry from pp_size steps ago and wait for its recv event,
        then filter out entries whose request was freed since `receive`.
        """
        if not self.queue:
            return None
        slot = self.queue.popleft()
        # Reserve this step's slot; `receive` overwrites it if applicable.
        self.queue.append(None)
        if slot is None:
            return None

        # Skip requests which did not need sampled output and/or those already
        # finished. The post_update kernel skips the -1 entries.
        freed = self.req_idx_gen_np[slot.idx_mapping_np] != slot.gen_at_receive_np
        exclude_mask = freed | ~slot.need_sampled_mask
        idx_mapping = slot.idx_mapping
        if exclude_mask.any():
            if exclude_mask.all():
                # No states require update anymore.
                return None
            # Filter excluded request indices.
            idx_mapping_np = np.where(exclude_mask, -1, slot.idx_mapping_np)
            idx_mapping = async_tensor_h2d(idx_mapping_np, device=self.device)

        self.main_stream.wait_event(slot.event)
        return dict(
            sampled_tokens=slot.sampled_tokens,
            num_sampled=slot.num_sampled,
            num_rejected=slot.num_rejected,
            idx_mapping=idx_mapping,
            draft_tokens=slot.draft_tokens,
        )

    def broadcast_drafts(self, draft_tokens: torch.Tensor) -> None:
        """Broadcast draft proposals so non-last ranks can embed real token ids."""
        assert self.is_last_rank
        if self.sampled_tokens_broadcast:
            with torch.cuda.stream(self.broadcast_stream):
                self.broadcast_stream.wait_stream(self.main_stream)
                torch.distributed.broadcast(
                    draft_tokens, src=self.last_rank, group=self.broadcast_group
                )
                draft_tokens.record_stream(self.broadcast_stream)

    def receive(self, input_batch: InputBatch, recv_drafts: bool) -> bool:
        """Returns True iff sampled tokens need to be gathered from *all*
        requests in the batch. `recv_drafts` must match whether the last rank
        calls `broadcast_drafts` for this step."""
        assert not self.is_last_rank
        need_sampled_mask = compute_need_sampled_mask(input_batch)
        if need_sampled_mask is None:
            # Leave this step's reserved slot as None.
            return False

        # Snapshot the per-slot generation counter so a later free of any of
        # these RequestStates request indices is detectable at consume time.
        gen_at_receive_np = self.req_idx_gen_np[input_batch.idx_mapping_np]

        num_reqs = input_batch.num_reqs
        with torch.cuda.stream(self.broadcast_stream):
            self.broadcast_stream.wait_stream(self.main_stream)
            sample_len = self._sample_len(input_batch)
            sampled_tokens = torch.empty(
                num_reqs, sample_len, dtype=torch.int64, device=self.device
            )
            combined = torch.empty(2, num_reqs, dtype=torch.int32, device=self.device)
            torch.distributed.broadcast(
                sampled_tokens, src=self.last_rank, group=self.broadcast_group
            )
            torch.distributed.broadcast(
                combined, src=self.last_rank, group=self.broadcast_group
            )
            draft_tokens = None
            if recv_drafts:
                draft_tokens = sampled_tokens.new_empty(
                    num_reqs, self.num_speculative_steps
                )
                torch.distributed.broadcast(
                    draft_tokens, src=self.last_rank, group=self.broadcast_group
                )
            event = self.broadcast_stream.record_event()
            num_sampled, num_rejected = combined.unbind(dim=0)
            # Must record_stream since these were allocated on broadcast stream but
            # later used on the main stream.
            sampled_tokens.record_stream(self.main_stream)
            combined.record_stream(self.main_stream)
            if draft_tokens is not None:
                draft_tokens.record_stream(self.main_stream)
        self.queue[-1] = PendingRecv(
            event,
            sampled_tokens,
            num_sampled,
            num_rejected,
            input_batch.idx_mapping,
            input_batch.idx_mapping_np,
            need_sampled_mask,
            gen_at_receive_np,
            draft_tokens,
        )
        return bool(need_sampled_mask.all())

    def broadcast(
        self,
        sampled_token_ids: torch.Tensor,
        num_sampled: torch.Tensor,
        num_rejected: torch.Tensor,
        input_batch: InputBatch,
    ) -> None:
        assert self.is_last_rank
        if compute_need_sampled_mask(input_batch) is None:
            # No request needs sampled outputs for a subsequent decode step.
            self.sampled_tokens_broadcast = False
            return

        assert sampled_token_ids.dtype == torch.int64
        self.sampled_tokens_broadcast = True

        if current_platform.is_xpu():
            self.main_stream.synchronize()

        with torch.cuda.stream(self.broadcast_stream):
            self.broadcast_stream.wait_stream(self.main_stream)
            # receive() allocates the same width on the other ranks.
            assert sampled_token_ids.shape[-1] == self._sample_len(input_batch)
            torch.distributed.broadcast(
                sampled_token_ids.contiguous(),
                src=self.last_rank,
                group=self.broadcast_group,
            )
            combined = torch.stack((num_sampled, num_rejected), dim=0)
            torch.distributed.broadcast(
                combined, src=self.last_rank, group=self.broadcast_group
            )
            for tensor in (sampled_token_ids, num_sampled, num_rejected):
                tensor.record_stream(self.broadcast_stream)

add_upstream_aux_hidden_states(intermediate_tensors, model_output)

Turn the last stage's model output, which carries only the aux hidden states this stage took, into the full list the speculator reads. Call it once per forward.

Source code in vllm/v1/worker/gpu/pp_utils.py

def add_upstream_aux_hidden_states(
    self,
    intermediate_tensors: IntermediateTensors | None,
    model_output: torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]],
) -> tuple[torch.Tensor, list[torch.Tensor]]:
    """Turn the last stage's model output, which carries only the aux hidden
    states this stage took, into the full list the speculator reads. Call it
    once per forward."""
    if isinstance(model_output, tuple):
        hidden_states, aux_hidden_states = model_output
    else:
        hidden_states, aux_hidden_states = model_output, []
    if self.upstream_aux_keys:
        assert intermediate_tensors is not None
        aux_hidden_states = [
            intermediate_tensors[key] for key in self.upstream_aux_keys
        ] + aux_hidden_states
    if len(aux_hidden_states) != self.num_aux_hidden_states:
        raise RuntimeError(
            f"The last PP stage has {len(aux_hidden_states)} aux hidden states "
            f"for {self.num_aux_hidden_states} requested layers; the model's "
            "forward must take them at global layer indices."
        )
    return hidden_states, aux_hidden_states

broadcast_drafts(draft_tokens)

Broadcast draft proposals so non-last ranks can embed real token ids.

Source code in vllm/v1/worker/gpu/pp_utils.py

def broadcast_drafts(self, draft_tokens: torch.Tensor) -> None:
    """Broadcast draft proposals so non-last ranks can embed real token ids."""
    assert self.is_last_rank
    if self.sampled_tokens_broadcast:
        with torch.cuda.stream(self.broadcast_stream):
            self.broadcast_stream.wait_stream(self.main_stream)
            torch.distributed.broadcast(
                draft_tokens, src=self.last_rank, group=self.broadcast_group
            )
            draft_tokens.record_stream(self.broadcast_stream)

get_prev_sampled_outputs()

Consume the entry from pp_size steps ago and wait for its recv event, then filter out entries whose request was freed since receive.

Source code in vllm/v1/worker/gpu/pp_utils.py

def get_prev_sampled_outputs(self) -> dict[str, torch.Tensor | None] | None:
    """Consume the entry from pp_size steps ago and wait for its recv event,
    then filter out entries whose request was freed since `receive`.
    """
    if not self.queue:
        return None
    slot = self.queue.popleft()
    # Reserve this step's slot; `receive` overwrites it if applicable.
    self.queue.append(None)
    if slot is None:
        return None

    # Skip requests which did not need sampled output and/or those already
    # finished. The post_update kernel skips the -1 entries.
    freed = self.req_idx_gen_np[slot.idx_mapping_np] != slot.gen_at_receive_np
    exclude_mask = freed | ~slot.need_sampled_mask
    idx_mapping = slot.idx_mapping
    if exclude_mask.any():
        if exclude_mask.all():
            # No states require update anymore.
            return None
        # Filter excluded request indices.
        idx_mapping_np = np.where(exclude_mask, -1, slot.idx_mapping_np)
        idx_mapping = async_tensor_h2d(idx_mapping_np, device=self.device)

    self.main_stream.wait_event(slot.event)
    return dict(
        sampled_tokens=slot.sampled_tokens,
        num_sampled=slot.num_sampled,
        num_rejected=slot.num_rejected,
        idx_mapping=idx_mapping,
        draft_tokens=slot.draft_tokens,
    )

receive(input_batch, recv_drafts)

Returns True iff sampled tokens need to be gathered from all requests in the batch. recv_drafts must match whether the last rank calls broadcast_drafts for this step.

Source code in vllm/v1/worker/gpu/pp_utils.py

def receive(self, input_batch: InputBatch, recv_drafts: bool) -> bool:
    """Returns True iff sampled tokens need to be gathered from *all*
    requests in the batch. `recv_drafts` must match whether the last rank
    calls `broadcast_drafts` for this step."""
    assert not self.is_last_rank
    need_sampled_mask = compute_need_sampled_mask(input_batch)
    if need_sampled_mask is None:
        # Leave this step's reserved slot as None.
        return False

    # Snapshot the per-slot generation counter so a later free of any of
    # these RequestStates request indices is detectable at consume time.
    gen_at_receive_np = self.req_idx_gen_np[input_batch.idx_mapping_np]

    num_reqs = input_batch.num_reqs
    with torch.cuda.stream(self.broadcast_stream):
        self.broadcast_stream.wait_stream(self.main_stream)
        sample_len = self._sample_len(input_batch)
        sampled_tokens = torch.empty(
            num_reqs, sample_len, dtype=torch.int64, device=self.device
        )
        combined = torch.empty(2, num_reqs, dtype=torch.int32, device=self.device)
        torch.distributed.broadcast(
            sampled_tokens, src=self.last_rank, group=self.broadcast_group
        )
        torch.distributed.broadcast(
            combined, src=self.last_rank, group=self.broadcast_group
        )
        draft_tokens = None
        if recv_drafts:
            draft_tokens = sampled_tokens.new_empty(
                num_reqs, self.num_speculative_steps
            )
            torch.distributed.broadcast(
                draft_tokens, src=self.last_rank, group=self.broadcast_group
            )
        event = self.broadcast_stream.record_event()
        num_sampled, num_rejected = combined.unbind(dim=0)
        # Must record_stream since these were allocated on broadcast stream but
        # later used on the main stream.
        sampled_tokens.record_stream(self.main_stream)
        combined.record_stream(self.main_stream)
        if draft_tokens is not None:
            draft_tokens.record_stream(self.main_stream)
    self.queue[-1] = PendingRecv(
        event,
        sampled_tokens,
        num_sampled,
        num_rejected,
        input_batch.idx_mapping,
        input_batch.idx_mapping_np,
        need_sampled_mask,
        gen_at_receive_np,
        draft_tokens,
    )
    return bool(need_sampled_mask.all())

relay_aux_hidden_states(intermediate_tensors, output_intermediate_tensors)

Forward earlier stages' aux hidden states next to this stage's.

Source code in vllm/v1/worker/gpu/pp_utils.py

def relay_aux_hidden_states(
    self,
    intermediate_tensors: IntermediateTensors | None,
    output_intermediate_tensors: IntermediateTensors,
) -> IntermediateTensors:
    """Forward earlier stages' aux hidden states next to this stage's."""
    missing = [
        key
        for key in self.local_aux_keys
        if key not in output_intermediate_tensors.tensors
    ]
    if missing:
        raise RuntimeError(
            f"The model did not send aux hidden states {missing} to the next "
            "PP stage; its forward must add pack_local_aux_hidden_states() "
            "to the IntermediateTensors it returns."
        )
    if not self.upstream_aux_keys:
        return output_intermediate_tensors
    assert intermediate_tensors is not None
    return IntermediateTensors(
        output_intermediate_tensors.tensors
        | {key: intermediate_tensors[key] for key in self.upstream_aux_keys}
    )

PendingRecv dataclass

Per-step slot data for a deferred postprocess on the main stream.

Source code in vllm/v1/worker/gpu/pp_utils.py

@dataclass
class PendingRecv:
    """Per-step slot data for a deferred postprocess on the main stream."""

    event: torch.cuda.Event

    sampled_tokens: torch.Tensor  # [num_reqs, sample_len]
    num_sampled: torch.Tensor  # [num_reqs]
    num_rejected: torch.Tensor  # [num_reqs]
    idx_mapping: torch.Tensor  # [num_reqs]
    idx_mapping_np: np.ndarray  # [num_reqs]
    # Records which rows need a deferred postprocess (bool).
    need_sampled_mask: np.ndarray  # [num_reqs]
    # Snapshot of slot generation counters at receive time, used to
    # detect requests aborted since then.
    gen_at_receive_np: np.ndarray  # [num_reqs]
    draft_tokens: torch.Tensor | None = None  # [num_reqs, num_speculative_steps]

compute_need_sampled_mask(input_batch)

Return a bool array of shape [input_batch.num_reqs] marking requests with outputs that might be needed in a subsequent (decode) step. Returns None if no sampled outputs are needed in the requests' next step.

Source code in vllm/v1/worker/gpu/pp_utils.py

def compute_need_sampled_mask(input_batch: InputBatch) -> np.ndarray | None:
    """Return a bool array of shape `[input_batch.num_reqs]` marking requests
    with outputs that might be needed in a subsequent (decode) step.
    Returns None if no sampled outputs are needed in the requests' next step."""
    old_computed = input_batch.num_computed_tokens_np
    prefill_len = input_batch.prefill_len_np
    # Exclude non-final prefill chunks (they don't produce a sample).
    need_sampled_mask = old_computed + input_batch.num_scheduled_tokens >= prefill_len
    if input_batch.max_seq_len_np is not None:
        # Also exclude final prefill chunks whose single sampled token reaches
        # the request's length cap.
        finished_prefill = prefill_len + 1 >= input_batch.max_seq_len_np
        finished_prefill &= input_batch.is_prefilling_np
        need_sampled_mask &= ~finished_prefill
    return need_sampled_mask if need_sampled_mask.any() else None