[Perf] Zero-copy torch.Tensor pickling in shm_broadcast MessageQueue (#48442)

Signed-off-by: Ruinan Ma <r7ma3088@gmail.com>
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Nick Hill <nickhill123@gmail.com>

tests/distributed/test_shm_broadcast.py

@@ -1,6 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import io
+import pickle
 import random
 import threading
 import time
@@ -10,12 +12,15 @@ from unittest import mock
 import multiprocess as mp
 import numpy as np
 import pytest
+import torch
 import torch.distributed as dist
 
 from vllm.distributed.device_communicators import shm_broadcast
 from vllm.distributed.device_communicators.shm_broadcast import (
     MessageQueue,
     ShmRingBuffer,
+    _rebuild_tensor,
+    _reduce_tensor,
     check_shm_free_space,
 )
 from vllm.distributed.utils import StatelessProcessGroup
@@ -354,6 +359,150 @@ def test_message_queue_busy_to_idle():
     distributed_run(worker_fn_test_busy_to_idle, 4)
 
 
+@worker_fn_wrapper
+def worker_fn_tensor_broadcast():
+    rank = dist.get_rank()
+    writer_rank = 0
+    message_queue = MessageQueue.create_from_process_group(
+        dist.group.WORLD, 8 * 1024 * 1024, 4, writer_rank
+    )
+
+    # Both ranks construct the identical reference payload.
+    torch.manual_seed(42)
+    payload = {
+        # 2MiB: rides the shm ring as an out-of-band buffer (the receiving
+        # side must copy out of the reusable ring chunk).
+        "mid": torch.randn(1024, 512),
+        # 16MiB > max_chunk_bytes: overflows to the zmq socket (the
+        # receiving side aliases the zmq.Frame zero-copy).
+        "big": torch.randn(4096, 2048, dtype=torch.bfloat16),
+        "nested": ["plain", 123, {"inner": torch.arange(5)}],
+    }
+
+    if rank == writer_rank:
+        with mock.patch(
+            "vllm.distributed.device_communicators.shm_broadcast._reduce_tensor",
+            wraps=_reduce_tensor,
+        ) as wrapped_reduce:
+            message_queue.enqueue(payload)
+        assert wrapped_reduce.call_count == 3
+        # Cycle the ring (max_chunks=4) several times over so that aliased
+        # ring chunks would be overwritten.
+        for i in range(16):
+            message_queue.enqueue({"junk": torch.full((1024, 512), float(i))})
+    else:
+        received = message_queue.dequeue(timeout=30)
+        for key in ("mid", "big"):
+            assert torch.equal(received[key], payload[key]), key
+            assert received[key].dtype == payload[key].dtype, key
+        assert torch.equal(received["nested"][2]["inner"], torch.arange(5))
+
+        snapshot = received["mid"].clone()
+        for i in range(16):
+            junk = message_queue.dequeue(timeout=30)
+            assert torch.equal(junk["junk"], torch.full((1024, 512), float(i)))
+        # Tensors received via the shm ring must not alias chunk memory
+        # that the writer has reused for subsequent messages.
+        assert torch.equal(received["mid"], snapshot)
+        # Rebuilt tensors must be writable, like regular tensors.
+        received["mid"] += 1.0
+        received["big"][0, 0] = 1.0
+
+    dist.barrier()
+    print(f"tensor broadcast passed the test! Rank {rank}")
+
+
+def test_tensor_broadcast():
+    distributed_run(worker_fn_tensor_broadcast, 2)
+
+
+def _dumps_oob(obj) -> tuple[bytes, list]:
+    """Pickle `obj` the same way `MessageQueue.enqueue` does: tensor
+    dispatch table + out-of-band buffers >= 1MiB."""
+    buffers = []
+
+    def callback(buf: pickle.PickleBuffer) -> bool:
+        raw = buf.raw()
+        if raw.nbytes < 1024 * 1024:
+            return True
+        buffers.append(raw)
+        return False
+
+    bio = io.BytesIO()
+    pickler = pickle.Pickler(
+        bio, protocol=pickle.HIGHEST_PROTOCOL, buffer_callback=callback
+    )
+    pickler.dispatch_table = {torch.Tensor: _reduce_tensor}
+    pickler.dump(obj)
+    return bio.getvalue(), buffers
+
+
+@pytest.mark.parametrize(
+    "case",
+    [
+        "small",
+        "mid",
+        "bf16",
+        "fp8",
+        "empty",
+        "scalar",
+        "noncontig",
+        "requires_grad",
+        "conj",
+        "param",
+    ],
+)
+def test_tensor_pickle_roundtrip(case: str):
+    tensor = {
+        # Inlined in-band (< 1MiB) and out-of-band (>= 1MiB) buffers.
+        "small": lambda: torch.randn(100, 10),
+        "mid": lambda: torch.randn(1024, 512),
+        # Dtypes numpy doesn't recognize.
+        "bf16": lambda: torch.randn(512, 512, dtype=torch.bfloat16),
+        "fp8": lambda: torch.randn(32, 32).to(torch.float8_e4m3fn),
+        # Shape edge cases.
+        "empty": lambda: torch.empty(0, 8),
+        "scalar": lambda: torch.tensor(3.14),
+        "noncontig": lambda: torch.randn(64, 64).t(),
+        # These fall back to torch's default reducer.
+        "requires_grad": lambda: torch.randn(8, 8, requires_grad=True),
+        "conj": lambda: torch.randn(4, dtype=torch.complex64).conj(),
+        "param": lambda: torch.nn.Parameter(torch.randn(4), requires_grad=False),
+    }[case]()
+
+    data, buffers = _dumps_oob({"tensor": tensor, "meta": list(range(10))})
+    received = pickle.loads(data, buffers=buffers)["tensor"]
+
+    assert received.shape == tensor.shape
+    assert received.dtype == tensor.dtype
+    if tensor.dtype == torch.float8_e4m3fn:
+        assert torch.equal(received.view(torch.uint8), tensor.view(torch.uint8))
+    else:
+        assert torch.equal(received, tensor)
+    assert received.requires_grad == tensor.requires_grad
+    assert isinstance(received, type(tensor))
+    if tensor.numel() and not tensor.requires_grad:
+        # Rebuilt tensors must be writable, like regular tensors.
+        received.view(-1)[0] = 1.0
+
+
+@pytest.mark.parametrize("case", ["cuda", "requires_grad", "conj"])
+def test_reduce_tensor_fallback(case: str):
+    """Tensors the zero-copy reducer can't safely alias must fall back to
+    torch's default reduction."""
+    if case == "cuda":
+        if not torch.cuda.is_available():
+            pytest.skip("requires CUDA")
+        tensor = torch.randn(4, device="cuda")
+    elif case == "requires_grad":
+        tensor = torch.randn(8, requires_grad=True)
+    else:
+        tensor = torch.randn(4, dtype=torch.complex64).conj()
+
+    reduced = _reduce_tensor(tensor)
+    assert reduced[0] is not _rebuild_tensor
+
+
 @pytest.mark.parametrize("should_warn", [False, True])
 def test_reader_timeout_caps_indefinite_waits(should_warn):
     with (

tools/pre_commit/check_forbidden_imports.py

@@ -48,6 +48,7 @@ CHECK_IMPORTS = {
             "vllm/distributed/device_communicators/shm_object_storage.py",
             "vllm/distributed/weight_transfer/ipc_engine.py",
             "vllm/distributed/weight_transfer/clients.py",
+            "tests/distributed/test_shm_broadcast.py",
             "tests/distributed/test_weight_transfer.py",
             "vllm/utils/hashing.py",
             "tests/multimodal/media/test_base.py",

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 copyreg
 import functools
+import io
 import os
 import pickle
 import shutil
@@ -385,6 +387,70 @@ class ShmRingBuffer:
             yield buf
 
 
+def _rebuild_tensor(buf: Any, shape: tuple[int, ...], dtype_str: str) -> torch.Tensor:
+    """Rebuild a tensor from an out-of-band pickle buffer.
+
+    Counterpart of `_reduce_tensor`. Note that pickle passes the original
+    buffer-providing object from `loads(buffers=...)` straight to this
+    function (no `PickleBuffer` wrapper on the receiving side), so `buf` is
+    a `zmq.Frame`, a `memoryview` of a shared-memory ring chunk, or `bytes`
+    if the buffer was serialized in-band.
+    """
+    dtype = getattr(torch, dtype_str)
+    assert isinstance(dtype, torch.dtype)
+    if isinstance(buf, zmq.Frame):
+        # ZMQ frames own their message memory independently of any context,
+        # so the tensor can safely alias it with zero copies. The tensor's
+        # storage keeps the frame (and thus its bytes) alive via a strong
+        # reference for as long as the tensor is.
+        try:
+            return torch.frombuffer(buf, dtype=torch.uint8).view(dtype).view(shape)
+        except ValueError:
+            # Empty or read-only frame buffer; fall through to the copy path.
+            pass
+    # Shared-memory ring buffer chunks are reused by the writer once all
+    # readers have marked them read, so we must copy out of them. bytearray
+    # (vs bytes) keeps the resulting tensor writable, matching normal tensor
+    # semantics.
+    raw = bytearray(buf)
+    if not raw:
+        assert 0 in shape
+        return torch.empty(shape, dtype=dtype)
+    return torch.frombuffer(raw, dtype=torch.uint8).view(dtype).view(shape)
+
+
+def _reduce_tensor(tensor: torch.Tensor):
+    """Reduce a CPU tensor to a `PickleBuffer` for out-of-band pickling.
+
+    `torch.Tensor.__reduce_ex__` copies the tensor bytes into the pickle
+    byte stream via `torch.serialization` and never emits a `PickleBuffer`,
+    which defeats the out-of-band buffer handling in `MessageQueue.enqueue`.
+    This reducer instead exposes the tensor's memory directly, so large
+    tensors (e.g. `prompt_embeds` in `SchedulerOutput`) traverse the queue
+    without being copied into and back out of the pickled message.
+    """
+    if (
+        tensor.device.type == "cpu"
+        and tensor.layout == torch.strided
+        and not tensor.requires_grad
+    ):
+        try:
+            # The uint8 view exposes the raw bytes via the buffer protocol,
+            # including for dtypes numpy doesn't recognize (bfloat16, fp8, ...).
+            # reshape(-1) first so that 0-dim tensors can be viewed as well.
+            raw = tensor.contiguous().reshape(-1).view(torch.uint8).numpy()
+        except RuntimeError:
+            # Exotic tensors (e.g. with the conjugate bit set) that don't
+            # support aliasing views; let torch handle them.
+            pass
+        else:
+            dtype_str = str(tensor.dtype).removeprefix("torch.")
+            return _rebuild_tensor, (PickleBuffer(raw), tuple(tensor.shape), dtype_str)
+
+    # Fall back to torch's default (copying) reduction.
+    return tensor.__reduce_ex__(pickle.HIGHEST_PROTOCOL)
+
+
 @dataclass
 class Handle:
     local_reader_ranks: list[int] = field(default_factory=list)
@@ -771,9 +837,23 @@ class MessageQueue:
             total_bytes += len(raw_buf) + 4
             return False
 
-        all_buffers[0] = pickle.dumps(
-            obj, protocol=pickle.HIGHEST_PROTOCOL, buffer_callback=oob_callback
-        )
+        # CPU tensors are routed through `_reduce_tensor` so that their
+        # bytes are emitted as out-of-band buffers instead of being
+        # copied into the pickle stream by torch's default reducer.
+        # Start from `copyreg.dispatch_table` to preserve globally
+        # registered reducers (e.g. `re.Pattern`); the per-pickler
+        # dispatch table would otherwise shadow them.
+        dispatch_table = dict(copyreg.dispatch_table)
+        dispatch_table[torch.Tensor] = _reduce_tensor
+        with io.BytesIO() as bio:
+            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()
         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: