| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Pipeline Parallelism utils for V2 Model Runner.""" |
| 4 | |
| 5 | from collections import deque |
| 6 | from dataclasses import dataclass |
| 7 | |
| 8 | import numpy as np |
| 9 | import torch |
| 10 | |
| 11 | import vllm.envs as envs |
| 12 | from vllm.distributed.parallel_state import get_pp_group |
| 13 | from vllm.platforms import current_platform |
| 14 | from vllm.sequence import IntermediateTensors |
| 15 | from vllm.utils.torch_utils import async_tensor_h2d |
| 16 | from vllm.v1.worker.gpu.input_batch import InputBatch |
| 17 | |
| 18 | |
| 19 | @dataclass |
| 20 | class PendingRecv: |
| 21 | """Per-step slot data for a deferred postprocess on the main stream.""" |
| 22 | |
| 23 | event: torch.cuda.Event |
| 24 | |
| 25 | sampled_tokens: torch.Tensor # [num_reqs, max_sample_len] |
| 26 | num_sampled: torch.Tensor # [num_reqs] |
| 27 | num_rejected: torch.Tensor # [num_reqs] |
| 28 | idx_mapping: torch.Tensor # [num_reqs] |
| 29 | idx_mapping_np: np.ndarray # [num_reqs] |
| 30 | # Records which rows need a deferred postprocess (bool). |
| 31 | need_sampled_mask: np.ndarray # [num_reqs] |
| 32 | # Snapshot of slot generation counters at receive time, used to |
| 33 | # detect requests aborted since then. |
| 34 | gen_at_receive_np: np.ndarray # [num_reqs] |
| 35 | draft_tokens: torch.Tensor | None = None # [num_reqs, num_speculative_steps] |
| 36 | |
| 37 | |
| 38 | def compute_need_sampled_mask(input_batch: InputBatch) -> np.ndarray | None: |
| 39 | """Return a bool array of shape `[input_batch.num_reqs]` marking requests |
| 40 | with outputs that might be needed in a subsequent (decode) step. |
| 41 | Returns None if no sampled outputs are needed in the requests' next step.""" |
| 42 | old_computed = input_batch.num_computed_tokens_np |
| 43 | prefill_len = input_batch.prefill_len_np |
| 44 | # Exclude non-final prefill chunks (they don't produce a sample). |
| 45 | need_sampled_mask = old_computed + input_batch.num_scheduled_tokens >= prefill_len |
| 46 | if input_batch.max_seq_len_np is not None: |
| 47 | # Also exclude final prefill chunks whose single sampled token reaches |
| 48 | # the request's length cap. |
| 49 | finished_prefill = prefill_len + 1 >= input_batch.max_seq_len_np |
| 50 | finished_prefill &= input_batch.is_prefilling_np |
| 51 | need_sampled_mask &= ~finished_prefill |
| 52 | return need_sampled_mask if need_sampled_mask.any() else None |
| 53 | |
| 54 | |
| 55 | class PPHandler: |
| 56 | """Runs the PP sampled-token broadcast/recv on a side stream so the |
| 57 | default stream isn't gated by the matching peer call. Step T's recv is |
| 58 | consumed at step T+pp_size via `get_prev_sampled_outputs`. |
| 59 | |
| 60 | Uses a dedicated NCCL communicator (sibling of the PP `device_group`) |
| 61 | for the broadcast so it does not serialize on the wire with the |
| 62 | inter-stage hidden-state p2p send/recv ops. |
| 63 | """ |
| 64 | |
| 65 | def __init__( |
| 66 | self, max_num_reqs: int, num_speculative_steps: int, device: torch.device |
| 67 | ): |
| 68 | self.is_last_rank = get_pp_group().is_last_rank |
| 69 | self.last_rank = get_pp_group().last_rank |
| 70 | self.max_sample_len = num_speculative_steps + 1 |
| 71 | self.num_speculative_steps = num_speculative_steps |
| 72 | self.device = device |
| 73 | self.main_stream = torch.cuda.current_stream(device) |
| 74 | self.broadcast_stream = torch.cuda.Stream(device) |
| 75 | |
| 76 | # On non-last ranks, a FIFO with one entry per in-flight step: the entry |
| 77 | # pushed by step T's `receive` is consumed pp_size steps later. Pre-seeded |
| 78 | # with pp_size None placeholders so the first pp_size consumes are no-ops. |
| 79 | # None means no postprocess is pending for that step (broadcast skipped). |
| 80 | # Only XPU can disable microbatching via VLLM_XPU_PP_MICROBATCH. |
| 81 | ring_depth = get_pp_group().world_size |
| 82 | if current_platform.is_xpu() and not envs.VLLM_XPU_PP_MICROBATCH: |
| 83 | ring_depth = 1 |
| 84 | self.queue: deque[PendingRecv | None] = ( |
| 85 | deque() if self.is_last_rank else deque([None] * ring_depth) |
| 86 | ) |
| 87 | |
| 88 | # Per req-index generation counter, incremented every time a request |
| 89 | # index is freed in RequestStats. Used for invalidating freed req data |
| 90 | # between PP decodes. |
| 91 | self.req_idx_gen_np = np.zeros(max_num_reqs, dtype=np.int32) |
| 92 | |
| 93 | # Dedicated subgroup for the sampled-token broadcast. |
| 94 | self.broadcast_group = get_pp_group().make_sibling_device_group( |
| 95 | group_desc="pp_broadcast" |
| 96 | ) |
| 97 | self.aux_hidden_state_relay_keys: tuple[str, ...] = () |
| 98 | |
| 99 | def on_req_idx_freed(self, req_idx: int) -> None: |
| 100 | self.req_idx_gen_np[req_idx] += 1 |
| 101 | |
| 102 | def configure_aux_hidden_state_relay(self, model: torch.nn.Module) -> None: |
| 103 | from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import ( |
| 104 | aux_hidden_state_relay_keys, |
| 105 | ) |
| 106 | |
| 107 | self.aux_hidden_state_relay_keys = aux_hidden_state_relay_keys(model) |
| 108 | |
| 109 | def relay_aux_hidden_states( |
| 110 | self, |
| 111 | intermediate_tensors: IntermediateTensors | None, |
| 112 | output_intermediate_tensors: IntermediateTensors, |
| 113 | ) -> IntermediateTensors: |
| 114 | if not self.aux_hidden_state_relay_keys: |
| 115 | return output_intermediate_tensors |
| 116 | assert intermediate_tensors is not None |
| 117 | return IntermediateTensors( |
| 118 | output_intermediate_tensors.tensors |
| 119 | | { |
| 120 | key: intermediate_tensors[key] |
| 121 | for key in self.aux_hidden_state_relay_keys |
| 122 | } |
| 123 | ) |
| 124 | |
| 125 | def get_prev_sampled_outputs( |
| 126 | self, draft_tokens_to_update: torch.Tensor | None = None |
| 127 | ) -> dict[str, torch.Tensor] | None: |
| 128 | """Consume the entry from pp_size steps ago and wait for its recv event, |
| 129 | then filter out entries whose request was freed since `receive`. |
| 130 | """ |
| 131 | if not self.queue: |
| 132 | return None |
| 133 | slot = self.queue.popleft() |
| 134 | # Reserve this step's slot; `receive` overwrites it if applicable. |
| 135 | self.queue.append(None) |
| 136 | if slot is None: |
| 137 | return None |
| 138 | |
| 139 | # Skip requests which did not need sampled output and/or those already |
| 140 | # finished. The post_update kernel skips the -1 entries. |
| 141 | freed = self.req_idx_gen_np[slot.idx_mapping_np] != slot.gen_at_receive_np |
| 142 | exclude_mask = freed | ~slot.need_sampled_mask |
| 143 | idx_mapping = slot.idx_mapping |
| 144 | if exclude_mask.any(): |
| 145 | if exclude_mask.all(): |
| 146 | # No states require update anymore. |
| 147 | return None |
| 148 | # Filter excluded request indices. |
| 149 | idx_mapping_np = np.where(exclude_mask, -1, slot.idx_mapping_np) |
| 150 | idx_mapping = async_tensor_h2d(idx_mapping_np, device=self.device) |
| 151 | |
| 152 | self.main_stream.wait_event(slot.event) |
| 153 | if slot.draft_tokens is not None and draft_tokens_to_update is not None: |
| 154 | draft_tokens = slot.draft_tokens |
| 155 | draft_idx_mapping = slot.idx_mapping |
| 156 | if exclude_mask.any(): |
| 157 | keep = ~exclude_mask |
| 158 | keep_t = torch.as_tensor(keep, device=self.device) |
| 159 | draft_tokens = draft_tokens[keep_t] |
| 160 | draft_idx_mapping = async_tensor_h2d( |
| 161 | slot.idx_mapping_np[keep], device=self.device |
| 162 | ) |
| 163 | draft_tokens_to_update[draft_idx_mapping] = draft_tokens |
| 164 | |
| 165 | return dict( |
| 166 | sampled_tokens=slot.sampled_tokens, |
| 167 | num_sampled=slot.num_sampled, |
| 168 | num_rejected=slot.num_rejected, |
| 169 | idx_mapping=idx_mapping, |
| 170 | ) |
| 171 | |
| 172 | def broadcast_drafts( |
| 173 | self, draft_tokens: torch.Tensor, input_batch: InputBatch |
| 174 | ) -> None: |
| 175 | """Broadcast draft proposals so non-last ranks can embed real token ids.""" |
| 176 | assert self.is_last_rank |
| 177 | if compute_need_sampled_mask(input_batch) is None: |
| 178 | return |
| 179 | with torch.cuda.stream(self.broadcast_stream): |
| 180 | self.broadcast_stream.wait_stream(self.main_stream) |
| 181 | send = draft_tokens[input_batch.idx_mapping].contiguous() |
| 182 | # Must record the idx_mapping tensor since it was allocated |
| 183 | # on the main stream. |
| 184 | input_batch.idx_mapping.record_stream(self.broadcast_stream) |
| 185 | torch.distributed.broadcast( |
| 186 | send, src=self.last_rank, group=self.broadcast_group |
| 187 | ) |
| 188 | |
| 189 | def receive(self, input_batch: InputBatch) -> bool: |
| 190 | """Returns True iff sampled tokens need to be gathered from *all* |
| 191 | requests in the batch.""" |
| 192 | assert not self.is_last_rank |
| 193 | need_sampled_mask = compute_need_sampled_mask(input_batch) |
| 194 | if need_sampled_mask is None: |
| 195 | # Leave this step's reserved slot as None. |
| 196 | return False |
| 197 | |
| 198 | # Snapshot the per-slot generation counter so a later free of any of |
| 199 | # these RequestStates request indices is detectable at consume time. |
| 200 | gen_at_receive_np = self.req_idx_gen_np[input_batch.idx_mapping_np] |
| 201 | |
| 202 | num_reqs = input_batch.num_reqs |
| 203 | with torch.cuda.stream(self.broadcast_stream): |
| 204 | self.broadcast_stream.wait_stream(self.main_stream) |
| 205 | sampled_tokens = torch.empty( |
| 206 | num_reqs, self.max_sample_len, dtype=torch.int64, device=self.device |
| 207 | ) |
| 208 | combined = torch.empty(2, num_reqs, dtype=torch.int32, device=self.device) |
| 209 | torch.distributed.broadcast( |
| 210 | sampled_tokens, src=self.last_rank, group=self.broadcast_group |
| 211 | ) |
| 212 | torch.distributed.broadcast( |
| 213 | combined, src=self.last_rank, group=self.broadcast_group |
| 214 | ) |
| 215 | draft_tokens = None |
| 216 | if self.num_speculative_steps > 0: |
| 217 | draft_tokens = torch.empty( |
| 218 | num_reqs, |
| 219 | self.num_speculative_steps, |
| 220 | dtype=torch.int64, |
| 221 | device=self.device, |
| 222 | ) |
| 223 | torch.distributed.broadcast( |
| 224 | draft_tokens, src=self.last_rank, group=self.broadcast_group |
| 225 | ) |
| 226 | event = self.broadcast_stream.record_event() |
| 227 | num_sampled, num_rejected = combined.unbind(dim=0) |
| 228 | # Must record_stream since these were allocated on broadcast stream but |
| 229 | # later used on the main stream. |
| 230 | sampled_tokens.record_stream(self.main_stream) |
| 231 | combined.record_stream(self.main_stream) |
| 232 | if draft_tokens is not None: |
| 233 | draft_tokens.record_stream(self.main_stream) |
| 234 | self.queue[-1] = PendingRecv( |
| 235 | event, |
| 236 | sampled_tokens, |
| 237 | num_sampled, |
| 238 | num_rejected, |
| 239 | input_batch.idx_mapping, |
| 240 | input_batch.idx_mapping_np, |
| 241 | need_sampled_mask, |
| 242 | gen_at_receive_np, |
| 243 | draft_tokens, |
| 244 | ) |
| 245 | return bool(need_sampled_mask.all()) |
| 246 | |
| 247 | def broadcast( |
| 248 | self, |
| 249 | sampled_token_ids: torch.Tensor, |
| 250 | num_sampled: torch.Tensor, |
| 251 | num_rejected: torch.Tensor, |
| 252 | input_batch: InputBatch, |
| 253 | ) -> None: |
| 254 | assert self.is_last_rank |
| 255 | if compute_need_sampled_mask(input_batch) is None: |
| 256 | # No request needs sampled outputs for a subsequent decode step. |
| 257 | return |
| 258 | |
| 259 | assert sampled_token_ids.dtype == torch.int64 |
| 260 | |
| 261 | if current_platform.is_xpu(): |
| 262 | self.main_stream.synchronize() |
| 263 | |
| 264 | with torch.cuda.stream(self.broadcast_stream): |
| 265 | self.broadcast_stream.wait_stream(self.main_stream) |
| 266 | send_tokens = torch.nn.functional.pad( |
| 267 | sampled_token_ids, |
| 268 | (0, self.max_sample_len - sampled_token_ids.shape[-1]), |
| 269 | ) |
| 270 | torch.distributed.broadcast( |
| 271 | send_tokens.contiguous(), |
| 272 | src=self.last_rank, |
| 273 | group=self.broadcast_group, |
| 274 | ) |
| 275 | combined = torch.stack((num_sampled, num_rejected), dim=0) |
| 276 | torch.distributed.broadcast( |
| 277 | combined, src=self.last_rank, group=self.broadcast_group |
| 278 | ) |
| 279 | for tensor in (sampled_token_ids, num_sampled, num_rejected): |
| 280 | tensor.record_stream(self.broadcast_stream) |