vllm.distributed.device_communicators.shm_tensor_arena

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.

Classes:

• ShmTensorArena –

  Slotted shared-memory arena: one writer, n_reader zero-copy readers.

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

Methods:

• flush_releases –

  Retire this reader's arena slots that are ready for writer reuse.

• get_tensor –

  Zero-copy view of a slot as a tensor. The slot is queued for

• write_tensor –

  Copy tensor bytes into a free slot; return slot idx or None

Source code in vllm/distributed/device_communicators/shm_tensor_arena.py

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

_ensure_pinned()

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.

Source code in vllm/distributed/device_communicators/shm_tensor_arena.py

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)

_mark_released(idxs)

Set THIS reader's done flag on the given slots; a slot becomes reusable once every reader has done so.

Source code in vllm/distributed/device_communicators/shm_tensor_arena.py

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()

_record_release_event()

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.

Source code in vllm/distributed/device_communicators/shm_tensor_arena.py

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

_unpin()

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.

Source code in vllm/distributed/device_communicators/shm_tensor_arena.py

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)

flush_releases()

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).

Source code in vllm/distributed/device_communicators/shm_tensor_arena.py

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))

get_tensor(idx, nbytes, dtype, shape)

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.

Source code in vllm/distributed/device_communicators/shm_tensor_arena.py

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)

write_tensor(t)

Copy tensor bytes into a free slot; return slot idx or None (caller must then fall back to the default pickle path).

Source code in vllm/distributed/device_communicators/shm_tensor_arena.py

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

_ArenaPickler

Bases: 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.

Source code in vllm/distributed/device_communicators/shm_tensor_arena.py

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

_rebuild_arena_parameter(arena_name, slot_idx, nbytes, dtype_str, shape)

Unpickle hook for torch.nn.Parameter (shares the slot's storage).

Source code in vllm/distributed/device_communicators/shm_tensor_arena.py

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)

_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.

Source code in vllm/distributed/device_communicators/shm_tensor_arena.py

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)