@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 589a5b884b

[PD][NixlPush][Bugfix] Fix blocking handshake call on writer thread (#49221)

Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Nicolò Lucchesi committed Jul 24, 2026 at 10:41 UTC 589a5b884bcc9a7662aac4a184529aa68da77436

4 files changed +196 -80

docs/design/nixl_kv_push_connector.md

+28 -14

@@ -47,7 +47,8 @@ sequenceDiagram

47 else only one side present

48 PWriter->>PWriter: stash and wait, self-poll only when blocks unmatched

49 end

50 - PWriter->>PWriter: ensure D handshake (one-time)

50 + PWriter->>PWriter: _ensure_handshake to D (async; defer WRITE)

51 + PWriter->>PWriter: handshake callback re-queues on _deferred_push_inbox, wake

52 PWriter->>DWriter: NIXL WRITE direct to D GPU + completion notif

54 note over DWorker,DWriter: D side - completion accounting

@@ -100,15 +101,20 @@ event:

101 D, completion notifs after a WRITE, late-arriving ``PUSH_REG``)

102 even when there is no new metadata to act on.

103 3. **Handshake-completion callback** (background handshake executor

103 - thread) — when a deferred D→P handshake finishes successfully, the

104 - future's done-callback re-enqueues the registration onto

105 - ``_reg_send_inbox`` and sets the wake so the corresponding

106 - ``send_notif`` runs on the writer (we never call ``send_notif`` from

107 - the executor thread). On this second pass ``_ensure_handshake``

108 - returns ``None`` (the agent is now connected), so the writer sends

109 - the ``PUSH_REG`` directly. If the handshake *failed*, the callback

110 - fails the request instead of re-enqueuing, so there is no retry

111 - loop.

104 + thread) — both handshakes run on the executor and never block the

105 + writer; their done-callbacks re-enqueue the deferred op and set the

106 + wake, since neither ``send_notif`` nor the NIXL WRITE may run off the

107 + writer thread:

108 + * the **D→P** handshake (before sending ``PUSH_REG``) re-enqueues the

109 + registration onto ``_reg_send_inbox``;

110 + * the **P→D** handshake (before a WRITE) re-enqueues the matched

111 + ``(req_id, blocks, reg_data)`` onto ``_deferred_push_inbox``.

112 +

113 + On this second pass ``_ensure_handshake`` returns ``None`` (the agent

114 + is now connected), so the writer sends the ``PUSH_REG`` / issues the

115 + WRITE directly. If a handshake *failed*, the callback fails or drops

116 + the request instead of re-enqueuing, so there is no retry loop (see

117 + Failure handling).

118

119 In addition to event-driven wakes, the writer self-polls at

120 ``_PUSH_WRITER_POLL_INTERVAL_MS = 1.0`` ms while there are P-side

@@ -227,6 +233,13 @@ Two per-request timers are armed on the scheduler:

233 * **D-side ``send_notif`` failure when shipping the PUSH_REG to P** —

234 identical handling: ``_handle_failed_transfer`` marks the recv as

235 failed.

236 +* **P-side handshake failure (P→D handshake before a WRITE)** — the

237 + future's done-callback logs ``push_handshake_failed`` and drops the

238 + request without re-queuing. It deliberately does *not* call

239 + ``_handle_failed_transfer`` (there is no ``_recving_metadata`` entry to

240 + invalidate on the producer side, same reasoning as the WRITE-submission

241 + failure below). P's blocks are reclaimed by the ``_kv_lease_duration``

242 + lease and D's stale registration by its watchdog.

243 * **P-side WRITE submission failure** — the WRITE handle (if any) is

244 released and ``xfer_stats.record_failed_transfer()`` bumps the

245 failure counter. We deliberately do not call

@@ -246,10 +259,11 @@ existing NIXL connector:

259 class — all subclasses of the existing base classes;

260 * one dedicated background thread per worker;

261 * a few cross-thread queues, each with a single consumer (the writer);

249 - most have one producer, except ``_reg_send_inbox``, which is fed both

250 - by the engine main thread (new registrations) and by the

251 - handshake-completion callback (registrations replayed after their

252 - D→P handshake finishes);

262 + most have one producer, except the two replay queues fed by both the

263 + engine main thread and a handshake-completion callback:

264 + ``_reg_send_inbox`` (registrations replayed after their D→P handshake)

265 + and ``_deferred_push_inbox`` (matched pushes replayed after their P→D

266 + handshake);

267 * one new notification type (`PUSH_REG:<msgpack>`).

268

269 Behavior on the engine main thread is otherwise unchanged. The writer

tests/v1/kv_connector/unit/test_nixl_push_connector.py

+109

@@ -25,6 +25,7 @@ import queue

25 import threading

26 import time

27 from collections import defaultdict

28 +from concurrent.futures import Future

29 from typing import Any

30 from unittest.mock import MagicMock, patch

@@ -320,6 +321,7 @@ class _StubWriterWorker(NixlPushConnectorWorker):

321 w._finished_blocks_inbox = queue.Queue()

322 w._pending_completion_notifs = queue.Queue()

323 w._evict_finished_inbox = queue.Queue()

324 + w._deferred_push_inbox = queue.Queue()

325 w._push_writer_wake = threading.Event()

326 w._push_writer_stop = threading.Event()

327 w._push_writer_thread = None

@@ -334,6 +336,7 @@ class _StubWriterWorker(NixlPushConnectorWorker):

336 w.world_size = 1

337 w.engine_id = "test-decode-engine"

338 w._remote_agents = {}

339 + w._physical_blocks_per_logical_kv_block = 1

340

341 # Track _do_start_push_kv invocations.

342 calls: list[tuple[str, Any, dict[str, Any]]] = []

@@ -506,6 +509,112 @@ class TestPushWriterStartLoadKv:

509 assert w.start_push_calls == []

510

511

512 +# The P→D handshake must run on the base worker's background executor, never

513 +# blocking the writer thread: ``_do_start_push_kv`` defers the WRITE until the

514 +# handshake resolves, then re-drives via ``_deferred_push_inbox``. These call

515 +# the *real* ``_do_start_push_kv`` (the stub overrides it for matching tests).

516 +def _real_do_start_push_kv(w, *args):

517 + return NixlPushConnectorWorker._do_start_push_kv(w, *args)

518 +

519 +

520 +def test_do_start_push_kv_defers_then_writes_when_handshake_ready():

521 + """Full happy-path lifecycle: an in-flight handshake defers the WRITE (no

522 + NIXL op from the writer or the executor callback); once it resolves the

523 + request is re-queued on ``_deferred_push_inbox`` with the wake set; and on

524 + re-drive with the handshake ready the WRITE is issued with a correct

525 + ReqMeta."""

526 + w = _StubWriterWorker.fresh()

527 + w._logical_to_kernel_block_ids = lambda x, ratio: x

528 + xfer_calls: list[dict[str, Any]] = []

529 + w._xfer_blocks_for_req = lambda **kw: xfer_calls.append(kw)

530 +

531 + fut: Future = Future()

532 + w._ensure_handshake = lambda *a, **k: fut

533 +

534 + rd = _registration_data("req-hs", decode_engine_id="decode-engine")

535 + _real_do_start_push_kv(w, "req-hs", ([1, 2, 3],), rd)

536 +

537 + # Handshake pending -> nothing issued, nothing queued, no wake.

538 + assert xfer_calls == []

539 + assert w._deferred_push_inbox.qsize() == 0

540 + assert not w._push_writer_wake.is_set()

541 +

542 + # Handshake completes: request re-queued for the writer, wake set, but no

543 + # *direct* WRITE from the callback (wrong thread for NIXL ops).

544 + fut.set_result(({(0, 0): "agent"}, 0.0))

545 + assert xfer_calls == []

546 + assert w._push_writer_wake.is_set()

547 + rid, blocks, reg = w._deferred_push_inbox.get_nowait()

548 + assert (rid, blocks, reg) == ("req-hs", ([1, 2, 3],), rd)

549 +

550 + # Re-drive on the writer with the handshake now ready -> WRITE issued.

551 + w._ensure_handshake = lambda *a, **k: None

552 + _real_do_start_push_kv(w, rid, blocks, reg)

553 + assert len(xfer_calls) == 1

554 + assert xfer_calls[0]["req_id"] == "req-hs"

555 + meta = xfer_calls[0]["meta"]

556 + assert meta.remote is not None

557 + assert meta.remote.engine_id == "decode-engine"

558 + # RemoteMeta.request_id is D's request id from the registration.

559 + assert meta.remote.request_id == "req-hs"

560 +

561 +

562 +def test_do_start_push_kv_drops_request_on_handshake_failure():

563 + """Handshake raises: the request is dropped (not re-queued, no WRITE) and

564 + the failure is logged. Blocks are reclaimed by the lease/watchdog, matching

565 + the old blocking behaviour."""

566 + w = _StubWriterWorker.fresh()

567 + w._logical_to_kernel_block_ids = lambda x: x

568 + xfer_calls: list[dict[str, Any]] = []

569 + w._xfer_blocks_for_req = lambda **kw: xfer_calls.append(kw)

570 + failures: list[dict[str, Any]] = []

571 + w._log_failure = lambda **kw: failures.append(kw)

572 +

573 + fut: Future = Future()

574 + w._ensure_handshake = lambda *a, **k: fut

575 +

576 + _real_do_start_push_kv(w, "req-fail", ([9],), _registration_data("req-fail"))

577 + fut.set_exception(RuntimeError("handshake boom"))

578 +

579 + assert w._deferred_push_inbox.qsize() == 0

580 + assert xfer_calls == []

581 + assert len(failures) == 1

582 + assert failures[0]["failure_type"] == "push_handshake_failed"

583 +

584 +

585 +def test_writer_loop_drains_deferred_push_inbox():

586 + """The writer loop drains ``_deferred_push_inbox`` and re-drives

587 + ``_do_start_push_kv`` for each entry (event-driven, no polling)."""

588 + w = _StubWriterWorker.fresh()

589 + w.nixl_wrapper = MagicMock()

590 + w.nixl_wrapper.get_new_notifs.return_value = {}

591 +

592 + processed = threading.Event()

593 +

594 + def _tracked(rid, blocks, rd):

595 + w.start_push_calls.append((rid, blocks, rd))

596 + processed.set()

597 +

598 + w._do_start_push_kv = _tracked # type: ignore[method-assign]

599 +

600 + w._deferred_push_inbox.put(

601 + ("req-retry", ([1, 2],), _registration_data("req-retry"))

602 + )

603 + w._push_writer_wake.set()

604 +

605 + t = threading.Thread(target=w._push_writer_loop, daemon=True)

606 + t.start()

607 + try:

608 + assert processed.wait(timeout=2.0), "writer did not drain deferred inbox"

609 + finally:

610 + w._push_writer_stop.set()

611 + w._push_writer_wake.set()

612 + t.join(timeout=2)

613 +

614 + assert len(w.start_push_calls) == 1

615 + assert w.start_push_calls[0][0] == "req-retry"

616 +

617 +

618 class TestPushWriterNotifs:

619 def test_get_new_notifs_processes_forwarded_completion_notif(self):

620 """Non-PUSH_REG notifs forwarded by the writer thread are drained

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

+1 -1

@@ -792,7 +792,7 @@ class NixlBaseConnectorWorker:

792 failure_type: str,

793 req_id: str | None,

794 msg: str = "",

795 - error: Exception | None = None,

795 + error: BaseException | None = None,

796 meta: ReqMeta | None = None,

797 **extra_context,

798 ):

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

+58 -65

@@ -8,10 +8,13 @@ notifs are forwarded to the engine main thread), sends PUSH_REG via

8 ``send_notif``, matches D registrations with P finished blocks, and

9 issues WRITE transfers via ``make_prepped_xfer`` / ``transfer``.

11 -The engine main thread feeds the writer through three queues:

11 +The engine main thread feeds the writer through queues:

12 ``_reg_send_inbox`` (D-side regs to send), ``_finished_blocks_inbox``

13 (P-side blocks from metadata) and ``_pending_completion_notifs``

14 (non-PUSH_REG notifs forwarded back for HB / completion accounting).

15 +The handshake-completion callback feeds ``_deferred_push_inbox`` with

16 +matched pushes whose P→D handshake has finished so the writer can

17 +(re-)issue the WRITE without ever blocking on the network.

19 Wake model: the writer self-polls every

20 ``_PUSH_WRITER_POLL_INTERVAL_MS`` only while it has unmatched

@@ -21,8 +24,8 @@ event-driven: the engine main thread sets ``_push_writer_wake`` from

24 ``start_load_kv`` (when handing it new work) and from ``get_finished``

25 (so each engine step gives the writer a chance to drain NIXL notifs);

26 the handshake-completion callback sets the same event after a deferred

24 -PUSH_REG send has been queued. When a request's lease expires (the base

25 -worker reports it via ``done_sending``) or the WRITE completes,

27 +PUSH_REG send or a deferred push WRITE has been queued. When a request's

28 +lease expires (the base worker reports it via ``done_sending``) or the WRITE completes,

29 ``get_finished`` enqueues an eviction onto ``_evict_finished_inbox`` so

30 the writer drops any leftover ``_push_finished_blocks`` /

31 ``_pending_d_registrations`` and stops self-polling.

@@ -109,6 +112,8 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):

112 # ``_push_finished_blocks`` so an unmatched entry doesn't keep the

113 # writer busy-polling forever.

114 self._evict_finished_inbox: queue.Queue[str] = queue.Queue()

115 + # Handshakes that have just completed and are ready for the WRITE on wthread

116 + self._deferred_push_inbox = queue.Queue[tuple[str, BlockIds, dict[str, Any]]]()

117

118 # Wake signal from engine main thread (start_load_kv / get_finished).

119 # Writer self-polls at _PUSH_WRITER_POLL_INTERVAL_MS while it has

@@ -206,7 +211,15 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):

211 break

212 self._send_registration_to_p(rid, rd)

213

209 - # 2. P-side finished blocks; match against pending regs.

214 + # 2. Deferred P→D pushes whose handshake just completed; do xfer now

215 + while True:

216 + try:

217 + rid, blocks, rd = self._deferred_push_inbox.get_nowait()

218 + except queue.Empty:

219 + break

220 + self._do_start_push_kv(rid, blocks, rd)

221 +

222 + # 3. P-side finished blocks; match against pending regs.

223 while True:

224 try:

225 rid, blocks = self._finished_blocks_inbox.get_nowait()

@@ -218,7 +231,7 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):

231 else:

232 self._push_finished_blocks[rid] = blocks

233

221 - # 2b. Evict finished blocks for requests that have either

234 + # 3b. Evict finished blocks for requests that have either

235 # completed (WRITE acknowledged) or whose lease expired

236 # without a D registration. Drop pending registrations

237 # for the same reason so we don't leak state.

@@ -230,7 +243,7 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):

243 self._push_finished_blocks.pop(rid, None)

244 self._pending_d_registrations.pop(rid, None)

245

233 - # 3. NIXL notifs: route PUSH_REG; forward the rest.

246 + # 4. NIXL notifs: route PUSH_REG; forward the rest.

247 for notifs in self.nixl_wrapper.get_new_notifs().values():

248 for notif in notifs:

249 if notif.startswith(PUSH_REG_NOTIF_PREFIX):

@@ -384,34 +397,53 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):

397 ) -> None:

398 """Start push-based KV transfer from P worker to D node.

399

387 - ``local_block_ids`` are P's *logical* block IDs (from the P

388 - scheduler's metadata). ``registration_data["local_block_ids"]``

389 - are D's *logical* block IDs (from D's scheduler, sent over the

390 - PUSH_REG notif). All conversion to physical block IDs is

391 - deferred to ``_xfer_blocks_for_req`` so each side uses its own

392 - physical-blocks-per-logical ratio (P uses

393 - ``self._physical_blocks_per_logical_kv_block``; D's ratio is

394 - learned during the NIXL handshake)."""

395 - decode_engine_id = registration_data["decode_engine_id"]

396 - remote_block_ids = registration_data["local_block_ids"]

397 - decode_host = registration_data["decode_host"]

398 - decode_port = registration_data["decode_port"]

399 - decode_request_id = registration_data["request_id"]

400 + The P→D handshake runs on the base worker's background executor.

401 + If it isn't ready yet we register a completion callback, defer the

402 + WRITE, and re-drive this request via ``_deferred_push_inbox`` once

403 + the handshake resolves -- so the writer thread never blocks on the

404 + network (mirrors ``_send_registration_to_p``).

405 + """

406 if not local_block_ids:

407 logger.warning("No local blocks to push for request %s", request_id)

408 return

409

404 - if not self._ensure_d_handshake(

410 + # ``local_block_ids`` are P's logical block IDs; ``remote_block_ids``

411 + # (D's, from the PUSH_REG notif) are also logical.

412 + decode_engine_id = registration_data["decode_engine_id"]

413 + remote_block_ids = registration_data["local_block_ids"]

414 + decode_request_id = registration_data["request_id"]

415 +

416 + # Runs on the background executor; defer the WRITE until it's ready.

417 + fut = self._ensure_handshake(

418 decode_engine_id,

406 - decode_host,

407 - decode_port,

419 + registration_data["decode_host"],

420 + registration_data["decode_port"],

421 registration_data["decode_tp_size"],

409 - request_id,

410 - ):

422 + )

423 + if fut is not None:

424 +

425 + def _on_handshake(

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

427 + rid: str = request_id,

428 + blocks: BlockIds = local_block_ids,

429 + rd: dict[str, Any] = registration_data,

430 + ) -> None:

431 + if (e := f.exception()) is not None:

432 + # The engine reclaims the blocks via the TTL so we dont free here

433 + self._log_failure(

434 + failure_type="push_handshake_failed", req_id=rid, error=e

435 + )

436 + return

437 + self._deferred_push_inbox.put((rid, blocks, rd))

438 + self._push_writer_wake.set()

439 +

440 + fut.add_done_callback(_on_handshake)

441 return

442

413 - # Both sides are kept in logical form here; ``_xfer_blocks_for_req``

414 - # expands each side using the appropriate ratio.

443 + # Both sides stay logical here; ``_xfer_blocks_for_req`` converts each

444 + # to physical with its own physical-blocks-per-logical ratio -- P uses

445 + # ``self._physical_blocks_per_logical_kv_block``, D's is learned during

446 + # the NIXL handshake.

447 logical_local = self._as_grouped_block_ids(local_block_ids)

448 logical_remote = self._as_grouped_block_ids(remote_block_ids)

449 physical_local = self._logical_to_kernel_block_ids(

@@ -441,45 +473,6 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):

473 elapsed_ms,

474 )

475

444 - def _ensure_d_handshake(

445 - self,

446 - decode_engine_id: str,

447 - decode_host: str,

448 - decode_port: int,

449 - decode_tp_size: int,

450 - request_id: str,

451 - ) -> bool:

452 - """First-time P→D handshake. Blocking call on the writer thread.

453 -

454 - Returns True iff the handshake succeeded (or had already been

455 - completed). Returns False if the handshake raised; the request is

456 - skipped in that case (the engine layer will reschedule or fail it

457 - via the standard lease/timeout path)."""

458 - if decode_engine_id in self._remote_agents:

459 - return True

460 - try:

461 - remote_agents, _ = self._nixl_handshake(

462 - decode_host,

463 - decode_port,

464 - decode_tp_size,

465 - decode_engine_id,

466 - )

467 - except Exception:

468 - logger.exception(

469 - "Failed handshake to D %s for push %s",

470 - decode_engine_id,

471 - request_id,

472 - )

473 - return False

474 - with self._handshake_lock:

475 - self._remote_agents[decode_engine_id] = remote_agents

476 - logger.info(

477 - "Push handshake to D %s done (%d agents)",

478 - decode_engine_id,

479 - len(remote_agents),

480 - )

481 - return True

482 -

476 @staticmethod

477 def _as_grouped_block_ids(block_ids: BlockIds) -> BlockIds:

478 """Normalise a sequence of block IDs to a tuple-of-groups shape.