From 7a74a9662bff4e5ef91f210f41b7f930de3df0bc Mon Sep 17 00:00:00 2001 From: tomerg-nvidia Date: Tue, 14 Jul 2026 15:03:41 +0200 Subject: [PATCH] [NIXL] Avoid reading expired blocks in bidirectional turn-2 read (#47021) Signed-off-by: Tomer Gilad Signed-off-by: NickLucche Co-authored-by: NickLucche --- docs/design/nixl_kv_cache_lease.md | 2 +- .../unit/test_bidirectional_kv_transfer.py | 165 +++++++++++++++++- .../kv_connector/unit/test_nixl_connector.py | 21 ++- .../unit/test_nixl_connector_hma.py | 1 + .../kv_connector/v1/nixl/base_scheduler.py | 7 +- .../kv_connector/v1/nixl/base_worker.py | 49 +++++- .../kv_connector/v1/nixl/metadata.py | 6 +- .../kv_connector/v1/nixl/pull_scheduler.py | 5 + .../kv_connector/v1/nixl/pull_worker.py | 26 +++ .../kv_connector/v1/nixl/push_worker.py | 4 +- 10 files changed, 264 insertions(+), 22 deletions(-) diff --git a/docs/design/nixl_kv_cache_lease.md b/docs/design/nixl_kv_cache_lease.md index aa7683bb9e1..35018aca22f 100644 --- a/docs/design/nixl_kv_cache_lease.md +++ b/docs/design/nixl_kv_cache_lease.md @@ -101,7 +101,7 @@ The heartbeat is deferred to the next step once the handshake completes --- the ## Bidirectional KV Transfer -For multi-turn conversations, [bidirectional KV transfer](../features/disagg_prefill.md) allows D to cache KV blocks that P can pull from on subsequent turns. Since the timing of the next conversational turn is **client-dependent** (not controlled by the system), the heartbeat-based lease mechanism does not apply here. Instead, a separate `decoder_kv_blocks_ttl` (default 480s) provides a simple fixed timeout for blocks cached on D. If the client takes too long to continue the conversation, the blocks expire and P recomputes. Future work may extend a symmetric heartbeat mechanism to this case. +For multi-turn conversations, [bidirectional KV transfer](../features/disagg_prefill.md) allows D to cache KV blocks that P can pull from on subsequent turns. Since the timing of the next conversational turn is **client-dependent** (not controlled by the system), the heartbeat-based lease mechanism does not apply here. Instead, a separate `decoder_kv_blocks_ttl` (default 480s) provides a simple fixed timeout for blocks cached on D. If the client takes too long to continue the conversation, the blocks expire. D communicates back the expiry time so P can know when blocks are expired and recompute. Because the deadline is a `perf_counter` value produced on D and the two engines run in separate processes (with unrelated clocks), P estimates the clock offset to D from the handshake round-trip and applies it before comparing the deadline against its own `perf_counter`. Future work may extend a symmetric heartbeat mechanism to this case. ## Key Design Decisions diff --git a/tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py b/tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py index 12831601cba..13183042ee9 100644 --- a/tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py +++ b/tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py @@ -87,7 +87,7 @@ def _make_connector_with_fake_worker( hand_shake_latency=0, cycles_before_done=0, do_handshake=True ): """Create a NixlConnector with FakeNixlConnectorWorker.""" - vllm_config = create_vllm_config() + vllm_config = create_vllm_config(kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG) kv_cache_config = make_kv_cache_config(block_size=16, num_blocks=2) connector = NixlConnector(vllm_config, KVConnectorRole.WORKER, kv_cache_config) connector.connector_worker = FakeNixlConnectorWorker( @@ -100,7 +100,7 @@ def _make_connector_with_fake_worker( assert isinstance(worker.nixl_wrapper, FakeNixlWrapper) worker.kv_cache_layout = "HND" if do_handshake: - remote_agents = worker._nixl_handshake( + remote_agents, _ = worker._nixl_handshake( host="localhost", port=1234, remote_tp_size=1, @@ -912,3 +912,164 @@ def test_remote_blocks_processed_flag_persists(): mro.kv_connector_output = KVConnectorOutput(finished_sending={req_id}) scheduler.update_from_output(so, mro) assert_scheduler_empty(scheduler) + + +# 6. Turn-2 deadline expiry check (P declines reading expired D blocks) + +_REMOTE = FakeNixlConnectorWorker.REMOTE_ENGINE_ID + + +@pytest.mark.parametrize( + ("offset", "expiry_delta", "expect_declined"), + [ + # offset known + deadline expired/near-expiry -> decline & recompute. + pytest.param(0.0, -10.0, True, id="expired"), + # offset applied before comparison: +50s in D's clock is expired locally + # when D is 100s ahead. + pytest.param(100.0, 50.0, True, id="offset_makes_expired"), + # within the default 5s safety margin -> treated as expired. + pytest.param(0.0, 2.0, True, id="near_expiry_within_margin"), + # far-future deadline -> read proceeds. + pytest.param(0.0, 1000.0, False, id="valid_far_deadline"), + # no deadline field (older peer / router did not forward) -> skipped. + pytest.param(0.0, None, False, id="missing_expiry_field"), + ], +) +@patch( + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", + FakeNixlWrapper, +) +def test_turn2_deadline_gate(dist_init, offset, expiry_delta, expect_declined): + """P declines on the turn-2 readback if the deadline is known, + near-expiry and a handshake clock offset is known""" + connector, worker = _make_connector_with_fake_worker() + worker._engine_clock_offset[_REMOTE] = offset + expiry_time = None if expiry_delta is None else time.perf_counter() + expiry_delta + meta = NixlConnectorMetadata() + params = { + "do_remote_prefill": False, + "do_remote_decode": True, + "remote_block_ids": ([20, 21],), + "remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID, + "remote_request_id": "decode-req", + "remote_host": "localhost", + "remote_port": 1234, + "remote_tp_size": 1, + } + if expiry_time is not None: + params["remote_blocks_expiry_time"] = expiry_time + meta.add_new_req_to_recv( + request_id="req", + local_block_ids=([10, 11],), + kv_transfer_params=params, + ) + _do_load_kv(connector, meta) + assert worker.xfer_stats.data["num_kv_expired_reqs"] == ( + [1] if expect_declined else [] + ) + _, done_recving = connector.get_finished(finished_req_ids=set()) + assert "req" in done_recving + + +@patch( + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", + FakeNixlWrapper, +) +def test_turn2_full_prefix_hit_with_expired_deadline_skips_gate(dist_init): + """A full local prefix hit issues no READ. The expiry gate must be skipped + even when D's deadline is expired.""" + connector, worker = _make_connector_with_fake_worker() + worker._engine_clock_offset[_REMOTE] = 0.0 + + meta = NixlConnectorMetadata() + meta.add_new_req_to_recv( + request_id="req", + local_block_ids=(), # full prefix hit -> zero local block groups + kv_transfer_params={ + "do_remote_prefill": False, + "do_remote_decode": True, + "remote_block_ids": ([20, 21],), + "remote_engine_id": _REMOTE, + "remote_request_id": "decode-req", + "remote_host": "localhost", + "remote_port": 1234, + "remote_tp_size": 1, + # Expired on D's clock; with offset 0 it is expired locally too. + "remote_blocks_expiry_time": time.perf_counter() - 10.0, + }, + ) + # Must not raise, and the gate must not fire for a no-read prefix hit. + _do_load_kv(connector, meta) + assert worker.xfer_stats.data["num_kv_expired_reqs"] == [] + assert worker.xfer_stats.data["num_failed_transfers"] == [] + + +def test_d_node_request_finished_exports_blocks_expiry_time(): + """D-node (do_remote_prefill request) exports a future float expiry time.""" + vllm_config = create_vllm_config(kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG) + scheduler = create_scheduler(vllm_config) + BS = vllm_config.cache_config.block_size + req = create_request( + request_id=600, block_size=BS, num_tokens=int(BS * 2.5), do_remote_prefill=True + ) + scheduler.add_request(req) + req_id = req.request_id + so = scheduler.schedule() + scheduler.update_from_output( + so, create_model_runner_output(reqs=[], finished_recving={req_id}) + ) + so = scheduler.schedule() + eco = scheduler.update_from_output( + so, create_model_runner_output(reqs=[req], use_eos=True) + ) + kv = eco[0].outputs[0].kv_transfer_params + assert kv["do_remote_decode"] is True + assert isinstance(kv["remote_blocks_expiry_time"], float) + assert kv["remote_blocks_expiry_time"] > time.perf_counter() + + +def test_handshake_listener_appends_perf_counter_frame(): + """The handshake reply carries a 2nd frame with the listener's live + perf_counter.""" + import threading + + import msgspec + import zmq + + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_scheduler import ( + NixlBaseConnectorScheduler, + ) + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import GET_META_MSG + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import zmq_ctx + from vllm.utils.network_utils import get_open_port, make_zmq_path + + encoded_data = {(0, 0): b"payload-rank-0"} + ready_event = threading.Event() + stop_event = threading.Event() + host = "127.0.0.1" + port = get_open_port() + listener = threading.Thread( + target=NixlBaseConnectorScheduler._nixl_handshake_listener, + args=(encoded_data, ready_event, stop_event, host, port), + daemon=True, + ) + listener.start() + try: + assert ready_event.wait(timeout=5) + path = make_zmq_path("tcp", host, port) + with zmq_ctx(zmq.REQ, path) as sock: # type: ignore[attr-defined] + sock.setsockopt(zmq.RCVTIMEO, 5000) # type: ignore[attr-defined] + t0 = time.perf_counter() + sock.send(msgspec.msgpack.encode((GET_META_MSG, 0, 0))) + parts = sock.recv_multipart() + t1 = time.perf_counter() + assert len(parts) == 2 + assert parts[0] == b"payload-rank-0" + remote_perf = msgspec.msgpack.decode(parts[1]) + assert isinstance(remote_perf, float) + # same-clock => offset is near zero. + offset = remote_perf - (t0 + t1) / 2 + assert abs(offset) < 1.0 + finally: + stop_event.set() + listener.join(timeout=5) diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 0724e776937..c1088f1c6e8 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -508,7 +508,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker): expected_engine_id: str, remote_pp_size: int = 1, notif_agents_only: bool = False, - ) -> dict[tuple[int, int], str]: + ) -> tuple[dict[tuple[int, int], str], float]: # Mimic slow _nixl_handshake, as well as bypass zmq communication. time.sleep(self._hand_shake_latency) # These should've been done in register_kv_caches(), called by @@ -560,7 +560,8 @@ class FakeNixlConnectorWorker(NixlConnectorWorker): remote_tp_size=remote_tp_size, ) remote_agents[(0, remote_tp_rank)] = remote_agent_name - return remote_agents + # Handshake bypasses zmq, so report a zero clock offset to the peer. + return remote_agents, 0.0 class TestNixlHandshake: @@ -768,7 +769,7 @@ class TestNixlHandshake: range(tp_ratio) ) - remote_agents = worker._nixl_handshake( + remote_agents, _ = worker._nixl_handshake( host="localhost", port=1234, remote_tp_size=4, @@ -780,7 +781,7 @@ class TestNixlHandshake: # discovered. This is not a scenario we actively support right now, but # the connector allows it. worker.REMOTE_ENGINE_ID = "remote_engine_2" - remote_agents = worker._nixl_handshake( + remote_agents, _ = worker._nixl_handshake( host="localhost", port=1234, remote_tp_size=6, @@ -2878,7 +2879,10 @@ def test_compatibility_hash_validation( # Mock ZMQ socket to return our handshake payload mock_socket = MagicMock() - mock_socket.recv.return_value = msgspec.msgpack.encode(handshake_payload) + mock_socket.recv_multipart.return_value = [ + msgspec.msgpack.encode(handshake_payload), + msgspec.msgpack.encode(time.perf_counter()), + ] # Mock add_remote_agent to avoid actual NIXL operations # Patch zmq_ctx to return our mock socket @@ -2897,7 +2901,7 @@ def test_compatibility_hash_validation( expected_engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID, ) else: - result = decode_worker._nixl_handshake( + result, _ = decode_worker._nixl_handshake( host="localhost", port=1234, remote_tp_size=1, @@ -2981,7 +2985,10 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario) raise AssertionError(f"{error_scenario} not a valid scenario") mock_socket = MagicMock() - mock_socket.recv.return_value = msg_bytes + mock_socket.recv_multipart.return_value = [ + msg_bytes, + msgspec.msgpack.encode(time.perf_counter()), + ] with ( patch.object(decode_worker, "add_remote_agent", return_value="fake_agent"), patch.object(nixl.base_worker, "zmq_ctx") as mock_zmq_ctx, diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index c6bcdc2896c..16c47962aef 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -209,6 +209,7 @@ def test_read_blocks_for_req_expands_remote_ids( worker = object.__new__(NixlConnectorWorker) worker._physical_blocks_per_logical_kv_block = local_physical_per_logical worker._engine_last_active = {} + worker._bidirectional_kv_xfer_enabled = False has_mamba = any(t is MambaSpec for t in resolved_types) has_swa = any(t is SlidingWindowSpec for t in resolved_types) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py index 942ef909263..7a3629c7229 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py @@ -322,8 +322,13 @@ class NixlBaseConnectorScheduler: ) if msg != GET_META_MSG: logger.warning("Connection listener got unexpected message %s", msg) + # Echo our perf_counter so P can estimate the clock offset. + # perf_counter is only comparable within a process, so this + # listener must run in the same process that stamps the block + # expiry deadline (`_reqs_need_send`). + ts = msgspec.msgpack.encode(time.perf_counter()) sock.send_multipart( - (identity, b"", encoded_data[(target_pp_rank, target_tp_rank)]) + (identity, b"", encoded_data[(target_pp_rank, target_tp_rank)], ts) ) def _get_remote_prefill_token_count(self, num_prompt_tokens: int) -> int: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 10554e7ede7..47243968936 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -266,6 +266,12 @@ class NixlBaseConnectorWorker: # NOTE (NickLucche): For now we use a hardcoded value for a simpler interface. self._lease_extension = kv_lease_duration * 2 // 3 + self._bidirectional_kv_xfer_enabled: bool = ( + vllm_config.kv_transfer_config.get_from_extra_config( + "bidirectional_kv_xfer", False + ) + ) + self._is_hma_required = ( not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager and any( @@ -340,6 +346,8 @@ class NixlBaseConnectorWorker: self._remote_agents: dict[EngineId, dict[tuple[int, int], str]] = defaultdict( dict ) + # Map of engine_id -> clock offset. + self._engine_clock_offset: dict[EngineId, float] = {} # Metadata. self.engine_id: EngineId = engine_id @@ -470,7 +478,9 @@ class NixlBaseConnectorWorker: thread_name_prefix="vllm-nixl-handshake-initiator", ) self._ready_requests = queue.Queue[tuple[ReqId, ReqMeta]]() - self._handshake_futures: dict[EngineId, Future[dict[tuple[int, int], str]]] = {} + self._handshake_futures: dict[ + EngineId, Future[tuple[dict[tuple[int, int], str], float]] + ] = {} # Protects _handshake_futures and _remote_agents. self._handshake_lock = threading.RLock() @@ -559,7 +569,7 @@ class NixlBaseConnectorWorker: expected_engine_id: str, remote_pp_size: int = 1, notif_agents_only: bool = False, - ) -> dict[tuple[int, int], str]: + ) -> tuple[dict[tuple[int, int], str], float]: """Do a NIXL handshake with a remote instance.""" # the first time we connect to a remote agent. @@ -583,6 +593,11 @@ class NixlBaseConnectorWorker: p_remote_ranks = self.transfer_topo.handshake_target_ranks(remote_tp_size) remote_rank_to_agent_name: dict[tuple[int, int], str] = {} path = make_zmq_path("tcp", host, port) + # Clock offset to the peer, estimated from the handshake round-trip. + # Keep the lowest-RTT sample: hop cost is ~uniform across ranks, so a + # higher RTT is just noise that skews the midpoint estimate. + best_rtt = float("inf") + best_offset: float | None = None with zmq_ctx(zmq.REQ, path) as sock: for remote_pp_rank, remote_rank in itertools.product( @@ -595,15 +610,24 @@ class NixlBaseConnectorWorker: remote_rank, ) - start_time = time.perf_counter() # Send query for the request. msg = msgspec.msgpack.encode( (GET_META_MSG, remote_pp_rank, remote_rank) ) # Set receive timeout to 5 seconds to avoid hanging on dead server sock.setsockopt(zmq.RCVTIMEO, 5000) # milliseconds + start_time = time.perf_counter() sock.send(msg) - handshake_bytes = sock.recv() + reply_parts = sock.recv_multipart() + recv_time = time.perf_counter() + assert len(reply_parts) == 2 + handshake_bytes = reply_parts[0] + + remote_perf = msgspec.msgpack.decode(reply_parts[1]) + rtt = recv_time - start_time + if rtt < best_rtt: + best_rtt = rtt + best_offset = remote_perf - (start_time + recv_time) / 2 # Decode handshake payload to get compatibility hash handshake_decoder = msgspec.msgpack.Decoder(NixlHandshakePayload) @@ -681,7 +705,9 @@ class NixlBaseConnectorWorker: ) remote_ranks = (remote_pp_rank, remote_rank) remote_rank_to_agent_name[remote_ranks] = remote_agent_name - return remote_rank_to_agent_name + + assert best_offset is not None + return remote_rank_to_agent_name, best_offset def _add_notif_only_remote_agent( self, metadata: NixlAgentMetadata, remote_tp_size: int @@ -818,7 +844,7 @@ class NixlBaseConnectorWorker: tp_size: int, pp_size: int = 1, notif_agents_only: bool = False, - ) -> Future[dict[tuple[int, int], str]] | None: + ) -> Future[tuple[dict[tuple[int, int], str], float]] | None: """ Ensure a handshake is in-flight (or already done) for *engine_id*. @@ -846,11 +872,16 @@ class NixlBaseConnectorWorker: ) self._handshake_futures[engine_id] = fut - def done_callback(f: Future[dict[tuple[int, int], str]], eid=engine_id): + def done_callback( + f: Future[tuple[dict[tuple[int, int], str], float]], + eid=engine_id, + ): with self._handshake_lock: del self._handshake_futures[eid] try: - self._remote_agents[eid] = f.result() + remote_agents, clock_offset = f.result() + self._remote_agents[eid] = remote_agents + self._engine_clock_offset[eid] = clock_offset self._engine_last_active[eid] = time.perf_counter() except Exception as e: self._log_failure( @@ -2421,6 +2452,8 @@ class NixlBaseConnectorWorker: if self.transfer_topo is not None: self.transfer_topo.unregister_remote_engine(engine_id) + # Drop the cached clock offset; it is re-measured on the next handshake. + self._engine_clock_offset.pop(engine_id, None) # Push P-side engines are tracked in _remote_agents but not in # _engine_last_active (they don't participate in stale eviction), so # tolerate a missing entry. diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py index 8ca10c2dc46..f6f3de0a1f2 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py @@ -39,8 +39,10 @@ PUSH_REG_NOTIF_PREFIX = b"PUSH_REG:" # 2: Add remote_request_id to kv_transfer_params # 3: Add physical_blocks_per_logical_kv_block to NixlAgentMetadata # 4: Add KV block lease renewal through heartbeats +# 5: Add remote_blocks_expiry_time to kv_transfer_params + handshake +# clock-sync timestamp # -NIXL_CONNECTOR_VERSION: int = 4 +NIXL_CONNECTOR_VERSION: int = 5 @dataclass @@ -157,6 +159,7 @@ class RemoteMeta: port: int engine_id: str request_id: str + blocks_expiry_time: float | None = None @dataclass @@ -225,5 +228,6 @@ class NixlConnectorMetadata(KVConnectorMetadata): request_id=kv_transfer_params["remote_request_id"], host=kv_transfer_params["remote_host"], port=kv_transfer_params["remote_port"], + blocks_expiry_time=kv_transfer_params.get("remote_blocks_expiry_time"), ) self.reqs_to_recv[request_id] = req diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py index 8c73461c5cd..51a54d06f07 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py @@ -238,6 +238,7 @@ class NixlPullConnectorScheduler(NixlBaseConnectorScheduler): # remove the conditional below delay_free_blocks = any(len(group) > 0 for group in block_ids) remote_num_tokens = 0 + blocks_expiry_time = None if delay_free_blocks: # Prefill request on remote. It will be read from D upon completion request_kv_blocks_ttl = self._kv_lease_duration @@ -254,6 +255,9 @@ class NixlPullConnectorScheduler(NixlBaseConnectorScheduler): self._reqs_need_send[request.request_id] = ( time.perf_counter() + request_kv_blocks_ttl ) + if is_d_node: + # D blocks expiry time exported for the turn-2 readback. + blocks_expiry_time = self._reqs_need_send[request.request_id] # NOTE HMA will "mark" empty/null blocks in groups with 0s (eg SWA ones), # trimming down after allocating for the whole sequence length. Empty # blocks are always at the start of the list. @@ -272,4 +276,5 @@ class NixlPullConnectorScheduler(NixlBaseConnectorScheduler): remote_port=self.side_channel_port, tp_size=self.vllm_config.parallel_config.tensor_parallel_size, remote_num_tokens=remote_num_tokens, + remote_blocks_expiry_time=blocks_expiry_time, ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py index 84e8d292854..4414bf60704 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py @@ -25,6 +25,10 @@ if TYPE_CHECKING: logger = init_logger(__name__) +# Slack (seconds) subtracted from D's exported block-expiry deadline on the turn-2 +# readback, absorbing clock-offset error and read latency. +_KV_BLOCKS_EXPIRY_SAFETY_MARGIN = 5.0 + class NixlPullConnectorWorker(NixlBaseConnectorWorker): """Pull-specific (READ) worker logic.""" @@ -98,12 +102,34 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker): # requests sit in the D scheduler WAITING queue. self._send_heartbeats(metadata) + def _is_turn2_read_expired(self, meta: ReqMeta) -> bool: + """Whether D's cached blocks for this turn-2 readback have (nearly) expired.""" + assert meta.remote is not None + blocks_expiry_time = meta.remote.blocks_expiry_time + # Deadline may be absent (router may not forward it) -> read as usual. + if blocks_expiry_time is None or not meta.local_physical_block_ids: + return False + clock_offset = self._engine_clock_offset[meta.remote.engine_id] + deadline = blocks_expiry_time - clock_offset + return time.perf_counter() + _KV_BLOCKS_EXPIRY_SAFETY_MARGIN >= deadline + def _read_blocks_for_req(self, req_id: str, meta: ReqMeta): assert meta.remote is not None and self.transfer_topo is not None engine_id = meta.remote.engine_id # Update last activity from this remote. Mind that cleanup is done on main # thread (this one), so we don't race on this structure. self._engine_last_active[engine_id] = time.perf_counter() + + if self._bidirectional_kv_xfer_enabled and self._is_turn2_read_expired(meta): + logger.warning( + "Declining expired remote read for %s from engine %s.", + req_id, + engine_id, + ) + self.xfer_stats.record_kv_expired_req() + self._handle_failed_transfer(req_id, None) + return + plan = self.tp_mappings[engine_id] remote_info = self.transfer_topo.get_engine_info(engine_id) tp_ratio = self.transfer_topo.tp_ratio(remote_info.remote_tp_size) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py index 0d1ab49767e..090c2074ddd 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py @@ -294,7 +294,7 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker): return def _on_handshake( - f: Future[dict[tuple[int, int], str]], + f: Future[tuple[dict[tuple[int, int], str], float]], rid: str = req_id, rd: dict[str, Any] = reg_data, ) -> None: @@ -456,7 +456,7 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker): if decode_engine_id in self._remote_agents: return True try: - remote_agents = self._nixl_handshake( + remote_agents, _ = self._nixl_handshake( decode_host, decode_port, decode_tp_size,