| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import contextlib |
| 4 | from collections.abc import Iterator |
| 5 | from dataclasses import dataclass, field |
| 6 | from typing import TYPE_CHECKING |
| 7 | |
| 8 | import numpy as np |
| 9 | import torch |
| 10 | |
| 11 | import vllm.envs as envs |
| 12 | from vllm.model_executor.layers.fused_moe.all2all_utils import get_ep_all2all_manager |
| 13 | from vllm.v1.outputs import ( |
| 14 | AsyncModelRunnerOutput, |
| 15 | LogprobsTensors, |
| 16 | ModelRunnerOutput, |
| 17 | PoolerOutput, |
| 18 | ) |
| 19 | from vllm.v1.worker.gpu.sample.output import SamplerOutput, SamplingMaskTensors |
| 20 | from vllm.v1.worker.utils import raise_if_nan_logits |
| 21 | |
| 22 | if TYPE_CHECKING: |
| 23 | from vllm.distributed.aux_output_connector.worker import PendingAuxOutput |
| 24 | from vllm.v1.worker.gpu.input_batch import InputBatch |
| 25 | |
| 26 | |
| 27 | @dataclass(frozen=True) |
| 28 | class StepTimingSample: |
| 29 | forward_ms: float |
| 30 | drafter_ms: float |
| 31 | num_target_tokens: int |
| 32 | num_reqs: int |
| 33 | full_cudagraph: bool |
| 34 | |
| 35 | |
| 36 | def _timing_event() -> torch.cuda.Event: |
| 37 | return torch.cuda.Event(enable_timing=True) |
| 38 | |
| 39 | |
| 40 | @dataclass |
| 41 | class StepTimingEvents: |
| 42 | forward_start: torch.cuda.Event = field(default_factory=_timing_event) |
| 43 | forward_end: torch.cuda.Event = field(default_factory=_timing_event) |
| 44 | drafter_start: torch.cuda.Event = field(default_factory=_timing_event) |
| 45 | drafter_end: torch.cuda.Event = field(default_factory=_timing_event) |
| 46 | |
| 47 | |
| 48 | class StepTimingCollector: |
| 49 | """Times the steps run inside ``collect``; record calls no-op outside it. |
| 50 | |
| 51 | Every step gets its own events, so the steps queue back-to-back and the |
| 52 | block resolves them all behind one sync on the way out. |
| 53 | """ |
| 54 | |
| 55 | def __init__(self): |
| 56 | self._collecting = False |
| 57 | self._step: StepTimingEvents | None = None |
| 58 | self._batch = (False, 0, 0) |
| 59 | self._timed: list[tuple[StepTimingEvents, tuple[bool, int, int]]] = [] |
| 60 | |
| 61 | @contextlib.contextmanager |
| 62 | def collect(self) -> Iterator[list[StepTimingSample]]: |
| 63 | """Time every step run in this block. |
| 64 | |
| 65 | The yielded list holds one sample per timed step once the block exits; |
| 66 | it stays empty inside the block, where the timings are still on device. |
| 67 | """ |
| 68 | samples: list[StepTimingSample] = [] |
| 69 | self._collecting = True |
| 70 | try: |
| 71 | yield samples |
| 72 | finally: |
| 73 | self._collecting = False |
| 74 | timed, self._timed, self._step = self._timed, [], None |
| 75 | if not timed: |
| 76 | return |
| 77 | # Same stream, issue order: once the last step is done, so are the rest. |
| 78 | timed[-1][0].drafter_end.synchronize() |
| 79 | samples.extend( |
| 80 | StepTimingSample( |
| 81 | events.forward_start.elapsed_time(events.forward_end), |
| 82 | events.drafter_start.elapsed_time(events.drafter_end), |
| 83 | num_target_tokens, |
| 84 | num_reqs, |
| 85 | full_cudagraph, |
| 86 | ) |
| 87 | for events, (full_cudagraph, num_target_tokens, num_reqs) in timed |
| 88 | ) |
| 89 | |
| 90 | def record_batch(self, input_batch: "InputBatch", full_cudagraph: bool) -> None: |
| 91 | """Costs from different execution modes must not share a cost curve.""" |
| 92 | self._batch = (full_cudagraph, input_batch.num_tokens, input_batch.num_reqs) |
| 93 | |
| 94 | def forward_start(self) -> None: |
| 95 | if self._collecting: |
| 96 | self._step = StepTimingEvents() |
| 97 | self._step.forward_start.record() |
| 98 | |
| 99 | def forward_end(self) -> None: |
| 100 | if self._step is not None: |
| 101 | self._step.forward_end.record() |
| 102 | |
| 103 | def drafter_start(self) -> None: |
| 104 | if self._step is not None: |
| 105 | self._step.drafter_start.record() |
| 106 | |
| 107 | def drafter_end(self) -> None: |
| 108 | """Ends the step: only steps that reach here have a draft cost.""" |
| 109 | if self._step is not None: |
| 110 | self._step.drafter_end.record() |
| 111 | self._timed.append((self._step, self._batch)) |
| 112 | self._step = None |
| 113 | |
| 114 | |
| 115 | class AsyncOutput(AsyncModelRunnerOutput): |
| 116 | def __init__( |
| 117 | self, |
| 118 | model_runner_output: ModelRunnerOutput, |
| 119 | sampler_output: SamplerOutput, |
| 120 | num_sampled_tokens: torch.Tensor, |
| 121 | main_stream: torch.cuda.Stream, |
| 122 | copy_stream: torch.cuda.Stream, |
| 123 | check_ep_fault: bool, |
| 124 | pending_aux_output: "PendingAuxOutput | None", |
| 125 | ): |
| 126 | # NOTE(woosuk): We must retain references to the GPU tensors, |
| 127 | # as the copy operations are performed on a different CUDA stream than |
| 128 | # the one where the tensors were created. |
| 129 | self.model_runner_output = model_runner_output |
| 130 | self.sampler_output = sampler_output |
| 131 | self.num_sampled_tokens = num_sampled_tokens |
| 132 | self.pending_aux_output = pending_aux_output |
| 133 | # Blocking (sleep) event to avoid busy-polling the CUDA driver lock. |
| 134 | self.copy_event = torch.cuda.Event(blocking=True) |
| 135 | self._has_fault: torch.Tensor | None = None |
| 136 | |
| 137 | with stream(copy_stream, main_stream): |
| 138 | copy_stream.wait_stream(main_stream) |
| 139 | |
| 140 | self.sampled_token_ids = async_copy_to_np(sampler_output.sampled_token_ids) |
| 141 | self.logprobs_tensors: LogprobsTensors | None = None |
| 142 | if sampler_output.logprobs_tensors is not None: |
| 143 | self.logprobs_tensors = ( |
| 144 | sampler_output.logprobs_tensors.to_cpu_nonblocking() |
| 145 | ) |
| 146 | self.num_nans: np.ndarray | None = None |
| 147 | if sampler_output.num_nans is not None: |
| 148 | self.num_nans = async_copy_to_np(sampler_output.num_nans) |
| 149 | self.num_sampled_tokens_np = async_copy_to_np(num_sampled_tokens) |
| 150 | self.sampling_mask_tensors: SamplingMaskTensors | None = None |
| 151 | if sampler_output.sampling_mask_tensors is not None: |
| 152 | self.sampling_mask_tensors = ( |
| 153 | sampler_output.sampling_mask_tensors.to_cpu_nonblocking() |
| 154 | ) |
| 155 | self.prompt_logprobs_dict = { |
| 156 | k: v.to_cpu_nonblocking() if v is not None else None |
| 157 | for k, v in self.model_runner_output.prompt_logprobs_dict.items() |
| 158 | } |
| 159 | token_id_logprobs = self.model_runner_output.prompt_token_id_logprobs_dict |
| 160 | self.prompt_token_id_logprobs_dict = { |
| 161 | k: v.to("cpu", non_blocking=True) for k, v in token_id_logprobs.items() |
| 162 | } |
| 163 | if self.pending_aux_output is not None: |
| 164 | self.pending_aux_output.enqueue_cpu_copy( |
| 165 | num_sampled=self.num_sampled_tokens_np, |
| 166 | num_rejected=async_copy_to_np(sampler_output.num_rejected), |
| 167 | ) |
| 168 | if check_ep_fault: |
| 169 | has_fault = get_ep_all2all_manager().query_fault() |
| 170 | self._has_fault = has_fault.to("cpu", non_blocking=True) |
| 171 | self.copy_event.record(copy_stream) |
| 172 | |
| 173 | def get_output(self) -> ModelRunnerOutput: |
| 174 | self.copy_event.synchronize() |
| 175 | |
| 176 | # NOTE(woosuk): The following code is to ensure compatibility with |
| 177 | # the existing model runner. |
| 178 | # Going forward, we should keep the data structures as NumPy arrays |
| 179 | # rather than Python lists. |
| 180 | sampled_token_ids: list[list[int]] = self.sampled_token_ids.tolist() |
| 181 | num_sampled_tokens: list[int] = self.num_sampled_tokens_np.tolist() |
| 182 | for token_ids, num_tokens in zip(sampled_token_ids, num_sampled_tokens): |
| 183 | del token_ids[num_tokens:] |
| 184 | self.model_runner_output.sampled_token_ids = sampled_token_ids |
| 185 | |
| 186 | if self.sampling_mask_tensors is not None: |
| 187 | self.model_runner_output.sampling_masks = ( |
| 188 | self.sampling_mask_tensors.tolists(self.num_sampled_tokens_np) |
| 189 | ) |
| 190 | |
| 191 | if self.num_nans is not None: |
| 192 | self.model_runner_output.num_nans_in_logits = dict( |
| 193 | zip(self.model_runner_output.req_ids, self.num_nans.tolist()) |
| 194 | ) |
| 195 | if envs.VLLM_RAISE_ON_LOGIT_NANS: |
| 196 | raise_if_nan_logits(self.model_runner_output.num_nans_in_logits) |
| 197 | |
| 198 | if self.logprobs_tensors is not None: |
| 199 | self.model_runner_output.logprobs = self.logprobs_tensors.tolists() |
| 200 | self.model_runner_output.prompt_logprobs_dict = self.prompt_logprobs_dict |
| 201 | self.model_runner_output.prompt_token_id_logprobs_dict = ( |
| 202 | self.prompt_token_id_logprobs_dict |
| 203 | ) |
| 204 | if self.pending_aux_output is not None: |
| 205 | self.model_runner_output.aux_output_connector_output = ( |
| 206 | self.pending_aux_output.process_output() |
| 207 | ) |
| 208 | |
| 209 | if self._has_fault is not None and self._has_fault.item(): |
| 210 | mask = get_ep_all2all_manager().query_active_mask() |
| 211 | raise RuntimeError( |
| 212 | "Fault detected in EP all2all communication: " |
| 213 | "one or more ranks timed out during dispatch/combine. " |
| 214 | f"Mask: {mask.cpu().tolist()}" |
| 215 | ) |
| 216 | |
| 217 | return self.model_runner_output |
| 218 | |
| 219 | |
| 220 | class AsyncPoolingOutput(AsyncModelRunnerOutput): |
| 221 | def __init__( |
| 222 | self, |
| 223 | model_runner_output: ModelRunnerOutput, |
| 224 | pooler_output: PoolerOutput, |
| 225 | finished_mask: list[bool], |
| 226 | main_stream: torch.cuda.Stream, |
| 227 | copy_stream: torch.cuda.Stream, |
| 228 | ): |
| 229 | self.model_runner_output = model_runner_output |
| 230 | self.pooler_output = pooler_output |
| 231 | # Blocking (sleep) event to avoid busy-polling the CUDA driver lock. |
| 232 | self.copy_event = torch.cuda.Event(blocking=True) |
| 233 | |
| 234 | with stream(copy_stream, main_stream): |
| 235 | copy_stream.wait_stream(main_stream) |
| 236 | if isinstance(self.pooler_output, torch.Tensor) and all(finished_mask): |
| 237 | self.pooler_output_cpu: PoolerOutput = self.pooler_output.to( |
| 238 | "cpu", non_blocking=True |
| 239 | ) |
| 240 | else: |
| 241 | outputs = ( |
| 242 | self.pooler_output.unbind() |
| 243 | if isinstance(self.pooler_output, torch.Tensor) |
| 244 | else self.pooler_output |
| 245 | ) |
| 246 | self.pooler_output_cpu = [ |
| 247 | None |
| 248 | if output is None or not is_finished |
| 249 | else output.to("cpu", non_blocking=True) |
| 250 | for output, is_finished in zip(outputs, finished_mask, strict=True) |
| 251 | ] |
| 252 | self.copy_event.record(copy_stream) |
| 253 | |
| 254 | def get_output(self) -> ModelRunnerOutput: |
| 255 | if isinstance(self.pooler_output_cpu, torch.Tensor): |
| 256 | pooler_output = list(self.pooler_output_cpu.unbind(dim=0)) |
| 257 | else: |
| 258 | pooler_output = self.pooler_output_cpu |
| 259 | self.copy_event.synchronize() |
| 260 | self.model_runner_output.pooler_output = pooler_output |
| 261 | return self.model_runner_output |
| 262 | |
| 263 | |
| 264 | def async_copy_to_np(x: torch.Tensor) -> np.ndarray: |
| 265 | return x.to("cpu", non_blocking=True).numpy() |
| 266 | |
| 267 | |
| 268 | @contextlib.contextmanager |
| 269 | def stream(to_stream: torch.cuda.Stream, from_stream: torch.cuda.Stream): |
| 270 | """Lightweight version of torch.cuda.stream() context manager which |
| 271 | avoids current_stream and device lookups. |
| 272 | """ |
| 273 | try: |
| 274 | torch.cuda.set_stream(to_stream) |
| 275 | yield |
| 276 | finally: |
| 277 | torch.cuda.set_stream(from_stream) |