@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 6453fc0b8c

[Bugfix] Don't reuse engine core payload buffer while zmq is sending it (#50053)

Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>

Nick Hill committed Jul 28, 2026 at 09:09 UTC 6453fc0b8cb50668dceaf5839033b0a2a8fa7f7f

4 files changed +179 -58

tests/v1/test_serial_utils.py

+136

@@ -423,3 +423,139 @@ def test_multiple_senders_single_receiver_ipc():

423 assert torch.allclose(decoded.prompt_embeds, original_tensor), (

424 f"Value mismatch for sender {sender_idx} msg {msg_idx}"

425 )

426 +

427 +

428 +def _logprobs_outputs(num_reqs: int, num_prompt_tokens: int):

429 + """An EngineCoreOutputs carrying prompt logprobs, as the engine core sends

430 + it: many requests, each with per-token tensors small enough that pyzmq

431 + copies their frames, while the accumulated payload frame is large enough

432 + that pyzmq sends it zero-copy."""

433 + from vllm.v1.engine import EngineCoreOutput, EngineCoreOutputs

434 + from vllm.v1.outputs import LogprobsTensors

435 +

436 + outputs = []

437 + for req in range(num_reqs):

438 + num_tokens = num_prompt_tokens + req % 4

439 + outputs.append(

440 + EngineCoreOutput(

441 + request_id=f"req-{req:08d}",

442 + new_token_ids=[req],

443 + new_prompt_logprobs_tensors=LogprobsTensors(

444 + logprob_token_ids=torch.arange(

445 + num_tokens * 2, dtype=torch.int64

446 + ).view(num_tokens, 2),

447 + logprobs=torch.zeros(num_tokens, 2, dtype=torch.float32),

448 + selected_token_ranks=torch.zeros(num_tokens, dtype=torch.int32),

449 + ),

450 + )

451 + )

452 + return EngineCoreOutputs(outputs=outputs)

453 +

454 +

455 +def test_payload_buffer_reuse_does_not_corrupt_in_flight_messages():

456 + """The engine core recycles the msgpack payload buffer across messages

457 + (`MsgpackEncoder.encode_into`). It may only do so once zmq has finished

458 + sending that buffer, otherwise a newer payload is delivered alongside the

459 + older message's zero-copy tensor frames.

460 +

461 + `Socket.send_multipart(track=True)` cannot be used to detect this: it

462 + returns a tracker for the last frame only, and pyzmq copies frames below

463 + `zmq.COPY_THRESHOLD` and reports them as already-sent.

464 + """

465 + import zmq

466 +

467 + from vllm.v1.engine import EngineCoreOutputs

468 + from vllm.v1.engine.core import EngineCoreProc

469 +

470 + num_msgs = 100

471 + encoder = MsgpackEncoder()

472 + decoder = MsgpackDecoder(EngineCoreOutputs)

473 + # Enough requests that the payload frame is zero-copied rather than copied

474 + # by pyzmq, which is what makes early reuse observable.

475 + messages = [_logprobs_outputs(300, 24 + i % 8) for i in range(num_msgs)]

476 + assert len(encoder.encode(messages[0])[0]) >= zmq.COPY_THRESHOLD

477 +

478 + reuse_buffers: list[bytearray] = []

479 + pending: list[tuple[zmq.MessageTracker, bytearray]] = []

480 + with zmq.Context() as ctx:

481 + push = ctx.socket(zmq.PUSH)

482 + push.bind("inproc://test-payload-reuse")

483 + pull = ctx.socket(zmq.PULL)

484 + pull.connect("inproc://test-payload-reuse")

485 +

486 + for outputs in messages:

487 + while pending and pending[0][0].done:

488 + reuse_buffers.append(pending.pop(0)[1])

489 + buffer = reuse_buffers.pop() if reuse_buffers else bytearray()

490 + buffers = encoder.encode_into(outputs, buffer)

491 + tracker = EngineCoreProc._send_msg_tracking_payload(push, buffers)

492 + if tracker.done:

493 + reuse_buffers.append(buffer)

494 + else:

495 + pending.append((tracker, buffer))

496 +

497 + for i, sent in enumerate(messages):

498 + received = decoder.decode(pull.recv_multipart(copy=False))

499 + assert len(received.outputs) == len(sent.outputs), f"message {i}"

500 + for expected, actual in zip(sent.outputs, received.outputs):

501 + sent_ids = expected.new_prompt_logprobs_tensors.logprob_token_ids

502 + got_ids = actual.new_prompt_logprobs_tensors.logprob_token_ids

503 + assert actual.request_id == expected.request_id, f"message {i}"

504 + assert torch.equal(got_ids, sent_ids), (

505 + f"message {i} request {actual.request_id}: corrupted "

506 + f"prompt logprobs, {got_ids.shape} vs {sent_ids.shape}"

507 + )

508 + push.close(linger=0)

509 + pull.close(linger=0)

510 +

511 +

512 +def test_zero_copy_frames_survive_without_caller_side_references():

513 + """Callers don't need to retain the encoded object until zmq has sent it:

514 + for a zero-copy frame, zmq holds its own reference to the backing buffer.

515 +

516 + The engine core clients rely on this when sending requests that carry

517 + tensors (e.g. prompt embeds) without tracking the messages.

518 +

519 + What makes that safe is that `tensor_data()` hands zmq a memoryview which

520 + transitively references the source tensor, so refcounting - not timing -

521 + keeps the memory from being freed and reused underneath zmq.

522 + """

523 + import gc

524 +

525 + import zmq

526 +

527 + from vllm.v1.utils import tensor_data

528 +

529 + num_elems = 100_000 # comfortably over zmq.COPY_THRESHOLD

530 + expected = torch.arange(num_elems, dtype=torch.int64)

531 + encoder = MsgpackEncoder()

532 + decoder = MsgpackDecoder(RequestWithTensor)

533 +

534 + # The buffer handed to zmq must keep the tensor's storage alive by itself.

535 + holder = tensor_data(expected).obj

536 + while getattr(holder, "base", None) is not None:

537 + holder = holder.base

538 + assert isinstance(holder, torch.Tensor)

539 + assert holder.data_ptr() == expected.data_ptr()

540 +

541 + with zmq.Context() as ctx:

542 + push = ctx.socket(zmq.PUSH)

543 + push.bind("inproc://test-zero-copy-lifetime")

544 + pull = ctx.socket(zmq.PULL)

545 + pull.connect("inproc://test-zero-copy-lifetime")

546 +

547 + request = RequestWithTensor(prompt_embeds=expected.clone(), data="req")

548 + buffers = encoder.encode(request)

549 + assert max(len(buf) for buf in buffers) >= zmq.COPY_THRESHOLD

550 + push.send_multipart(buffers, copy=False)

551 +

552 + # Drop every reference the sender holds, then churn the allocator.

553 + del request, buffers

554 + gc.collect()

555 + torch.arange(num_elems * 4, dtype=torch.int64)

556 +

557 + decoded = decoder.decode(pull.recv_multipart(copy=False))

558 + assert decoded.prompt_embeds is not None

559 + assert torch.equal(decoded.prompt_embeds, expected)

560 + push.close(linger=0)

561 + pull.close(linger=0)

vllm/envs.py

+1 -1

@@ -1547,7 +1547,7 @@ environment_variables: dict[str, Callable[[], Any]] = {

1547 # tensors above will instead be sent via a separate message.

1548 # While the sending side still actually copies the tensor

1549 # in all cases, on the receiving side, tensors above this

1550 - # limit will actually be zero-copy decoded.

1550 + # limit will actually be zero-copy decoded. The unit is bytes.

1551 "VLLM_MSGPACK_ZERO_COPY_THRESHOLD": lambda: int(

1552 os.getenv("VLLM_MSGPACK_ZERO_COPY_THRESHOLD", "256")

1553 ),

vllm/v1/engine/core.py

+30 -11

@@ -7,7 +7,7 @@ import signal

7 import threading

8 import time

9 from collections import defaultdict, deque

10 -from collections.abc import Callable, Generator

10 +from collections.abc import Callable, Generator, Sequence

11 from concurrent.futures import Future

12 from contextlib import ExitStack, contextmanager

13 from enum import IntEnum

@@ -87,7 +87,7 @@ from vllm.v1.kv_cache_interface import KVCacheConfig, get_kv_cache_spec_kind

87 from vllm.v1.metrics.stats import SchedulerIterationDetails, SchedulerStats

88 from vllm.v1.outputs import ModelRunnerOutput

89 from vllm.v1.request import Request, RequestStatus

90 -from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder

90 +from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder, bytestr

91 from vllm.v1.structured_output import StructuredOutputManager

92 from vllm.v1.utils import compute_iteration_details

93 from vllm.version import __version__ as VLLM_VERSION

@@ -1748,10 +1748,11 @@ class EngineCoreProc(EngineCore):

1748 encoder = MsgpackEncoder()

1749 # Send buffers to reuse.

1750 reuse_buffers: list[bytearray] = []

1751 - # Keep references to outputs and buffers until zmq is finished

1752 - # with them (outputs may contain tensors/np arrays whose

1753 - # backing buffers were extracted for zero-copy send).

1754 - pending = deque[tuple[zmq.MessageTracker, Any, bytearray]]()

1751 + # Payload buffers that can't be reused yet because zmq may still be

1752 + # sending them.

1753 + # Buffers of the zero-copy tensor/ndarray frames don't need tracking

1754 + # here: zmq itself holds a reference to each until it's done with it.

1755 + pending = deque[tuple[zmq.MessageTracker, bytearray]]()

1756

1757 # We must set linger to ensure the ENGINE_CORE_DEAD

1758 # message is sent prior to closing the socket.

@@ -1792,20 +1793,38 @@ class EngineCoreProc(EngineCore):

1793

1794 # Reclaim buffers that zmq is finished with.

1795 while pending and pending[-1][0].done:

1795 - reuse_buffers.append(pending.pop()[2])

1796 + reclaimed = pending.pop()[1]

1797 + if len(reuse_buffers) < max_reuse_bufs:

1798 + reuse_buffers.append(reclaimed)

1799

1800 buffer = reuse_buffers.pop() if reuse_buffers else bytearray()

1801 buffers = encoder.encode_into(outputs, buffer)

1799 - tracker = sockets[client_index].send_multipart(

1800 - buffers, copy=False, track=True

1802 + tracker = self._send_msg_tracking_payload(

1803 + sockets[client_index], buffers

1804 )

1805 if not tracker.done:

1803 - ref = outputs if len(buffers) > 1 else None

1804 - pending.appendleft((tracker, ref, buffer))

1806 + pending.appendleft((tracker, buffer))

1807 elif len(reuse_buffers) < max_reuse_bufs:

1808 # Limit the number of buffers to reuse.

1809 reuse_buffers.append(buffer)

1810

1811 + @staticmethod

1812 + def _send_msg_tracking_payload(

1813 + socket: zmq.Socket, buffers: Sequence[bytestr]

1814 + ) -> zmq.MessageTracker:

1815 + """Send `buffers` as a zero-copy multipart message, returning a tracker

1816 + for the *first* frame.

1817 +

1818 + Used instead of `Socket.send_multipart()` because we reuse the buffer

1819 + passed to `MsgpackEncoder.encode_into()`: `send_multipart()` returns a

1820 + tracker for the last frame only.

1821 + """

1822 + more_flag = zmq.SNDMORE if len(buffers) > 1 else 0

1823 + tracker = socket.send(buffers[0], more_flag, copy=False, track=True)

1824 + if more_flag:

1825 + socket.send_multipart(buffers[1:], copy=False)

1826 + return tracker

1827 +

1828 def _handle_request_preproc_error(self, request: EngineCoreRequest) -> None:

1829 """Log and return a request-scoped error response for exceptions raised

1830 from the add request preprocessing in the input socket processing thread.

vllm/v1/engine/core_client.py

+12 -46

@@ -7,7 +7,7 @@ import sys

7 import uuid

8 import weakref

9 from abc import ABC, abstractmethod

10 -from collections import Counter, defaultdict, deque

10 +from collections import Counter, defaultdict

11 from collections.abc import Awaitable, Callable, Sequence

12 from concurrent.futures import Future

13 from dataclasses import dataclass

@@ -671,11 +671,6 @@ class MPClient(EngineCoreClient):

671 self.core_engine: EngineIdentity = self.core_engines[0]

672 self.utility_results: dict[int, AnyFuture] = {}

673

674 - # Request objects which may contain pytorch-allocated tensors

675 - # that we need to keep references to until zmq is done with the

676 - # underlying data.

677 - self.pending_messages = deque[tuple[zmq.MessageTracker, Any]]()

678 -

674 # Start monitoring engine core processes for unexpected failures

675 self.start_engine_core_monitor()

676

@@ -707,14 +702,6 @@ class MPClient(EngineCoreClient):

702 if self.resources.engine_dead:

703 raise EngineDeadError()

704

710 - def add_pending_message(self, tracker: zmq.MessageTracker, msg: Any):

711 - if not tracker.done:

712 - self.pending_messages.appendleft((tracker, msg))

713 -

714 - def free_pending_messages(self):

715 - while self.pending_messages and self.pending_messages[-1][0].done:

716 - self.pending_messages.pop()

717 -

705 def dp_engines_running(self) -> bool:

706 return self.engines_running

707

@@ -896,17 +883,12 @@ class SyncMPClient(MPClient):

883

884 def _send_input(self, request_type: EngineCoreRequestType, request: Any):

885 self.ensure_alive()

899 - self.free_pending_messages()

886 # (Identity, RequestType, SerializedRequest)

887 msg = (self.core_engine, request_type.value, *self.encoder.encode(request))

902 -

903 - if len(msg) <= 3:

904 - # No auxiliary buffers => no tensor backing buffers in request.

905 - self.input_socket.send_multipart(msg, copy=False)

906 - return

907 -

908 - tracker = self.input_socket.send_multipart(msg, copy=False, track=True)

909 - self.add_pending_message(tracker, request)

888 + # Any zero-copy tensor/ndarray frames are kept alive by zmq itself

889 + # until it's finished sending them (there is a ref chain from the underlying

890 + # memoryview back to the original owning tensor/ndarray).

891 + self.input_socket.send_multipart(msg, copy=False)

892

893 def call_utility(self, method: str, *args) -> Any:

894 call_id = uuid.uuid1().int >> 64

@@ -1129,32 +1111,16 @@ class AsyncMPClient(MPClient):

1111 engine = self.core_engine

1112

1113 message = (request_type.value, *self.encoder.encode(request))

1132 - return self._send_input_message(message, engine, request)

1114 + return self._send_input_message(message, engine)

1115

1116 def _send_input_message(

1135 - self, message: tuple[bytestr, ...], engine: EngineIdentity, objects: Any

1117 + self, message: tuple[bytestr, ...], engine: EngineIdentity

1118 ) -> Awaitable[Any]:

1137 - """

1138 - objects is a reference to retain until zmq is finished with the

1139 - buffers, in case they were extracted from tensors in the request.

1140 - """

1119 self.ensure_alive()

1142 - self.free_pending_messages()

1143 -

1144 - msg = (engine,) + message

1145 - if not objects or len(msg) <= 3:

1146 - # No auxiliary buffers => no tensor backing buffers in request.

1147 - return self.input_socket.send_multipart(msg, copy=False)

1148 -

1149 - future: asyncio.Future[zmq.MessageTracker]

1150 - future = self.input_socket.send_multipart(msg, copy=False, track=True)

1151 -

1152 - def add_pending(f: asyncio.Future[zmq.MessageTracker]):

1153 - with contextlib.suppress(BaseException):

1154 - self.add_pending_message(f.result(), objects)

1155 -

1156 - future.add_done_callback(add_pending)

1157 - return future

1120 + # Any zero-copy tensor/ndarray frames are kept alive by zmq itself

1121 + # until it's finished sending them (there is a ref chain from the underlying

1122 + # memoryview back to the original owning tensor/ndarray).

1123 + return self.input_socket.send_multipart((engine,) + message, copy=False)

1124

1125 async def call_utility_async(self, method: str, *args) -> Any:

1126 return await self._call_utility_async(method, *args, engine=self.core_engine)

@@ -1169,7 +1135,7 @@ class AsyncMPClient(MPClient):

1135 EngineCoreRequestType.UTILITY.value,

1136 *self.encoder.encode((self.client_index, call_id, method, args)),

1137 )

1172 - await self._send_input_message(message, engine, args)

1138 + await self._send_input_message(message, engine)

1139 self._ensure_output_queue_task()

1140 return await future

1141