| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Scheduler-side control plane for execution auxiliary outputs.""" |
| 4 | |
| 5 | from __future__ import annotations |
| 6 | |
| 7 | from collections.abc import Iterator, Sequence |
| 8 | from dataclasses import dataclass |
| 9 | from typing import TYPE_CHECKING |
| 10 | |
| 11 | import numpy as np |
| 12 | |
| 13 | if TYPE_CHECKING: |
| 14 | from vllm.v1.core.sched.output import SchedulerOutput |
| 15 | from vllm.v1.request import Request |
| 16 | |
| 17 | |
| 18 | @dataclass |
| 19 | class PackedBlockHashes: |
| 20 | """Contiguous, self-contained block hashes for scheduler-worker IPC.""" |
| 21 | |
| 22 | data: bytes |
| 23 | item_size: int |
| 24 | |
| 25 | def __iter__(self) -> Iterator[bytes]: |
| 26 | for start in range(0, len(self.data), self.item_size): |
| 27 | yield self.data[start : start + self.item_size] |
| 28 | |
| 29 | |
| 30 | @dataclass |
| 31 | class AuxOutputConnectorMetadata: |
| 32 | generation: int |
| 33 | requests: dict[str, int] |
| 34 | block_hashes: dict[str, PackedBlockHashes] |
| 35 | finished_requests: tuple[str, ...] |
| 36 | |
| 37 | |
| 38 | @dataclass |
| 39 | class AuxRequestOutput: |
| 40 | token_start: int |
| 41 | rows: np.ndarray |
| 42 | |
| 43 | |
| 44 | class AuxOutputSchedulerConnector: |
| 45 | """Build worker metadata without owning auxiliary output payloads or stores.""" |
| 46 | |
| 47 | def __init__(self) -> None: |
| 48 | # Number of hashes already sent to the worker for each active request. |
| 49 | self._sent_hash_counts: dict[str, int] = {} |
| 50 | # Terminal events are delivered with the next connector metadata. |
| 51 | self._finished_requests: dict[str, PackedBlockHashes | None] = {} |
| 52 | self._generation = 0 |
| 53 | |
| 54 | def build_connector_meta( |
| 55 | self, |
| 56 | scheduler_output: SchedulerOutput, |
| 57 | requests: dict[str, Request], |
| 58 | ) -> AuxOutputConnectorMetadata: |
| 59 | """Build one step's incremental worker metadata.""" |
| 60 | scheduled_requests: dict[str, int] = {} |
| 61 | block_hashes_by_request: dict[str, PackedBlockHashes] = {} |
| 62 | for request_id in scheduler_output.num_scheduled_tokens: |
| 63 | num_sent = self._sent_hash_counts.setdefault(request_id, 0) |
| 64 | request = requests[request_id] |
| 65 | packed = self._pack_new_hashes(request.block_hashes, num_sent) |
| 66 | if packed is not None: |
| 67 | block_hashes_by_request[request_id] = packed |
| 68 | self._sent_hash_counts[request_id] = len(request.block_hashes) |
| 69 | assert request.sampling_params is not None |
| 70 | scheduled_requests[request_id] = max( |
| 71 | request.sampling_params.routed_experts_prompt_start, |
| 72 | 0 if request.num_output_tokens == 0 else request.num_tokens - 1, |
| 73 | ) |
| 74 | # A settled token can complete a hash block after the next async schedule |
| 75 | # was built. Send a hash-only update if the request was not rescheduled. |
| 76 | if len(self._sent_hash_counts) != len(scheduled_requests): |
| 77 | for request_id, num_sent in self._sent_hash_counts.items(): |
| 78 | if request_id in scheduled_requests: |
| 79 | continue |
| 80 | request = requests[request_id] |
| 81 | packed = self._pack_new_hashes(request.block_hashes, num_sent) |
| 82 | if packed is None: |
| 83 | continue |
| 84 | block_hashes_by_request[request_id] = packed |
| 85 | self._sent_hash_counts[request_id] = len(request.block_hashes) |
| 86 | # Sending transfers ownership of these one-shot events. |
| 87 | finished_requests = tuple(self._finished_requests) |
| 88 | block_hashes_by_request.update( |
| 89 | (request_id, block_hashes) |
| 90 | for request_id, block_hashes in self._finished_requests.items() |
| 91 | if block_hashes is not None |
| 92 | ) |
| 93 | self._finished_requests = {} |
| 94 | return AuxOutputConnectorMetadata( |
| 95 | self._generation, |
| 96 | scheduled_requests, |
| 97 | block_hashes_by_request, |
| 98 | finished_requests, |
| 99 | ) |
| 100 | |
| 101 | def take_output( |
| 102 | self, |
| 103 | request: Request, |
| 104 | output: dict[str, AuxRequestOutput] | None, |
| 105 | ) -> np.ndarray | None: |
| 106 | """Return the accepted R3 rows for one scheduled request.""" |
| 107 | request_id = request.request_id |
| 108 | assert output is not None and request_id in output, ( |
| 109 | f"auxiliary output worker output is missing {request_id}" |
| 110 | ) |
| 111 | request_output = output[request_id] |
| 112 | token_end = request.num_tokens - 1 |
| 113 | local_end = token_end - request_output.token_start |
| 114 | if local_end < 0: |
| 115 | assert not request.is_finished(), ( |
| 116 | "finished auxiliary output output has no accepted token range: " |
| 117 | f"request={request_id}, token_end={token_end}, " |
| 118 | f"output_start={request_output.token_start}, " |
| 119 | "output_end=" |
| 120 | f"{request_output.token_start + len(request_output.rows)}" |
| 121 | ) |
| 122 | return None |
| 123 | assert local_end <= len(request_output.rows), ( |
| 124 | "auxiliary output worker output has an invalid token range: " |
| 125 | f"request={request_id}, token_end={token_end}, " |
| 126 | f"output_start={request_output.token_start}, " |
| 127 | f"output_end={request_output.token_start + len(request_output.rows)}" |
| 128 | ) |
| 129 | return request_output.rows[:local_end] |
| 130 | |
| 131 | def request_finished(self, request: Request) -> None: |
| 132 | """Queue a request's terminal event and final block hashes.""" |
| 133 | request_id = request.request_id |
| 134 | num_sent = self._sent_hash_counts.pop(request_id, None) |
| 135 | if num_sent is None: |
| 136 | return |
| 137 | # The next metadata delivers this terminal event to the worker. |
| 138 | self._finished_requests[request_id] = self._pack_new_hashes( |
| 139 | request.block_hashes, num_sent |
| 140 | ) |
| 141 | |
| 142 | @staticmethod |
| 143 | def _pack_new_hashes( |
| 144 | block_hashes: Sequence[bytes], num_sent: int |
| 145 | ) -> PackedBlockHashes | None: |
| 146 | assert num_sent <= len(block_hashes), "KV block-hash history shrank" |
| 147 | new_hashes = block_hashes[num_sent:] |
| 148 | if not new_hashes: |
| 149 | return None |
| 150 | return PackedBlockHashes(b"".join(new_hashes), len(new_hashes[0])) |
| 151 | |
| 152 | def reset(self) -> None: |
| 153 | """Start a new auxiliary output namespace after a prefix-cache reset.""" |
| 154 | # The worker drops temporary state on generation changes; resend hashes. |
| 155 | self._sent_hash_counts.clear() |
| 156 | self._finished_requests.clear() |
| 157 | self._generation += 1 |