[Core] shm tensor arena for efficient cpu->gpu worker broadcast (#51207)

Signed-off-by: Bolin Sun <bolins@nvidia.com>
Signed-off-by: <bolins@nvidia.com>
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Nick Hill <nickhill123@gmail.com>

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.

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"),
+    )

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",

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",

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:

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

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,

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