@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 7a74a9662b

[NIXL] Avoid reading expired blocks in bidirectional turn-2 read (#47021)

Signed-off-by: Tomer Gilad <tgilad@nvidia.com> Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai> Co-authored-by: NickLucche <nicolo.lucchesi@mistral.ai>

tomerg-nvidia committed Jul 14, 2026 at 15:03 UTC 7a74a9662bff4e5ef91f210f41b7f930de3df0bc

10 files changed +264 -22

docs/design/nixl_kv_cache_lease.md

+1 -1

@@ -101,7 +101,7 @@ The heartbeat is deferred to the next step once the handshake completes --- the

101

102 ## Bidirectional KV Transfer

103

104 -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.

104 +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.

105

106 ## Key Design Decisions

107

tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py

+163 -2

@@ -87,7 +87,7 @@ def _make_connector_with_fake_worker(

87 hand_shake_latency=0, cycles_before_done=0, do_handshake=True

88 ):

89 """Create a NixlConnector with FakeNixlConnectorWorker."""

90 - vllm_config = create_vllm_config()

90 + vllm_config = create_vllm_config(kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG)

91 kv_cache_config = make_kv_cache_config(block_size=16, num_blocks=2)

92 connector = NixlConnector(vllm_config, KVConnectorRole.WORKER, kv_cache_config)

93 connector.connector_worker = FakeNixlConnectorWorker(

@@ -100,7 +100,7 @@ def _make_connector_with_fake_worker(

100 assert isinstance(worker.nixl_wrapper, FakeNixlWrapper)

101 worker.kv_cache_layout = "HND"

102 if do_handshake:

103 - remote_agents = worker._nixl_handshake(

103 + remote_agents, _ = worker._nixl_handshake(

104 host="localhost",

105 port=1234,

106 remote_tp_size=1,

@@ -912,3 +912,164 @@ def test_remote_blocks_processed_flag_persists():

912 mro.kv_connector_output = KVConnectorOutput(finished_sending={req_id})

913 scheduler.update_from_output(so, mro)

914 assert_scheduler_empty(scheduler)

915 +

916 +

917 +# 6. Turn-2 deadline expiry check (P declines reading expired D blocks)

918 +

919 +_REMOTE = FakeNixlConnectorWorker.REMOTE_ENGINE_ID

920 +

921 +

922 +@pytest.mark.parametrize(

923 + ("offset", "expiry_delta", "expect_declined"),

924 + [

925 + # offset known + deadline expired/near-expiry -> decline & recompute.

926 + pytest.param(0.0, -10.0, True, id="expired"),

927 + # offset applied before comparison: +50s in D's clock is expired locally

928 + # when D is 100s ahead.

929 + pytest.param(100.0, 50.0, True, id="offset_makes_expired"),

930 + # within the default 5s safety margin -> treated as expired.

931 + pytest.param(0.0, 2.0, True, id="near_expiry_within_margin"),

932 + # far-future deadline -> read proceeds.

933 + pytest.param(0.0, 1000.0, False, id="valid_far_deadline"),

934 + # no deadline field (older peer / router did not forward) -> skipped.

935 + pytest.param(0.0, None, False, id="missing_expiry_field"),

936 + ],

937 +)

938 +@patch(

939 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",

940 + FakeNixlWrapper,

941 +)

942 +def test_turn2_deadline_gate(dist_init, offset, expiry_delta, expect_declined):

943 + """P declines on the turn-2 readback if the deadline is known,

944 + near-expiry and a handshake clock offset is known"""

945 + connector, worker = _make_connector_with_fake_worker()

946 + worker._engine_clock_offset[_REMOTE] = offset

947 + expiry_time = None if expiry_delta is None else time.perf_counter() + expiry_delta

948 + meta = NixlConnectorMetadata()

949 + params = {

950 + "do_remote_prefill": False,

951 + "do_remote_decode": True,

952 + "remote_block_ids": ([20, 21],),

953 + "remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID,

954 + "remote_request_id": "decode-req",

955 + "remote_host": "localhost",

956 + "remote_port": 1234,

957 + "remote_tp_size": 1,

958 + }

959 + if expiry_time is not None:

960 + params["remote_blocks_expiry_time"] = expiry_time

961 + meta.add_new_req_to_recv(

962 + request_id="req",

963 + local_block_ids=([10, 11],),

964 + kv_transfer_params=params,

965 + )

966 + _do_load_kv(connector, meta)

967 + assert worker.xfer_stats.data["num_kv_expired_reqs"] == (

968 + [1] if expect_declined else []

969 + )

970 + _, done_recving = connector.get_finished(finished_req_ids=set())

971 + assert "req" in done_recving

972 +

973 +

974 +@patch(

975 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",

976 + FakeNixlWrapper,

977 +)

978 +def test_turn2_full_prefix_hit_with_expired_deadline_skips_gate(dist_init):

979 + """A full local prefix hit issues no READ. The expiry gate must be skipped

980 + even when D's deadline is expired."""

981 + connector, worker = _make_connector_with_fake_worker()

982 + worker._engine_clock_offset[_REMOTE] = 0.0

983 +

984 + meta = NixlConnectorMetadata()

985 + meta.add_new_req_to_recv(

986 + request_id="req",

987 + local_block_ids=(), # full prefix hit -> zero local block groups

988 + kv_transfer_params={

989 + "do_remote_prefill": False,

990 + "do_remote_decode": True,

991 + "remote_block_ids": ([20, 21],),

992 + "remote_engine_id": _REMOTE,

993 + "remote_request_id": "decode-req",

994 + "remote_host": "localhost",

995 + "remote_port": 1234,

996 + "remote_tp_size": 1,

997 + # Expired on D's clock; with offset 0 it is expired locally too.

998 + "remote_blocks_expiry_time": time.perf_counter() - 10.0,

999 + },

1000 + )

1001 + # Must not raise, and the gate must not fire for a no-read prefix hit.

1002 + _do_load_kv(connector, meta)

1003 + assert worker.xfer_stats.data["num_kv_expired_reqs"] == []

1004 + assert worker.xfer_stats.data["num_failed_transfers"] == []

1005 +

1006 +

1007 +def test_d_node_request_finished_exports_blocks_expiry_time():

1008 + """D-node (do_remote_prefill request) exports a future float expiry time."""

1009 + vllm_config = create_vllm_config(kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG)

1010 + scheduler = create_scheduler(vllm_config)

1011 + BS = vllm_config.cache_config.block_size

1012 + req = create_request(

1013 + request_id=600, block_size=BS, num_tokens=int(BS * 2.5), do_remote_prefill=True

1014 + )

1015 + scheduler.add_request(req)

1016 + req_id = req.request_id

1017 + so = scheduler.schedule()

1018 + scheduler.update_from_output(

1019 + so, create_model_runner_output(reqs=[], finished_recving={req_id})

1020 + )

1021 + so = scheduler.schedule()

1022 + eco = scheduler.update_from_output(

1023 + so, create_model_runner_output(reqs=[req], use_eos=True)

1024 + )

1025 + kv = eco[0].outputs[0].kv_transfer_params

1026 + assert kv["do_remote_decode"] is True

1027 + assert isinstance(kv["remote_blocks_expiry_time"], float)

1028 + assert kv["remote_blocks_expiry_time"] > time.perf_counter()

1029 +

1030 +

1031 +def test_handshake_listener_appends_perf_counter_frame():

1032 + """The handshake reply carries a 2nd frame with the listener's live

1033 + perf_counter."""

1034 + import threading

1035 +

1036 + import msgspec

1037 + import zmq

1038 +

1039 + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_scheduler import (

1040 + NixlBaseConnectorScheduler,

1041 + )

1042 + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import GET_META_MSG

1043 + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import zmq_ctx

1044 + from vllm.utils.network_utils import get_open_port, make_zmq_path

1045 +

1046 + encoded_data = {(0, 0): b"payload-rank-0"}

1047 + ready_event = threading.Event()

1048 + stop_event = threading.Event()

1049 + host = "127.0.0.1"

1050 + port = get_open_port()

1051 + listener = threading.Thread(

1052 + target=NixlBaseConnectorScheduler._nixl_handshake_listener,

1053 + args=(encoded_data, ready_event, stop_event, host, port),

1054 + daemon=True,

1055 + )

1056 + listener.start()

1057 + try:

1058 + assert ready_event.wait(timeout=5)

1059 + path = make_zmq_path("tcp", host, port)

1060 + with zmq_ctx(zmq.REQ, path) as sock: # type: ignore[attr-defined]

1061 + sock.setsockopt(zmq.RCVTIMEO, 5000) # type: ignore[attr-defined]

1062 + t0 = time.perf_counter()

1063 + sock.send(msgspec.msgpack.encode((GET_META_MSG, 0, 0)))

1064 + parts = sock.recv_multipart()

1065 + t1 = time.perf_counter()

1066 + assert len(parts) == 2

1067 + assert parts[0] == b"payload-rank-0"

1068 + remote_perf = msgspec.msgpack.decode(parts[1])

1069 + assert isinstance(remote_perf, float)

1070 + # same-clock => offset is near zero.

1071 + offset = remote_perf - (t0 + t1) / 2

1072 + assert abs(offset) < 1.0

1073 + finally:

1074 + stop_event.set()

1075 + listener.join(timeout=5)

tests/v1/kv_connector/unit/test_nixl_connector.py

+14 -7

@@ -508,7 +508,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):

508 expected_engine_id: str,

509 remote_pp_size: int = 1,

510 notif_agents_only: bool = False,

511 - ) -> dict[tuple[int, int], str]:

511 + ) -> tuple[dict[tuple[int, int], str], float]:

512 # Mimic slow _nixl_handshake, as well as bypass zmq communication.

513 time.sleep(self._hand_shake_latency)

514 # These should've been done in register_kv_caches(), called by

@@ -560,7 +560,8 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):

560 remote_tp_size=remote_tp_size,

561 )

562 remote_agents[(0, remote_tp_rank)] = remote_agent_name

563 - return remote_agents

563 + # Handshake bypasses zmq, so report a zero clock offset to the peer.

564 + return remote_agents, 0.0

565

566

567 class TestNixlHandshake:

@@ -768,7 +769,7 @@ class TestNixlHandshake:

769 range(tp_ratio)

770 )

771

771 - remote_agents = worker._nixl_handshake(

772 + remote_agents, _ = worker._nixl_handshake(

773 host="localhost",

774 port=1234,

775 remote_tp_size=4,

@@ -780,7 +781,7 @@ class TestNixlHandshake:

781 # discovered. This is not a scenario we actively support right now, but

782 # the connector allows it.

783 worker.REMOTE_ENGINE_ID = "remote_engine_2"

783 - remote_agents = worker._nixl_handshake(

784 + remote_agents, _ = worker._nixl_handshake(

785 host="localhost",

786 port=1234,

787 remote_tp_size=6,

@@ -2878,7 +2879,10 @@ def test_compatibility_hash_validation(

2879

2880 # Mock ZMQ socket to return our handshake payload

2881 mock_socket = MagicMock()

2881 - mock_socket.recv.return_value = msgspec.msgpack.encode(handshake_payload)

2882 + mock_socket.recv_multipart.return_value = [

2883 + msgspec.msgpack.encode(handshake_payload),

2884 + msgspec.msgpack.encode(time.perf_counter()),

2885 + ]

2886

2887 # Mock add_remote_agent to avoid actual NIXL operations

2888 # Patch zmq_ctx to return our mock socket

@@ -2897,7 +2901,7 @@ def test_compatibility_hash_validation(

2901 expected_engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID,

2902 )

2903 else:

2900 - result = decode_worker._nixl_handshake(

2904 + result, _ = decode_worker._nixl_handshake(

2905 host="localhost",

2906 port=1234,

2907 remote_tp_size=1,

@@ -2981,7 +2985,10 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario)

2985 raise AssertionError(f"{error_scenario} not a valid scenario")

2986

2987 mock_socket = MagicMock()

2984 - mock_socket.recv.return_value = msg_bytes

2988 + mock_socket.recv_multipart.return_value = [

2989 + msg_bytes,

2990 + msgspec.msgpack.encode(time.perf_counter()),

2991 + ]

2992 with (

2993 patch.object(decode_worker, "add_remote_agent", return_value="fake_agent"),

2994 patch.object(nixl.base_worker, "zmq_ctx") as mock_zmq_ctx,

tests/v1/kv_connector/unit/test_nixl_connector_hma.py

@@ -209,6 +209,7 @@ def test_read_blocks_for_req_expands_remote_ids(

209 worker = object.__new__(NixlConnectorWorker)

210 worker._physical_blocks_per_logical_kv_block = local_physical_per_logical

211 worker._engine_last_active = {}

212 + worker._bidirectional_kv_xfer_enabled = False

213

214 has_mamba = any(t is MambaSpec for t in resolved_types)

215 has_swa = any(t is SlidingWindowSpec for t in resolved_types)

vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py

+6 -1

@@ -322,8 +322,13 @@ class NixlBaseConnectorScheduler:

322 )

323 if msg != GET_META_MSG:

324 logger.warning("Connection listener got unexpected message %s", msg)

325 + # Echo our perf_counter so P can estimate the clock offset.

326 + # perf_counter is only comparable within a process, so this

327 + # listener must run in the same process that stamps the block

328 + # expiry deadline (`_reqs_need_send`).

329 + ts = msgspec.msgpack.encode(time.perf_counter())

330 sock.send_multipart(

326 - (identity, b"", encoded_data[(target_pp_rank, target_tp_rank)])

331 + (identity, b"", encoded_data[(target_pp_rank, target_tp_rank)], ts)

332 )

333

334 def _get_remote_prefill_token_count(self, num_prompt_tokens: int) -> int:

vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py

+41 -8

@@ -266,6 +266,12 @@ class NixlBaseConnectorWorker:

266 # NOTE (NickLucche): For now we use a hardcoded value for a simpler interface.

267 self._lease_extension = kv_lease_duration * 2 // 3

268

269 + self._bidirectional_kv_xfer_enabled: bool = (

270 + vllm_config.kv_transfer_config.get_from_extra_config(

271 + "bidirectional_kv_xfer", False

272 + )

273 + )

274 +

275 self._is_hma_required = (

276 not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager

277 and any(

@@ -340,6 +346,8 @@ class NixlBaseConnectorWorker:

346 self._remote_agents: dict[EngineId, dict[tuple[int, int], str]] = defaultdict(

347 dict

348 )

349 + # Map of engine_id -> clock offset.

350 + self._engine_clock_offset: dict[EngineId, float] = {}

351

352 # Metadata.

353 self.engine_id: EngineId = engine_id

@@ -470,7 +478,9 @@ class NixlBaseConnectorWorker:

478 thread_name_prefix="vllm-nixl-handshake-initiator",

479 )

480 self._ready_requests = queue.Queue[tuple[ReqId, ReqMeta]]()

473 - self._handshake_futures: dict[EngineId, Future[dict[tuple[int, int], str]]] = {}

481 + self._handshake_futures: dict[

482 + EngineId, Future[tuple[dict[tuple[int, int], str], float]]

483 + ] = {}

484 # Protects _handshake_futures and _remote_agents.

485 self._handshake_lock = threading.RLock()

486

@@ -559,7 +569,7 @@ class NixlBaseConnectorWorker:

569 expected_engine_id: str,

570 remote_pp_size: int = 1,

571 notif_agents_only: bool = False,

562 - ) -> dict[tuple[int, int], str]:

572 + ) -> tuple[dict[tuple[int, int], str], float]:

573 """Do a NIXL handshake with a remote instance."""

574

575 # the first time we connect to a remote agent.

@@ -583,6 +593,11 @@ class NixlBaseConnectorWorker:

593 p_remote_ranks = self.transfer_topo.handshake_target_ranks(remote_tp_size)

594 remote_rank_to_agent_name: dict[tuple[int, int], str] = {}

595 path = make_zmq_path("tcp", host, port)

596 + # Clock offset to the peer, estimated from the handshake round-trip.

597 + # Keep the lowest-RTT sample: hop cost is ~uniform across ranks, so a

598 + # higher RTT is just noise that skews the midpoint estimate.

599 + best_rtt = float("inf")

600 + best_offset: float | None = None

601

602 with zmq_ctx(zmq.REQ, path) as sock:

603 for remote_pp_rank, remote_rank in itertools.product(

@@ -595,15 +610,24 @@ class NixlBaseConnectorWorker:

610 remote_rank,

611 )

612

598 - start_time = time.perf_counter()

613 # Send query for the request.

614 msg = msgspec.msgpack.encode(

615 (GET_META_MSG, remote_pp_rank, remote_rank)

616 )

617 # Set receive timeout to 5 seconds to avoid hanging on dead server

618 sock.setsockopt(zmq.RCVTIMEO, 5000) # milliseconds

619 + start_time = time.perf_counter()

620 sock.send(msg)

606 - handshake_bytes = sock.recv()

621 + reply_parts = sock.recv_multipart()

622 + recv_time = time.perf_counter()

623 + assert len(reply_parts) == 2

624 + handshake_bytes = reply_parts[0]

625 +

626 + remote_perf = msgspec.msgpack.decode(reply_parts[1])

627 + rtt = recv_time - start_time

628 + if rtt < best_rtt:

629 + best_rtt = rtt

630 + best_offset = remote_perf - (start_time + recv_time) / 2

631

632 # Decode handshake payload to get compatibility hash

633 handshake_decoder = msgspec.msgpack.Decoder(NixlHandshakePayload)

@@ -681,7 +705,9 @@ class NixlBaseConnectorWorker:

705 )

706 remote_ranks = (remote_pp_rank, remote_rank)

707 remote_rank_to_agent_name[remote_ranks] = remote_agent_name

684 - return remote_rank_to_agent_name

708 +

709 + assert best_offset is not None

710 + return remote_rank_to_agent_name, best_offset

711

712 def _add_notif_only_remote_agent(

713 self, metadata: NixlAgentMetadata, remote_tp_size: int

@@ -818,7 +844,7 @@ class NixlBaseConnectorWorker:

844 tp_size: int,

845 pp_size: int = 1,

846 notif_agents_only: bool = False,

821 - ) -> Future[dict[tuple[int, int], str]] | None:

847 + ) -> Future[tuple[dict[tuple[int, int], str], float]] | None:

848 """

849 Ensure a handshake is in-flight (or already done) for *engine_id*.

850

@@ -846,11 +872,16 @@ class NixlBaseConnectorWorker:

872 )

873 self._handshake_futures[engine_id] = fut

874

849 - def done_callback(f: Future[dict[tuple[int, int], str]], eid=engine_id):

875 + def done_callback(

876 + f: Future[tuple[dict[tuple[int, int], str], float]],

877 + eid=engine_id,

878 + ):

879 with self._handshake_lock:

880 del self._handshake_futures[eid]

881 try:

853 - self._remote_agents[eid] = f.result()

882 + remote_agents, clock_offset = f.result()

883 + self._remote_agents[eid] = remote_agents

884 + self._engine_clock_offset[eid] = clock_offset

885 self._engine_last_active[eid] = time.perf_counter()

886 except Exception as e:

887 self._log_failure(

@@ -2421,6 +2452,8 @@ class NixlBaseConnectorWorker:

2452 if self.transfer_topo is not None:

2453 self.transfer_topo.unregister_remote_engine(engine_id)

2454

2455 + # Drop the cached clock offset; it is re-measured on the next handshake.

2456 + self._engine_clock_offset.pop(engine_id, None)

2457 # Push P-side engines are tracked in _remote_agents but not in

2458 # _engine_last_active (they don't participate in stale eviction), so

2459 # tolerate a missing entry.

vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py

+5 -1

@@ -39,8 +39,10 @@ PUSH_REG_NOTIF_PREFIX = b"PUSH_REG:"

39 # 2: Add remote_request_id to kv_transfer_params

40 # 3: Add physical_blocks_per_logical_kv_block to NixlAgentMetadata

41 # 4: Add KV block lease renewal through heartbeats

42 +# 5: Add remote_blocks_expiry_time to kv_transfer_params + handshake

43 +# clock-sync timestamp

44 #

43 -NIXL_CONNECTOR_VERSION: int = 4

45 +NIXL_CONNECTOR_VERSION: int = 5

48 @dataclass

@@ -157,6 +159,7 @@ class RemoteMeta:

159 port: int

160 engine_id: str

161 request_id: str

162 + blocks_expiry_time: float | None = None

163

164

165 @dataclass

@@ -225,5 +228,6 @@ class NixlConnectorMetadata(KVConnectorMetadata):

228 request_id=kv_transfer_params["remote_request_id"],

229 host=kv_transfer_params["remote_host"],

230 port=kv_transfer_params["remote_port"],

231 + blocks_expiry_time=kv_transfer_params.get("remote_blocks_expiry_time"),

232 )

233 self.reqs_to_recv[request_id] = req

vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py

@@ -238,6 +238,7 @@ class NixlPullConnectorScheduler(NixlBaseConnectorScheduler):

238 # remove the conditional below

239 delay_free_blocks = any(len(group) > 0 for group in block_ids)

240 remote_num_tokens = 0

241 + blocks_expiry_time = None

242 if delay_free_blocks:

243 # Prefill request on remote. It will be read from D upon completion

244 request_kv_blocks_ttl = self._kv_lease_duration

@@ -254,6 +255,9 @@ class NixlPullConnectorScheduler(NixlBaseConnectorScheduler):

255 self._reqs_need_send[request.request_id] = (

256 time.perf_counter() + request_kv_blocks_ttl

257 )

258 + if is_d_node:

259 + # D blocks expiry time exported for the turn-2 readback.

260 + blocks_expiry_time = self._reqs_need_send[request.request_id]

261 # NOTE HMA will "mark" empty/null blocks in groups with 0s (eg SWA ones),

262 # trimming down after allocating for the whole sequence length. Empty

263 # blocks are always at the start of the list.

@@ -272,4 +276,5 @@ class NixlPullConnectorScheduler(NixlBaseConnectorScheduler):

276 remote_port=self.side_channel_port,

277 tp_size=self.vllm_config.parallel_config.tensor_parallel_size,

278 remote_num_tokens=remote_num_tokens,

279 + remote_blocks_expiry_time=blocks_expiry_time,

280 )

vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py

+26

@@ -25,6 +25,10 @@ if TYPE_CHECKING:

26 logger = init_logger(__name__)

28 +# Slack (seconds) subtracted from D's exported block-expiry deadline on the turn-2

29 +# readback, absorbing clock-offset error and read latency.

30 +_KV_BLOCKS_EXPIRY_SAFETY_MARGIN = 5.0

31 +

33 class NixlPullConnectorWorker(NixlBaseConnectorWorker):

34 """Pull-specific (READ) worker logic."""

@@ -98,12 +102,34 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):

102 # requests sit in the D scheduler WAITING queue.

103 self._send_heartbeats(metadata)

104

105 + def _is_turn2_read_expired(self, meta: ReqMeta) -> bool:

106 + """Whether D's cached blocks for this turn-2 readback have (nearly) expired."""

107 + assert meta.remote is not None

108 + blocks_expiry_time = meta.remote.blocks_expiry_time

109 + # Deadline may be absent (router may not forward it) -> read as usual.

110 + if blocks_expiry_time is None or not meta.local_physical_block_ids:

111 + return False

112 + clock_offset = self._engine_clock_offset[meta.remote.engine_id]

113 + deadline = blocks_expiry_time - clock_offset

114 + return time.perf_counter() + _KV_BLOCKS_EXPIRY_SAFETY_MARGIN >= deadline

115 +

116 def _read_blocks_for_req(self, req_id: str, meta: ReqMeta):

117 assert meta.remote is not None and self.transfer_topo is not None

118 engine_id = meta.remote.engine_id

119 # Update last activity from this remote. Mind that cleanup is done on main

120 # thread (this one), so we don't race on this structure.

121 self._engine_last_active[engine_id] = time.perf_counter()

122 +

123 + if self._bidirectional_kv_xfer_enabled and self._is_turn2_read_expired(meta):

124 + logger.warning(

125 + "Declining expired remote read for %s from engine %s.",

126 + req_id,

127 + engine_id,

128 + )

129 + self.xfer_stats.record_kv_expired_req()

130 + self._handle_failed_transfer(req_id, None)

131 + return

132 +

133 plan = self.tp_mappings[engine_id]

134 remote_info = self.transfer_topo.get_engine_info(engine_id)

135 tp_ratio = self.transfer_topo.tp_ratio(remote_info.remote_tp_size)

vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py

+2 -2

@@ -294,7 +294,7 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):

294 return

295

296 def _on_handshake(

297 - f: Future[dict[tuple[int, int], str]],

297 + f: Future[tuple[dict[tuple[int, int], str], float]],

298 rid: str = req_id,

299 rd: dict[str, Any] = reg_data,

300 ) -> None:

@@ -456,7 +456,7 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):

456 if decode_engine_id in self._remote_agents:

457 return True

458 try:

459 - remote_agents = self._nixl_handshake(

459 + remote_agents, _ = self._nixl_handshake(

460 decode_host,

461 decode_port,

462 decode_tp_size,