[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>
Ruinan Ma committed
Jul 28, 2026 at 16:28 UTC
a07fac758f2cae58f6ebad29da726e9432012c13
3 files changed
+233
-3
tests/distributed/test_shm_broadcast.py
+149
@@ -1,6 +1,8 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import io
5
+import pickle
6
import random
7
import threading
8
import time
@@ -10,12 +12,15 @@ from unittest import mock
12
import multiprocess as mp
13
import numpy as np
14
import pytest
15
+import torch
16
import torch.distributed as dist
17
18
from vllm.distributed.device_communicators import shm_broadcast
19
from vllm.distributed.device_communicators.shm_broadcast import (
20
MessageQueue,
21
ShmRingBuffer,
22
+ _rebuild_tensor,
23
+ _reduce_tensor,
24
check_shm_free_space,
25
)
26
from vllm.distributed.utils import StatelessProcessGroup
@@ -354,6 +359,150 @@ def test_message_queue_busy_to_idle():
359
distributed_run(worker_fn_test_busy_to_idle, 4)
360
361
362
+@worker_fn_wrapper
363
+def worker_fn_tensor_broadcast():
364
+ rank = dist.get_rank()
365
+ writer_rank = 0
366
+ message_queue = MessageQueue.create_from_process_group(
367
+ dist.group.WORLD, 8 * 1024 * 1024, 4, writer_rank
368
+ )
369
+
370
+ # Both ranks construct the identical reference payload.
371
+ torch.manual_seed(42)
372
+ payload = {
373
+ # 2MiB: rides the shm ring as an out-of-band buffer (the receiving
374
+ # side must copy out of the reusable ring chunk).
375
+ "mid": torch.randn(1024, 512),
376
+ # 16MiB > max_chunk_bytes: overflows to the zmq socket (the
377
+ # receiving side aliases the zmq.Frame zero-copy).
378
+ "big": torch.randn(4096, 2048, dtype=torch.bfloat16),
379
+ "nested": ["plain", 123, {"inner": torch.arange(5)}],
380
+ }
381
+
382
+ if rank == writer_rank:
383
+ with mock.patch(
384
+ "vllm.distributed.device_communicators.shm_broadcast._reduce_tensor",
385
+ wraps=_reduce_tensor,
386
+ ) as wrapped_reduce:
387
+ message_queue.enqueue(payload)
388
+ assert wrapped_reduce.call_count == 3
389
+ # Cycle the ring (max_chunks=4) several times over so that aliased
390
+ # ring chunks would be overwritten.
391
+ for i in range(16):
392
+ message_queue.enqueue({"junk": torch.full((1024, 512), float(i))})
393
+ else:
394
+ received = message_queue.dequeue(timeout=30)
395
+ for key in ("mid", "big"):
396
+ assert torch.equal(received[key], payload[key]), key
397
+ assert received[key].dtype == payload[key].dtype, key
398
+ assert torch.equal(received["nested"][2]["inner"], torch.arange(5))
399
+
400
+ snapshot = received["mid"].clone()
401
+ for i in range(16):
402
+ junk = message_queue.dequeue(timeout=30)
403
+ assert torch.equal(junk["junk"], torch.full((1024, 512), float(i)))
404
+ # Tensors received via the shm ring must not alias chunk memory
405
+ # that the writer has reused for subsequent messages.
406
+ assert torch.equal(received["mid"], snapshot)
407
+ # Rebuilt tensors must be writable, like regular tensors.
408
+ received["mid"] += 1.0
409
+ received["big"][0, 0] = 1.0
410
+
411
+ dist.barrier()
412
+ print(f"tensor broadcast passed the test! Rank {rank}")
413
+
414
+
415
+def test_tensor_broadcast():
416
+ distributed_run(worker_fn_tensor_broadcast, 2)
417
+
418
+
419
+def _dumps_oob(obj) -> tuple[bytes, list]:
420
+ """Pickle `obj` the same way `MessageQueue.enqueue` does: tensor
421
+ dispatch table + out-of-band buffers >= 1MiB."""
422
+ buffers = []
423
+
424
+ def callback(buf: pickle.PickleBuffer) -> bool:
425
+ raw = buf.raw()
426
+ if raw.nbytes < 1024 * 1024:
427
+ return True
428
+ buffers.append(raw)
429
+ return False
430
+
431
+ bio = io.BytesIO()
432
+ pickler = pickle.Pickler(
433
+ bio, protocol=pickle.HIGHEST_PROTOCOL, buffer_callback=callback
434
+ )
435
+ pickler.dispatch_table = {torch.Tensor: _reduce_tensor}
436
+ pickler.dump(obj)
437
+ return bio.getvalue(), buffers
438
+
439
+
440
+@pytest.mark.parametrize(
441
+ "case",
442
+ [
443
+ "small",
444
+ "mid",
445
+ "bf16",
446
+ "fp8",
447
+ "empty",
448
+ "scalar",
449
+ "noncontig",
450
+ "requires_grad",
451
+ "conj",
452
+ "param",
453
+ ],
454
+)
455
+def test_tensor_pickle_roundtrip(case: str):
456
+ tensor = {
457
+ # Inlined in-band (< 1MiB) and out-of-band (>= 1MiB) buffers.
458
+ "small": lambda: torch.randn(100, 10),
459
+ "mid": lambda: torch.randn(1024, 512),
460
+ # Dtypes numpy doesn't recognize.
461
+ "bf16": lambda: torch.randn(512, 512, dtype=torch.bfloat16),
462
+ "fp8": lambda: torch.randn(32, 32).to(torch.float8_e4m3fn),
463
+ # Shape edge cases.
464
+ "empty": lambda: torch.empty(0, 8),
465
+ "scalar": lambda: torch.tensor(3.14),
466
+ "noncontig": lambda: torch.randn(64, 64).t(),
467
+ # These fall back to torch's default reducer.
468
+ "requires_grad": lambda: torch.randn(8, 8, requires_grad=True),
469
+ "conj": lambda: torch.randn(4, dtype=torch.complex64).conj(),
470
+ "param": lambda: torch.nn.Parameter(torch.randn(4), requires_grad=False),
471
+ }[case]()
472
+
473
+ data, buffers = _dumps_oob({"tensor": tensor, "meta": list(range(10))})
474
+ received = pickle.loads(data, buffers=buffers)["tensor"]
475
+
476
+ assert received.shape == tensor.shape
477
+ assert received.dtype == tensor.dtype
478
+ if tensor.dtype == torch.float8_e4m3fn:
479
+ assert torch.equal(received.view(torch.uint8), tensor.view(torch.uint8))
480
+ else:
481
+ assert torch.equal(received, tensor)
482
+ assert received.requires_grad == tensor.requires_grad
483
+ assert isinstance(received, type(tensor))
484
+ if tensor.numel() and not tensor.requires_grad:
485
+ # Rebuilt tensors must be writable, like regular tensors.
486
+ received.view(-1)[0] = 1.0
487
+
488
+
489
+@pytest.mark.parametrize("case", ["cuda", "requires_grad", "conj"])
490
+def test_reduce_tensor_fallback(case: str):
491
+ """Tensors the zero-copy reducer can't safely alias must fall back to
492
+ torch's default reduction."""
493
+ if case == "cuda":
494
+ if not torch.cuda.is_available():
495
+ pytest.skip("requires CUDA")
496
+ tensor = torch.randn(4, device="cuda")
497
+ elif case == "requires_grad":
498
+ tensor = torch.randn(8, requires_grad=True)
499
+ else:
500
+ tensor = torch.randn(4, dtype=torch.complex64).conj()
501
+
502
+ reduced = _reduce_tensor(tensor)
503
+ assert reduced[0] is not _rebuild_tensor
504
+
505
+
506
@pytest.mark.parametrize("should_warn", [False, True])
507
def test_reader_timeout_caps_indefinite_waits(should_warn):
508
with (
tools/pre_commit/check_forbidden_imports.py
+1
@@ -48,6 +48,7 @@ CHECK_IMPORTS = {
48
"vllm/distributed/device_communicators/shm_object_storage.py",
49
"vllm/distributed/weight_transfer/ipc_engine.py",
50
"vllm/distributed/weight_transfer/clients.py",
51
+ "tests/distributed/test_shm_broadcast.py",
52
"tests/distributed/test_weight_transfer.py",
53
"vllm/utils/hashing.py",
54
"tests/multimodal/media/test_base.py",
vllm/distributed/device_communicators/shm_broadcast.py
+83
-3
@@ -1,6 +1,8 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import copyreg
4
import functools
5
+import io
6
import os
7
import pickle
8
import shutil
@@ -385,6 +387,70 @@ class ShmRingBuffer:
387
yield buf
388
389
390
+def _rebuild_tensor(buf: Any, shape: tuple[int, ...], dtype_str: str) -> torch.Tensor:
391
+ """Rebuild a tensor from an out-of-band pickle buffer.
392
+
393
+ Counterpart of `_reduce_tensor`. Note that pickle passes the original
394
+ buffer-providing object from `loads(buffers=...)` straight to this
395
+ function (no `PickleBuffer` wrapper on the receiving side), so `buf` is
396
+ a `zmq.Frame`, a `memoryview` of a shared-memory ring chunk, or `bytes`
397
+ if the buffer was serialized in-band.
398
+ """
399
+ dtype = getattr(torch, dtype_str)
400
+ assert isinstance(dtype, torch.dtype)
401
+ if isinstance(buf, zmq.Frame):
402
+ # ZMQ frames own their message memory independently of any context,
403
+ # so the tensor can safely alias it with zero copies. The tensor's
404
+ # storage keeps the frame (and thus its bytes) alive via a strong
405
+ # reference for as long as the tensor is.
406
+ try:
407
+ return torch.frombuffer(buf, dtype=torch.uint8).view(dtype).view(shape)
408
+ except ValueError:
409
+ # Empty or read-only frame buffer; fall through to the copy path.
410
+ pass
411
+ # Shared-memory ring buffer chunks are reused by the writer once all
412
+ # readers have marked them read, so we must copy out of them. bytearray
413
+ # (vs bytes) keeps the resulting tensor writable, matching normal tensor
414
+ # semantics.
415
+ raw = bytearray(buf)
416
+ if not raw:
417
+ assert 0 in shape
418
+ return torch.empty(shape, dtype=dtype)
419
+ return torch.frombuffer(raw, dtype=torch.uint8).view(dtype).view(shape)
420
+
421
+
422
+def _reduce_tensor(tensor: torch.Tensor):
423
+ """Reduce a CPU tensor to a `PickleBuffer` for out-of-band pickling.
424
+
425
+ `torch.Tensor.__reduce_ex__` copies the tensor bytes into the pickle
426
+ byte stream via `torch.serialization` and never emits a `PickleBuffer`,
427
+ which defeats the out-of-band buffer handling in `MessageQueue.enqueue`.
428
+ This reducer instead exposes the tensor's memory directly, so large
429
+ tensors (e.g. `prompt_embeds` in `SchedulerOutput`) traverse the queue
430
+ without being copied into and back out of the pickled message.
431
+ """
432
+ if (
433
+ tensor.device.type == "cpu"
434
+ and tensor.layout == torch.strided
435
+ and not tensor.requires_grad
436
+ ):
437
+ try:
438
+ # The uint8 view exposes the raw bytes via the buffer protocol,
439
+ # including for dtypes numpy doesn't recognize (bfloat16, fp8, ...).
440
+ # reshape(-1) first so that 0-dim tensors can be viewed as well.
441
+ raw = tensor.contiguous().reshape(-1).view(torch.uint8).numpy()
442
+ except RuntimeError:
443
+ # Exotic tensors (e.g. with the conjugate bit set) that don't
444
+ # support aliasing views; let torch handle them.
445
+ pass
446
+ else:
447
+ dtype_str = str(tensor.dtype).removeprefix("torch.")
448
+ return _rebuild_tensor, (PickleBuffer(raw), tuple(tensor.shape), dtype_str)
449
+
450
+ # Fall back to torch's default (copying) reduction.
451
+ return tensor.__reduce_ex__(pickle.HIGHEST_PROTOCOL)
452
+
453
+
454
@dataclass
455
class Handle:
456
local_reader_ranks: list[int] = field(default_factory=list)
@@ -771,9 +837,23 @@ class MessageQueue:
837
total_bytes += len(raw_buf) + 4
838
return False
839
774
- all_buffers[0] = pickle.dumps(
775
- obj, protocol=pickle.HIGHEST_PROTOCOL, buffer_callback=oob_callback
776
- )
840
+ # CPU tensors are routed through `_reduce_tensor` so that their
841
+ # bytes are emitted as out-of-band buffers instead of being
842
+ # copied into the pickle stream by torch's default reducer.
843
+ # Start from `copyreg.dispatch_table` to preserve globally
844
+ # registered reducers (e.g. `re.Pattern`); the per-pickler
845
+ # dispatch table would otherwise shadow them.
846
+ dispatch_table = dict(copyreg.dispatch_table)
847
+ dispatch_table[torch.Tensor] = _reduce_tensor
848
+ with io.BytesIO() as bio:
849
+ pickler = pickle.Pickler(
850
+ bio,
851
+ protocol=pickle.HIGHEST_PROTOCOL,
852
+ buffer_callback=oob_callback,
853
+ )
854
+ pickler.dispatch_table = dispatch_table
855
+ pickler.dump(obj)
856
+ all_buffers[0] = bio.getvalue()
857
if self.n_local_reader > 0:
858
if total_bytes + len(all_buffers[0]) >= self.buffer.max_chunk_bytes:
859
with self.acquire_write(timeout) as buf: