[Core] Cheapen the spin while waiting for model output

tests/distributed/test_shm_broadcast.py

@@ -757,15 +757,50 @@ def test_readable_reports_pending_message_without_consuming():
         reader.shutdown()
 
 
-def test_readable_sleeps_on_remote_socket():
+def test_spin_until_readable_returns_on_each_condition():
+    writer = MessageQueue(
+        n_reader=1,
+        n_local_reader=1,
+        max_chunk_bytes=1024 * 1024,
+        max_chunks=2,
+    )
+    reader = MessageQueue.create_from_handle(writer.export_handle(), rank=0)
+    try:
+        writer.wait_until_ready()
+        reader.wait_until_ready()
+        far_future = time.monotonic() + 60
+
+        # None of these would return before the deadline if the condition failed.
+        reader.spin_until_readable(["input"], far_future)
+        reader.spin_until_readable([], time.monotonic())
+        assert not reader.readable()
+
+        threading.Timer(0.05, writer.enqueue, args=("message",)).start()
+        reader.spin_until_readable([], far_future)
+        assert reader.readable()
+        # Doesn't consume the message.
+        assert reader.dequeue(timeout=5) == "message"
+    finally:
+        writer.shutdown()
+        reader.shutdown()
+
+
+def test_spin_until_readable_sleeps_on_remote_socket():
+    reader = mock.Mock(_is_local_reader=False, _is_remote_reader=True)
+    MessageQueue.spin_until_readable(reader, [], time.monotonic() + 60)
+    reader.remote_socket.poll.assert_called_once_with(
+        timeout=shm_broadcast.REMOTE_POLL_WAIT_MS
+    )
+    reader.buffer.get_metadata.assert_not_called()
+
+
+def test_readable_does_not_wait_on_remote_socket():
     reader = mock.Mock(
         shutting_down=False, _is_local_reader=False, _is_remote_reader=True
     )
     reader.remote_socket.poll.return_value = 0
     assert MessageQueue.readable(reader) is False
-    reader.remote_socket.poll.assert_called_once_with(
-        timeout=shm_broadcast.REMOTE_READABLE_WAIT_MS
-    )
+    reader.remote_socket.poll.assert_called_once_with(timeout=0)
 
 
 def test_reader_in_reference_cycle_is_collected_without_blocking():

tests/v1/engine/test_inputs_during_step.py

@@ -35,6 +35,7 @@ def __init__(
         self.ready_after = ready_after
         self.on_poll = on_poll
         self.polls = 0
+        self.waits = 0
 
     def poll(self) -> bool:
         self.polls += 1
@@ -44,6 +45,9 @@ def poll(self) -> bool:
             self.set_result("output")
         return self.done()
 
+    def wait_briefly(self, pending, deadline) -> None:
+        self.waits += 1
+
     def result(self, timeout=None):
         # Stands in for blocking until the output arrives.
         if not self.done():
@@ -128,11 +132,8 @@ def test_blocks_for_output_after_spin_period():
     engine, handled = make_engine()
     engine.input_queue.put((ADD, "a"))
     future = PollableFuture(ready_after=1000)
-    with (
-        mock.patch.object(engine_core, "_INPUT_SPIN_DURATION_S", 0.0),
-        mock.patch.object(engine_core, "sched_yield") as sched_yield,
-    ):
+    with mock.patch.object(engine_core, "_INPUT_SPIN_DURATION_S", 0.0):
         assert engine._wait_for_output(future) == "output"
-    sched_yield.assert_not_called()
+    assert future.waits == 0
     assert future.polls == 2
     assert handled == [(ADD, "a")]

tests/v1/executor/test_future_wrapper_poll.py

@@ -82,3 +82,21 @@ def fail():
     assert future.poll()
     with pytest.raises(ValueError, match="worker failed"):
         future.result()
+
+
+def test_wait_briefly_waits_for_oldest_response():
+    queue: deque[FutureWrapper] = deque()
+    waited: list[str] = []
+
+    def waiter(name: str):
+        return lambda pending, deadline: waited.append(name)
+
+    f1 = FutureWrapper(queue, lambda: "a", wait_for_response=waiter("first"))
+    f2 = FutureWrapper(queue, lambda: "b", wait_for_response=waiter("second"))
+
+    f2.wait_briefly([], 0.0)
+    assert waited == ["first"]
+
+    assert f1.poll()
+    f2.wait_briefly([], 0.0)
+    assert waited == ["first", "second"]

vllm/distributed/device_communicators/shm_broadcast.py

@@ -10,6 +10,7 @@
 import threading
 import time
 import weakref
+from collections.abc import Sized
 from contextlib import contextmanager
 from dataclasses import dataclass, field
 from multiprocessing import shared_memory
@@ -56,7 +57,8 @@
             "spinloop extension could not be loaded, disabling VLLM_USE_SPINLOOP_EXT!"
         )
 SPINLOOP_TIMEOUT_SECONDS = 0.1
-REMOTE_READABLE_WAIT_MS = 1
+# How long MessageQueue.spin_until_readable() sleeps on a remote reader's socket.
+REMOTE_POLL_WAIT_MS = 1
 
 if TYPE_CHECKING:
     from _typeshed import SizedBuffer
@@ -884,19 +886,36 @@ def _is_readable(self, metadata_buffer) -> bool:
         return bool(written_flag) and not read_flag
 
     def readable(self) -> bool:
-        """Check whether a message is ready to be dequeued, without consuming it.
-        A local reader doesn't wait; a remote reader waits up to
-        REMOTE_READABLE_WAIT_MS for a message to arrive.
+        """Non-blocking check for whether a message is ready to be dequeued.
         Raises RuntimeError if the queue has been shut down."""
         if self.shutting_down:
             raise RuntimeError("cancelled")
         if self._is_local_reader:
             with self.buffer.get_metadata(self.current_idx) as metadata_buffer:
                 return self._is_readable(metadata_buffer)
         if self._is_remote_reader:
-            return bool(self.remote_socket.poll(timeout=REMOTE_READABLE_WAIT_MS))
+            return bool(self.remote_socket.poll(timeout=0))
         raise RuntimeError("Only readers can check readability")
 
+    def spin_until_readable(self, pending: Sized, deadline: float) -> None:
+        """Spin until a message is ready to be dequeued, `pending` is non-empty or
+        time.monotonic() reaches `deadline`. Callers must recheck readable().
+
+        Iterations are kept minimal, since each holds the GIL until it yields. A
+        remote reader instead sleeps briefly on its socket, without the GIL.
+        """
+        if self._is_remote_reader:
+            self.remote_socket.poll(timeout=REMOTE_POLL_WAIT_MS)
+            return
+        read_flag_idx = self.local_reader_rank + 1
+        monotonic = time.monotonic
+        with self.buffer.get_metadata(self.current_idx) as metadata_buffer:
+            while not pending and monotonic() < deadline:
+                memory_fence()
+                if metadata_buffer[0] and not metadata_buffer[read_flag_idx]:
+                    return
+                sched_yield()
+
     @contextmanager
     def acquire_read(
         self,

vllm/v1/engine/core.py

@@ -7,7 +7,7 @@
 import threading
 import time
 from collections import defaultdict, deque
-from collections.abc import Callable, Generator, Sequence
+from collections.abc import Callable, Generator, Sequence, Sized
 from concurrent.futures import Future
 from contextlib import ExitStack, contextmanager
 from enum import IntEnum
@@ -32,7 +32,6 @@
 from vllm.distributed.ec_transfer.ec_connector.encoder_output_width import (
     measure_encoder_output_widths,
 )
-from vllm.distributed.utils import sched_yield
 from vllm.envs import enable_envs_cache
 from vllm.exceptions import VLLMValidationError
 from vllm.logger import configure_logging, init_logger
@@ -1900,7 +1899,10 @@ def _wait_for_output(self, future: Future[Any]) -> Any:
         """Wait for a step's model output, meanwhile handling new requests and
         aborts as they arrive."""
         poll: Callable[[], bool] | None = getattr(future, "poll", None)
-        if poll is not None:
+        wait_briefly: Callable[[Sized, float], None] | None = getattr(
+            future, "wait_briefly", None
+        )
+        if poll is not None and wait_briefly is not None:
             pending_inputs = self.input_queue.queue
             spin_deadline = time.monotonic() + _INPUT_SPIN_DURATION_S
             while not self._deferred_inputs and not poll():
@@ -1912,7 +1914,9 @@ def _wait_for_output(self, future: Future[Any]) -> Any:
                     else:
                         self._deferred_inputs.append(req)
                 elif time.monotonic() < spin_deadline:
-                    sched_yield()
+                    # Spin on something cheaper than poll(), to leave the GIL
+                    # free more of the time for the input and output threads.
+                    wait_briefly(pending_inputs, spin_deadline)
                 else:
                     # Long-running step: stop busy-looping and block until the
                     # output arrives. Later inputs are handled after the step.

vllm/v1/executor/multiproc_executor.py

@@ -10,7 +10,7 @@
 import traceback
 import weakref
 from collections import deque
-from collections.abc import Callable, Sequence
+from collections.abc import Callable, Sequence, Sized
 from concurrent.futures import Future, InvalidStateError, ThreadPoolExecutor
 from contextlib import suppress
 from dataclasses import dataclass
@@ -96,13 +96,17 @@ def __init__(
         get_response: Callable[[], Any],
         aggregate: Callable = lambda x: x,
         response_ready: Callable[[], bool] | None = None,
+        wait_for_response: Callable[[Sized, float], None] | None = None,
     ):
         self.futures_queue = futures_queue
         self.get_response = get_response
         self.aggregate = aggregate
-        # Check, waiting at most briefly, for whether get_response() can return
-        # without waiting. If None, poll() treats the response as ready (and blocks).
+        # Non-blocking check for whether get_response() can return without
+        # waiting. If None, poll() treats the response as ready (and blocks).
         self.response_ready = response_ready
+        # Spins until get_response() is ready, the given collection is non-empty
+        # or time.monotonic() reaches the given deadline.
+        self.wait_for_response = wait_for_response
         super().__init__()
         self.futures_queue.appendleft(self)
 
@@ -125,6 +129,13 @@ def poll(self) -> bool:
             future._wait_for_response()
         return True
 
+    def wait_briefly(self, pending: Sized, deadline: float) -> None:
+        """Wait for the oldest uncollected response, until `pending` is non-empty
+        or time.monotonic() reaches `deadline`. Call after poll() returns False."""
+        wait_for_response = self.futures_queue[-1].wait_for_response
+        if wait_for_response is not None:
+            wait_for_response(pending, deadline)
+
     def _wait_for_response(self):
         try:
             response = self.aggregate(self.get_response())
@@ -562,11 +573,18 @@ def get_response():
 
             return responses[0] if output_rank is not None else responses
 
+        def wait_for_response(pending: Sized, deadline: float) -> None:
+            for mq in response_mqs:
+                if not mq.readable():
+                    mq.spin_until_readable(pending, deadline)
+                    return
+
         future = FutureWrapper(
             self.futures_queue,
             get_response=get_response,
             aggregate=aggregate,
             response_ready=lambda: all(mq.readable() for mq in response_mqs),
+            wait_for_response=wait_for_response,
         )
 
         return future if non_block else future.result()