| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Worker-side auxiliary-output data plane.""" |
| 4 | |
| 5 | from __future__ import annotations |
| 6 | |
| 7 | from collections.abc import Iterable, Sequence |
| 8 | from dataclasses import dataclass, field |
| 9 | from threading import Lock |
| 10 | from typing import TYPE_CHECKING, Any |
| 11 | |
| 12 | import numpy as np |
| 13 | import torch |
| 14 | |
| 15 | from vllm.config import VllmConfig |
| 16 | from vllm.distributed.aux_output_connector.connector import ( |
| 17 | AuxOutputConnectorMetadata, |
| 18 | AuxRequestOutput, |
| 19 | ) |
| 20 | from vllm.distributed.aux_output_connector.routed_experts import ( |
| 21 | RoutedExpertsBuffer, |
| 22 | materialize_routed_experts, |
| 23 | publish_routed_experts, |
| 24 | routed_experts_keys, |
| 25 | ) |
| 26 | from vllm.distributed.aux_output_connector.store import ( |
| 27 | BackgroundBlockObjectStore, |
| 28 | BlockObjectStore, |
| 29 | ) |
| 30 | from vllm.distributed.parallel_state import get_tp_group |
| 31 | from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( |
| 32 | RoutedExpertsCapturer, |
| 33 | bind_routed_experts_capturer, |
| 34 | ) |
| 35 | from vllm.v1.core.kv_cache_utils import resolve_kv_cache_block_sizes |
| 36 | |
| 37 | if TYPE_CHECKING: |
| 38 | from vllm.v1.kv_cache_interface import KVCacheConfig |
| 39 | from vllm.v1.worker.gpu.input_batch import InputBatch |
| 40 | |
| 41 | |
| 42 | @dataclass |
| 43 | class _WorkerRequestState: |
| 44 | aux_output_keys: list[str] = field(default_factory=list) |
| 45 | pending_blocks: list[tuple[int, np.ndarray]] = field(default_factory=list) |
| 46 | capture_cursor: int | None = None |
| 47 | scheduled_cursor: int = 0 |
| 48 | emit_cursor: int = 0 |
| 49 | pending_outputs: int = 0 |
| 50 | # Terminal event received; teardown waits for in-flight step outputs. |
| 51 | finished: bool = False |
| 52 | |
| 53 | |
| 54 | @dataclass |
| 55 | class PendingAuxOutput: |
| 56 | """Own one step's R3 snapshot until its asynchronous copy is consumed.""" |
| 57 | |
| 58 | connector: AuxOutputWorkerConnector |
| 59 | request_ids: list[str] |
| 60 | token_starts: np.ndarray |
| 61 | query_start_loc: np.ndarray |
| 62 | # GPU snapshot taken on the main stream. |
| 63 | routed_experts_gpu: torch.Tensor |
| 64 | # Filled by enqueue_cpu_copy on the output copy stream. |
| 65 | routed_experts: np.ndarray | None = None |
| 66 | num_sampled: np.ndarray | None = None |
| 67 | num_rejected: np.ndarray | None = None |
| 68 | |
| 69 | def enqueue_cpu_copy( |
| 70 | self, num_sampled: np.ndarray, num_rejected: np.ndarray |
| 71 | ) -> None: |
| 72 | """Enqueue asynchronous D2H copies; call on the output copy stream.""" |
| 73 | self.num_sampled = num_sampled |
| 74 | self.num_rejected = num_rejected |
| 75 | self.routed_experts = self.routed_experts_gpu.to( |
| 76 | "cpu", non_blocking=True |
| 77 | ).numpy() |
| 78 | |
| 79 | def process_output(self) -> dict[str, AuxRequestOutput]: |
| 80 | return self.connector.process_output(self) |
| 81 | |
| 82 | |
| 83 | class AuxOutputWorkerConnector: |
| 84 | """Own capture, request tails, and backend resources on the output worker.""" |
| 85 | |
| 86 | def __init__( |
| 87 | self, |
| 88 | *, |
| 89 | vllm_config: VllmConfig, |
| 90 | model: torch.nn.Module, |
| 91 | kv_cache_config: KVCacheConfig, |
| 92 | ) -> None: |
| 93 | max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens |
| 94 | capturer = RoutedExpertsCapturer( |
| 95 | max_num_batched_tokens=max_num_batched_tokens, vllm_config=vllm_config |
| 96 | ) |
| 97 | bind_routed_experts_capturer(model, capturer) |
| 98 | self._capturer = capturer |
| 99 | self._store: BackgroundBlockObjectStore | None = None |
| 100 | self._buffer: RoutedExpertsBuffer | None = None |
| 101 | self._requests: dict[str, _WorkerRequestState] = {} |
| 102 | self._generation = 0 |
| 103 | self._step_metadata: AuxOutputConnectorMetadata | None = None |
| 104 | # Steps whose asynchronous copy has not been consumed yet. The engine |
| 105 | # consumes step outputs in order, so at most max_concurrent_batches |
| 106 | # are outstanding; guarded by _lock against the output thread. |
| 107 | self._pending_outputs: list[PendingAuxOutput] = [] |
| 108 | self._lock = Lock() |
| 109 | self._max_concurrent_batches = vllm_config.max_concurrent_batches |
| 110 | # Every TP rank participates in capture collectives, but only the |
| 111 | # executor output rank owns the auxiliary output data plane. |
| 112 | if not get_tp_group().is_first_rank: |
| 113 | return |
| 114 | |
| 115 | shape_per_token = self._capturer.shape_per_token |
| 116 | dtype: np.dtype[Any] = np.dtype(self._capturer.output_dtype_name) |
| 117 | scheduler_block_size, hash_block_size = resolve_kv_cache_block_sizes( |
| 118 | kv_cache_config, vllm_config |
| 119 | ) |
| 120 | hashes_per_kv_block = scheduler_block_size // hash_block_size |
| 121 | block_nbytes = hash_block_size * int(np.prod(shape_per_token)) * dtype.itemsize |
| 122 | max_bytes = vllm_config.aux_output_config.max_bytes |
| 123 | if max_bytes is None: |
| 124 | max_bytes = kv_cache_config.num_blocks * hashes_per_kv_block * block_nbytes |
| 125 | self._store = BackgroundBlockObjectStore( |
| 126 | BlockObjectStore(max_bytes=max_bytes, object_nbytes=block_nbytes), |
| 127 | max_pending_batches=2 * vllm_config.scheduler_config.max_num_seqs, |
| 128 | ) |
| 129 | self._buffer = RoutedExpertsBuffer( |
| 130 | dtype, |
| 131 | shape_per_token, |
| 132 | hash_block_size, |
| 133 | vllm_config.scheduler_config.max_num_seqs, |
| 134 | max_num_batched_tokens, |
| 135 | vllm_config.max_concurrent_batches, |
| 136 | ) |
| 137 | |
| 138 | def prepare_output(self, input_batch: InputBatch) -> PendingAuxOutput | None: |
| 139 | """Snapshot one step's R3 tensor for asynchronous CPU transfer.""" |
| 140 | if self._buffer is None or self._step_metadata is None: |
| 141 | return None |
| 142 | |
| 143 | request_ids = list(input_batch.req_ids) |
| 144 | query_start_loc = input_batch.query_start_loc_np[: len(request_ids) + 1] |
| 145 | num_rows = int(query_start_loc[-1]) |
| 146 | pending_output = PendingAuxOutput( |
| 147 | connector=self, |
| 148 | request_ids=request_ids, |
| 149 | token_starts=input_batch.num_computed_tokens_np, |
| 150 | query_start_loc=query_start_loc, |
| 151 | routed_experts_gpu=self._capturer.snapshot_routing_data(num_rows), |
| 152 | ) |
| 153 | with self._lock: |
| 154 | assert len(self._pending_outputs) < self._max_concurrent_batches, ( |
| 155 | "auxiliary output step outputs are not consumed in order" |
| 156 | ) |
| 157 | self._pending_outputs.append(pending_output) |
| 158 | for request_id in pending_output.request_ids: |
| 159 | self._requests[request_id].pending_outputs += 1 |
| 160 | return pending_output |
| 161 | |
| 162 | def process_output(self, pending: PendingAuxOutput) -> dict[str, AuxRequestOutput]: |
| 163 | """Commit one consumed R3 snapshot and build request outputs. |
| 164 | |
| 165 | Request teardown deferred by begin_step is completed here once no |
| 166 | in-flight step covers the request. |
| 167 | """ |
| 168 | with self._lock: |
| 169 | teardown = [] |
| 170 | release_keys: list[str] = [] |
| 171 | try: |
| 172 | outputs = self._commit_output(pending) |
| 173 | finally: |
| 174 | self._pending_outputs.remove(pending) |
| 175 | for request_id in pending.request_ids: |
| 176 | state = self._requests[request_id] |
| 177 | state.pending_outputs -= 1 |
| 178 | if state.pending_outputs == 0 and state.finished: |
| 179 | teardown.append(request_id) |
| 180 | release_keys.extend(reversed(state.aux_output_keys)) |
| 181 | if teardown: |
| 182 | self._publish_blocks([], release_keys=release_keys) |
| 183 | self._teardown(teardown) |
| 184 | return outputs |
| 185 | |
| 186 | def _commit_output(self, pending: PendingAuxOutput) -> dict[str, AuxRequestOutput]: |
| 187 | buffer = self._buffer |
| 188 | store = self._store |
| 189 | assert buffer is not None and store is not None |
| 190 | block_size = buffer.block_size |
| 191 | |
| 192 | routed_experts = pending.routed_experts |
| 193 | num_sampled = pending.num_sampled |
| 194 | num_rejected = pending.num_rejected |
| 195 | assert ( |
| 196 | routed_experts is not None |
| 197 | and num_sampled is not None |
| 198 | and num_rejected is not None |
| 199 | ), "auxiliary output CPU copy was not enqueued" |
| 200 | |
| 201 | # Publish the whole batch before materializing any consumer output. |
| 202 | materialize_outputs: list[tuple[str, int, int]] = [] |
| 203 | block_batches = [] |
| 204 | outputs: dict[str, AuxRequestOutput] = {} |
| 205 | |
| 206 | # Use the ModelRunner's actual batch boundaries rather than rebuilding them. |
| 207 | for request_id, token_start, start, end, sampled, rejected in zip( |
| 208 | pending.request_ids, |
| 209 | pending.token_starts, |
| 210 | pending.query_start_loc[:-1], |
| 211 | pending.query_start_loc[1:], |
| 212 | num_sampled, |
| 213 | num_rejected, |
| 214 | strict=True, |
| 215 | ): |
| 216 | request_num_tokens = end - start |
| 217 | assert request_num_tokens > 0, ( |
| 218 | "auxiliary output request token count must be positive" |
| 219 | ) |
| 220 | state = self._requests[request_id] |
| 221 | |
| 222 | # Capture precedes speculative acceptance, so discard the rejected |
| 223 | # suffix. Batch boundaries still span the full executed range. |
| 224 | rejected = int(rejected) |
| 225 | assert 0 <= rejected <= request_num_tokens, ( |
| 226 | "auxiliary output rejected-token count is invalid" |
| 227 | ) |
| 228 | rows = routed_experts[start : end - rejected] |
| 229 | |
| 230 | capture_start = token_start |
| 231 | capture_cursor = state.capture_cursor |
| 232 | if capture_cursor is None: |
| 233 | capture_cursor = capture_start |
| 234 | |
| 235 | assert capture_start >= capture_cursor, ( |
| 236 | "auxiliary output capture moved backwards" |
| 237 | ) |
| 238 | if capture_start > capture_cursor: |
| 239 | # Reattach after an optimistically scheduled suffix was rejected. |
| 240 | assert capture_cursor < state.scheduled_cursor, ( |
| 241 | "auxiliary output capture has an unbacked token gap" |
| 242 | ) |
| 243 | capture_start = capture_cursor |
| 244 | |
| 245 | emit_start = state.emit_cursor |
| 246 | # Complete blocks without keys remain pending until a hash update. |
| 247 | completed = buffer.capture(request_id, capture_start, rows) |
| 248 | token_end = capture_start + len(rows) |
| 249 | state.capture_cursor = token_end |
| 250 | state.scheduled_cursor = token_start + request_num_tokens |
| 251 | block_batches.append((state, completed)) |
| 252 | |
| 253 | if sampled > 0 and emit_start <= token_end: |
| 254 | if emit_start >= capture_start: |
| 255 | outputs[request_id] = AuxRequestOutput( |
| 256 | emit_start, rows[emit_start - capture_start :] |
| 257 | ) |
| 258 | state.emit_cursor = token_end |
| 259 | else: |
| 260 | materialize_outputs.append((request_id, emit_start, token_end)) |
| 261 | |
| 262 | # A consumer may reuse a block produced earlier in the same batch. |
| 263 | self._publish_blocks(block_batches) |
| 264 | |
| 265 | for request_id, emit_start, token_end in materialize_outputs: |
| 266 | state = self._requests[request_id] |
| 267 | stored_end = ( |
| 268 | min(token_end // block_size, len(state.aux_output_keys)) * block_size |
| 269 | ) |
| 270 | if emit_start < stored_end: |
| 271 | first_block = emit_start // block_size |
| 272 | stored = materialize_routed_experts( |
| 273 | store, |
| 274 | state.aux_output_keys[first_block : stored_end // block_size], |
| 275 | shape_per_token=buffer.shape_per_token, |
| 276 | dtype=buffer.dtype, |
| 277 | ) |
| 278 | local_start = emit_start % block_size |
| 279 | rows = stored[local_start : local_start + stored_end - emit_start] |
| 280 | if stored_end < token_end: |
| 281 | rows = np.concatenate( |
| 282 | (rows, buffer.read(request_id, stored_end, token_end)) |
| 283 | ) |
| 284 | else: |
| 285 | rows = buffer.read(request_id, emit_start, token_end) |
| 286 | outputs[request_id] = AuxRequestOutput(emit_start, rows) |
| 287 | state.emit_cursor = token_end |
| 288 | return outputs |
| 289 | |
| 290 | def _publish_blocks( |
| 291 | self, |
| 292 | batches: list[tuple[_WorkerRequestState, list[tuple[int, np.ndarray]]]], |
| 293 | retain_keys: Sequence[str] = (), |
| 294 | release_keys: Sequence[str] = (), |
| 295 | ) -> None: |
| 296 | store = self._store |
| 297 | buffer = self._buffer |
| 298 | assert store is not None and buffer is not None |
| 299 | ready_batches = [] |
| 300 | for state, completed in batches: |
| 301 | blocks = state.pending_blocks + completed |
| 302 | keyed_end = len(state.aux_output_keys) * buffer.block_size |
| 303 | ready = [(start, rows) for start, rows in blocks if start < keyed_end] |
| 304 | state.pending_blocks = [ |
| 305 | (start, buffer.retain_block(rows)) |
| 306 | for start, rows in blocks |
| 307 | if start >= keyed_end |
| 308 | ] |
| 309 | if ready: |
| 310 | ready_batches.append((state.aux_output_keys, ready)) |
| 311 | if ready_batches or retain_keys or release_keys: |
| 312 | publish_routed_experts( |
| 313 | store, |
| 314 | batches=ready_batches, |
| 315 | block_size=buffer.block_size, |
| 316 | retain_keys=retain_keys, |
| 317 | release_keys=release_keys, |
| 318 | ) |
| 319 | for _, blocks in ready_batches: |
| 320 | for _, rows in blocks: |
| 321 | buffer.release_block(rows) |
| 322 | |
| 323 | def _teardown(self, request_ids: Iterable[str]) -> None: |
| 324 | """Drop finished requests' temporary state after publishing key releases.""" |
| 325 | buffer = self._buffer |
| 326 | assert buffer is not None |
| 327 | for request_id in request_ids: |
| 328 | state = self._requests.pop(request_id) |
| 329 | for _, rows in state.pending_blocks: |
| 330 | buffer.release_block(rows) |
| 331 | buffer.discard(request_id) |
| 332 | |
| 333 | def begin_step(self, metadata: AuxOutputConnectorMetadata | None) -> None: |
| 334 | """Apply one scheduler step's request and block-hash updates. |
| 335 | |
| 336 | While an unconsumed step covers a finished request, its teardown is |
| 337 | deferred to that step's process_output. |
| 338 | """ |
| 339 | self._step_metadata = metadata |
| 340 | if self._buffer is None or metadata is None: |
| 341 | return |
| 342 | assert not metadata.requests.keys() & metadata.finished_requests, ( |
| 343 | "auxiliary output request cannot run and finish in one step" |
| 344 | ) |
| 345 | assert metadata.generation >= self._generation, ( |
| 346 | "auxiliary output metadata generation moved backwards" |
| 347 | ) |
| 348 | with self._lock: |
| 349 | release_keys: list[str] = [] |
| 350 | if metadata.generation > self._generation: |
| 351 | # A generation change follows a prefix-cache reset, which the |
| 352 | # scheduler only performs once all model output is consumed. |
| 353 | assert not self._pending_outputs, ( |
| 354 | "auxiliary output generation changed with output in flight" |
| 355 | ) |
| 356 | release_keys.extend( |
| 357 | key |
| 358 | for state in self._requests.values() |
| 359 | for key in reversed(state.aux_output_keys) |
| 360 | ) |
| 361 | self._buffer.reset() |
| 362 | self._requests.clear() |
| 363 | self._generation = metadata.generation |
| 364 | for request_id, emit_start in metadata.requests.items(): |
| 365 | state = self._requests.setdefault( |
| 366 | request_id, _WorkerRequestState(emit_cursor=emit_start) |
| 367 | ) |
| 368 | if state.pending_outputs == 0: |
| 369 | # In-flight steps leave the worker cursor behind the |
| 370 | # scheduler's optimistic view; only check settled requests. |
| 371 | assert emit_start <= state.emit_cursor, ( |
| 372 | "auxiliary output Scheduler emit cursor moved ahead" |
| 373 | ) |
| 374 | block_batches: list[ |
| 375 | tuple[_WorkerRequestState, list[tuple[int, np.ndarray]]] |
| 376 | ] = [] |
| 377 | retained_keys: list[str] = [] |
| 378 | for request_id, block_hashes in metadata.block_hashes.items(): |
| 379 | state = self._requests[request_id] |
| 380 | keys = routed_experts_keys(block_hashes, str(self._generation)) |
| 381 | state.aux_output_keys.extend(keys) |
| 382 | retained_keys.extend(keys) |
| 383 | block_batches.append((state, [])) |
| 384 | finished_now: list[str] = [] |
| 385 | for request_id in metadata.finished_requests: |
| 386 | state = self._requests[request_id] |
| 387 | state.finished = True |
| 388 | if state.pending_outputs: |
| 389 | continue |
| 390 | finished_now.append(request_id) |
| 391 | release_keys.extend(reversed(state.aux_output_keys)) |
| 392 | self._publish_blocks(block_batches, retained_keys, release_keys) |
| 393 | self._teardown(finished_now) |
| 394 | |
| 395 | def close(self) -> None: |
| 396 | with self._lock: |
| 397 | self._pending_outputs.clear() |
| 398 | if self._store is not None: |
| 399 | try: |
| 400 | self._store.close() |
| 401 | finally: |
| 402 | self._store = None |
| 403 | |
| 404 | |
| 405 | def get_aux_output_connector( |
| 406 | model: torch.nn.Module, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig |
| 407 | ) -> AuxOutputWorkerConnector: |
| 408 | return AuxOutputWorkerConnector( |
| 409 | model=model, |
| 410 | kv_cache_config=kv_cache_config, |
| 411 | vllm_config=vllm_config, |
| 412 | ) |