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

Signed-off-by: Nick Hill <[email protected]>
Co-authored-by: Andreas Karatzas <[email protected]>

tests/v1/test_serial_utils.py

@@ -423,3 +423,139 @@ def test_multiple_senders_single_receiver_ipc():
         assert torch.allclose(decoded.prompt_embeds, original_tensor), (
             f"Value mismatch for sender {sender_idx} msg {msg_idx}"
         )
+
+
+def _logprobs_outputs(num_reqs: int, num_prompt_tokens: int):
+    """An EngineCoreOutputs carrying prompt logprobs, as the engine core sends
+    it: many requests, each with per-token tensors small enough that pyzmq
+    copies their frames, while the accumulated payload frame is large enough
+    that pyzmq sends it zero-copy."""
+    from vllm.v1.engine import EngineCoreOutput, EngineCoreOutputs
+    from vllm.v1.outputs import LogprobsTensors
+
+    outputs = []
+    for req in range(num_reqs):
+        num_tokens = num_prompt_tokens + req % 4
+        outputs.append(
+            EngineCoreOutput(
+                request_id=f"req-{req:08d}",
+                new_token_ids=[req],
+                new_prompt_logprobs_tensors=LogprobsTensors(
+                    logprob_token_ids=torch.arange(
+                        num_tokens * 2, dtype=torch.int64
+                    ).view(num_tokens, 2),
+                    logprobs=torch.zeros(num_tokens, 2, dtype=torch.float32),
+                    selected_token_ranks=torch.zeros(num_tokens, dtype=torch.int32),
+                ),
+            )
+        )
+    return EngineCoreOutputs(outputs=outputs)
+
+
+def test_payload_buffer_reuse_does_not_corrupt_in_flight_messages():
+    """The engine core recycles the msgpack payload buffer across messages
+    (`MsgpackEncoder.encode_into`). It may only do so once zmq has finished
+    sending that buffer, otherwise a newer payload is delivered alongside the
+    older message's zero-copy tensor frames.
+
+    `Socket.send_multipart(track=True)` cannot be used to detect this: it
+    returns a tracker for the last frame only, and pyzmq copies frames below
+    `zmq.COPY_THRESHOLD` and reports them as already-sent.
+    """
+    import zmq
+
+    from vllm.v1.engine import EngineCoreOutputs
+    from vllm.v1.engine.core import EngineCoreProc
+
+    num_msgs = 100
+    encoder = MsgpackEncoder()
+    decoder = MsgpackDecoder(EngineCoreOutputs)
+    # Enough requests that the payload frame is zero-copied rather than copied
+    # by pyzmq, which is what makes early reuse observable.
+    messages = [_logprobs_outputs(300, 24 + i % 8) for i in range(num_msgs)]
+    assert len(encoder.encode(messages[0])[0]) >= zmq.COPY_THRESHOLD
+
+    reuse_buffers: list[bytearray] = []
+    pending: list[tuple[zmq.MessageTracker, bytearray]] = []
+    with zmq.Context() as ctx:
+        push = ctx.socket(zmq.PUSH)
+        push.bind("inproc://test-payload-reuse")
+        pull = ctx.socket(zmq.PULL)
+        pull.connect("inproc://test-payload-reuse")
+
+        for outputs in messages:
+            while pending and pending[0][0].done:
+                reuse_buffers.append(pending.pop(0)[1])
+            buffer = reuse_buffers.pop() if reuse_buffers else bytearray()
+            buffers = encoder.encode_into(outputs, buffer)
+            tracker = EngineCoreProc._send_msg_tracking_payload(push, buffers)
+            if tracker.done:
+                reuse_buffers.append(buffer)
+            else:
+                pending.append((tracker, buffer))
+
+        for i, sent in enumerate(messages):
+            received = decoder.decode(pull.recv_multipart(copy=False))
+            assert len(received.outputs) == len(sent.outputs), f"message {i}"
+            for expected, actual in zip(sent.outputs, received.outputs):
+                sent_ids = expected.new_prompt_logprobs_tensors.logprob_token_ids
+                got_ids = actual.new_prompt_logprobs_tensors.logprob_token_ids
+                assert actual.request_id == expected.request_id, f"message {i}"
+                assert torch.equal(got_ids, sent_ids), (
+                    f"message {i} request {actual.request_id}: corrupted "
+                    f"prompt logprobs, {got_ids.shape} vs {sent_ids.shape}"
+                )
+        push.close(linger=0)
+        pull.close(linger=0)
+
+
+def test_zero_copy_frames_survive_without_caller_side_references():
+    """Callers don't need to retain the encoded object until zmq has sent it:
+    for a zero-copy frame, zmq holds its own reference to the backing buffer.
+
+    The engine core clients rely on this when sending requests that carry
+    tensors (e.g. prompt embeds) without tracking the messages.
+
+    What makes that safe is that `tensor_data()` hands zmq a memoryview which
+    transitively references the source tensor, so refcounting - not timing -
+    keeps the memory from being freed and reused underneath zmq.
+    """
+    import gc
+
+    import zmq
+
+    from vllm.v1.utils import tensor_data
+
+    num_elems = 100_000  # comfortably over zmq.COPY_THRESHOLD
+    expected = torch.arange(num_elems, dtype=torch.int64)
+    encoder = MsgpackEncoder()
+    decoder = MsgpackDecoder(RequestWithTensor)
+
+    # The buffer handed to zmq must keep the tensor's storage alive by itself.
+    holder = tensor_data(expected).obj
+    while getattr(holder, "base", None) is not None:
+        holder = holder.base
+    assert isinstance(holder, torch.Tensor)
+    assert holder.data_ptr() == expected.data_ptr()
+
+    with zmq.Context() as ctx:
+        push = ctx.socket(zmq.PUSH)
+        push.bind("inproc://test-zero-copy-lifetime")
+        pull = ctx.socket(zmq.PULL)
+        pull.connect("inproc://test-zero-copy-lifetime")
+
+        request = RequestWithTensor(prompt_embeds=expected.clone(), data="req")
+        buffers = encoder.encode(request)
+        assert max(len(buf) for buf in buffers) >= zmq.COPY_THRESHOLD
+        push.send_multipart(buffers, copy=False)
+
+        # Drop every reference the sender holds, then churn the allocator.
+        del request, buffers
+        gc.collect()
+        torch.arange(num_elems * 4, dtype=torch.int64)
+
+        decoded = decoder.decode(pull.recv_multipart(copy=False))
+        assert decoded.prompt_embeds is not None
+        assert torch.equal(decoded.prompt_embeds, expected)
+        push.close(linger=0)
+        pull.close(linger=0)

vllm/envs.py

@@ -1547,7 +1547,7 @@ environment_variables: dict[str, Callable[[], Any]] = {
     # tensors above will instead be sent via a separate message.
     # While the sending side still actually copies the tensor
     # in all cases, on the receiving side, tensors above this
-    # limit will actually be zero-copy decoded.
+    # limit will actually be zero-copy decoded. The unit is bytes.
     "VLLM_MSGPACK_ZERO_COPY_THRESHOLD": lambda: int(
         os.getenv("VLLM_MSGPACK_ZERO_COPY_THRESHOLD", "256")
     ),

vllm/v1/engine/core.py

@@ -7,7 +7,7 @@ import signal
 import threading
 import time
 from collections import defaultdict, deque
-from collections.abc import Callable, Generator
+from collections.abc import Callable, Generator, Sequence
 from concurrent.futures import Future
 from contextlib import ExitStack, contextmanager
 from enum import IntEnum
@@ -87,7 +87,7 @@ from vllm.v1.kv_cache_interface import KVCacheConfig, get_kv_cache_spec_kind
 from vllm.v1.metrics.stats import SchedulerIterationDetails, SchedulerStats
 from vllm.v1.outputs import ModelRunnerOutput
 from vllm.v1.request import Request, RequestStatus
-from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder
+from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder, bytestr
 from vllm.v1.structured_output import StructuredOutputManager
 from vllm.v1.utils import compute_iteration_details
 from vllm.version import __version__ as VLLM_VERSION
@@ -1748,10 +1748,11 @@ class EngineCoreProc(EngineCore):
         encoder = MsgpackEncoder()
         # Send buffers to reuse.
         reuse_buffers: list[bytearray] = []
-        # Keep references to outputs and buffers until zmq is finished
-        # with them (outputs may contain tensors/np arrays whose
-        # backing buffers were extracted for zero-copy send).
-        pending = deque[tuple[zmq.MessageTracker, Any, bytearray]]()
+        # Payload buffers that can't be reused yet because zmq may still be
+        # sending them.
+        # Buffers of the zero-copy tensor/ndarray frames don't need tracking
+        # here: zmq itself holds a reference to each until it's done with it.
+        pending = deque[tuple[zmq.MessageTracker, bytearray]]()
 
         # We must set linger to ensure the ENGINE_CORE_DEAD
         # message is sent prior to closing the socket.
@@ -1792,20 +1793,38 @@ class EngineCoreProc(EngineCore):
 
                 # Reclaim buffers that zmq is finished with.
                 while pending and pending[-1][0].done:
-                    reuse_buffers.append(pending.pop()[2])
+                    reclaimed = pending.pop()[1]
+                    if len(reuse_buffers) < max_reuse_bufs:
+                        reuse_buffers.append(reclaimed)
 
                 buffer = reuse_buffers.pop() if reuse_buffers else bytearray()
                 buffers = encoder.encode_into(outputs, buffer)
-                tracker = sockets[client_index].send_multipart(
-                    buffers, copy=False, track=True
+                tracker = self._send_msg_tracking_payload(
+                    sockets[client_index], buffers
                 )
                 if not tracker.done:
-                    ref = outputs if len(buffers) > 1 else None
-                    pending.appendleft((tracker, ref, buffer))
+                    pending.appendleft((tracker, buffer))
                 elif len(reuse_buffers) < max_reuse_bufs:
                     # Limit the number of buffers to reuse.
                     reuse_buffers.append(buffer)
 
+    @staticmethod
+    def _send_msg_tracking_payload(
+        socket: zmq.Socket, buffers: Sequence[bytestr]
+    ) -> zmq.MessageTracker:
+        """Send `buffers` as a zero-copy multipart message, returning a tracker
+        for the *first* frame.
+
+        Used instead of `Socket.send_multipart()` because we reuse the buffer
+        passed to `MsgpackEncoder.encode_into()`: `send_multipart()` returns a
+        tracker for the last frame only.
+        """
+        more_flag = zmq.SNDMORE if len(buffers) > 1 else 0
+        tracker = socket.send(buffers[0], more_flag, copy=False, track=True)
+        if more_flag:
+            socket.send_multipart(buffers[1:], copy=False)
+        return tracker
+
     def _handle_request_preproc_error(self, request: EngineCoreRequest) -> None:
         """Log and return a request-scoped error response for exceptions raised
         from the add request preprocessing in the input socket processing thread.

vllm/v1/engine/core_client.py

@@ -7,7 +7,7 @@ import sys
 import uuid
 import weakref
 from abc import ABC, abstractmethod
-from collections import Counter, defaultdict, deque
+from collections import Counter, defaultdict
 from collections.abc import Awaitable, Callable, Sequence
 from concurrent.futures import Future
 from dataclasses import dataclass
@@ -671,11 +671,6 @@ class MPClient(EngineCoreClient):
             self.core_engine: EngineIdentity = self.core_engines[0]
             self.utility_results: dict[int, AnyFuture] = {}
 
-            # Request objects which may contain pytorch-allocated tensors
-            # that we need to keep references to until zmq is done with the
-            # underlying data.
-            self.pending_messages = deque[tuple[zmq.MessageTracker, Any]]()
-
             # Start monitoring engine core processes for unexpected failures
             self.start_engine_core_monitor()
 
@@ -707,14 +702,6 @@ class MPClient(EngineCoreClient):
         if self.resources.engine_dead:
             raise EngineDeadError()
 
-    def add_pending_message(self, tracker: zmq.MessageTracker, msg: Any):
-        if not tracker.done:
-            self.pending_messages.appendleft((tracker, msg))
-
-    def free_pending_messages(self):
-        while self.pending_messages and self.pending_messages[-1][0].done:
-            self.pending_messages.pop()
-
     def dp_engines_running(self) -> bool:
         return self.engines_running
 
@@ -896,17 +883,12 @@ class SyncMPClient(MPClient):
 
     def _send_input(self, request_type: EngineCoreRequestType, request: Any):
         self.ensure_alive()
-        self.free_pending_messages()
         # (Identity, RequestType, SerializedRequest)
         msg = (self.core_engine, request_type.value, *self.encoder.encode(request))
-
-        if len(msg) <= 3:
-            # No auxiliary buffers => no tensor backing buffers in request.
-            self.input_socket.send_multipart(msg, copy=False)
-            return
-
-        tracker = self.input_socket.send_multipart(msg, copy=False, track=True)
-        self.add_pending_message(tracker, request)
+        # Any zero-copy tensor/ndarray frames are kept alive by zmq itself
+        # until it's finished sending them (there is a ref chain from the underlying
+        # memoryview back to the original owning tensor/ndarray).
+        self.input_socket.send_multipart(msg, copy=False)
 
     def call_utility(self, method: str, *args) -> Any:
         call_id = uuid.uuid1().int >> 64
@@ -1129,32 +1111,16 @@ class AsyncMPClient(MPClient):
             engine = self.core_engine
 
         message = (request_type.value, *self.encoder.encode(request))
-        return self._send_input_message(message, engine, request)
+        return self._send_input_message(message, engine)
 
     def _send_input_message(
-        self, message: tuple[bytestr, ...], engine: EngineIdentity, objects: Any
+        self, message: tuple[bytestr, ...], engine: EngineIdentity
     ) -> Awaitable[Any]:
-        """
-        objects is a reference to retain until zmq is finished with the
-        buffers, in case they were extracted from tensors in the request.
-        """
         self.ensure_alive()
-        self.free_pending_messages()
-
-        msg = (engine,) + message
-        if not objects or len(msg) <= 3:
-            # No auxiliary buffers => no tensor backing buffers in request.
-            return self.input_socket.send_multipart(msg, copy=False)
-
-        future: asyncio.Future[zmq.MessageTracker]
-        future = self.input_socket.send_multipart(msg, copy=False, track=True)
-
-        def add_pending(f: asyncio.Future[zmq.MessageTracker]):
-            with contextlib.suppress(BaseException):
-                self.add_pending_message(f.result(), objects)
-
-        future.add_done_callback(add_pending)
-        return future
+        # Any zero-copy tensor/ndarray frames are kept alive by zmq itself
+        # until it's finished sending them (there is a ref chain from the underlying
+        # memoryview back to the original owning tensor/ndarray).
+        return self.input_socket.send_multipart((engine,) + message, copy=False)
 
     async def call_utility_async(self, method: str, *args) -> Any:
         return await self._call_utility_async(method, *args, engine=self.core_engine)
@@ -1169,7 +1135,7 @@ class AsyncMPClient(MPClient):
             EngineCoreRequestType.UTILITY.value,
             *self.encoder.encode((self.client_index, call_id, method, args)),
         )
-        await self._send_input_message(message, engine, args)
+        await self._send_input_message(message, engine)
         self._ensure_output_queue_task()
         return await future