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"
)
self.aux_hidden_state_relay_keys: tuple[str, ...] = ()
def on_req_idx_freed(self, req_idx: int) -> None:
self.req_idx_gen_np[req_idx] += 1
def configure_aux_hidden_state_relay(self, model: torch.nn.Module) -> None:
from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import (
aux_hidden_state_relay_keys,
)
self.aux_hidden_state_relay_keys = aux_hidden_state_relay_keys(model)
def relay_aux_hidden_states(
self,
intermediate_tensors: IntermediateTensors | None,
output_intermediate_tensors: IntermediateTensors,
) -> IntermediateTensors:
if not self.aux_hidden_state_relay_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.aux_hidden_state_relay_keys
}
)
def get_prev_sampled_outputs(
self, draft_tokens_to_update: torch.Tensor | None = None
) -> dict[str, torch.Tensor] | 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)
if slot.draft_tokens is not None and draft_tokens_to_update is not None:
draft_tokens = slot.draft_tokens
draft_idx_mapping = slot.idx_mapping
if exclude_mask.any():
keep = ~exclude_mask
keep_t = torch.as_tensor(keep, device=self.device)
draft_tokens = draft_tokens[keep_t]
draft_idx_mapping = async_tensor_h2d(
slot.idx_mapping_np[keep], device=self.device
)
draft_tokens_to_update[draft_idx_mapping] = draft_tokens
return dict(
sampled_tokens=slot.sampled_tokens,
num_sampled=slot.num_sampled,
num_rejected=slot.num_rejected,
idx_mapping=idx_mapping,
)
def broadcast_drafts(
self, draft_tokens: torch.Tensor, input_batch: InputBatch
) -> None:
"""Broadcast draft proposals so non-last ranks can embed real token ids."""
assert self.is_last_rank
if compute_need_sampled_mask(input_batch) is None:
return
with torch.cuda.stream(self.broadcast_stream):
self.broadcast_stream.wait_stream(self.main_stream)
send = draft_tokens[input_batch.idx_mapping].contiguous()
# Must record the idx_mapping tensor since it was allocated
# on the main stream.
input_batch.idx_mapping.record_stream(self.broadcast_stream)
torch.distributed.broadcast(
send, src=self.last_rank, group=self.broadcast_group
)
def receive(self, input_batch: InputBatch) -> bool:
"""Returns True iff sampled tokens need to be gathered from *all*
requests in the batch."""
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)
sampled_tokens = torch.empty(
num_reqs, self.max_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 self.num_speculative_steps > 0:
draft_tokens = torch.empty(
num_reqs,
self.num_speculative_steps,
dtype=torch.int64,
device=self.device,
)
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.
return
assert sampled_token_ids.dtype == torch.int64
if current_platform.is_xpu():
self.main_stream.synchronize()
with torch.cuda.stream(self.broadcast_stream):
self.broadcast_stream.wait_stream(self.main_stream)
send_tokens = torch.nn.functional.pad(
sampled_token_ids,
(0, self.max_sample_len - sampled_token_ids.shape[-1]),
)
torch.distributed.broadcast(
send_tokens.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)