A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Perf] Exploit out-of-band buffers in shm_broadcast (#26961)
Signed-off-by: Nick Hill <nhill@redhat.com>
Nick Hill committed
Oct 16, 2025 at 20:08 UTC
ab81379ea6b0135c6a89c7deb97b49c962dabb4a
1 file changed
+54
-16
vllm/distributed/device_communicators/shm_broadcast.py
+54
-16
@@ -1,13 +1,14 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
3
+import functools
4
import pickle
5
import time
6
from contextlib import contextmanager
7
from dataclasses import dataclass, field
8
from multiprocessing import shared_memory
9
+from pickle import PickleBuffer
10
from threading import Event
10
-from typing import Any
11
+from typing import TYPE_CHECKING, Any
12
from unittest.mock import patch
13
14
import torch
@@ -33,8 +34,18 @@ from vllm.utils import (
34
is_valid_ipv6_address,
35
)
36
37
+if TYPE_CHECKING:
38
+ from _typeshed import SizedBuffer
39
+
40
VLLM_RINGBUFFER_WARNING_INTERVAL = envs.VLLM_RINGBUFFER_WARNING_INTERVAL
41
42
+from_bytes_big = functools.partial(int.from_bytes, byteorder="big")
43
+
44
+
45
+def to_bytes_big(value: int, size: int) -> bytes:
46
+ return value.to_bytes(size, byteorder="big")
47
+
48
+
49
logger = init_logger(__name__)
50
51
@@ -225,7 +236,7 @@ class MessageQueue:
236
n_reader, # number of all readers
237
n_local_reader, # number of local readers through shared memory
238
local_reader_ranks: list[int] | None = None,
228
- max_chunk_bytes: int = 1024 * 1024 * 10,
239
+ max_chunk_bytes: int = 1024 * 1024 * 24, # 24MiB
240
max_chunks: int = 10,
241
connect_ip: str | None = None,
242
):
@@ -505,18 +516,41 @@ class MessageQueue:
516
def enqueue(self, obj, timeout: float | None = None):
517
"""Write to message queue with optional timeout (in seconds)"""
518
assert self._is_writer, "Only writers can enqueue"
508
- serialized_obj = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL)
519
+ all_buffers: list[SizedBuffer] = [b""]
520
+ total_bytes = 6 # 2 bytes for oob buffer count, 4 for main buffer size
521
+
522
+ def oob_callback(buf: PickleBuffer) -> bool:
523
+ raw_buf = buf.raw()
524
+ if len(raw_buf) < 1024 * 1024:
525
+ # In-line buffers smaller than 1MiB.
526
+ return True
527
+ all_buffers.append(raw_buf)
528
+ nonlocal total_bytes
529
+ total_bytes += len(raw_buf) + 4
530
+ return False
531
+
532
+ all_buffers[0] = pickle.dumps(
533
+ obj, protocol=pickle.HIGHEST_PROTOCOL, buffer_callback=oob_callback
534
+ )
535
if self.n_local_reader > 0:
510
- if len(serialized_obj) >= self.buffer.max_chunk_bytes:
536
+ if total_bytes + len(all_buffers[0]) >= self.buffer.max_chunk_bytes:
537
with self.acquire_write(timeout) as buf:
538
buf[0] = 1 # overflow
513
- self.local_socket.send(serialized_obj)
539
+ self.local_socket.send_multipart(all_buffers, copy=False)
540
else:
541
with self.acquire_write(timeout) as buf:
542
buf[0] = 0 # not overflow
517
- buf[1 : len(serialized_obj) + 1] = serialized_obj
543
+ offset = 3
544
+ buf[1:offset] = to_bytes_big(len(all_buffers), 2) # oob buf count
545
+ for buffer in all_buffers:
546
+ buf_len = len(buffer)
547
+ # prepend each buffer with 4 bytes containing its size.
548
+ buf_offset = offset + 4
549
+ buf[offset:buf_offset] = to_bytes_big(buf_len, 4)
550
+ buf[buf_offset : (offset := buf_offset + buf_len)] = buffer
551
+
552
if self.n_remote_reader > 0:
519
- self.remote_socket.send(serialized_obj)
553
+ self.remote_socket.send_multipart(all_buffers, copy=False)
554
555
def dequeue(
556
self,
@@ -529,10 +563,15 @@ class MessageQueue:
563
with self.acquire_read(timeout, cancel, indefinite) as buf:
564
overflow = buf[0] == 1
565
if not overflow:
532
- # no need to know the size of serialized object
533
- # pickle format contains the size information internally
534
- # see https://docs.python.org/3/library/pickle.html
535
- obj = pickle.loads(buf[1:])
566
+ offset = 3
567
+ buf_count = from_bytes_big(buf[1:offset])
568
+ all_buffers = []
569
+ for i in range(buf_count):
570
+ buf_offset = offset + 4
571
+ buf_len = from_bytes_big(buf[offset:buf_offset])
572
+ offset = buf_offset + buf_len
573
+ all_buffers.append(buf[buf_offset:offset])
574
+ obj = pickle.loads(all_buffers[0], buffers=all_buffers[1:])
575
if overflow:
576
obj = MessageQueue.recv(self.local_socket, timeout)
577
elif self._is_remote_reader:
@@ -546,15 +585,14 @@ class MessageQueue:
585
timeout_ms = None if timeout is None else int(timeout * 1000)
586
if not socket.poll(timeout=timeout_ms):
587
raise TimeoutError
549
- recv = socket.recv(copy=False)
550
- return pickle.loads(recv.buffer)
588
+ recv, *recv_oob = socket.recv_multipart(copy=False)
589
+ return pickle.loads(recv, buffers=recv_oob)
590
591
def broadcast_object(self, obj=None):
592
if self._is_writer:
593
self.enqueue(obj)
594
return obj
556
- else:
557
- return self.dequeue()
595
+ return self.dequeue()
596
597
@staticmethod
598
def create_from_process_group(