| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Shared pieces of the sharded-RDT backend: the op-chain allowlist, buffer |
| 4 | sizing and the minimum Ray version, plus ``RdtRouter`` — consumer-only, since |
| 5 | routing is a consumer decision. |
| 6 | |
| 7 | The gather-group partition itself is ``base.layerwise_groups``: it defines what a |
| 8 | group index means for any ``WeightSource``, not just this transport. |
| 9 | """ |
| 10 | |
| 11 | from collections.abc import Callable |
| 12 | from typing import Any |
| 13 | |
| 14 | import torch |
| 15 | |
| 16 | # The op-chain contract, with two enforcers. The consumer's ``FakeRDTTensor`` |
| 17 | # intercepts exactly these methods during the bake and records the string name; |
| 18 | # anything else reaches ``__torch_dispatch__`` and raises, so a loader needing |
| 19 | # real data fails at init rather than transferring wrong bytes. The producer |
| 20 | # replays a chain as ``getattr(tensor, op)(...)`` and refuses anything outside |
| 21 | # ``ALLOWED_OPS``, so a spoofed consumer cannot invoke arbitrary methods. |
| 22 | # |
| 23 | # Every entry must be a pure view; ``to`` in particular must never be added, as |
| 24 | # it is the dtype/device escape the bake exists to reject. Both sets derive from |
| 25 | # this one table so they cannot disagree. |
| 26 | SUPPORTED_OPS: dict[Callable, str] = { |
| 27 | torch.Tensor.narrow: "narrow", |
| 28 | torch.Tensor.view: "view", |
| 29 | torch.Tensor.reshape: "reshape", |
| 30 | torch.Tensor.__getitem__: "__getitem__", |
| 31 | torch.Tensor.unsqueeze: "unsqueeze", |
| 32 | torch.Tensor.squeeze: "squeeze", |
| 33 | torch.Tensor.transpose: "transpose", |
| 34 | torch.Tensor.t: "t", |
| 35 | torch.Tensor.permute: "permute", |
| 36 | torch.Tensor.flatten: "flatten", |
| 37 | torch.Tensor.contiguous: "contiguous", |
| 38 | torch.Tensor.chunk: "chunk", |
| 39 | # Multi-return like chunk: one child per output, trailing __getitem__(i). |
| 40 | torch.Tensor.unbind: "unbind", |
| 41 | } |
| 42 | |
| 43 | ALLOWED_OPS = frozenset(SUPPORTED_OPS.values()) |
| 44 | |
| 45 | # Ray Direct Transport landed in halves: the actor opt-in |
| 46 | # ``enable_tensor_transport`` in 2.49, and ``ray.experimental``'s |
| 47 | # ``register_nixl_memory`` / ``set_target_for_ref`` — which this transport calls |
| 48 | # directly — in 2.55. The floor is the tested version rather than the earliest |
| 49 | # one that could import; 2.55 has never been run against. |
| 50 | RDT_MIN_RAY_VERSION = "2.56.0" |
| 51 | |
| 52 | |
| 53 | def check_ray_rdt_version() -> None: |
| 54 | """Refuse an installed Ray older than the one this backend is tested on. |
| 55 | |
| 56 | vLLM does not depend on Ray, so there is no pin to carry this. Without the |
| 57 | check the failure is an opaque option-validation error out of |
| 58 | ``.options(enable_tensor_transport=True)`` (below 2.49) or an ImportError |
| 59 | raised deep in the first pull, long after init reported success (below 2.55). |
| 60 | |
| 61 | Raises: |
| 62 | ValueError: the installed Ray predates ``RDT_MIN_RAY_VERSION``. |
| 63 | |
| 64 | """ |
| 65 | import importlib.metadata |
| 66 | |
| 67 | from packaging import version |
| 68 | |
| 69 | required = version.parse(RDT_MIN_RAY_VERSION) |
| 70 | current = version.parse(importlib.metadata.version("ray")) |
| 71 | if current < required: |
| 72 | raise ValueError( |
| 73 | f"The 'sharded_rdt' weight transfer backend requires Ray " |
| 74 | f">= {required}, the version it is tested against; Ray Direct " |
| 75 | f"Transport (ray.experimental.register_nixl_memory / " |
| 76 | f"set_target_for_ref) needs at least 2.55. Found {current}. " |
| 77 | f"Run `pip install -U 'ray>={required}'`." |
| 78 | ) |
| 79 | |
| 80 | |
| 81 | def register_nixl_memory(tensor: torch.Tensor) -> None: |
| 82 | """Register a tensor with Ray using the platform's NIXL package.""" |
| 83 | from vllm.distributed.nixl_utils import alias_nixl_for_ray |
| 84 | |
| 85 | alias_nixl_for_ray() |
| 86 | from ray.experimental import register_nixl_memory as ray_register_nixl_memory |
| 87 | |
| 88 | ray_register_nixl_memory(tensor) |
| 89 | |
| 90 | |
| 91 | def assign_producer_indices( |
| 92 | num_producers: int, num_consumers: int, consumer_idx: int |
| 93 | ) -> list[int]: |
| 94 | """Producers (global indices) that consumer ``consumer_idx`` binds.""" |
| 95 | p = max(1, num_producers) |
| 96 | c = max(1, num_consumers) |
| 97 | if p >= c: |
| 98 | return list(range(consumer_idx * p // c, (consumer_idx + 1) * p // c)) |
| 99 | return [consumer_idx * p // c] |
| 100 | |
| 101 | |
| 102 | class RdtRouter: |
| 103 | """Decides which producer serves each weight name, and — once bound — sends. |
| 104 | |
| 105 | Ownership is per NAME. ``owner_sets`` holds the few distinct producer sets |
| 106 | that occur, and ``name_owner_class[i]`` indexes it for ``names[i]``; a name |
| 107 | determines both its owner set and its gather group, so neither appears in the |
| 108 | routing API. Empty tables mean every producer holds everything. |
| 109 | |
| 110 | Any placement is expressible this way: a pipeline stage is a set of names |
| 111 | sharing one owner set, an expert is a name whose owner set is a single rank, |
| 112 | and a group produced by two stages is just two classes inside one group. |
| 113 | |
| 114 | Both engines derive the same tables from the same wire data, so they agree on |
| 115 | who serves what. Disagreement is not a wrong answer but a hang or a loud |
| 116 | misroute: a pull sent to a producer that never gathered the name trips its |
| 117 | served-names guard. |
| 118 | |
| 119 | Routing is consumer-only. The trainer publishes what it holds and answers |
| 120 | whatever arrives; it never asks who serves what. |
| 121 | """ |
| 122 | |
| 123 | def __init__( |
| 124 | self, |
| 125 | num_producers: int, |
| 126 | num_consumers: int, |
| 127 | owner_sets: list[list[int]] | None = None, |
| 128 | name_owner_class: list[int] | None = None, |
| 129 | names: list[str] | None = None, |
| 130 | group_lens: list[int] | None = None, |
| 131 | workers_per_replica: int = 0, |
| 132 | ) -> None: |
| 133 | """Build the routing tables from the wire data both engines receive. |
| 134 | |
| 135 | The four table arguments travel together and default together: with all |
| 136 | of them empty the router degrades to "every producer holds everything". |
| 137 | |
| 138 | Args: |
| 139 | num_producers: Trainer ranks. Owner indices are positions in this |
| 140 | range, and ``validate`` rejects any that fall outside it. |
| 141 | num_consumers: Inference workers across the whole fleet. Fixes |
| 142 | the id space; the block carve uses ``workers_per_replica`` |
| 143 | of it. |
| 144 | owner_sets: The distinct producer sets that occur, one row per |
| 145 | owner class; each row is deduplicated and sorted here. Empty |
| 146 | means a single class owning every producer. |
| 147 | name_owner_class: Parallel to ``names`` — ``name_owner_class[i]`` |
| 148 | indexes ``owner_sets`` for ``names[i]``. A name with no entry |
| 149 | falls back to class 0. |
| 150 | names: Every parameter name, in group-major order: concatenating |
| 151 | the gather groups reproduces this list exactly. The other two |
| 152 | tables are keyed by this order. |
| 153 | group_lens: Length of each gather group, consecutive over |
| 154 | ``names``, so they sum to ``len(names)``. Fixes name -> group |
| 155 | and the per-group owner union the free barrier fans out to. |
| 156 | workers_per_replica: Consumers per inference DEPLOYMENT. The block |
| 157 | carve spreads this many consumers over an owner set and every |
| 158 | deployment reuses that carve, so the same worker of each |
| 159 | deployment resolves one producer (see ``producer_for``). 0 |
| 160 | means one deployment: carve over the whole fleet. |
| 161 | |
| 162 | Raises: |
| 163 | ValueError: ``workers_per_replica`` does not divide |
| 164 | ``num_consumers``, so the deployments are not uniform and two |
| 165 | workers of one deployment would share a block index. |
| 166 | |
| 167 | """ |
| 168 | self.num_producers = max(1, num_producers) |
| 169 | self.num_consumers = max(1, num_consumers) |
| 170 | # Width of the block carve: one deployment's worth of consumers, so |
| 171 | # deployments overlay rather than spread. See ``producer_for``. |
| 172 | self._block_consumers = min( |
| 173 | self.num_consumers, max(1, workers_per_replica or self.num_consumers) |
| 174 | ) |
| 175 | if self.num_consumers % self._block_consumers: |
| 176 | raise ValueError( |
| 177 | f"workers_per_replica={workers_per_replica} does not divide " |
| 178 | f"num_consumers={self.num_consumers}; the inference deployments " |
| 179 | f"must be uniform for the per-deployment block carve to be " |
| 180 | f"well defined." |
| 181 | ) |
| 182 | self._owner_sets = ( |
| 183 | [sorted(set(row)) for row in owner_sets] |
| 184 | if owner_sets |
| 185 | else [list(range(self.num_producers))] |
| 186 | ) |
| 187 | names = list(names or []) |
| 188 | classes = list(name_owner_class or []) |
| 189 | self._class_of = { |
| 190 | n: (classes[i] if i < len(classes) else 0) for i, n in enumerate(names) |
| 191 | } |
| 192 | # name -> gather group, and the per-group owner union the free barrier |
| 193 | # fans out to. Both are pure functions of the tables, so precompute once. |
| 194 | self._group_of: dict[str, int] = {} |
| 195 | self._group_owners: list[list[int]] = [] |
| 196 | pos = 0 |
| 197 | for gi, glen in enumerate(group_lens or []): |
| 198 | union: set[int] = set() |
| 199 | for n in names[pos : pos + glen]: |
| 200 | self._group_of[n] = gi |
| 201 | c = self._class_of.get(n, 0) |
| 202 | # Bounds-checked: a bad class must surface from validate() as a |
| 203 | # named ValueError, not as an IndexError from this precompute. |
| 204 | if 0 <= c < len(self._owner_sets): |
| 205 | union.update(self._owner_sets[c]) |
| 206 | pos += glen |
| 207 | self._group_owners.append(sorted(union)) |
| 208 | # Bound at init on the consumer; the rule methods never touch these. |
| 209 | self._actors: list[Any] = [] |
| 210 | self._produce_methods: list[Any] = [] |
| 211 | self.consumer_id = 0 |
| 212 | |
| 213 | # ---------------- rules (pure; safe to use unbound) ---------------- |
| 214 | |
| 215 | def class_of(self, name: str) -> int: |
| 216 | """The name's owner class — the planner's bucketing key, since all names |
| 217 | of a chunk must share one producer.""" |
| 218 | return self._class_of.get(name, 0) |
| 219 | |
| 220 | def owners(self, name: str) -> list[int]: |
| 221 | """Every producer holding ``name``.""" |
| 222 | return list(self._owner_sets[self.class_of(name)]) |
| 223 | |
| 224 | def group_owners(self, group_idx: int) -> list[int]: |
| 225 | """Every producer holding any name of ``group_idx``. |
| 226 | |
| 227 | The one group-keyed rule, because the free barrier is per group: a |
| 228 | consumer signals ``free_group(gi)`` at each of these and the producer |
| 229 | counts signals against its live-consumer total. Names route; groups free. |
| 230 | """ |
| 231 | if not self._group_owners: |
| 232 | return list(range(self.num_producers)) |
| 233 | return list(self._group_owners[group_idx]) |
| 234 | |
| 235 | def producer_for(self, consumer_id: int, name: str) -> int: |
| 236 | """The single producer ``consumer_id`` pulls ``name`` from. |
| 237 | |
| 238 | Blocks consumers across the name's owner set with the same rule that |
| 239 | binds producers globally, then rotates by the name's GROUP index so a |
| 240 | consumer spreads its groups over its block instead of hammering one NIC. |
| 241 | Rotating per group (not per name) keeps every name of a chunk on one |
| 242 | producer, which is what lets a chunk be a single pull. |
| 243 | |
| 244 | The carve is over ONE DEPLOYMENT (this consumer's index within its own |
| 245 | deployment, over ``workers_per_replica`` of them), so the same worker of |
| 246 | every deployment resolves the same producer for every name. Carving over |
| 247 | the whole fleet instead spreads the multi-owner names of different |
| 248 | deployments onto different producers while single-owner names still |
| 249 | route by ownership, so a producer serves several distinct worker |
| 250 | indices, each needing its own serve ring, and none of them can share a |
| 251 | slot. With one deployment the width is the fleet, so this is the plain |
| 252 | block rule. |
| 253 | """ |
| 254 | own = self.owners(name) |
| 255 | if not own: |
| 256 | raise ValueError(f"{name!r} has no owner") |
| 257 | block = assign_producer_indices( |
| 258 | len(own), self._block_consumers, consumer_id % self._block_consumers |
| 259 | ) |
| 260 | return own[block[self._group_of.get(name, 0) % len(block)]] |
| 261 | |
| 262 | def validate(self) -> None: |
| 263 | """Check the ownership tables can be served. |
| 264 | |
| 265 | Raises: |
| 266 | ValueError: an owner set is empty or out of range, or a class index |
| 267 | does not resolve. |
| 268 | |
| 269 | """ |
| 270 | for c, row in enumerate(self._owner_sets): |
| 271 | if not row: |
| 272 | raise ValueError(f"owner set {c} is empty: no producer holds it") |
| 273 | bad = [p for p in row if not 0 <= p < self.num_producers] |
| 274 | if bad: |
| 275 | raise ValueError(f"owner set {c} out of range: {bad}") |
| 276 | for name, c in self._class_of.items(): |
| 277 | if not 0 <= c < len(self._owner_sets): |
| 278 | raise ValueError( |
| 279 | f"{name!r} has owner class {c}, but only " |
| 280 | f"{len(self._owner_sets)} owner set(s) were shipped" |
| 281 | ) |
| 282 | |
| 283 | # ---------------- send surface (needs the binding) ---------------- |
| 284 | |
| 285 | def bind(self, actors: list, produce_methods: list, consumer_id: int) -> None: |
| 286 | """Attach this consumer's producer handles. Rebuilt wholesale per init, |
| 287 | never appended to: every owner index is a position in these lists, so a |
| 288 | rejoining engine that re-inits must not shift them.""" |
| 289 | if len(actors) != len(produce_methods): |
| 290 | raise ValueError("actors and produce_methods must be parallel") |
| 291 | self._actors = list(actors) |
| 292 | self._produce_methods = list(produce_methods) |
| 293 | self.consumer_id = consumer_id |
| 294 | |
| 295 | def _bound(self) -> None: |
| 296 | if not self._produce_methods: |
| 297 | raise RuntimeError("RdtRouter has no producers bound; call bind() first.") |
| 298 | |
| 299 | def pull(self, owner: int, keys: list, seq: int) -> Any: |
| 300 | """Issue one packed pull to ``owner``. |
| 301 | |
| 302 | The consumer id keys the producer's serve ring; without it every worker |
| 303 | is served out of ring 0 and concurrent pulls overwrite each other's blob. |
| 304 | ``seq`` is this call's index in the stream to ``owner`` (see ``_Chunk``): |
| 305 | it selects the serve slot in ISSUE order, so a slot is never repacked |
| 306 | while the read of its previous contents is still in flight. |
| 307 | """ |
| 308 | self._bound() |
| 309 | return self._produce_methods[owner].remote( |
| 310 | keys, consumer_id=self.consumer_id, seq=seq |
| 311 | ) |
| 312 | |
| 313 | def free_group(self, group_idx: int) -> list: |
| 314 | """Signal the group done at every producer holding any of its names.""" |
| 315 | self._bound() |
| 316 | return [ |
| 317 | self._actors[p].free_group.remote(group_idx) |
| 318 | for p in self.group_owners(group_idx) |
| 319 | ] |
| 320 | |
| 321 | def reserve_serve_buffers( |
| 322 | self, bytes_by_producer: list[int], plan_digests: list[str] | None = None |
| 323 | ) -> list: |
| 324 | """Ask each producer to pre-register a serve ring sized to the most this |
| 325 | consumer will pull from it. |
| 326 | |
| 327 | ``plan_digests[p]`` describes the chunks this consumer pulls from |
| 328 | producer ``p``, in pull order. A producer that shares one serve ring |
| 329 | across the consumers of several deployments compares it across them, so |
| 330 | a fleet whose deployments are not identical fails at init instead of |
| 331 | stalling mid-sync. A producer that shares nothing ignores it. |
| 332 | """ |
| 333 | self._bound() |
| 334 | digests: list = ( |
| 335 | list(plan_digests) if plan_digests else [None] * len(bytes_by_producer) |
| 336 | ) |
| 337 | return [ |
| 338 | self._actors[p].reserve_serve_buffer.remote( |
| 339 | self.consumer_id, nb, digests[p] |
| 340 | ) |
| 341 | for p, nb in enumerate(bytes_by_producer) |
| 342 | if nb > 0 |
| 343 | ] |
| 344 | |
| 345 | |
| 346 | def buffer_alloc_bytes(nbytes: int, presize: int = 0) -> int: |
| 347 | """Size a NIXL buffer / ring slot for ``nbytes``: the max of the request, an |
| 348 | optional ``presize`` floor, and a coarse 256MB round-up, so the buffer is |
| 349 | allocated ONCE and never regrows. Regrowth is a correctness hazard, not just a |
| 350 | perf one -- see ``buffer_presize_gb``. Shared by the consumer's receive buffers |
| 351 | and the producer's serve rings.""" |
| 352 | return max(nbytes, presize, -(-nbytes // (256 << 20)) * (256 << 20)) |