From 60cae66f2e218007d1f8942ad0b8a69a4a2bf690 Mon Sep 17 00:00:00 2001 From: Bolin Sun Date: Tue, 7 Jul 2026 15:18:30 -0700 Subject: [PATCH 1/3] [Core] Zero-copy shared-memory arena for large CPU tensors in MessageQueue Large multimodal pixel_values tensors (100-200MB) were pickle-serialized into the broadcast payload on the engine (THPStorage_writeFileRaw) and deserialized on every local TP reader (THPStorage_readFileRaw), blocking the EngineCore step loop ~1s per large image with all GPUs idle. Add ShmTensorArena: a slotted shared-memory region (per-slot reader flags, same protocol as ShmRingBuffer). A Pickler.reducer_override diverts large contiguous CPU tensors into a free slot (single memcpy) and pickles only a (slot, nbytes, dtype, shape) stub; readers rebuild the tensor as a zero-copy torch.frombuffer view of the mapped slot. Slots are released lazily at the reader's next dequeue (worker loop is sequential, so the previous step's HtoD has completed). The writer never blocks: no free slot or oversize tensor falls back to the default in-band pickling. Env knobs: VLLM_SHM_TENSOR_ARENA (default on), _SLOTS (8), _SLOT_MB (256), _MIN_MB (8). Local readers only; disabled when remote readers exist. Unit test: 199MB bf16 roundtrip to 2 forked readers = 66.7ms enqueue (vs ~1275ms pickled), checksums exact, slot reuse + fallback verified. Co-Authored-By: Claude Fable 5 --- .../device_communicators/shm_broadcast.py | 275 +++++++++++++++++- 1 file changed, 272 insertions(+), 3 deletions(-) diff --git a/vllm/distributed/device_communicators/shm_broadcast.py b/vllm/distributed/device_communicators/shm_broadcast.py index dc7e6d151a48..565ee6c63806 100644 --- a/vllm/distributed/device_communicators/shm_broadcast.py +++ b/vllm/distributed/device_communicators/shm_broadcast.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import functools +import io +import os import pickle import sys import threading @@ -344,6 +346,235 @@ def get_metadata(self, current_idx: int): yield buf +# --------------------------------------------------------------------------- +# Zero-copy shared-memory arena for large CPU tensors (multimodal fast path). +# +# Problem: `MessageQueue.enqueue` pickles the whole payload. For multimodal +# models the payload contains raw `pixel_values` tensors (up to hundreds of +# MB), and torch's storage pickling copies every byte into the pickle stream +# on the writer (THPStorage_writeFileRaw) and back out on EVERY local reader +# (THPStorage_readFileRaw). On a TP=4 engine this serialize/deserialize chain +# blocks the EngineCore step loop for ~1s per large image with all GPUs idle. +# +# Fast path: intercept large contiguous CPU tensors during pickling +# (`_ArenaPickler.reducer_override`), memcpy them ONCE into a free slot of a +# shared-memory arena, and pickle only a tiny (slot, nbytes, dtype, shape) +# stub. Readers rebuild the tensor as a zero-copy view of the mapped slot +# (`torch.frombuffer`) — no byte-copy on either side. +# +# Slot lifecycle: single writer, n_reader readers, per-slot metadata +# [written_flag, reader0_done, ..., readerN_done] (same protocol as +# ShmRingBuffer). A reader does NOT release the slot when the tensor is +# rebuilt — the tensor is still consumed (HtoD copy) while the worker +# executes that step. Instead releases are deferred and flushed at the +# reader's NEXT `dequeue` call: the worker loop is sequential, so by then the +# previous step (and its pageable-source HtoD, which completes staging before +# `cudaMemcpyAsync` returns) has finished. +# +# The writer NEVER blocks on the arena: if no slot is free (or the tensor is +# larger than a slot), it falls back to the default pickle path for that +# tensor. Worst case is the status-quo behavior, and deadlock is impossible. +# +# Physical memory: slots are allocated lazily by the kernel (pages are backed +# on first write), so arenas on queues that never carry big tensors cost ~0. +# --------------------------------------------------------------------------- + +VLLM_SHM_TENSOR_ARENA = os.getenv("VLLM_SHM_TENSOR_ARENA", "1") != "0" +VLLM_SHM_TENSOR_ARENA_SLOTS = int(os.getenv("VLLM_SHM_TENSOR_ARENA_SLOTS", "8")) +VLLM_SHM_TENSOR_ARENA_SLOT_MB = int(os.getenv("VLLM_SHM_TENSOR_ARENA_SLOT_MB", "256")) +VLLM_SHM_TENSOR_ARENA_MIN_MB = int(os.getenv("VLLM_SHM_TENSOR_ARENA_MIN_MB", "8")) + +# Reader-side registry: arena shm name -> attached ShmTensorArena. Populated +# by MessageQueue.create_from_handle; consumed by _rebuild_arena_tensor when +# unpickling a tensor stub. +_TENSOR_ARENAS: dict[str, "ShmTensorArena"] = {} + + +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, + n_slots: int, + 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 + # slots whose tensors were handed out by THIS reader and not yet + # released (flushed at the next dequeue). + self._pending_release: list[int] = [] + + if name is None: + self.is_creator = True + 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 + 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 + src = t.detach().reshape(-1).view(torch.uint8) + 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 in-band pickling. Consider raising " + "VLLM_SHM_TENSOR_ARENA_SLOTS/SLOT_MB.", + nbytes, + self._fallbacks, + ) + return None + + # ---- reader side ---- + + 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 released at this + reader's next flush_releases() (called from MessageQueue.dequeue).""" + # NOT a context-manager view: the tensor must keep the mapping alive. + slot_mv = self._slot(idx, nbytes) + t8 = torch.frombuffer(slot_mv, dtype=torch.uint8, count=nbytes) + t = t8.view(dtype).view(shape) + self._pending_release.append(idx) + return t + + def flush_releases(self): + if not self._pending_release: + return + for idx in self._pending_release: + with self._meta(idx) as meta: + meta[1 + self.reader_rank] = 1 + memory_fence() + self._pending_release.clear() + + def __del__(self): + if hasattr(self, "shared_memory"): + 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) + + +class _ArenaPickler(pickle.Pickler): + """Pickler that diverts large contiguous CPU tensors into the arena.""" + + 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): + if ( + isinstance(obj, torch.Tensor) + and obj.device.type == "cpu" + and obj.layout == torch.strided + and obj.is_contiguous() + and obj.numel() * obj.element_size() + >= VLLM_SHM_TENSOR_ARENA_MIN_MB * 1024 * 1024 + ): + idx = self.arena.write_tensor(obj) + if idx is not None: + return ( + _rebuild_arena_tensor, + ( + self.arena.shared_memory.name, + idx, + obj.numel() * obj.element_size(), + str(obj.dtype).removeprefix("torch."), + tuple(obj.shape), + ), + ) + return NotImplemented + + @dataclass class Handle: local_reader_ranks: list[int] = field(default_factory=list) @@ -353,6 +584,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 class MessageQueue: @@ -377,12 +609,24 @@ def __init__( 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 VLLM_SHM_TENSOR_ARENA and n_remote_reader == 0: + self.tensor_arena = ShmTensorArena( + n_local_reader, + VLLM_SHM_TENSOR_ARENA_SLOT_MB * 1024 * 1024, + VLLM_SHM_TENSOR_ARENA_SLOTS, + ) + # XPUB is very similar to PUB, # except that it can receive subscription messages # to confirm the number of subscribers @@ -443,6 +687,9 @@ def __init__( 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 + ), ) logger.debug("vLLM message queue communication handle: %s", self.handle) @@ -458,6 +705,7 @@ def create_from_handle(handle: Handle, rank) -> "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) @@ -466,6 +714,15 @@ def create_from_handle(handle: Handle, rank) -> "MessageQueue": self._is_local_reader = True self._is_remote_reader = False + arena_handle = getattr(handle, "tensor_arena_handle", None) + if arena_handle is not None: + self.tensor_arena = ShmTensorArena( + *arena_handle, reader_rank=self.local_reader_rank + ) + _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 @@ -733,9 +990,15 @@ def oob_callback(buf: PickleBuffer) -> bool: total_bytes += len(raw_buf) + 4 return False - all_buffers[0] = pickle.dumps( - obj, protocol=pickle.HIGHEST_PROTOCOL, buffer_callback=oob_callback - ) + arena = getattr(self, "tensor_arena", None) + if arena is not None: + bio = io.BytesIO() + _ArenaPickler(bio, arena, buffer_callback=oob_callback).dump(obj) + all_buffers[0] = bio.getbuffer() + else: + all_buffers[0] = pickle.dumps( + obj, protocol=pickle.HIGHEST_PROTOCOL, buffer_callback=oob_callback + ) if self.n_local_reader > 0: if total_bytes + len(all_buffers[0]) >= self.buffer.max_chunk_bytes: with self.acquire_write(timeout) as buf: @@ -769,6 +1032,12 @@ def dequeue( ): """Read from message queue with optional timeout (in seconds)""" if self._is_local_reader: + # Release arena slots consumed by the PREVIOUS message: the worker + # loop is sequential, so the previous step (incl. its HtoD of any + # zero-copy tensors) has completed by the time we dequeue again. + arena = getattr(self, "tensor_arena", None) + if arena is not None: + arena.flush_releases() with self.acquire_read(timeout, indefinite) as buf: overflow = buf[0] == 1 if not overflow: From 8bbb960b0e9bdffb05a4cecb548568fca62d517b Mon Sep 17 00:00:00 2001 From: Bolin Sun Date: Tue, 7 Jul 2026 15:18:30 -0700 Subject: [PATCH 2/3] [Core] Pin the ShmTensorArena mapping via cudaHostRegister on readers The zero-copy tensor views the tmpfs mapping; without pinning, the HtoD pays first-touch page faults + pageable staging (~600-800ms observed for a 192MB image). Lazy one-time cudaHostRegister at first get_tensor makes every later HtoD a true DMA. Falls back gracefully when the process has no CUDA context. Co-Authored-By: Claude Fable 5 --- .../device_communicators/shm_broadcast.py | 31 +++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/vllm/distributed/device_communicators/shm_broadcast.py b/vllm/distributed/device_communicators/shm_broadcast.py index 565ee6c63806..ef27e6682c7f 100644 --- a/vllm/distributed/device_communicators/shm_broadcast.py +++ b/vllm/distributed/device_communicators/shm_broadcast.py @@ -418,6 +418,8 @@ def __init__( self.total_bytes = (slot_bytes + self.metadata_size) * n_slots self._next_slot = 0 self._fallbacks = 0 + self._pin_attempted = False + self._pinned = False # slots whose tensors were handed out by THIS reader and not yet # released (flushed at the next dequeue). self._pending_release: list[int] = [] @@ -501,11 +503,40 @@ def write_tensor(self, t: torch.Tensor) -> int | None: # ---- reader side ---- + def _ensure_pinned(self): + """cudaHostRegister the whole arena mapping 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 torch.cuda.is_available(): + return + import ctypes + + buf = self.shared_memory.buf + assert buf is not None + ptr = ctypes.addressof(ctypes.c_char.from_buffer(buf)) + ret = torch.cuda.cudart().cudaHostRegister(ptr, self.total_bytes, 0) + self._pinned = int(ret) == 0 + 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 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 released at this reader's next flush_releases() (called from MessageQueue.dequeue).""" + self._ensure_pinned() # NOT a context-manager view: the tensor must keep the mapping alive. slot_mv = self._slot(idx, nbytes) t8 = torch.frombuffer(slot_mv, dtype=torch.uint8, count=nbytes) From 29d30fb16ad9a50fdb59f371f4cacb44a34a6987 Mon Sep 17 00:00:00 2001 From: Bolin Sun Date: Tue, 7 Jul 2026 15:18:30 -0700 Subject: [PATCH 3/3] [Docs] Design doc for the zero-copy shared-memory tensor arena Documents the multimodal engine->worker transport bottleneck (in-band pickle of large pixel_values tensors in MessageQueue.enqueue, per-rank deserialize, engine-loop head-of-line blocking) and the ShmTensorArena zero-copy design that removes it, including the pinning rationale, configuration knobs, validation results, and known limitations. Co-Authored-By: Claude Fable 5 --- docs/design/shm_tensor_arena.md | 207 ++++++++++++++++++++++++++++++++ 1 file changed, 207 insertions(+) create mode 100644 docs/design/shm_tensor_arena.md diff --git a/docs/design/shm_tensor_arena.md b/docs/design/shm_tensor_arena.md new file mode 100644 index 000000000000..9d5952ecd5cc --- /dev/null +++ b/docs/design/shm_tensor_arena.md @@ -0,0 +1,207 @@ +# Zero-Copy Shared-Memory Tensor Arena for `MessageQueue` + +This document describes a performance problem in vLLM's engine→worker IPC path that +severely penalizes multimodal serving under tensor parallelism, and the fix implemented +on this branch: a zero-copy shared-memory tensor arena inside +`vllm/distributed/device_communicators/shm_broadcast.py`. + +## TL;DR + +With TP > 1, every `execute_model` call is broadcast from the EngineCore process to the +TP worker processes through `MessageQueue.enqueue`, which serializes the whole payload +with `pickle` — **including raw multimodal `pixel_values` tensors**. Torch tensors do +not participate in pickle's out-of-band buffer protocol, so a large image is byte-copied +into the pickle stream on the engine side and byte-copied back out **once per TP rank** +on the reader side. For a ~200 MB pixel tensor on a TP=4 worker that is ~1.8 GB of host +byte-copying, most of it serial and **inside the engine's step loop** — the scheduler +stops, and all GPUs on the worker go idle for **0.9–1.3 s per large image**, including +decode for unrelated in-flight requests (head-of-line blocking). The vision-encoder +compute for the same images is 50–90 ms; transport dominates compute ~15:1. + +The fix routes large CPU tensors around the pickle stream entirely, through a slotted +shared-memory arena: one memcpy into a slot on the writer, a `torch.frombuffer` view +(zero bytes copied) on each reader, with the mapping `cudaHostRegister`-pinned so the +subsequent H2D copy is a true DMA. Measured end to end on an interactive multimodal +workload (uncapped input images, identical seeded request sequence): TTFT p99 +1321 ms → **862 ms**, TTFT max 2139 ms → **1375 ms**, requests over 1.5 s: 9 → **0**, +TTFT p50 unchanged (no overhead on the small-request path). + +## 1. The problem + +### 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 — with the multimodal kwargs, i.e. the raw preprocessed pixel +tensors — from the EngineCore to every TP worker on each step. Image inputs at native +resolution easily reach tens to hundreds of MB of `pixel_values` (a ~4k×4k image is +~200 MB as bf16 patches). Latency-sensitive (interactive) serving makes the resulting +tail visible in TTFT and, through head-of-line blocking, in the end-to-end latency of +*other* requests. + +### The mechanism, at source level + +`MessageQueue.enqueue` (before this branch): + +```python +serialized_obj = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL, + buffer_callback=oob_callback) # ← the cost lives here +if total_bytes >= self.buffer.max_chunk_bytes: # 16 MB default + # write a 1-byte overflow marker into the shm ring... + self.local_socket.send_multipart(all_buffers) # ...and ZMQ to every local reader +else: + # copy the bytes into a shm ring chunk +``` + +Three compounding costs: + +1. **Serialize.** Torch tensors do not implement pickle protocol-5 out-of-band + buffers, so `pickle.dumps` drives `THPStorage_writeFileRaw`, which byte-copies the + entire tensor into the pickle stream through Python-driven small writes. Measured at + ~6.5 ms/MB (~155 MB/s) on the critical path — ~1.3 s for a 200 MB tensor. +2. **It blocks the engine.** The serialize runs inside the EngineCore step loop + (`collective_rpc → enqueue`), so no new GPU work is scheduled anywhere on the worker + until it finishes. CUDA-level profiling shows all GPUs of the TP group completely + silent for the duration — decode graph launches for other requests included. +3. **Deserialize, per rank.** Payloads above `max_chunk_bytes` bypass the ring and go + through a local ZMQ socket to each reader; every rank then runs `pickle.loads` → + `THPStorage_readFileRaw` — another full byte-copy per rank — before the (cheap, + ~20 ms) H2D copy. + +Per 200 MB image on TP=4: ~200 MB × (1 serialize + 4 socket transfers + 4 deserializes) +≈ **1.8 GB of host data motion**, where the model itself only needs one 200 MB H2D. + +### What does *not* fix it + +Raising `VLLM_MQ_MAX_CHUNK_BYTES_MB` so the payload takes the shm-ring path instead of +the ZMQ socket was tested first (same seeded workload, per-image A/B): **no +improvement**. The path switch happens (zmq frames disappear from the stall +callchains), but `pickle.dumps` runs *before* the ring-vs-socket branch — the serialize +and the ×N deserialize copies are common to both paths and are the actual cost. This +negative result pins the fix target: the pickle byte-copies themselves must go. + +## 2. The fix + +All changes are contained in `shm_broadcast.py`. + +### 2.1 `ShmTensorArena` + +A second shared-memory region created alongside the existing `ShmRingBuffer`: + +- **N slots × slot_bytes** (default 8 × 256 MB, env-tunable), 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. + +### 2.2 Writer path — `_ArenaPickler.reducer_override` + +```python +class _ArenaPickler(pickle.Pickler): + def reducer_override(self, obj): + if (isinstance(obj, torch.Tensor) and obj.device.type == "cpu" + and obj.layout is torch.strided and obj.is_contiguous() + 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_arena_tensor, + (arena_name, idx, nbytes, dtype_str, shape)) + return NotImplemented # everything else: default pickling +``` + +`reducer_override` is consulted for every object before its `__reduce_ex__`, so a +diverted tensor never reaches `THPStorage_writeFileRaw`: the pickle stream carries a +~100-byte rebuild stub instead of the tensor bytes. The one remaining copy — +`torch.frombuffer(slot).copy_(t.view(torch.uint8))` — is a multithreaded memcpy at +memory bandwidth (~50–70 ms for 200 MB, vs ~1.3 s for the pickle serialize). + +**The writer never blocks.** If no slot is free or the tensor exceeds the slot size, +`write_tensor` returns `None` and the pickler falls through to the default path — the +worst case is exactly today's 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 attached arena + return arena.get_tensor(slot_idx, nbytes, getattr(torch, dtype_str), shape) + # get_tensor: torch.frombuffer over the mapped slot — zero bytes copied +``` + +The rebuilt tensor *is* the shared memory — no deserialize on any rank. + +**Slot lifecycle.** The rebuilt tensor is consumed (H2D'd) while the worker executes +that step, so the reader must not release the slot at unpickle time. Releases are +deferred and flushed at the reader's **next `dequeue`**: the worker loop is strictly +sequential (dequeue step N → execute → dequeue step N+1), so by the next dequeue the +previous step's H2D has completed. The writer requires all readers' done flags before +reusing a slot. + +### 2.4 Pinning — the zero-copy trap + +The first version of the patch only halved the stall. Profiling showed the residual +entirely inside `cudaMemcpyAsync`: the zero-copy view references tmpfs pages the reader +process has never touched, so the H2D pays **first-touch page faults on ~50k pages plus +pageable-copy staging**. (The old deserialize, for all its waste, incidentally left the +bytes in hot process-local heap pages — which is why the baseline H2D was fast.) + +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. + +## 3. Configuration + +| Env var | Default | Meaning | +|---|---|---| +| `VLLM_SHM_TENSOR_ARENA` | `1` | Enable the arena (`0` disables; behavior reverts to stock). | +| `VLLM_SHM_TENSOR_ARENA_SLOTS` | `8` | Number of slots. | +| `VLLM_SHM_TENSOR_ARENA_SLOT_MB` | `256` | Slot size; tensors larger than this fall back to pickle. | +| `VLLM_SHM_TENSOR_ARENA_MIN_MB` | `8` | Minimum tensor size to divert; smaller tensors pickle as before. | + +## 4. Validation + +1. **Unit test**: a 199 MB bf16 tensor pushed through a real `MessageQueue` to two + forked reader processes — byte-exact checksums, arena path confirmed, deferred + release, slot reuse, and exhaustion/oversize fallbacks exercised. Enqueue of the + 199 MB payload: **66.7 ms** (vs ~1275 ms observed in production for the same size). +2. **Same-seed A/B**, 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 | + |---|---:|---:|---:|---:|---:|---:| + | stock | 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** (worst offenders, CUDA-level all-GPU idle gap around the pixel + H2D): 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. **Load scaling**: on the same workload the stock build failed a 1.5 s e2e-p99 target + even at very low arrival rates (the stall is a fixed per-image cost); with the patch + the passing arrival rate is bounded by prefill/decode compute interference instead — + transport no longer appears in the profile at the tail. + +## 5. Limitations and future work + +- **Slot release granularity**: releases are flushed at the reader's next dequeue. A + CUDA-event-based release would close the theoretical reuse window under extreme + multi-image bursts deeper than the slot count (today such bursts safely fall back to + pickle when the arena is exhausted). +- **Fallback observability**: arena exhaustion / oversize fallbacks are rate-limited + log lines today; a counter metric would be better. +- **Scope**: the arena activates only when every queue reader is node-local. Remote + readers (multi-node PP/TP) keep the existing socket path. +- **Generality**: the problem is generic to any multimodal model on TP > 1 with the + multiproc executor; upstreaming the approach (or an equivalent out-of-band tensor + channel) is worth discussing.