| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Routed-experts auxiliary output keys, publication, and materialization.""" |
| 4 | |
| 5 | from __future__ import annotations |
| 6 | |
| 7 | from collections.abc import Hashable, Iterable, Sequence |
| 8 | from dataclasses import dataclass |
| 9 | from typing import Any |
| 10 | |
| 11 | import numpy as np |
| 12 | |
| 13 | from vllm.distributed.aux_output_connector.store import ( |
| 14 | BackgroundBlockObjectStore, |
| 15 | BlockObject, |
| 16 | BlockObjectStore, |
| 17 | ) |
| 18 | |
| 19 | |
| 20 | @dataclass |
| 21 | class _RequestTail: |
| 22 | slot: int |
| 23 | block_start: int |
| 24 | length: int = 0 |
| 25 | |
| 26 | |
| 27 | class RoutedExpertsBuffer: |
| 28 | """Own incomplete and unkeyed full routed-experts blocks.""" |
| 29 | |
| 30 | def __init__( |
| 31 | self, |
| 32 | dtype: np.dtype[Any], |
| 33 | shape_per_token: tuple[int, ...], |
| 34 | block_size: int, |
| 35 | max_num_seqs: int, |
| 36 | max_num_batched_tokens: int, |
| 37 | max_concurrent_batches: int, |
| 38 | ) -> None: |
| 39 | self.dtype = dtype |
| 40 | self.shape_per_token = shape_per_token |
| 41 | self.block_size = block_size |
| 42 | max_step_blocks = ( |
| 43 | max_num_batched_tokens + block_size - 1 |
| 44 | ) // block_size + max_num_seqs |
| 45 | # In-flight full blocks plus one incomplete tail per active request. |
| 46 | max_blocks = max_concurrent_batches * max_step_blocks + max_num_seqs |
| 47 | self._rows: np.ndarray = np.empty( |
| 48 | (max_blocks, block_size, *shape_per_token), dtype=dtype |
| 49 | ) |
| 50 | self._free_slots = list(range(len(self._rows) - 1, -1, -1)) |
| 51 | self._owned_slots: dict[int, int] = {} |
| 52 | self._requests: dict[Hashable, _RequestTail] = {} |
| 53 | |
| 54 | def _tail(self, request_id: Hashable, block_start: int) -> _RequestTail: |
| 55 | tail = self._requests.get(request_id) |
| 56 | if tail is not None: |
| 57 | assert tail.block_start == block_start, ( |
| 58 | "auxiliary output capture skipped an incomplete block: " |
| 59 | f"request={request_id}, expected={tail.block_start}, " |
| 60 | f"actual={block_start}" |
| 61 | ) |
| 62 | return tail |
| 63 | assert self._free_slots, "auxiliary output block pool is exhausted" |
| 64 | tail = _RequestTail(self._free_slots.pop(), block_start) |
| 65 | self._requests[request_id] = tail |
| 66 | return tail |
| 67 | |
| 68 | def capture( |
| 69 | self, request_id: Hashable, token_start: int, rows: np.ndarray |
| 70 | ) -> list[tuple[int, np.ndarray]]: |
| 71 | """Stage rows and return completed blocks without retaining them.""" |
| 72 | rows = np.asarray(rows) |
| 73 | assert rows.shape[1:] == self.shape_per_token and rows.dtype == self.dtype, ( |
| 74 | "routed-experts capture profile changed" |
| 75 | ) |
| 76 | if token_start < 0: |
| 77 | raise ValueError("auxiliary output token start must be non-negative") |
| 78 | |
| 79 | completed: list[tuple[int, np.ndarray]] = [] |
| 80 | offset = 0 |
| 81 | while offset < len(rows): |
| 82 | position = token_start + offset |
| 83 | block_start = position // self.block_size * self.block_size |
| 84 | local_start = position - block_start |
| 85 | |
| 86 | # Full aligned input blocks do not need tail staging. |
| 87 | if ( |
| 88 | request_id not in self._requests |
| 89 | and local_start == 0 |
| 90 | and len(rows) - offset >= self.block_size |
| 91 | ): |
| 92 | completed.append((block_start, rows[offset : offset + self.block_size])) |
| 93 | offset += self.block_size |
| 94 | continue |
| 95 | |
| 96 | tail = self._tail(request_id, block_start) |
| 97 | assert local_start == tail.length, ( |
| 98 | "auxiliary output capture is not contiguous: " |
| 99 | f"request={request_id}, expected={block_start + tail.length}, " |
| 100 | f"actual={position}" |
| 101 | ) |
| 102 | count = min(self.block_size - local_start, len(rows) - offset) |
| 103 | if count == 1: |
| 104 | self._rows[tail.slot, local_start] = rows[offset] |
| 105 | else: |
| 106 | self._rows[tail.slot, local_start : local_start + count] = rows[ |
| 107 | offset : offset + count |
| 108 | ] |
| 109 | tail.length += count |
| 110 | offset += count |
| 111 | if tail.length == self.block_size: |
| 112 | block = self._rows[tail.slot] |
| 113 | del self._requests[request_id] |
| 114 | self._owned_slots[id(block)] = tail.slot |
| 115 | completed.append((block_start, block)) |
| 116 | |
| 117 | return completed |
| 118 | |
| 119 | def read( |
| 120 | self, request_id: Hashable, token_start: int, token_end: int |
| 121 | ) -> np.ndarray: |
| 122 | tail = self._requests.get(request_id) |
| 123 | assert tail is not None, ( |
| 124 | f"auxiliary output buffer is missing request {request_id}" |
| 125 | ) |
| 126 | local_start = token_start - tail.block_start |
| 127 | local_end = token_end - tail.block_start |
| 128 | assert 0 <= local_start < local_end <= tail.length, ( |
| 129 | "auxiliary output range is unavailable: " |
| 130 | f"request={request_id}, range=[{token_start}, {token_end}), " |
| 131 | f"available=[{tail.block_start}, {tail.block_start + tail.length})" |
| 132 | ) |
| 133 | return self._rows[tail.slot, local_start:local_end].copy() |
| 134 | |
| 135 | def retain_block(self, rows: np.ndarray) -> np.ndarray: |
| 136 | """Retain one unkeyed block after the current capture call.""" |
| 137 | if id(rows) in self._owned_slots: |
| 138 | return rows |
| 139 | assert self._free_slots, "auxiliary output block pool is exhausted" |
| 140 | slot = self._free_slots.pop() |
| 141 | retained = self._rows[slot] |
| 142 | retained[...] = rows |
| 143 | self._owned_slots[id(retained)] = slot |
| 144 | return retained |
| 145 | |
| 146 | def release_block(self, rows: np.ndarray) -> None: |
| 147 | slot = self._owned_slots.pop(id(rows), None) |
| 148 | if slot is not None: |
| 149 | self._free_slots.append(slot) |
| 150 | |
| 151 | def discard(self, request_id: Hashable) -> None: |
| 152 | tail = self._requests.pop(request_id, None) |
| 153 | if tail is not None: |
| 154 | self._free_slots.append(tail.slot) |
| 155 | |
| 156 | def reset(self) -> None: |
| 157 | self._requests.clear() |
| 158 | self._owned_slots.clear() |
| 159 | self._free_slots[:] = range(len(self._rows) - 1, -1, -1) |
| 160 | |
| 161 | |
| 162 | def routed_experts_keys( |
| 163 | block_hashes: Iterable[bytes], aux_output_namespace: str |
| 164 | ) -> list[str]: |
| 165 | prefix = f"vllm-artifact/{aux_output_namespace}/" |
| 166 | return [prefix + block_hash.hex() for block_hash in block_hashes] |
| 167 | |
| 168 | |
| 169 | def materialize_routed_experts( |
| 170 | store: BackgroundBlockObjectStore | BlockObjectStore, |
| 171 | aux_output_keys: list[str], |
| 172 | *, |
| 173 | shape_per_token: tuple[int, ...], |
| 174 | dtype: np.dtype[Any], |
| 175 | ) -> np.ndarray: |
| 176 | if not aux_output_keys: |
| 177 | raise ValueError("routed-experts auxiliary output key list must not be empty") |
| 178 | payload = store.get_concatenated(aux_output_keys) |
| 179 | return np.frombuffer(payload, dtype=dtype).reshape((-1, *shape_per_token)) |
| 180 | |
| 181 | |
| 182 | def publish_routed_experts( |
| 183 | store: BackgroundBlockObjectStore | BlockObjectStore, |
| 184 | *, |
| 185 | batches: Sequence[tuple[Sequence[str], list[tuple[int, np.ndarray]]]], |
| 186 | block_size: int, |
| 187 | retain_keys: Sequence[str] = (), |
| 188 | release_keys: Sequence[str] = (), |
| 189 | ) -> None: |
| 190 | """Publish immutable full R3 blocks.""" |
| 191 | objects = [] |
| 192 | for aux_output_keys, blocks in batches: |
| 193 | for block_start, array in blocks: |
| 194 | if block_start < 0 or block_start % block_size: |
| 195 | raise ValueError( |
| 196 | "auxiliary output block start is not hash-block aligned" |
| 197 | ) |
| 198 | block_index = block_start // block_size |
| 199 | if block_index >= len(aux_output_keys): |
| 200 | raise ValueError( |
| 201 | "auxiliary output block has no corresponding key: " |
| 202 | f"start={block_start}, index={block_index}, " |
| 203 | f"keys={len(aux_output_keys)}, block_size={block_size}" |
| 204 | ) |
| 205 | if len(array) != block_size: |
| 206 | raise ValueError( |
| 207 | "auxiliary output block length does not match hash block size" |
| 208 | ) |
| 209 | objects.append( |
| 210 | BlockObject( |
| 211 | key=aux_output_keys[block_index], |
| 212 | payload=array.tobytes(order="C"), |
| 213 | ) |
| 214 | ) |
| 215 | store.put( |
| 216 | objects, |
| 217 | retain_keys=retain_keys, |
| 218 | release_keys=release_keys, |
| 219 | ) |