From 04b02fc8a56fc0140d8559b3a07d84b2fdc46ae8 Mon Sep 17 00:00:00 2001 From: Bolin Sun Date: Fri, 9 Oct 2026 19:13:24 -0400 Subject: [PATCH] [Core] shm tensor arena for efficient cpu->gpu worker broadcast (#51207) Signed-off-by: Bolin Sun Signed-off-by: Signed-off-by: Nick Hill Co-authored-by: Claude Co-authored-by: Nick Hill --- docs/design/shm_tensor_arena.md | 302 ++++++++++ tests/distributed/test_shm_broadcast.py | 522 +++++++++++++++++- tools/pre_commit/check_forbidden_imports.py | 1 + vllm/config/parallel.py | 10 + .../device_communicators/shm_broadcast.py | 59 +- .../device_communicators/shm_tensor_arena.py | 397 +++++++++++++ vllm/engine/arg_utils.py | 6 + vllm/v1/executor/multiproc_executor.py | 1 + 8 files changed, 1290 insertions(+), 8 deletions(-) create mode 100644 docs/design/shm_tensor_arena.md create mode 100644 vllm/distributed/device_communicators/shm_tensor_arena.py diff --git a/docs/design/shm_tensor_arena.md b/docs/design/shm_tensor_arena.md new file mode 100644 index 0000000000..c37ac89e36 --- /dev/null +++ b/docs/design/shm_tensor_arena.md @@ -0,0 +1,302 @@ +# Zero-Copy Shared-Memory Tensor Arena for `MessageQueue` + +This document describes an optimization to vLLM's engine→worker IPC path for +multimodal serving under tensor parallelism, implemented in +`vllm/distributed/device_communicators/shm_tensor_arena.py` and wired into +`shm_broadcast.py`'s `MessageQueue`. It **builds on the existing out-of-band +tensor channel** (`_reduce_tensor` / `_rebuild_tensor`, still in +`shm_broadcast.py`), which already keeps CPU-tensor bytes out of the pickle +stream. + +> **Status: experimental, opt-in.** This lands as an increment over the +> existing out-of-band channel, not a replacement for it. vLLM also has an +> existing multimodal-tensor shm caching mechanism; reconciling the two is +> intentionally left as follow-up work rather than blocking this PR. + +## TL;DR + +`MessageQueue.enqueue` already avoids copying CPU tensor bytes into the pickle +stream: `_reduce_tensor` emits each tensor as a protocol-5 out-of-band +`PickleBuffer`, and large payloads are published to the local readers over a +ZMQ socket. That removed the dominant cost of the older *in-band* path — the +per-tensor `pickle` serialize (see [History](#history-the-in-band-pickle-era)). + +Two costs remain for a large multimodal `pixel_values` tensor on a TP=N worker: + +1. **N transport copies.** On the overflow (large-payload) path the out-of-band + buffer is published to each of the N node-local readers via `send_multipart`, + so the tensor bytes are moved once **per reader** by the transport — ~N×200 MB + of host data motion for a payload the model needs exactly once. +2. **A pageable H2D.** Each reader zero-copies a view over its received ZMQ + frame, but that frame is ordinary pageable memory, so the reader's subsequent + `cudaMemcpyAsync` to the GPU pays pageable-copy staging plus first-touch + faults. + +This change adds a slotted shared-memory **arena**: the writer does **one** +memcpy of the tensor into a free slot; every reader takes a **zero-copy view** +of the same slot (no per-reader transport copy); and the mapping is +`cudaHostRegister`-pinned so the H2D is a true DMA. Net for a 200 MB image on +TP=4: **1 copy + a pinned DMA** instead of **4 transport copies + a pageable +H2D**. Small tensors, non-contiguous tensors, and arena exhaustion all fall +through to the existing `_reduce_tensor` path unchanged. + +## Scope of the win (and what still needs measuring) + +The arena is an **incremental** improvement layered on the out-of-band channel, +not a from-scratch fix — the out-of-band channel already eliminated the single +largest cost (the serialize copy). The benefit that remains is: + +- **N→1 copies**, which grows with TP degree (no gain at TP=1), and +- the **pinned vs. pageable H2D**, which shows up as reader-side H2D latency. + +> **The benchmarks in [§4](#4-validation) were measured against the *in-band +> pickle* baseline that predates `_reduce_tensor`**, where the dominant cost was +> the serialize copy. They are retained as motivation and as an upper bound. The +> incremental improvement of the arena *over the current out-of-band path on +> `main`* is smaller and has **not yet been re-measured**; that A/B (arena vs. +> oob, same seeded workload) is the outstanding validation item before this is +> merge-ready. + +## 1. Where it bites + +Any multimodal model served with `--tensor-parallel-size > 1` through the +multiproc executor. The `rpc_broadcast_mq` `MessageQueue` +(`v1/executor/multiproc_executor.py`) carries scheduler output — including the +raw preprocessed `pixel_values` — from the EngineCore to every TP worker each +step. Images at native resolution reach tens to hundreds of MB (a ~4k×4k image +is ~200 MB as bf16 patches). The out-of-band channel keeps those bytes out of +the pickle stream, but on the overflow path they are still published **once per +node-local reader**, and each reader's H2D is from **pageable** memory. The +per-reader copy cost scales with TP degree; the pinning cost is a fixed per-image +tax on the reader H2D. Latency-sensitive (interactive) serving makes the residual +visible in TTFT and, via head-of-line blocking, in the end-to-end latency of +*other* requests. + +### History: the in-band-pickle era + +Before `_reduce_tensor`, `MessageQueue.enqueue` pickled the whole payload — +including tensor bytes — because `torch.Tensor.__reduce_ex__` does not emit +protocol-5 out-of-band buffers. A ~200 MB tensor was byte-copied into the pickle +stream on the engine (`THPStorage_writeFileRaw`, ~6.5 ms/MB → ~1.3 s), *inside* +the EngineCore step loop, so all GPUs on the worker went idle for **0.9–1.3 s per +large image** (the vision-encoder compute for the same image is 50–90 ms → +transport dominated compute ~15:1); each reader then paid a full `pickle.loads` +byte-copy. Per 200 MB image on TP=4 that was ~1.8 GB of host data motion. + +The `_reduce_tensor` out-of-band channel removed the serialize copy (and the +per-reader `pickle.loads` copy on the zero-copy ZMQ-frame path). Note that +merely raising `VLLM_MQ_MAX_CHUNK_BYTES_MB` (ring vs. socket path) did **not** +help in the in-band era: the serialize ran before the ring-vs-socket branch, so +it was common to both — a negative result that pinned the real cost on the +byte-copies, not the transport choice. + +## 2. The arena + +The arena itself (`ShmTensorArena`, `_ArenaPickler`, `_rebuild_arena_tensor`) +lives in `shm_tensor_arena.py`; `shm_broadcast.py`'s `MessageQueue` creates +and wires it in but otherwise treats it as an internal implementation detail. + +### 2.1 `ShmTensorArena` + +A second shared-memory region created alongside the existing `ShmRingBuffer`: + +- **N slots × slot_bytes** (default 8 × 256 MB, internal constants -- see + §3), plus per-slot metadata `[written_flag, reader0_done … readerN_done]`. +- Concurrency uses the **same lock-free single-writer/N-reader protocol as + `ShmRingBuffer`** (memory fences, per-reader done flags), so the model is one + the codebase already trusts. +- Created by the queue **writer** in `MessageQueue.__init__` when all readers are + node-local; readers attach via a new `tensor_arena_handle` field on the queue + `Handle`. Queues with remote readers get no arena and keep today's behavior. +- Pages are allocated lazily by the kernel on first write, so arenas on queues + that never carry big tensors (e.g. worker→engine response queues) cost + approximately nothing beyond the (virtual) reservation. + +### 2.2 Writer path — arena diversion, composing with `_reduce_tensor` + +`enqueue` builds one `dispatch_table` that routes CPU tensors through +`_reduce_tensor` (out-of-band `PickleBuffer`). When a node-local arena exists, +the pickler is an `_ArenaPickler` whose `reducer_override` *additionally* diverts +**large contiguous** CPU tensors into the arena: + +```python +class _ArenaPickler(pickle.Pickler): + def reducer_override(self, obj): + rebuild_fn = _ARENA_REBUILD_FNS.get(type(obj)) # {Tensor: ..., Parameter: ...} + if (rebuild_fn is not None and obj.device.type == "cpu" + and obj.layout is torch.strided and obj.is_contiguous() + and not obj.requires_grad + and obj.numel() * obj.element_size() >= MIN_BYTES): + idx = self.arena.write_tensor(obj) # ONE memcpy into a free slot + if idx is not None: + return (rebuild_fn, + (self.arena.shared_memory.name, idx, nbytes, dtype_str, shape)) + return NotImplemented # fall through to dispatch_table → _reduce_tensor +``` + +Note the exact-type dict lookup (not `isinstance`), matching how +`dispatch_table` itself dispatches: an unrecognized `Tensor` subclass must +decline here too, or it would come back as a plain `Tensor`, silently +losing its subclass identity. `torch.nn.Parameter` is the one subclass this +*does* divert, via its own rebuild function (§2.3). + +`reducer_override` is consulted before an object's normal reduction, and +returning `NotImplemented` falls through to the `dispatch_table` — so a diverted +tensor becomes a ~100-byte rebuild stub, and **everything the arena declines** +(too small, non-contiguous, or arena full) is handled by `_reduce_tensor` +exactly as on `main`. The single arena copy — +`torch.frombuffer(slot).copy_(t.view(torch.uint8))` — is a multithreaded memcpy +at memory bandwidth (~50–70 ms for 200 MB), and, unlike the oob path, the *same* +slot is then read zero-copy by all N readers rather than transported to each. + +**The writer never blocks.** If no slot is free or the tensor exceeds the slot +size, `write_tensor` returns `None`; the pickler falls through to `_reduce_tensor` +— the worst case is exactly the current out-of-band behavior, and deadlock is +structurally impossible. + +### 2.3 Reader path — zero copies + +The stub unpickles through a module-level rebuild function: + +```python +def _rebuild_arena_tensor(arena_name, slot_idx, nbytes, dtype_str, shape): + arena = _TENSOR_ARENAS[arena_name] # this process's registry of attached arenas + # get_tensor: a view over the mapped slot — zero bytes copied + return arena.get_tensor(slot_idx, nbytes, getattr(torch, dtype_str), shape) +``` + +`_rebuild_arena_parameter` wraps the same tensor in +`torch.nn.Parameter(t, requires_grad=False)`, which shares its storage. +`_ARENA_REBUILD_FNS = {torch.Tensor: _rebuild_arena_tensor, torch.nn. +Parameter: _rebuild_arena_parameter}` is the full set of types the arena +diverts; anything else (any other `Tensor` subclass) declines in +`reducer_override` and takes the out-of-band path. + +The registry is keyed by arena shm name (`_TENSOR_ARENAS: +weakref.WeakValueDictionary[str, ShmTensorArena]`) rather than a single +per-process slot, since nothing guarantees a process only ever attaches to +one arena-bearing queue. It holds only a *weak* reference — a strong one +would keep the arena (and its pinned mapping) alive for the process +lifetime — and needs no explicit removal on `MessageQueue.shutdown()` +(unsafe anyway, since shutdown can run on a different thread than an +in-flight `dequeue()`); an entry just disappears once nothing else +references that arena. + +The rebuilt tensor *is* the shared memory — no transport copy and no deserialize +on any rank. + +**Slot lifecycle.** The rebuilt tensor is the *source* of an async H2D while +the worker executes that step, so the reader must not release the slot at +unpickle time. Release is tied to **the lifetime of the slot's tensor +storage**, not a fixed "next dequeue" schedule — needed for callers like +`prompt_embeds` that retain the tensor across many `dequeue` calls during +chunked prefill, where a fixed schedule would let the writer reclaim (and +corrupt) the slot while the worker was still reading it. `get_tensor` builds +the tensor via `torch.from_numpy` over a numpy array wrapping the slot; the +tensor's storage holds the only reference to that array, so a +`weakref.finalize` on it fires only once every tensor, view/slice and wrapper +(e.g. `Parameter`) sharing the storage has been freed. Once queued, the +pinned fast path (§2.4) adds a second gate: the H2D is a true async DMA that +can outlive `execute_model`, so "the storage was freed" alone isn't +sufficient. +`flush_releases` records a CUDA event on the compute stream after each +step's H2D and only marks a slot done once that event completes +(non-blocking `event.query()`; not-yet-done just waits one more dequeue). +Unpinned, `cudaMemcpyAsync` from pageable memory already stages +synchronously, so the slot releases immediately. Either way, the writer +requires every reader's done flag before reusing a slot. + +> **Note.** Assumes the multimodal H2D is issued on the worker's +> current/default compute stream (true today: mm inputs are copied eagerly, outside the decode CUDA +> graph). If a future vLLM issues that copy on a dedicated side stream, the event +> must be recorded at the copy site rather than at `flush_releases`. + +### 2.4 Pinning — the zero-copy trap + +A zero-copy view over shared memory the reader process has never touched makes +the subsequent `cudaMemcpyAsync` pay **first-touch page faults on ~50k pages plus +pageable-copy staging** — the same tax the current oob path pays on its pageable +ZMQ frame. Fix: each reader lazily `cudaHostRegister`s the whole arena mapping +once (~1 s, first use), after which every H2D from the arena is a **pinned-memory +DMA** (~10 ms for 192 MB). Processes without a CUDA context skip registration +silently. (Because the arena mapping is stable and reused, this one-time pin +amortizes across all future images — an option the transient per-message ZMQ +frame does not have.) + +## 3. Configuration + +A single CLI flag toggles the arena; the slot count/size and divert threshold +are internal constants (no env vars). + +| Flag (`ParallelConfig` field) | Default | Meaning | +| --- | --- | --- | +| `--enable-shm-tensor-arena` / `--no-enable-shm-tensor-arena` (`enable_shm_tensor_arena`) | **off** | Opt in to route large CPU tensors through the arena (reserves slots in `/dev/shm`). Off = the out-of-band `_reduce_tensor` path only, identical to stock behavior. | + +Internal constants in `shm_tensor_arena.py`: **8 slots × 256 MB**; tensors larger +than a slot, or smaller than the **8 MB** divert threshold, take the out-of-band +`_reduce_tensor` path. + +## 4. Validation + +> **Baseline caveat.** The A/B numbers below were collected against the *in-band +> pickle* path that predates `_reduce_tensor`. They quantify the arena vs. that +> older baseline (an upper bound), **not** the arena vs. the current out-of-band +> path. The current-`main` oob baseline would land between the "stock" and +> "arena + pinning" rows; measuring exactly where is the open item. + +1. **Unit tests** (`tests/distributed/test_shm_broadcast.py`): byte-exact + zero-copy round-trips (incl. bf16/fp8), slot lifecycle (no reuse until every + reader releases; exhaustion / oversize / non-contiguous fall back to the + out-of-band path), event-gated release on the pinned path, `_ArenaPickler` ⊕ + `_reduce_tensor` composition, and an end-to-end `MessageQueue` broadcast + through forked processes. (Historical microbenchmark: enqueue of a 199 MB + payload took **66.7 ms** vs ~1275 ms for the old in-band serialize.) +2. **Same-seed A/B** (baseline = in-band pickle, pre-`_reduce_tensor`), interactive + multimodal workload, uncapped input images, TP=4 ×2 workers on one 8-GPU node, + low qps so individual images decompose cleanly (~1.2k aligned requests): + + | TTFT (ms) | p50 | p90 | p99 | max | >1 s | >1.5 s | + | --- | ---: | ---: | ---: | ---: | ---: | ---: | + | in-band pickle (old baseline) | 93 | 401 | 1321 | 2139 | 30 | 9 | + | + bigger MQ chunk (config lever) | 91 | 368 | 1287 | 2027 | 22 | 5 | + | **arena + pinning** | **89** | **241** | **862** | **1375** | **7** | **0** | + + p50 unchanged → `reducer_override` adds no measurable overhead when no large + tensor is present. +3. **Per-image stall** (CUDA-level all-GPU idle gap around the pixel H2D), vs the + in-band baseline: 192 MB image 1318 ms → 681 ms (arena only) → **351 ms class** + (arena + pinning); vision-encoder compute for the same image is ~80 ms. +4. **Outstanding — arena vs. out-of-band on `main`.** Re-run the same seeded A/B + with the baseline set to current `main` (i.e. `_reduce_tensor` on, arena off + via `--no-enable-shm-tensor-arena`) vs. arena on, to quantify the incremental N→1 + copy + pinning gain at TP=4 (and ideally TP=2/8 to show the copy-count scaling). + +## 5. Limitations and future work + +- **Value is TP- and size-dependent.** The arena helps only when a queue has + ≥2 node-local readers *and* carries tensors ≥ `MIN_MB`; at TP=1, or for small + payloads, it adds nothing and everything takes the out-of-band path. +- **Shared-memory reservation.** The arena reserves `slots × slot_bytes` of + `/dev/shm` (default 2 GB, lazily paged). Like `ShmRingBuffer`, creation is + guarded by a free-space check (`check_shm_free_space`), so an undersized + `/dev/shm` fails fast with a clear error instead of a `SIGBUS` the first + time a tensor is copied into pages beyond tmpfs capacity. +- **Relationship to `_reduce_tensor`.** This is strictly additive: the arena is + an opt-in fast path for the large-multimodal-tensor case; declining it (or + disabling via env) reverts to the merged out-of-band behavior. +- **Relationship to existing mm tensor shm caching.** vLLM already has a + separate shared-memory caching path for multimodal tensors; this arena is + not yet unified with it (see the experimental note in §1). Reconciling the + two is follow-up work. +- **Slot release granularity**: releases are tied to the slot's tensor + storage being freed and, on the pinned path, further gated on a per-slot + CUDA event recorded after the consuming H2D (§2.3), which closes both the + multi-step-retention and the async-DMA reuse windows. Bursts deeper than + the slot count safely fall back to the out-of-band path when the arena is + exhausted. +- **Fallback observability**: arena exhaustion (no free slot) is a + rate-limited log line today; a counter metric would be better. An + oversize tensor (bigger than a slot) falls back silently, with no log at + all -- worth adding if that path turns out to matter in practice. +- **Scope**: the arena activates only when every queue reader is node-local. + Remote readers (multi-node PP/TP) keep the existing socket path. diff --git a/tests/distributed/test_shm_broadcast.py b/tests/distributed/test_shm_broadcast.py index d28732261a..4984a2480a 100644 --- a/tests/distributed/test_shm_broadcast.py +++ b/tests/distributed/test_shm_broadcast.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import dataclasses import gc import io import pickle @@ -17,7 +18,7 @@ import pytest import torch import torch.distributed as dist -from vllm.distributed.device_communicators import shm_broadcast +from vllm.distributed.device_communicators import shm_broadcast, shm_tensor_arena from vllm.distributed.device_communicators.shm_broadcast import ( MessageQueue, ShmRingBuffer, @@ -25,6 +26,10 @@ from vllm.distributed.device_communicators.shm_broadcast import ( _reduce_tensor, check_shm_free_space, ) +from vllm.distributed.device_communicators.shm_tensor_arena import ( + ShmTensorArena, + _ArenaPickler, +) from vllm.distributed.utils import StatelessProcessGroup from vllm.utils.network_utils import get_open_port from vllm.utils.system_utils import update_environment_variables @@ -38,15 +43,18 @@ def get_arrays(n: int, seed: int = 0) -> list[np.ndarray]: return [np.random.randint(1, 100, i) for i in sizes] -def distributed_run(fn, world_size, timeout=60): +def distributed_run(fn, world_size, timeout=60, mp_context=None): """Run a function in multiple processes with proper error handling. Args: fn: Function to run in each process world_size: Number of processes to spawn timeout: Maximum time in seconds to wait for processes (default: 60) + mp_context: Optional multiprocess context (e.g. ``mp.get_context("spawn")``) + for workers that must not inherit the parent's thread/CUDA state. """ + ctx = mp_context or mp number_of_processes = world_size processes = [] for i in range(number_of_processes): @@ -57,7 +65,7 @@ def distributed_run(fn, world_size, timeout=60): env["LOCAL_WORLD_SIZE"] = str(number_of_processes) env["MASTER_ADDR"] = "localhost" env["MASTER_PORT"] = "12345" - p = mp.Process(target=fn, args=(env,)) + p = ctx.Process(target=fn, args=(env,)) processes.append(p) p.start() @@ -765,6 +773,20 @@ def test_shm_ring_buffer_creation_checks_free_space(): ShmRingBuffer(n_reader=1, max_chunk_bytes=24 * 1024 * 1024, max_chunks=10) +def test_shm_tensor_arena_creation_checks_free_space(): + """Like ShmRingBuffer, arena creation must fail fast with a clear error + on an undersized /dev/shm instead of discovering the shortfall via a + SIGBUS the first time a tensor is copied into pages past tmpfs capacity.""" + with ( + mock.patch.object( + shm_broadcast.shutil, "disk_usage", return_value=_fake_disk_usage(1 << 20) + ), + mock.patch.object(shm_broadcast.os.path, "isdir", return_value=True), + pytest.raises(RuntimeError, match="Insufficient space"), + ): + ShmTensorArena(n_reader=1, slot_bytes=256 << 20, n_slots=8) + + def test_remote_subscribe_addr_unique_concurrent_writers( monkeypatch: pytest.MonkeyPatch, ): @@ -828,3 +850,497 @@ def test_remote_subscribe_addr_unique_concurrent_writers( for q in queues: q.remote_socket.close(linger=0) + + +# --------------- ShmTensorArena (zero-copy tensor arena) tests --------------- + + +def _make_arena(n_reader: int = 2, slot_bytes: int = 4 << 20, n_slots: int = 3): + """Writer arena plus attached per-reader arenas (same process).""" + writer = ShmTensorArena(n_reader, slot_bytes, n_slots) + readers = [ShmTensorArena(*writer.handle(), reader_rank=i) for i in range(n_reader)] + return writer, readers + + +def _get_view(reader: ShmTensorArena, idx: int, ref: torch.Tensor) -> torch.Tensor: + return reader.get_tensor( + idx, ref.numel() * ref.element_size(), ref.dtype, tuple(ref.shape) + ) + + +def _drain(reader: ShmTensorArena) -> None: + """Flush a reader's releases to completion. On the pinned path the first + flush defers the release behind an H2D-completion event; the second flush + retires it once the event has fired.""" + reader.flush_releases() + if reader._deferred_releases: + torch.accelerator.synchronize() + reader.flush_releases() + assert not reader._pending_release + assert not reader._deferred_releases + + +def test_arena_zero_copy_roundtrip(): + writer, (r0, r1) = _make_arena(n_reader=2) + src = torch.randn(400, 512) + idx = writer.write_tensor(src) + assert idx is not None + a = _get_view(r0, idx, src) + b = _get_view(r1, idx, src) + assert torch.equal(a, src) + assert torch.equal(b, src) + assert a.dtype == src.dtype and a.shape == src.shape + # Readers alias one shared mapping: a write through one reader's view is + # visible through the other's (this is what "zero-copy" means here). + a[0, 0] = 12345.0 + assert b[0, 0].item() == 12345.0 + del a, b + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn]) +def test_arena_dtype_roundtrip(dtype): + """Dtypes numpy doesn't recognize traverse the arena's uint8 view path.""" + writer, (reader,) = _make_arena(n_reader=1) + src = torch.randn(256, 256).to(dtype) + idx = writer.write_tensor(src) + assert idx is not None + got = _get_view(reader, idx, src) + assert got.dtype == dtype + assert torch.equal(got.view(torch.uint8), src.view(torch.uint8)) + del got + + +def test_arena_slot_lifecycle(): + """A slot must not be reusable until EVERY reader has released it — + the writer overwriting a slot a reader still consumes would corrupt data. + Release is tied to the returned tensor's storage lifetime (`get_tensor`): + a reader that still holds its view keeps the slot reserved even across + repeated `flush_releases()` calls, not just until the "next" one.""" + writer, readers = _make_arena(n_reader=2, n_slots=3) + t = torch.ones(1000) + first = writer.write_tensor(t) + views = [_get_view(r, first, t) for r in readers] + # Fill the remaining slots; the arena is now exhausted. + assert all(writer.write_tensor(t) is not None for _ in range(2)) + assert writer.write_tensor(t) is None # exhausted -> caller falls back + + # Flushing while both readers still hold their view releases nothing: + # the slot is only queued once the returned tensor's storage is freed. + readers[0].flush_releases() + readers[1].flush_releases() + assert writer.write_tensor(t) is None + + # Dropping one reader's view queues its release; still not enough. + views[0] = None + _drain(readers[0]) + assert writer.write_tensor(t) is None + # Once every reader has dropped (and released) its view, it's reused. + views[1] = None + _drain(readers[1]) + assert writer.write_tensor(t) == first + + +def test_arena_slot_not_released_while_tensor_retained_across_multiple_flushes(): + """A caller retaining the tensor across many `flush_releases()` calls + (e.g. `prompt_embeds` across chunked-prefill steps) must keep its slot + reserved the whole time, not just until the reader's next flush.""" + writer, (reader,) = _make_arena(n_reader=1, n_slots=2) + t = torch.ones(1000) + idx = writer.write_tensor(t) + retained = _get_view(reader, idx, t) + # Several flush cycles pass while the caller still holds `retained`. + for _ in range(5): + reader.flush_releases() + assert reader._pending_release == [] + # The slot is still reserved: only the arena's other slot is available... + assert writer.write_tensor(t) is not None + assert writer.write_tensor(t) is None # ...and now the arena is exhausted. + # ...and the bytes the caller is reading are unchanged throughout. + assert torch.equal(retained, t) + # Only once the caller drops its reference is the slot released. + del retained + _drain(reader) + assert writer.write_tensor(t) == idx + + +def test_arena_slot_not_released_while_view_retained(): + """A view/slice (or wrapper) sharing the returned tensor's storage keeps + the slot reserved even after the original tensor object is dropped.""" + writer, (reader,) = _make_arena(n_reader=1, n_slots=1) + t = torch.arange(1000, dtype=torch.float32) + idx = writer.write_tensor(t) + view = _get_view(reader, idx, t)[10:20].unsqueeze(0) + param = torch.nn.Parameter(view, requires_grad=False) + del view + reader.flush_releases() + assert reader._pending_release == [] + assert writer.write_tensor(t) is None + assert torch.equal(param[0], t[10:20]) + del param + assert reader._pending_release == [idx] + _drain(reader) + assert writer.write_tensor(t) == idx + + +def test_arena_release_survives_intervening_flush_from_another_tensor(): + """`flush_releases()` must clear `_pending_release` in place, never + rebind it: a finalizer registered before a rebind would append into the + now-orphaned old list and permanently leak its slot. Exercises the + interleaving a long-lived `prompt_embeds` tensor sees once a + shorter-lived tensor through the same reader triggers a flush first.""" + writer, (reader,) = _make_arena(n_reader=1, n_slots=2) + t = torch.ones(1000) + # `a` is long-lived: its finalizer is registered first, while + # _pending_release is still the arena's original (empty) list object. + idx_a = writer.write_tensor(t) + a = _get_view(reader, idx_a, t) + # `b` is short-lived and its finalizer is registered against that SAME + # list object (nothing has flushed yet). + idx_b = writer.write_tensor(t) + b = _get_view(reader, idx_b, t) + del b + assert reader._pending_release == [idx_b] + # This flush must not orphan `a`'s still-pending finalizer. `_drain`, + # not a bare `flush_releases()` call: on the pinned fast path (real + # CUDA present) the release is gated behind an event and may need a + # second flush once it fires -- see `flush_releases`. + _drain(reader) + assert writer.write_tensor(t) == idx_b # b's slot correctly reused. + # `a` is dropped well after that intervening flush. + del a + assert reader._pending_release == [idx_a] + _drain(reader) + assert writer.write_tensor(t) == idx_a # a's slot must NOT be leaked. + + +def test_arena_oversize_falls_back(): + writer, _ = _make_arena(n_reader=1, slot_bytes=1 << 20, n_slots=2) + big = torch.empty((1 << 20) + 4096, dtype=torch.uint8) + assert writer.write_tensor(big) is None + + +def _dumps_arena(obj, arena: ShmTensorArena) -> tuple[bytes, list, list[int]]: + """Pickle `obj` the same way `MessageQueue.enqueue` does when an arena is + attached: arena diversion first (reducer_override), then the tensor + dispatch table, with out-of-band buffers >= 1MiB. Also returns the slot + indices the arena actually wrote to, in write order.""" + buffers = [] + + def callback(buf: pickle.PickleBuffer) -> bool: + raw = buf.raw() + if raw.nbytes < 1024 * 1024: + return True + buffers.append(raw) + return False + + written: list[int] = [] + orig_write_tensor = arena.write_tensor + + def _spy_write_tensor(t): + idx = orig_write_tensor(t) + if idx is not None: + written.append(idx) + return idx + + bio = io.BytesIO() + pickler = _ArenaPickler(bio, arena, buffer_callback=callback) + pickler.dispatch_table = {torch.Tensor: _reduce_tensor} + with mock.patch.object(arena, "write_tensor", side_effect=_spy_write_tensor): + pickler.dump(obj) + return bio.getvalue(), buffers, written + + +def test_arena_pickler_composes(monkeypatch): + """Large contiguous tensors are diverted into the arena; everything the + arena declines falls through to `_reduce_tensor` unchanged.""" + writer, (reader,) = _make_arena(n_reader=1) + monkeypatch.setattr(shm_tensor_arena, "_ARENA_MIN_BYTES", 1 << 20) + monkeypatch.setitem( + shm_tensor_arena._TENSOR_ARENAS, reader.shared_memory.name, reader + ) + big = torch.randn(1024, 1024) # 4MiB -> diverted into the arena + small = torch.randn(16, 16) # 1KiB -> falls through to _reduce_tensor + data, buffers, written = _dumps_arena({"big": big, "small": small}, writer) + # The diverted tensor's bytes are in the arena, not the pickle stream. + assert len(data) + sum(b.nbytes for b in buffers) < big.numel() * 4 + out = pickle.loads(data, buffers=buffers) + assert torch.equal(out["big"], big) + assert torch.equal(out["small"], small) + # "big" is a zero-copy view of the reader's slot; "small" is not. + (idx,) = written + nbytes = big.numel() * big.element_size() + slot_ptr = torch.frombuffer(reader._slot(idx, nbytes), dtype=torch.uint8).data_ptr() + assert out["big"].data_ptr() == slot_ptr + assert out["small"].data_ptr() != slot_ptr + # Not yet queued for release: `out["big"]` is still alive (see get_tensor). + assert reader._pending_release == [] + del out + assert reader._pending_release == [idx] + _drain(reader) + + +def test_arena_pickler_noncontig_falls_through(monkeypatch): + writer, (reader,) = _make_arena(n_reader=1) + monkeypatch.setattr(shm_tensor_arena, "_ARENA_MIN_BYTES", 1 << 20) + monkeypatch.setitem( + shm_tensor_arena._TENSOR_ARENAS, reader.shared_memory.name, reader + ) + nc = torch.randn(2048, 1024)[:, ::2] # non-contiguous, above threshold + assert not nc.is_contiguous() + data, buffers, written = _dumps_arena(nc, writer) + out = pickle.loads(data, buffers=buffers) + assert torch.equal(out, nc) + # The arena never saw it: no slot consumed on the reader. + assert written == [] + assert reader._pending_release == [] + del out + + +@pytest.mark.parametrize("case", ["requires_grad", "conj"]) +def test_arena_pickler_excludes_requires_grad_and_conj(monkeypatch, case: str): + """Tensors `_reduce_tensor` excludes (requires_grad, or an exotic + conjugate-bit view that can't alias as uint8) must fall through the + arena the same way, even when every other criterion qualifies.""" + writer, (reader,) = _make_arena(n_reader=1) + monkeypatch.setattr(shm_tensor_arena, "_ARENA_MIN_BYTES", 1 << 20) + monkeypatch.setitem( + shm_tensor_arena._TENSOR_ARENAS, reader.shared_memory.name, reader + ) + if case == "requires_grad": + tensor = torch.randn(1024, 1024, requires_grad=True) # 4MiB, above threshold + else: + tensor = torch.randn(600, 600, dtype=torch.complex64).conj() # ~2.9MiB + assert tensor.is_contiguous() + assert tensor.numel() * tensor.element_size() >= shm_tensor_arena._ARENA_MIN_BYTES + + data, buffers, written = _dumps_arena(tensor, writer) + # Declined by the arena -- no slot consumed -- and handled by + # _reduce_tensor's identical exclusion instead. + assert written == [] + out = pickle.loads(data, buffers=buffers) + assert torch.equal(out, tensor) + assert out.requires_grad == tensor.requires_grad + + +def test_arena_pickler_excludes_tensor_subclasses(monkeypatch): + """`torch.nn.Parameter` is diverted via its own `_rebuild_arena_parameter` + hook, which wraps an intermediate tensor it keeps no Python reference to. + The slot must stay reserved while the Parameter is alive -- the + `_pending_release` assertions below check this even though + `type(out) is torch.nn.Parameter` alone would not.""" + writer, (reader,) = _make_arena(n_reader=1) + monkeypatch.setattr(shm_tensor_arena, "_ARENA_MIN_BYTES", 1 << 20) + monkeypatch.setitem( + shm_tensor_arena._TENSOR_ARENAS, reader.shared_memory.name, reader + ) + # requires_grad=False to isolate the subclass check from the (already + # separately tested) requires_grad exclusion. + param = torch.nn.Parameter(torch.randn(1024, 1024), requires_grad=False) + assert type(param) is not torch.Tensor + assert param.is_contiguous() + assert param.numel() * param.element_size() >= shm_tensor_arena._ARENA_MIN_BYTES + + data, buffers, written = _dumps_arena(param, writer) + # The Parameter itself is diverted directly (_ARENA_REBUILD_FNS routes + # it to _rebuild_arena_parameter) -- Parameter.__reduce_ex__ never runs, + # so there's no separate "inner" tensor diversion to speak of. + assert written != [] + out = pickle.loads(data, buffers=buffers) + # The outer Parameter's type identity survives. + assert type(out) is torch.nn.Parameter + assert torch.equal(out, param) + assert out.requires_grad == param.requires_grad + # Not yet queued for release: `out` still shares the slot's storage. + assert reader._pending_release == [] + del out + assert reader._pending_release == written + _drain(reader) + + +def test_arena_unregisters_pinned_memory_on_del(monkeypatch): + """A reader that cudaHostRegister-ed its mapping must cudaHostUnregister + the same pointer before the mapping is closed; an unpinned arena must not + unregister memory it never registered.""" + writer, (reader,) = _make_arena(n_reader=1) + cudart = mock.MagicMock() + cudart.cudaHostRegister.return_value = 0 + cudart.cudaHostUnregister.return_value = 0 + monkeypatch.setattr( + shm_tensor_arena, + "current_platform", + SimpleNamespace(is_cuda_alike=lambda: True, cudart=lambda: cudart), + ) + reader._ensure_pinned() + assert reader._pinned + ptr = cudart.cudaHostRegister.call_args.args[0] + del reader + cudart.cudaHostUnregister.assert_called_once_with(ptr) + del writer # never pinned -> no unregister + cudart.cudaHostUnregister.assert_called_once() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_arena_event_gated_release(): + """On the pinned path a slot release is gated on an H2D-completion CUDA + event: the writer must not see the reader's done flag until the async DMA + sourced from the slot has been retired. Release is first queued once the + caller drops the CPU-side view -- the realistic pattern: the view is only + needed as the H2D source, while the GPU tensor is what's retained for the + rest of the step.""" + writer, (reader,) = _make_arena(n_reader=1) + reader._ensure_pinned() + if not reader._pinned: + pytest.skip("cudaHostRegister unavailable in this environment") + src = torch.randn(512, 1024) + idx = writer.write_tensor(src) + view = _get_view(reader, idx, src) + dev = view.to("cuda", non_blocking=True) + # Not yet queued: this test still holds `view`. + assert reader._pending_release == [] + del view + assert reader._pending_release == [idx] + reader.flush_releases() + # Deferred behind the event, not applied eagerly. + assert len(reader._deferred_releases) == 1 + with reader._meta(idx) as meta: + assert meta[1] == 0 + torch.accelerator.synchronize() + reader.flush_releases() + assert reader._deferred_releases == [] + with reader._meta(idx) as meta: + assert meta[1] == 1 + assert torch.equal(dev.cpu(), src) + del dev + + +def test_tensor_arena_registry_is_weak_and_survives_shutdown(): + """`_TENSOR_ARENAS` must not keep an arena alive past its owning + MessageQueue, but also must not be cleared by `shutdown()` -- which can + run on a different thread than an in-flight `dequeue()` still + reconstructing a tensor from that arena. A weak-value registry gets + both: the entry survives `shutdown()` and vanishes only once + unreferenced.""" + writer_mq = MessageQueue(n_reader=1, n_local_reader=1, enable_shm_tensor_arena=True) + handle = writer_mq.export_handle() + reader_mq = MessageQueue.create_from_handle(handle, rank=0) + + name = reader_mq.tensor_arena.shared_memory.name + assert shm_tensor_arena._TENSOR_ARENAS[name] is reader_mq.tensor_arena + + # shutdown() must NOT remove the entry: that would race with a + # concurrent dequeue() still reconstructing a tensor from this arena. + reader_mq.shutdown() + assert shm_tensor_arena._TENSOR_ARENAS[name] is reader_mq.tensor_arena + + # Once nothing else references the arena, the registry entry vanishes + # on its own. + reader_mq.tensor_arena = None + assert name not in shm_tensor_arena._TENSOR_ARENAS + + writer_mq.shutdown() + + +def test_create_from_handle_tolerates_arena_attach_failure(): + """If a reader's `ShmTensorArena.__init__` can't attach to the writer's + segment (`FileNotFoundError`, e.g. deserialized on a different node), + `create_from_handle` must not crash -- same deferred-failure convention + as the sibling `ShmRingBuffer` path.""" + writer_mq = MessageQueue(n_reader=1, n_local_reader=1, enable_shm_tensor_arena=True) + handle = writer_mq.export_handle() + n_reader, slot_bytes, n_slots, _real_name = handle.tensor_arena_handle + bogus_handle = dataclasses.replace( + handle, + tensor_arena_handle=( + n_reader, + slot_bytes, + n_slots, + "vllm-arena-does-not-exist", + ), + ) + + reader_mq = MessageQueue.create_from_handle(bogus_handle, rank=0) + + # Attach failed silently: the arena object exists but never got its + # shared_memory attribute, and was never registered. + assert reader_mq.tensor_arena is not None + assert not hasattr(reader_mq.tensor_arena, "shared_memory") + assert "vllm-arena-does-not-exist" not in shm_tensor_arena._TENSOR_ARENAS + + reader_mq.shutdown() + writer_mq.shutdown() + + +@worker_fn_wrapper +def worker_fn_arena_broadcast(): + rank = dist.get_rank() + writer_rank = 0 + if rank == writer_rank: + message_queue = MessageQueue( + 1, + 1, + local_reader_ranks=[1], + max_chunk_bytes=8 * 1024 * 1024, + enable_shm_tensor_arena=True, + ) + handles = [message_queue.export_handle()] + dist.broadcast_object_list(handles, src=writer_rank) + else: + handles = [None] + dist.broadcast_object_list(handles, src=writer_rank) + message_queue = MessageQueue.create_from_handle(handles[0], rank) + message_queue.wait_until_ready() + + torch.manual_seed(42) + payload = { + # 16MiB: above the arena divert threshold (8MiB) -> arena slot. + "huge": torch.randn(2048, 2048), + # 2MiB: declined by the arena -> out-of-band _reduce_tensor path. + "mid": torch.randn(1024, 512), + } + + if rank == writer_rank: + assert message_queue.tensor_arena is not None + with mock.patch( + "vllm.distributed.device_communicators.shm_broadcast._reduce_tensor", + wraps=_reduce_tensor, + ) as wrapped_reduce: + message_queue.enqueue(payload) + # "huge" was diverted into the arena before the dispatch table was + # consulted; only "mid" went through _reduce_tensor. + assert wrapped_reduce.call_count == 1 + else: + received = message_queue.dequeue(timeout=30) + assert torch.equal(received["huge"], payload["huge"]) + assert torch.equal(received["mid"], payload["mid"]) + # The huge tensor is a zero-copy view into the arena's shared memory + # (not a copy received over the transport). Its slot is not yet + # queued for release while `received` still holds the reference. + arena = message_queue.tensor_arena + assert arena is not None + assert arena._pending_release == [] + huge_ptr = received["huge"].data_ptr() + arena_base = torch.frombuffer( + arena.shared_memory.buf, dtype=torch.uint8 + ).data_ptr() + assert arena_base <= huge_ptr < arena_base + arena.total_bytes + # Dropping the last reference to the slot's storage queues it for + # release (see get_tensor). + del received + assert len(arena._pending_release) == 1 + + dist.barrier() + print(f"arena broadcast passed the test! Rank {rank}") + + +def test_arena_broadcast(): + # Spawn (not fork): by the time this test runs, the pytest process may be + # multi-threaded / CUDA-initialized (earlier GPU tests), and a forked + # child running the arena's slot memcpy can deadlock on inherited lock + # state. Spawned workers start clean; the larger timeout absorbs their + # interpreter + import startup. + distributed_run( + worker_fn_arena_broadcast, + 2, + timeout=180, + mp_context=mp.get_context("spawn"), + ) diff --git a/tools/pre_commit/check_forbidden_imports.py b/tools/pre_commit/check_forbidden_imports.py index 531f6fb8a2..d80a8e1fd2 100644 --- a/tools/pre_commit/check_forbidden_imports.py +++ b/tools/pre_commit/check_forbidden_imports.py @@ -56,6 +56,7 @@ CHECK_IMPORTS = { "vllm/distributed/device_communicators/all_reduce_utils.py", "vllm/distributed/device_communicators/shm_broadcast.py", "vllm/distributed/device_communicators/shm_object_storage.py", + "vllm/distributed/device_communicators/shm_tensor_arena.py", "vllm/distributed/weight_transfer/ipc_engine.py", "vllm/distributed/weight_transfer/clients.py", "vllm/model_executor/model_loader/weight_cache/protocol.py", diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index 37affbaf9e..6960661e0c 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -235,6 +235,15 @@ class ParallelConfig: disable_custom_all_reduce: bool = False """Disable the custom all-reduce kernel and fall back to NCCL.""" + enable_shm_tensor_arena: bool = False + """Route large CPU tensors (e.g. multimodal ``pixel_values``) in the + engine→worker broadcast through a zero-copy shared-memory arena instead of + sending a copy to every reader. **Experimental**: may be reconciled with + vLLM's existing multimodal shm tensor caching in a future release. Opt-in; + most beneficial for multimodal serving with large images and tensor + parallelism. Active only when all queue readers are node-local; reserves + slots in ``/dev/shm``.""" + enable_elastic_ep: bool = False """Enable elastic expert parallelism with stateless NCCL groups for DP/EP.""" elastic_ep_max_dp_size: int = Field(default=None, ge=1) # type: ignore[assignment] @@ -901,6 +910,7 @@ class ParallelConfig: "nnodes", "max_parallel_loading_workers", "disable_custom_all_reduce", + "enable_shm_tensor_arena", "ray_workers_use_nsight", "ray_runtime_env", "placement_group", diff --git a/vllm/distributed/device_communicators/shm_broadcast.py b/vllm/distributed/device_communicators/shm_broadcast.py index da2288f77b..5656e8aaa1 100644 --- a/vllm/distributed/device_communicators/shm_broadcast.py +++ b/vllm/distributed/device_communicators/shm_broadcast.py @@ -61,6 +61,10 @@ SPINLOOP_TIMEOUT_SECONDS = 0.1 if TYPE_CHECKING: from _typeshed import SizedBuffer + from vllm.distributed.device_communicators.shm_tensor_arena import ( + ShmTensorArena, + ) + VLLM_RINGBUFFER_WARNING_INTERVAL = envs.VLLM_RINGBUFFER_WARNING_INTERVAL # Cap on how long an idle reader parks before re-reading the authoritative SHM # written-flag. Bounds lost-notify recovery latency to ~5s while the periodic @@ -471,6 +475,7 @@ class Handle: local_notify_addr: str | None = None remote_subscribe_addr: str | None = None remote_addr_ipv6: bool = False + tensor_arena_handle: tuple[int, int, int, str] | None = None def _close_zmq(context: zmq.Context, sockets: list[zmq.Socket]) -> None: @@ -490,6 +495,7 @@ class MessageQueue: max_chunk_bytes: int = 1024 * 1024 * 24, max_chunks: int = 10, connect_ip: str | None = None, + enable_shm_tensor_arena: bool = False, ): if local_reader_ranks is None: local_reader_ranks = list(range(n_local_reader)) @@ -501,12 +507,22 @@ class MessageQueue: self.shutting_down = False context = Context() + self.tensor_arena: ShmTensorArena | None = None if n_local_reader > 0: # for local readers, we will: # 1. create a shared memory ring buffer to communicate small data # 2. create a publish-subscribe socket to communicate large data self.buffer = ShmRingBuffer(n_local_reader, max_chunk_bytes, max_chunks) + # Zero-copy arena for large CPU tensors (see ShmTensorArena). + # Local readers only: remote readers receive the pickled bytes + # over a socket and cannot map the arena, so the substitution + # would break them. + if enable_shm_tensor_arena and n_remote_reader == 0: + from vllm.distributed.device_communicators import shm_tensor_arena + + self.tensor_arena = shm_tensor_arena.ShmTensorArena(n_local_reader) + # XPUB is very similar to PUB, # except that it can receive subscription messages # to confirm the number of subscribers @@ -567,6 +583,9 @@ class MessageQueue: local_notify_addr=local_notify_addr, remote_subscribe_addr=remote_subscribe_addr, remote_addr_ipv6=remote_addr_ipv6, + tensor_arena_handle=( + self.tensor_arena.handle() if self.tensor_arena is not None else None + ), ) self._close_zmq_on_collect(context) @@ -599,6 +618,7 @@ class MessageQueue: context = Context() + self.tensor_arena = None if rank in handle.local_reader_ranks: assert handle.buffer_handle is not None self.buffer = ShmRingBuffer(*handle.buffer_handle) @@ -607,6 +627,23 @@ class MessageQueue: self._is_local_reader = True self._is_remote_reader = False + arena_handle = handle.tensor_arena_handle + if arena_handle is not None: + from vllm.distributed.device_communicators import shm_tensor_arena + + self.tensor_arena = shm_tensor_arena.ShmTensorArena( + *arena_handle, reader_rank=self.local_reader_rank + ) + # `shared_memory` is unset (not just None) if attaching to + # the writer's segment raised FileNotFoundError -- see + # ShmTensorArena.__init__'s reader branch. Registering a + # broken arena is a no-op deferred failure (same convention + # ShmRingBuffer uses above), not a crash here. + if hasattr(self.tensor_arena, "shared_memory"): + shm_tensor_arena._TENSOR_ARENAS[ + self.tensor_arena.shared_memory.name + ] = self.tensor_arena + self.local_socket = context.socket(SUB) self.local_socket.setsockopt_string(SUBSCRIBE, "") socket_addr = handle.local_subscribe_addr @@ -880,12 +917,21 @@ class MessageQueue: # dispatch table would otherwise shadow them. dispatch_table = dict(copyreg.dispatch_table) dispatch_table[torch.Tensor] = _reduce_tensor + arena = self.tensor_arena with io.BytesIO() as bio: - pickler = pickle.Pickler( - bio, - protocol=pickle.HIGHEST_PROTOCOL, - buffer_callback=oob_callback, - ) + pickler: pickle.Pickler + if arena is not None: + from vllm.distributed.device_communicators.shm_tensor_arena import ( + _ArenaPickler, + ) + + pickler = _ArenaPickler(bio, arena, buffer_callback=oob_callback) + else: + pickler = pickle.Pickler( + bio, + protocol=pickle.HIGHEST_PROTOCOL, + buffer_callback=oob_callback, + ) pickler.dispatch_table = dispatch_table pickler.dump(obj) all_buffers[0] = bio.getvalue() @@ -922,6 +968,9 @@ class MessageQueue: ): """Read from message queue with optional timeout (in seconds)""" if self._is_local_reader: + # See ShmTensorArena.flush_releases for what this retires and why. + if self.tensor_arena is not None: + self.tensor_arena.flush_releases() with self.acquire_read(timeout, indefinite) as buf: overflow = buf[0] == 1 if not overflow: diff --git a/vllm/distributed/device_communicators/shm_tensor_arena.py b/vllm/distributed/device_communicators/shm_tensor_arena.py new file mode 100644 index 0000000000..257cce06b2 --- /dev/null +++ b/vllm/distributed/device_communicators/shm_tensor_arena.py @@ -0,0 +1,397 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Zero-copy shared-memory arena for large CPU tensors (multimodal fast path). + +Problem: `MessageQueue.enqueue` (see `shm_broadcast.py`) pickles the whole +payload, and torch's default tensor pickling copies bytes into the pickle +stream on the writer and back out on every local reader. For a large +`pixel_values` tensor (up to hundreds of MB) this blocks the EngineCore step +loop for ~1s with all GPUs idle. + +Fast path: `_ArenaPickler.reducer_override` intercepts qualifying CPU +tensors during pickling, memcpys them ONCE into a free arena slot, and +pickles a tiny `(arena_name, slot, nbytes, dtype, shape)` stub. Readers +rebuild a zero-copy view of the mapped slot (`torch.frombuffer`). + +Slot lifecycle: single writer, n_reader readers, per-slot metadata +[written_flag, reader0_done, ..., readerN_done] (same protocol as +ShmRingBuffer). A slot is not released when rebuilt: the rebuilt tensor is +the SOURCE of an async H2D copy, and callers like chunked-prefill +`prompt_embeds` retain it across many steps. Release is queued only once the +tensor's storage is freed, i.e. once no tensor or view of the slot remains +(see `get_tensor`), and is further gated on a CUDA event on the pinned fast +path, since the H2D can outlive the step that issued it (see +`flush_releases`). + +The writer NEVER blocks: no free slot (or an oversized tensor) falls back to +the default pickle path, so worst case matches today's behavior. + +Physical memory: slots are paged in lazily, so an arena on a queue that +never carries big tensors costs ~0. Creation is guarded by +`check_shm_free_space` (like `ShmRingBuffer`) so an undersized `/dev/shm` +fails fast instead of a `SIGBUS` mid-copy. +""" + +import pickle +import weakref +from contextlib import suppress +from multiprocessing import shared_memory +from unittest.mock import patch + +import numpy as np +import torch + +from vllm.distributed.device_communicators.shm_broadcast import ( + check_shm_free_space, + memory_fence, +) +from vllm.logger import init_logger +from vllm.platforms import current_platform + +logger = init_logger(__name__) + +# Arena sizing — internal constants (kept off the env/CLI surface). +_ARENA_SLOTS = 8 +_ARENA_SLOT_BYTES = 256 * 1024 * 1024 +# Minimum contiguous CPU-tensor size to divert into the arena; smaller tensors +# take the out-of-band pickle-buffer path. +_ARENA_MIN_BYTES = 8 * 1024 * 1024 + +# Reader-side registry: arena shm name -> attached ShmTensorArena, keyed by +# name rather than a single slot since a process isn't guaranteed to attach +# to only one. Populated by `MessageQueue.create_from_handle`; consumed by +# the `_rebuild_arena_*` unpickle hooks. WEAK-value: a strong reference here +# would keep the arena (and its pinned mapping) alive for the process +# lifetime, and would need `MessageQueue.shutdown` to remove it explicitly -- +# unsafe since shutdown() can run on a different thread than an in-flight +# `dequeue`. Instead an entry just disappears once nothing else references +# that arena. +_TENSOR_ARENAS: weakref.WeakValueDictionary[str, "ShmTensorArena"] = ( + weakref.WeakValueDictionary() +) + + +class ShmTensorArena: + """Slotted shared-memory arena: one writer, n_reader zero-copy readers. + + Memory layout: [slot0 | slot1 | ... | slotN-1 | meta0 | meta1 | ... ] + where each meta is (1 + n_reader) bytes: [written_flag, reader_done...]. + Slot states mirror ShmRingBuffer: + written=0 -> free (never written / being written) + written=1, some reader_done=0 -> in use, cannot reuse + written=1, all reader_done=1 -> consumed, can reuse + """ + + def __init__( + self, + n_reader: int, + slot_bytes: int = _ARENA_SLOT_BYTES, + n_slots: int = _ARENA_SLOTS, + name: str | None = None, + reader_rank: int = -1, + ): + self.n_reader = n_reader + self.slot_bytes = slot_bytes + self.n_slots = n_slots + self.reader_rank = reader_rank # -1 for the writer + self.metadata_size = 1 + n_reader + self.metadata_offset = slot_bytes * n_slots + self.total_bytes = (slot_bytes + self.metadata_size) * n_slots + self._next_slot = 0 + self._fallbacks = 0 + self._pin_attempted = False + self._pinned = False + self._pinned_ptr = 0 + # slots queued for release once the storage of the tensor `get_tensor` + # handed out is freed; flushed at the next dequeue. + self._pending_release: list[int] = [] + # slots whose release is gated on an H2D-completion CUDA event: each + # entry is (event, [slot_idx, ...]). Only populated on the pinned fast + # path, where the source-side DMA is asynchronous (see flush_releases). + self._deferred_releases: list[tuple[torch.cuda.Event, list[int]]] = [] + + if name is None: + self.is_creator = True + # Same guard as ShmRingBuffer: SharedMemory/ftruncate can succeed + # on an undersized /dev/shm, so check up front instead of SIGBUS + # on first write past tmpfs capacity. + check_shm_free_space(self.total_bytes, allocation_name="ShmTensorArena") + self.shared_memory = shared_memory.SharedMemory( + create=True, size=self.total_bytes + ) + assert self.shared_memory.buf is not None + with self.shared_memory.buf[self.metadata_offset :] as meta: + torch.frombuffer(meta, dtype=torch.uint8).fill_(0) + else: + self.is_creator = False + # same resource_tracker workaround as ShmRingBuffer + with patch( + "multiprocessing.resource_tracker.register", + lambda *args, **kwargs: None, + ): + try: + self.shared_memory = shared_memory.SharedMemory(name=name) + assert self.shared_memory.size >= self.total_bytes + except FileNotFoundError: + # deserialized on a different node; arena unused there + pass + + def handle(self): + return (self.n_reader, self.slot_bytes, self.n_slots, self.shared_memory.name) + + def _meta(self, idx: int) -> memoryview: + start = self.metadata_offset + idx * self.metadata_size + assert self.shared_memory.buf is not None + return self.shared_memory.buf[start : start + self.metadata_size] + + def _slot(self, idx: int, nbytes: int) -> memoryview: + start = idx * self.slot_bytes + assert self.shared_memory.buf is not None + return self.shared_memory.buf[start : start + nbytes] + + # ---- writer side ---- + + def write_tensor(self, t: torch.Tensor) -> int | None: + """Copy tensor bytes into a free slot; return slot idx or None + (caller must then fall back to the default pickle path).""" + nbytes = t.numel() * t.element_size() + if nbytes > self.slot_bytes: + return None + try: + # Computed before claiming a slot: exotic tensors (e.g. conjugate + # bit set) raise RuntimeError here, same as `_reduce_tensor`'s + # identical view -- fail before claiming, not after. + src = t.detach().reshape(-1).view(torch.uint8) + except RuntimeError: + return None + memory_fence() + for probe in range(self.n_slots): + idx = (self._next_slot + probe) % self.n_slots + with self._meta(idx) as meta: + free = meta[0] == 0 or sum(meta[1:]) == self.n_reader + if not free: + continue + meta[0] = 0 # claim + slot_mv = self._slot(idx, nbytes) + try: + dst = torch.frombuffer(slot_mv, dtype=torch.uint8, count=nbytes) + dst.copy_(src) + finally: + slot_mv.release() + with self._meta(idx) as meta: + for i in range(1, self.n_reader + 1): + meta[i] = 0 + memory_fence() + meta[0] = 1 + memory_fence() + self._next_slot = (idx + 1) % self.n_slots + return idx + self._fallbacks += 1 + if self._fallbacks == 1 or self._fallbacks % 100 == 0: + logger.info( + "ShmTensorArena: no free slot (%d bytes, %d fallbacks so far); " + "falling back to the out-of-band pickle path.", + nbytes, + self._fallbacks, + ) + return None + + # ---- reader side ---- + + def _ensure_pinned(self): + """Pin the whole arena mapping via cudaHostRegister in THIS process + (lazy, once). Without it the HtoD of a zero-copy tensor pays + first-touch page faults on the tmpfs mapping plus pageable staging + (~hundreds of ms for 200MB); registration allocates+pins the pages + once, making every later HtoD a true DMA. Failure (no CUDA in this + process, etc.) is fine — the copy still works, just slower.""" + if self._pin_attempted: + return + self._pin_attempted = True + try: + if not current_platform.is_cuda_alike(): + return + import ctypes + + buf = self.shared_memory.buf + assert buf is not None + ptr = ctypes.addressof(ctypes.c_char.from_buffer(buf)) + ret = current_platform.cudart().cudaHostRegister(ptr, self.total_bytes, 0) + self._pinned = int(ret) == 0 + if self._pinned: + self._pinned_ptr = ptr + logger.info( + "ShmTensorArena: cudaHostRegister(%d MB) -> %s", + self.total_bytes >> 20, + "pinned" if self._pinned else f"error {int(ret)}", + ) + except Exception as e: + logger.info("ShmTensorArena: host-register skipped: %s", e) + + def _unpin(self): + """Undo the pin from _ensure_pinned via cudaHostUnregister; must run + before the mapping is closed. Failures are ignored — at interpreter + shutdown the CUDA context may already be gone, and the registration + dies with the process anyway.""" + if not self._pinned: + return + self._pinned = False + with suppress(Exception): + current_platform.cudart().cudaHostUnregister(self._pinned_ptr) + + def get_tensor( + self, idx: int, nbytes: int, dtype: torch.dtype, shape: tuple[int, ...] + ) -> torch.Tensor: + """Zero-copy view of a slot as a tensor. The slot is queued for + release once the returned tensor's storage is freed, i.e. once the + caller has dropped it along with any views/slices or wrappers (such + as `torch.nn.Parameter`) sharing that storage.""" + self._ensure_pinned() + # The tensor's storage holds the only reference to `holder` and drops + # it when freed, firing the finalizer. (`weakref.finalize` can't + # target the tensor itself, which views and wrappers don't keep alive, + # nor the memoryview, which doesn't support weakrefs.) + holder = np.frombuffer(self._slot(idx, nbytes), dtype=np.uint8) + # Binds to the current list OBJECT: `flush_releases` must mutate it in + # place, never rebind `self._pending_release`. + weakref.finalize(holder, self._pending_release.append, idx) + return torch.from_numpy(holder).view(dtype).view(shape) + + def _mark_released(self, idxs: list[int]): + """Set THIS reader's done flag on the given slots; a slot becomes + reusable once every reader has done so.""" + for idx in idxs: + with self._meta(idx) as meta: + meta[1 + self.reader_rank] = 1 + memory_fence() + + def _record_release_event(self): + """Record a CUDA event on the current (compute) stream, or return None + if this process has no usable CUDA context. Called from flush_releases, + which runs after the previous step's H2D was enqueued on that same + stream, so the event is ordered strictly after that H2D.""" + try: + event = current_platform.Event() + event.record() + return event + except Exception: + return None + + def flush_releases(self): + """Retire this reader's arena slots that are ready for writer reuse. + + Called on every `MessageQueue.dequeue`. `_pending_release` holds + slots whose tensor storage was just freed (see `get_tensor`). On + the pinned fast path that tensor was the source of an async + `non_blocking=True` H2D that can outlive the collection, so instead + of marking released immediately we record a CUDA event (ordered + after that H2D on the same stream) and move the slot to + `_deferred_releases` until `event.query()` says it's done -- a + not-yet-complete slot just waits one more `dequeue`. Unpinned (or no + CUDA context), the H2D already staged synchronously, so the slot is + safe to release right away. + + A slot only becomes free for the writer (`write_tensor`) once every + reader has marked it released (`_mark_released`). + """ + if self._deferred_releases: + still_pending = [] + for deferred_event, idxs in self._deferred_releases: + if deferred_event.query(): + self._mark_released(idxs) + else: + still_pending.append((deferred_event, idxs)) + self._deferred_releases = still_pending + + if not self._pending_release: + return + # Drain in place with atomic pops, never rebind (see get_tensor): a + # finalizer may append from another thread concurrently. + idxs = [] + while self._pending_release: + idxs.append(self._pending_release.pop()) + + event = self._record_release_event() if self._pinned else None + if event is None: + self._mark_released(idxs) + else: + self._deferred_releases.append((event, idxs)) + + def __del__(self): + if hasattr(self, "shared_memory"): + self._unpin() + try: + self.shared_memory.close() + if self.is_creator: + self.shared_memory.unlink() + except BufferError: + # zero-copy tensor views may still hold exported pointers at + # interpreter shutdown; the mapping dies with the process. + pass + + +def _rebuild_arena_tensor(arena_name, slot_idx, nbytes, dtype_str, shape): + """Unpickle hook: rebuild a tensor as a zero-copy view of an arena slot.""" + arena = _TENSOR_ARENAS[arena_name] + return arena.get_tensor(slot_idx, nbytes, getattr(torch, dtype_str), shape) + + +def _rebuild_arena_parameter(arena_name, slot_idx, nbytes, dtype_str, shape): + """Unpickle hook for `torch.nn.Parameter` (shares the slot's storage).""" + t = _rebuild_arena_tensor(arena_name, slot_idx, nbytes, dtype_str, shape) + return torch.nn.Parameter(t, requires_grad=False) + + +# Types `_ArenaPickler.reducer_override` diverts, mapped to the rebuild hook +# that reconstructs them. Any other type -- including other `torch.Tensor` +# subclasses -- falls through to the normal pickle path, since an arbitrary +# subclass can't be faithfully rebuilt from just the slot's bytes. +_ARENA_REBUILD_FNS = { + torch.Tensor: _rebuild_arena_tensor, + torch.nn.Parameter: _rebuild_arena_parameter, +} + + +class _ArenaPickler(pickle.Pickler): + """Pickler that diverts large contiguous CPU tensors into the arena. + + `reducer_override` is consulted before an object's normal reduction. For + a qualifying tensor (exact type in `_ARENA_REBUILD_FNS`, CPU, strided, + contiguous, not `requires_grad`, above the divert-size threshold -- the + same criteria `_reduce_tensor` uses, plus size) it does one memcpy into a + free slot and returns a tiny rebuild stub. Anything that doesn't qualify, + or finds the arena full, returns `NotImplemented` and falls through to + the normal out-of-band pickle path, exactly as if no arena were attached. + """ + + def __init__(self, file, arena: ShmTensorArena, buffer_callback=None): + super().__init__( + file, + protocol=pickle.HIGHEST_PROTOCOL, + buffer_callback=buffer_callback, + ) + self.arena = arena + + def reducer_override(self, obj): + rebuild_fn = _ARENA_REBUILD_FNS.get(type(obj)) + if ( + rebuild_fn is not None + and obj.device.type == "cpu" + and obj.layout == torch.strided + and obj.is_contiguous() + and not obj.requires_grad + and obj.numel() * obj.element_size() >= _ARENA_MIN_BYTES + ): + idx = self.arena.write_tensor(obj) + if idx is not None: + return ( + rebuild_fn, + ( + self.arena.shared_memory.name, + idx, + obj.numel() * obj.element_size(), + str(obj.dtype).removeprefix("torch."), + tuple(obj.shape), + ), + ) + return NotImplemented diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 5f32ec8deb..31777e266f 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -622,6 +622,7 @@ class EngineArgs: allow_deprecated_quantization: bool = ModelConfig.allow_deprecated_quantization enforce_eager: bool = ModelConfig.enforce_eager disable_custom_all_reduce: bool = ParallelConfig.disable_custom_all_reduce + enable_shm_tensor_arena: bool = ParallelConfig.enable_shm_tensor_arena language_model_only: bool = MultiModalConfig.language_model_only limit_mm_per_prompt: dict[str, int | dict[str, int]] = get_field( MultiModalConfig, "limit_per_prompt" @@ -1274,6 +1275,10 @@ class EngineArgs: "--all2all-backend", **parallel_kwargs["all2all_backend"] ) parallel_group.add_argument("--enable-dbo", **parallel_kwargs["enable_dbo"]) + parallel_group.add_argument( + "--enable-shm-tensor-arena", + **parallel_kwargs["enable_shm_tensor_arena"], + ) parallel_group.add_argument( "--ubatch-size", **parallel_kwargs["ubatch_size"], @@ -2534,6 +2539,7 @@ class EngineArgs: enable_elastic_ep=self.enable_elastic_ep, elastic_ep_max_dp_size=self.elastic_ep_max_dp_size, enable_dbo=self.enable_dbo, + enable_shm_tensor_arena=self.enable_shm_tensor_arena, ubatch_size=self.ubatch_size, dbo_decode_token_threshold=self.dbo_decode_token_threshold, dp_sync_interval=self.dp_sync_interval, diff --git a/vllm/v1/executor/multiproc_executor.py b/vllm/v1/executor/multiproc_executor.py index cad3436854..749f46b428 100644 --- a/vllm/v1/executor/multiproc_executor.py +++ b/vllm/v1/executor/multiproc_executor.py @@ -166,6 +166,7 @@ class MultiprocExecutor(Executor): self.local_world_size, max_chunk_bytes=max_chunk_bytes, connect_ip=mq_connect_ip, + enable_shm_tensor_arena=(self.parallel_config.enable_shm_tensor_arena), ) scheduler_output_handle = self.rpc_broadcast_mq.export_handle() # Create workers