From 7c36d63be58f332253ace4e181396a93e559bebb Mon Sep 17 00:00:00 2001 From: nilig Date: Wed, 16 Sep 2026 14:01:15 +0300 Subject: [PATCH 1/2] [KV Offload][P2P] Register peers off the scheduler thread P2PSession._on_connect registered the peer's NIXL metadata synchronously before sending ConnectAck: add_remote_agent plus one descriptor per CPU block (19,599 at a 64 GiB tier). Measured with trace stamps on both peers this is 14-40 ms per side, so it is not the source of the multi-second stalls in #55179, but it still holds a scheduling iteration on both peers for every new session and delays the acknowledgement the peer's first lookup waits behind. NixlTransport runs add_remote_peer on a single-worker executor and opens the agent in NIXL_THREAD_SYNC_STRICT (nixl_agent_config.sync_mode, NIXL 1.4.1) so registration may overlap transfer calls from the scheduler thread. The session sends ConnectAck once the future completes, rejects the peer if it fails, and reports the pending registration as pending work so an idle engine keeps ticking until the ack goes out. Signed-off-by: nilig --- .../tiering/p2p/test_data_transport.py | 36 +++++++- .../v1/kv_offload/tiering/p2p/test_manager.py | 10 +++ .../kv_offload/tiering/p2p/test_sessions.py | 59 +++++++++++++ vllm/v1/kv_offload/tiering/p2p/data/base.py | 29 +++++++ vllm/v1/kv_offload/tiering/p2p/data/nixl.py | 82 ++++++++++++++++++- .../kv_offload/tiering/p2p/session/session.py | 48 +++++++++-- 6 files changed, 249 insertions(+), 15 deletions(-) diff --git a/tests/v1/kv_offload/tiering/p2p/test_data_transport.py b/tests/v1/kv_offload/tiering/p2p/test_data_transport.py index d3bacc8f326f..15aefb10f2c2 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_data_transport.py +++ b/tests/v1/kv_offload/tiering/p2p/test_data_transport.py @@ -324,7 +324,9 @@ def test_non_ucx_backends_passes_backends_kwarg(self): ): NixlTransport("test:1", self._make_view(), backends=["MOONCAKE"]) - config_fn.assert_called_once_with(backends=["MOONCAKE"], capture_telemetry=True) + config_fn.assert_called_once_with( + backends=["MOONCAKE"], capture_telemetry=True, sync_mode=None + ) # num_threads must NOT be passed on the non-UCX branch. assert "num_threads" not in config_fn.call_args.kwargs @@ -340,7 +342,9 @@ def test_ucx_only_passes_num_threads(self): ): NixlTransport("test:1", self._make_view(), num_threads=8) - config_fn.assert_called_once_with(num_threads=8, capture_telemetry=True) + config_fn.assert_called_once_with( + num_threads=8, capture_telemetry=True, sync_mode=None + ) assert "backends" not in config_fn.call_args.kwargs def test_default_backends_is_ucx_only(self): @@ -356,4 +360,30 @@ def test_default_backends_is_ucx_only(self): NixlTransport("test:1", self._make_view()) # Default num_threads=4, no backends kwarg. - config_fn.assert_called_once_with(num_threads=4, capture_telemetry=True) + config_fn.assert_called_once_with( + num_threads=4, capture_telemetry=True, sync_mode=None + ) + + def test_agent_runs_in_strict_thread_sync_when_available(self): + """With a NIXL that exposes the sync enum, the agent serializes its + API so peer registration may run on the worker thread.""" + import sys + import types + + fake_module = types.ModuleType("fake_nixl_api") + fake_module.nixl_thread_sync_t = types.SimpleNamespace( + NIXL_THREAD_SYNC_STRICT="strict" + ) + config_fn = MagicMock(return_value=MagicMock(name="cfg")) + config_fn.__module__ = fake_module.__name__ + agent_cls = MagicMock() + with ( + patch.dict(sys.modules, {fake_module.__name__: fake_module}), + patch("vllm.v1.kv_offload.tiering.p2p.data.nixl._NixlAgent", agent_cls), + patch( + "vllm.v1.kv_offload.tiering.p2p.data.nixl._NixlAgentConfig", config_fn + ), + ): + NixlTransport("test:1", self._make_view()) + + assert config_fn.call_args.kwargs["sync_mode"] == "strict" diff --git a/tests/v1/kv_offload/tiering/p2p/test_manager.py b/tests/v1/kv_offload/tiering/p2p/test_manager.py index aceec1cd8fa1..8190a9f8b1d0 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_manager.py +++ b/tests/v1/kv_offload/tiering/p2p/test_manager.py @@ -1073,6 +1073,16 @@ def add_remote_peer( "block_len": block_len, } + def add_remote_peer_async( + self, peer_id, agent_metadata, base_addr, num_blocks, block_len + ): + from concurrent.futures import Future + + self.add_remote_peer(peer_id, agent_metadata, base_addr, num_blocks, block_len) + future: Future[None] = Future() + future.set_result(None) + return future + def remove_remote_peer(self, peer_id: str) -> None: self._remote_peers.pop(peer_id, None) diff --git a/tests/v1/kv_offload/tiering/p2p/test_sessions.py b/tests/v1/kv_offload/tiering/p2p/test_sessions.py index 2e4415fbb75d..be83dd4c71c5 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_sessions.py +++ b/tests/v1/kv_offload/tiering/p2p/test_sessions.py @@ -14,6 +14,7 @@ import time from collections.abc import Sequence +from concurrent.futures import Future import numpy as np import pytest @@ -86,6 +87,10 @@ def __init__( self._poll_failed: list[int] = [] self._cancel_still_inflight: set[int] = set() self._cancel_calls: list[tuple[list[int], str]] = [] + # When True, add_remote_peer_async parks the registration in + # pending_registrations for the test to complete or fail. + self.defer_registration = False + self.pending_registrations: dict[str, Future[None]] = {} @property def base_addr(self) -> int: @@ -116,6 +121,25 @@ def add_remote_peer( "block_len": block_len, } + def add_remote_peer_async( + self, peer_id, agent_metadata, base_addr, num_blocks, block_len + ) -> Future[None]: + future: Future[None] = Future() + if self.defer_registration: + self.pending_registrations[peer_id] = future + future.add_done_callback( + lambda f: ( + f.exception() is None + and self.add_remote_peer( + peer_id, agent_metadata, base_addr, num_blocks, block_len + ) + ) + ) + return future + self.add_remote_peer(peer_id, agent_metadata, base_addr, num_blocks, block_len) + future.set_result(None) + return future + def remove_remote_peer(self, peer_id: str) -> None: self._remote_peers.pop(peer_id, None) @@ -390,6 +414,41 @@ def test_peer_connect_triggers_add_remote_and_ack(self): ack = next(m for m in conn._sent if m[TYPE_KEY] == ConnectAckMsg.TYPE) assert ack[ConnectAckMsg.PEER_ID] == "local:9000" + def test_connect_ack_waits_for_peer_registration(self): + """ConnectAck is deferred until the transport has registered the peer.""" + transport = FakeDataTransport() + transport.defer_registration = True + session, conn, _ = _make_session(transport=transport) + conn.enqueue(_peer_connect_msg()) + session.poll() + assert not any(m[TYPE_KEY] == ConnectAckMsg.TYPE for m in conn._sent) + assert "peer:8000" not in transport._remote_peers + assert session.has_pending_work + session.poll() + assert not any(m[TYPE_KEY] == ConnectAckMsg.TYPE for m in conn._sent) + + transport.pending_registrations["peer:8000"].set_result(None) + session.poll() + assert "peer:8000" in transport._remote_peers + ack = next(m for m in conn._sent if m[TYPE_KEY] == ConnectAckMsg.TYPE) + assert ack[ConnectAckMsg.PEER_ID] == "local:9000" + assert not session.has_pending_work + + def test_failed_peer_registration_rejects_peer(self): + """A registration error rejects the peer like a validation failure.""" + transport = FakeDataTransport() + transport.defer_registration = True + session, conn, _ = _make_session(transport=transport) + conn.enqueue(_peer_connect_msg()) + session.poll() + transport.pending_registrations["peer:8000"].set_exception( + RuntimeError("agent metadata rejected") + ) + session.poll() + assert not session.alive + assert not any(m[TYPE_KEY] == ConnectAckMsg.TYPE for m in conn._sent) + assert not session.has_pending_work + def test_connect_ack_makes_session_ready(self): """Session.ready becomes True after ConnectAckMsg.""" session, conn, _ = _make_session() diff --git a/vllm/v1/kv_offload/tiering/p2p/data/base.py b/vllm/v1/kv_offload/tiering/p2p/data/base.py index ffe609d71c1e..d47f55584965 100644 --- a/vllm/v1/kv_offload/tiering/p2p/data/base.py +++ b/vllm/v1/kv_offload/tiering/p2p/data/base.py @@ -75,6 +75,7 @@ import hashlib import json from abc import ABC, abstractmethod +from concurrent.futures import Future from collections.abc import Iterable, Sequence from typing import Literal, NamedTuple @@ -177,6 +178,34 @@ def add_remote_peer( """ ... + def add_remote_peer_async( + self, + peer_id: str, + agent_metadata: bytes, + base_addr: int, + num_blocks: int, + block_len: int, + ) -> Future[None]: + """Register a remote peer without blocking the caller. + + Returns a future that completes once the peer is usable by + write_blocks(); a failed registration surfaces as the future's + exception. This default runs add_remote_peer() inline and returns + a completed future. Transports whose registration is slow enough + to hold up a scheduling iteration override it to run the work on + a worker thread. + """ + future: Future[None] = Future() + try: + self.add_remote_peer( + peer_id, agent_metadata, base_addr, num_blocks, block_len + ) + except Exception as exc: + future.set_exception(exc) + else: + future.set_result(None) + return future + @abstractmethod def remove_remote_peer(self, peer_id: str) -> None: """Unregister a remote peer and release associated resources. diff --git a/vllm/v1/kv_offload/tiering/p2p/data/nixl.py b/vllm/v1/kv_offload/tiering/p2p/data/nixl.py index bf3adf4b604b..ef237e764501 100644 --- a/vllm/v1/kv_offload/tiering/p2p/data/nixl.py +++ b/vllm/v1/kv_offload/tiering/p2p/data/nixl.py @@ -6,8 +6,11 @@ from __future__ import annotations +import contextlib import itertools +import sys from collections.abc import Iterable +from concurrent.futures import Future, ThreadPoolExecutor from typing import Any, NamedTuple import numpy as np @@ -23,6 +26,21 @@ logger = init_logger(__name__) + +def _strict_thread_sync() -> Any: + """The agent sync mode that serializes NIXL API calls across threads. + + Registration runs on a worker thread while the engine thread submits + and polls transfers, so the agent must serialize its API. Returns None + when the installed NIXL exposes no sync-mode enum; the agent then keeps + its default and registration still runs on the worker. + """ + module = sys.modules.get(getattr(_NixlAgentConfig, "__module__", "")) + enum = getattr(module, "nixl_thread_sync_t", None) + if enum is None: + return None + return enum.NIXL_THREAD_SYNC_STRICT + # Shared sentinel returned by poll() in the steady state (no inflight, or # no transfer changed state since the last poll). Tuples make it immutable; # callers only iterate / membership-test / equality-check. @@ -68,6 +86,12 @@ def __init__( # transfer_id → _Inflight(peer_id, handle). self._inflight: dict[int, _Inflight] = {} self._next_id = itertools.count() + # peer_id → registration future while add_remote_peer_async runs + # on the worker thread; the agent runs in strict sync mode so the + # worker's registration calls and the caller's transfer calls may + # overlap. + self._registrations: dict[str, Future[None]] = {} + self._registration_executor: ThreadPoolExecutor | None = None self._init(view) @@ -79,9 +103,14 @@ def _init(self, view: memoryview) -> None: if _NixlAgent is None: return + sync_mode = _strict_thread_sync() non_ucx_backends = [b for b in self._backends if b != "UCX"] if non_ucx_backends: - cfg = _NixlAgentConfig(backends=self._backends, capture_telemetry=True) + cfg = _NixlAgentConfig( + backends=self._backends, + capture_telemetry=True, + sync_mode=sync_mode, + ) logger.info( "NixlTransport %s: NIXL backends=%s", self._agent_name, @@ -89,7 +118,9 @@ def _init(self, view: memoryview) -> None: ) else: cfg = _NixlAgentConfig( - num_threads=self._num_threads, capture_telemetry=True + num_threads=self._num_threads, + capture_telemetry=True, + sync_mode=sync_mode, ) logger.info( "NixlTransport %s: NIXL backends=[UCX] num_threads=%d", @@ -137,7 +168,49 @@ def add_remote_peer( self._peer_nixl_names[peer_id] = nixl_name self._remote_dlists[peer_id] = remote_dlist + def add_remote_peer_async( + self, + peer_id: str, + agent_metadata: bytes, + base_addr: int, + num_blocks: int, + block_len: int, + ) -> Future[None]: + """Run add_remote_peer() on the registration worker. + + Loading the peer's agent metadata and preparing its descriptor list + scale with the peer's block count, so the caller must not hold a + scheduling iteration for them. Peers register one at a time in + arrival order. + """ + if self._registration_executor is None: + self._registration_executor = ThreadPoolExecutor( + max_workers=1, thread_name_prefix="p2p-peer-registration" + ) + future = self._registration_executor.submit( + self.add_remote_peer, + peer_id, + agent_metadata, + base_addr, + num_blocks, + block_len, + ) + self._registrations[peer_id] = future + future.add_done_callback( + lambda f, peer_id=peer_id: self._registrations.pop(peer_id, None) + ) + if future.done(): + # Completed before it was recorded; the callback found nothing. + self._registrations.pop(peer_id, None) + return future + def remove_remote_peer(self, peer_id: str) -> None: + pending = self._registrations.pop(peer_id, None) + if pending is not None: + # A registration still in flight would otherwise complete after + # the removal and leave the peer registered with the agent. + with contextlib.suppress(Exception): + pending.result() nixl_name = self._peer_nixl_names.pop(peer_id, None) dlist = self._remote_dlists.pop(peer_id, None) if self._agent is not None: @@ -292,8 +365,11 @@ def close(self) -> None: return self._release_handles([entry.handle for entry in self._inflight.values()]) self._inflight.clear() - for peer_id in list(self._remote_dlists): + for peer_id in list(self._registrations) + list(self._remote_dlists): self.remove_remote_peer(peer_id) + if self._registration_executor is not None: + self._registration_executor.shutdown(wait=True) + self._registration_executor = None if self._local_dlist is not None: self._agent.release_dlist_handle(self._local_dlist) self._local_dlist = None diff --git a/vllm/v1/kv_offload/tiering/p2p/session/session.py b/vllm/v1/kv_offload/tiering/p2p/session/session.py index ba2aca6438c6..5c262eb4cbbb 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/session.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/session.py @@ -44,6 +44,8 @@ ) if TYPE_CHECKING: + from concurrent.futures import Future + from vllm.v1.kv_offload.base import ReqContext from vllm.v1.kv_offload.tiering.base import JobId, ParentManager from vllm.v1.kv_offload.tiering.p2p.data import DataTransport @@ -125,6 +127,10 @@ def __init__( self._send_ready = False # True after the peer acked our ConnectMsg # Msgs waiting to be sent on connection establishment self._queued: list[dict] = [] + # Registration of the peer's NIXL metadata, started by its ConnectMsg. + # ConnectAck goes out once it completes, so the peer never sends a + # fetch this side cannot serve yet. + self._pending_registration: Future[None] | None = None # Consecutive non-protocol dispatch errors. Reset on success. self._dispatch_error_count: int = 0 @@ -166,8 +172,13 @@ def ready(self) -> bool: @property def has_pending_work(self) -> bool: - """True while inbound loads or outbound transfers are outstanding.""" - return self._client.has_active_loads or self._server.has_inflight_transfers + """True while inbound loads, outbound transfers, or a peer + registration are outstanding.""" + return ( + self._client.has_active_loads + or self._server.has_inflight_transfers + or self._pending_registration is not None + ) # ------------------------------------------------------------------ # Connection lifecycle @@ -260,6 +271,7 @@ def poll(self) -> SessionPollResult: for msg in self._conn.recv(): self._on_message(msg) + self._maybe_finish_registration() loads = self._client.collect_results() stores = self._server.collect_results() @@ -285,6 +297,7 @@ def close(self) -> SessionCloseResult: ``parent.on_request_finished`` (the manager flushes these on its next ``serve_external_requests``). """ + self._pending_registration = None client_result = self._client.close() failed_stores, failed_serves = self._server.close() @@ -463,19 +476,36 @@ def _on_connect(self, msg: dict) -> None: f"local={self._local_hash_seed!r}. Ensure PYTHONHASHSEED " "(if set) matches on all P2P peers." ) - self._transport.add_remote_peer( - self.peer_id, - agent_metadata=msg[ConnectMsg.AGENT_METADATA], - base_addr=msg[ConnectMsg.BASE_ADDR], - num_blocks=msg[ConnectMsg.NUM_BLOCKS], - block_len=msg[ConnectMsg.BLOCK_LEN], - ) except ValueError as exc: logger.error("P2PSession %s: rejecting peer connect: %s", self.peer_id, exc) if self._conn is not None: self._conn.mark_dead() return + self._pending_registration = self._transport.add_remote_peer_async( + self.peer_id, + agent_metadata=msg[ConnectMsg.AGENT_METADATA], + base_addr=msg[ConnectMsg.BASE_ADDR], + num_blocks=msg[ConnectMsg.NUM_BLOCKS], + block_len=msg[ConnectMsg.BLOCK_LEN], + ) + + def _maybe_finish_registration(self) -> None: + """Send ConnectAck once the peer's registration has completed. + + A failed registration rejects the peer the same way a ConnectMsg + validation failure does. + """ + future = self._pending_registration + if future is None or not future.done(): + return + self._pending_registration = None + exc = future.exception() + if exc is not None: + logger.error("P2PSession %s: rejecting peer connect: %s", self.peer_id, exc) + if self._conn is not None: + self._conn.mark_dead() + return if self._conn is not None: self._conn.send( { From 94ac8f3748adfeb8bfa3927afd00dae1e0d9c8e0 Mon Sep 17 00:00:00 2001 From: nilig Date: Wed, 16 Sep 2026 14:01:15 +0300 Subject: [PATCH 2/2] [Core][KV Connector] Service connector work while waiting on the model The offloading connector drains its secondary tiers and serves peer lookups only from the per-step scheduler hooks. On a source rank busy with long prefills every P2P control-plane round trip therefore waits for a step boundary: measured on GLM-5.2-FP8 (DP8, 8xH200 prefill pods, 64 GiB CPU tier), a pull from a source loaded with three 115K-token prefills spends 2.5-2.9 s in lookup delay against 0.1-0.2 s of data movement, the LookupMsg sitting in the server queue until the next step (0.55-0.61 s) and the FetchMsg dispatched at the step after (0.45-0.50 s). KVConnectorBase_V1.poll_pending_work is a scheduler-side hook that the tiering manager implements as the same sweep on_schedule_end runs. EngineCore drives it from a ConnectorPoller thread that is active only while the engine core thread is blocked in future.result(), with a lock that hands the connector between the two threads, so the connector is never entered concurrently and the executor futures need no timed wait. With it the same busy-source pulls measure 0.59-0.69 s of lookup delay. The dummy-batch wait of the data-parallel loop (execute_dummy_batch) is not routed through the poller yet; a rank holding only a deferred request still services the control plane once per dummy batch. Signed-off-by: nilig --- tests/v1/engine/test_connector_poller.py | 84 +++++++++++++++++++ .../tiering/test_tiering_offloading.py | 40 +++++++++ .../kv_transfer/kv_connector/v1/base.py | 12 +++ .../kv_connector/v1/multi_connector.py | 4 + .../kv_connector/v1/offloading/scheduler.py | 3 + .../kv_connector/v1/offloading_connector.py | 4 + vllm/v1/engine/core.py | 82 +++++++++++++++++- vllm/v1/kv_offload/base.py | 9 ++ vllm/v1/kv_offload/tiering/manager.py | 13 +++ vllm/v1/kv_offload/tiering/p2p/manager.py | 6 +- .../kv_offload/tiering/p2p/session/server.py | 4 +- 11 files changed, 255 insertions(+), 6 deletions(-) create mode 100644 tests/v1/engine/test_connector_poller.py diff --git a/tests/v1/engine/test_connector_poller.py b/tests/v1/engine/test_connector_poller.py new file mode 100644 index 000000000000..d3947e4c19f2 --- /dev/null +++ b/tests/v1/engine/test_connector_poller.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""ConnectorPoller: services a connector only while the engine waits.""" + +import threading +import time +from concurrent.futures import Future + +import pytest + +from vllm.v1.engine.core import ConnectorPoller + + +class _CountingConnector: + def __init__(self, fail: bool = False) -> None: + self.calls = 0 + self.fail = fail + self.thread_names: set[str] = set() + + def poll_pending_work(self) -> None: + self.calls += 1 + self.thread_names.add(threading.current_thread().name) + if self.fail: + raise RuntimeError("sweep failed") + + +def _resolve_after(future: Future, delay_s: float, value: object = "out") -> None: + def run(): + time.sleep(delay_s) + future.set_result(value) + + threading.Thread(target=run, daemon=True).start() + + +def test_poller_sweeps_only_while_waiting(): + connector = _CountingConnector() + poller = ConnectorPoller(connector, interval_s=0.001) + try: + time.sleep(0.02) + assert connector.calls == 0 + + future: Future = Future() + _resolve_after(future, 0.05) + assert poller.wait(future) == "out" + assert connector.calls > 0 + assert connector.thread_names == {"kv-connector-poller"} + + settled = connector.calls + time.sleep(0.02) + assert connector.calls == settled + finally: + poller.close() + + +def test_poller_reraises_sweep_error_on_engine_thread(): + connector = _CountingConnector(fail=True) + poller = ConnectorPoller(connector, interval_s=0.001) + try: + future: Future = Future() + _resolve_after(future, 0.03) + with pytest.raises(RuntimeError, match="sweep failed"): + poller.wait(future) + assert connector.calls == 1 + + # A later wait runs sweeps again; the stored error was consumed. + connector.fail = False + future = Future() + _resolve_after(future, 0.03) + assert poller.wait(future) == "out" + assert connector.calls > 1 + finally: + poller.close() + + +def test_poller_propagates_future_exception(): + connector = _CountingConnector() + poller = ConnectorPoller(connector, interval_s=0.001) + try: + future: Future = Future() + future.set_exception(ValueError("model failed")) + with pytest.raises(ValueError, match="model failed"): + poller.wait(future) + finally: + poller.close() diff --git a/tests/v1/kv_offload/tiering/test_tiering_offloading.py b/tests/v1/kv_offload/tiering/test_tiering_offloading.py index 771153ec4734..5fb78ba29422 100644 --- a/tests/v1/kv_offload/tiering/test_tiering_offloading.py +++ b/tests/v1/kv_offload/tiering/test_tiering_offloading.py @@ -148,6 +148,46 @@ def get_stats(self) -> OffloadingConnectorStats | None: return stats +class ServeRecordingSecondaryTierManager(MetricsSecondaryTierManager): + """Test-only secondary tier that records service sweeps.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.finished_polls = 0 + self.serve_calls = 0 + + def get_finished_jobs(self) -> Iterable[JobResult]: + self.finished_polls += 1 + return () + + def serve_external_requests(self, parent) -> None: + self.serve_calls += 1 + + +def test_poll_pending_work_services_tiers_between_steps(): + """poll_pending_work() collects finished jobs and serves every tier + without a scheduler step and without consuming the per-step gate.""" + mock_region = _mock_mmap_region(5) + primary = CPUPrimaryTierOffloadingManager(num_chunks=5, mmap_region=mock_region) + tier = ServeRecordingSecondaryTierManager( + offloading_spec=_MOCK_OFFLOADING_SPEC, + primary_kv_view=mock_region.create_kv_memoryview(), + tier_type="recording", + ) + manager = TieringOffloadingManager(primary_tier=primary, secondary_tiers=[tier]) + + manager.poll_pending_work() + manager.poll_pending_work() + assert tier.finished_polls == 2 + assert tier.serve_calls == 2 + assert manager._processed_jobs_this_step is False + + ctx = ScheduleEndContext(new_req_ids=[], preempted_req_ids=()) + manager.on_schedule_end(ctx) + assert tier.finished_polls == 3 + assert tier.serve_calls == 3 + + def test_tiering_spec_collects_secondary_metric_definitions(monkeypatch): monkeypatch.setitem( SecondaryTierFactory._registry, diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/base.py b/vllm/distributed/kv_transfer/kv_connector/v1/base.py index c97374832ffe..bd9c11744fd4 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/base.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/base.py @@ -638,6 +638,18 @@ def has_pending_push_work(self) -> bool: # scheduler alive (e.g. extend has_unfinished_requests). return False + def poll_pending_work(self) -> None: + """Advance scheduler-side work while the engine waits on a model step. + + The engine core calls this at a bounded interval between scheduling + iterations, never concurrently with any other scheduler-side method, + so a connector that services peers from the scheduler side (control + messages, transfer completions) is not limited to one service window + per step. Implementations must not touch scheduler state and must + return promptly when there is nothing to do. + """ + return + @classmethod def get_required_kvcache_layout(cls, vllm_config: "VllmConfig") -> str | None: """ diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py index d291d3852804..858fb69f1ffa 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py @@ -588,6 +588,10 @@ def has_pending_block_frees(self) -> bool: def has_pending_push_work(self) -> bool: return any(c.has_pending_push_work() for c in self._connectors) + def poll_pending_work(self) -> None: + for c in self._connectors: + c.poll_pending_work() + @classmethod def get_required_kvcache_layout(cls, vllm_config: "VllmConfig") -> str | None: """ diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index b253a2399867..00ccdabfa737 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -1839,6 +1839,9 @@ def has_pending_push_work(self) -> bool: """ return bool(self._jobs) or self.manager.has_pending_work() + def poll_pending_work(self) -> None: + self.manager.poll_pending_work() + def update_connector_output(self, connector_output: KVConnectorOutput): """ Update KVConnector state from worker-side connectors output. diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py index 73570881753d..6385099cf222 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py @@ -202,6 +202,10 @@ def has_pending_push_work(self) -> bool: assert self.connector_scheduler is not None return self.connector_scheduler.has_pending_push_work() + def poll_pending_work(self) -> None: + assert self.connector_scheduler is not None + self.connector_scheduler.poll_pending_work() + def update_connector_output(self, connector_output: KVConnectorOutput): assert self.connector_scheduler is not None self.connector_scheduler.update_connector_output(connector_output) diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index b3d94de3e7db..915f4745960e 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -27,6 +27,7 @@ cleanup_dist_env_and_memory, stateless_destroy_torch_distributed_process_group, ) +from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorBase_V1 from vllm.envs import enable_envs_cache from vllm.logger import init_logger from vllm.logging_utils.dump_input import dump_engine_exception @@ -102,6 +103,73 @@ _R = TypeVar("_R") # Return type for collective_rpc +# Interval between connector service sweeps while the engine core waits +# for a model step. +KV_CONNECTOR_POLL_INTERVAL_S = 0.005 + + +class ConnectorPoller: + """Services a KV connector's scheduler-side work while the engine core + thread waits on a model step. + + The connector is touched by one thread at a time: the engine core + thread outside ``wait`` and the poller thread inside it. ``wait`` hands + the connector over by raising ``_waiting`` under the lock and takes it + back the same way, so a sweep in progress when the model output arrives + completes before the engine core continues. An exception raised by a + sweep is re-raised from ``wait`` on the engine core thread. + """ + + def __init__( + self, + connector: KVConnectorBase_V1, + interval_s: float = KV_CONNECTOR_POLL_INTERVAL_S, + ) -> None: + self._connector = connector + self._interval_s = interval_s + self._lock = threading.Lock() + self._waiting = False + self._wake = threading.Event() + self._closed = False + self._error: BaseException | None = None + self._thread = threading.Thread( + target=self._run, name="kv-connector-poller", daemon=True + ) + self._thread.start() + + def wait(self, future: Future[Any]) -> Any: + with self._lock: + self._waiting = True + self._wake.set() + try: + return future.result() + finally: + with self._lock: + self._waiting = False + error, self._error = self._error, None + self._wake.clear() + if error is not None: + raise error + + def close(self) -> None: + self._closed = True + self._wake.set() + self._thread.join(timeout=1.0) + + def _run(self) -> None: + while not self._closed: + self._wake.wait() + if self._closed: + return + with self._lock: + if self._waiting and self._error is None: + try: + self._connector.poll_pending_work() + except BaseException as exc: + self._error = exc + time.sleep(self._interval_s) + + class EngineCore: """Inner loop of vLLM's Engine.""" @@ -187,6 +255,9 @@ def __init__( # handshake metadata from all workers so the connector in the scheduler # will have the full context kv_connector = self.scheduler.get_kv_connector() + self.connector_poller: ConnectorPoller | None = ( + ConnectorPoller(kv_connector) if kv_connector is not None else None + ) if kv_connector is not None: # Collect and store KV connector xfer metadata from workers # (after KV cache registration) @@ -620,7 +691,7 @@ def step(self) -> tuple[dict[int, EngineCoreOutputs], bool]: self.capture_iteration_details(scheduler_output) as iteration_details, self.log_error_detail(scheduler_output), ): - model_output = future.result() + model_output = self._wait_for_model_output(future) if model_output is None: model_output = self.model_executor.sample_tokens(grammar_output) @@ -634,6 +705,11 @@ def step(self) -> tuple[dict[int, EngineCoreOutputs], bool]: return engine_core_outputs, scheduler_output.total_num_scheduled_tokens > 0 + def _wait_for_model_output(self, future: Future[Any]) -> Any: + if self.connector_poller is None: + return future.result() + return self.connector_poller.wait(future) + def post_step(self, model_executed: bool) -> None: # When using async scheduling we can't get draft token ids in advance, # so we update draft token ids in the worker process and don't @@ -719,7 +795,7 @@ def step_with_batch_queue( self.capture_iteration_details(scheduler_output) as iteration_details, self.log_error_detail(scheduler_output), ): - model_output = future.result() + model_output = self._wait_for_model_output(future) if model_output is None: # None from sample_tokens() implies that the original execute_model() # call failed - raise that exception. @@ -771,6 +847,8 @@ def _process_aborts_queue(self): def shutdown(self): logger.debug_once("[shutdown] EngineCore: tearing down local resources") + if self.connector_poller is not None: + self.connector_poller.close() self.structured_output_manager.clear_backend() if self.model_executor: self.model_executor.shutdown() diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index 988747e06e1e..763aa84421ba 100644 --- a/vllm/v1/kv_offload/base.py +++ b/vllm/v1/kv_offload/base.py @@ -396,6 +396,15 @@ def has_pending_work(self) -> bool: """ return False + def poll_pending_work(self) -> None: + """Service background work between scheduler steps. + + Called on the scheduler side while the engine waits for a model + step, never concurrently with the per-step methods. The default + does nothing. + """ + return + def reset_cache(self) -> None: """Evict all tracked blocks and reset internal state.""" return diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index f19616828daa..6bce2493e0ab 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -852,6 +852,19 @@ def has_pending_work(self) -> bool: or any(tier.has_pending_work() for tier in self.secondary_tiers) ) + @override + def poll_pending_work(self) -> None: + """Service the secondary tiers between scheduler steps. + + The same sweep on_schedule_end() runs at the end of a step: finished + jobs are collected and every tier serves its external requests. The + per-step gate is untouched, so the next step still polls once before + its own lookups. + """ + self._process_finished_jobs() + for tier in self.secondary_tiers: + tier.serve_external_requests(self._tier_parents[tier]) + @override def take_events(self) -> Iterable[OffloadingEvent]: """Yield events owned by the primary and secondary tiers. diff --git a/vllm/v1/kv_offload/tiering/p2p/manager.py b/vllm/v1/kv_offload/tiering/p2p/manager.py index ba1973b76784..8453518bffca 100644 --- a/vllm/v1/kv_offload/tiering/p2p/manager.py +++ b/vllm/v1/kv_offload/tiering/p2p/manager.py @@ -201,8 +201,10 @@ class P2PSecondaryTierManager(SecondaryTierManager): blocks from the peer) and server-role (serving blocks to the peer) over the same control connection. - Single-threaded: every public method runs on the scheduler thread, and - the engine drives polling via ``get_finished_jobs()`` once per step. + Never re-entered: every public method runs on the engine core thread + or, while the engine waits on a model step, on the connector poller, + and the two never overlap. The engine drives polling via + ``get_finished_jobs()`` at least once per step. ``has_pending_work()`` keeps the engine ticking so the control transport and existing sessions are polled even when no requests are scheduled. """ diff --git a/vllm/v1/kv_offload/tiering/p2p/session/server.py b/vllm/v1/kv_offload/tiering/p2p/session/server.py index c78918323266..9883856cd7d1 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/server.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/server.py @@ -449,8 +449,8 @@ def on_lookup( def serve_external_requests(self, parent: ParentManager) -> None: """Resolve inbound peer lookups against the tiering manager. - Called once per scheduler step with a ``parent`` handle valid - only for this call. Drains newly-enqueued LookupMsgs, re-polls + Called at least once per scheduler step with a ``parent`` handle + valid only for this call. Drains newly-enqueued LookupMsgs, re-polls any parked HIT_PENDING / RETRY keys, and releases the bookkeeping for lookups closed since the last serve. """