@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 5672d100ed

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 3, 22:45 UTC

[KV Connector][NIXL][Bugfix] Fix NIXL handshake failures not honoring kv_load_failure_policy (#40364)

When NIXL handshake fails (e.g., due to compatibility hash mismatch between prefill and decode instances), requests fail with "engine dead" error instead of gracefully falling back to local recomputation as configured by kv_load_failure_policy='recompute'. Signed-off-by: NickLucche <nlucches@redhat.com>

Nicolò Lucchesi committed May 11, 2026 at 11:37 UTC 5672d100ed66490b84ea54b030b0785a6d20d111

2 files changed +162 -25

tests/v1/kv_connector/unit/test_nixl_connector.py

+121

@@ -2267,6 +2267,127 @@ def test_transfer_setup_failure_returns_finished(default_vllm_config, dist_init)

2267 assert request_id in done_recving

2268

2269

2270 +@patch(

2271 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker.NixlWrapper",

2272 + FailingNixlWrapper,

2273 +)

2274 +@pytest.mark.parametrize(

2275 + "failure_mode",

2276 + [

2277 + "handshake",

2278 + "transfer_setup",

2279 + "transfer_failed",

2280 + "transfer_exception",

2281 + ],

2282 +)

2283 +def test_failed_request_skips_kv_postprocessing(

2284 + default_vllm_config, dist_init, failure_mode

2285 +):

2286 + """Test that failed requests skip KV sync and post-processing in

2287 + get_finished().

2288 +

2289 + This is the core safety behavior: when a KV transfer fails at any stage,

2290 + the request must still appear in done_recving (so the scheduler can apply

2291 + kv_load_failure_policy), but sync_recved_kv_to_device and post-processing

2292 + must NOT be called since no valid KV data was received.

2293 +

2294 + Covers all failure paths that involve an actual (attempted) KV transfer:

2295 + - handshake: add_remote_agent raises during async handshake

2296 + - transfer_setup: make_prepped_xfer raises before handle is in transfers

2297 + - transfer_failed: check_xfer_state returns bad state ("ERR") in

2298 + _pop_done_transfers — this is the path that previously had the bug

2299 + where post-processing was NOT skipped

2300 + - transfer_exception: check_xfer_state raises in _pop_done_transfers

2301 +

2302 + Note: notification_failed (send_notif raises on the full-cache-hit path)

2303 + is intentionally excluded. That path is a best-effort D→P courtesy

2304 + notification; the blocks are already in D's cache, so no KV transfer

2305 + was attempted and done_recving is correctly empty.

2306 + """

2307 + # Map each failure mode to the FailingNixlWrapper attribute to set.

2308 + _WRAPPER_CONFIG: dict[str, str] = {

2309 + "handshake": "fail_handshake",

2310 + "transfer_setup": "fail_transfer_setup",

2311 + "transfer_failed": "fail_transfer_state",

2312 + "transfer_exception": "fail_transfer_exception",

2313 + }

2314 +

2315 + # Use enable_permute_local_kv=True so that

2316 + # post_process_device_kv_on_receive would be called on the success path,

2317 + # making the assertion meaningful (not trivially true).

2318 + vllm_config = create_vllm_config(enable_permute_local_kv=True)

2319 +

2320 + connector = NixlConnector(

2321 + vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)

2322 + )

2323 + connector.connector_worker = FakeNixlConnectorWorker(

2324 + vllm_config,

2325 + connector.engine_id,

2326 + hand_shake_latency=0.1 if failure_mode == "handshake" else 0,

2327 + )

2328 + worker = connector.connector_worker

2329 + setattr(worker.nixl_wrapper, _WRAPPER_CONFIG[failure_mode], True)

2330 +

2331 + request_id = f"test_{failure_mode}_skip_postprocess"

2332 + metadata = NixlConnectorMetadata()

2333 + metadata.add_new_req_to_recv(

2334 + request_id=request_id,

2335 + local_block_ids=([1, 2, 3],),

2336 + kv_transfer_params={

2337 + "remote_block_ids": ([4, 5, 6],),

2338 + "remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID,

2339 + "remote_request_id": f"prefill-{request_id}",

2340 + "remote_host": "localhost",

2341 + "remote_port": 1234,

2342 + "remote_tp_size": 1,

2343 + },

2344 + )

2345 + connector.bind_connector_metadata(metadata)

2346 +

2347 + dummy_ctx = ForwardContext(

2348 + no_compile_layers={},

2349 + attn_metadata={},

2350 + slot_mapping={},

2351 + )

2352 + connector.start_load_kv(dummy_ctx)

2353 +

2354 + if failure_mode == "handshake":

2355 + # Wait for async handshake to fail.

2356 + time.sleep(0.3)

2357 + else:

2358 + # All other modes: let the handshake complete, then process the

2359 + # ready_requests queue. For transfer_failed / transfer_exception the

2360 + # handle ends up in _recving_transfers; the failure surfaces in

2361 + # get_finished() via _pop_done_transfers below.

2362 + connector.bind_connector_metadata(NixlConnectorMetadata())

2363 + time.sleep(0.1)

2364 + connector.start_load_kv(dummy_ctx)

2365 +

2366 + # Spy on sync_recved_kv_to_device and post_process_device_kv_on_receive

2367 + # to verify they are NOT called for the failed request.

2368 + with (

2369 + patch.object(worker, "sync_recved_kv_to_device") as mock_sync,

2370 + patch.object(worker, "post_process_device_kv_on_receive") as mock_postprocess,

2371 + ):

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

2373 +

2374 + # The failed request must appear in done_recving so the scheduler

2375 + # can handle it (e.g., trigger recompute via kv_load_failure_policy).

2376 + assert request_id in done_recving

2377 +

2378 + # Critical: KV sync and post-processing must NOT have been called

2379 + # since no valid KV data was received for the failed request.

2380 + mock_sync.assert_not_called()

2381 + mock_postprocess.assert_not_called()

2382 +

2383 + # Metadata for the request should have been cleaned up.

2384 + assert request_id not in worker._recving_metadata

2385 +

2386 + # Blocks should have been marked as invalid.

2387 + invalid_blocks = connector.get_block_ids_with_load_errors()

2388 + assert invalid_blocks == {1, 2, 3}

2389 +

2390 +

2391 @pytest.mark.parametrize(

2392 "mismatch_type,config_overrides,version_override,should_fail,enforce_handshake_compat",

2393 [

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

+41 -25

@@ -388,10 +388,12 @@ class NixlConnectorWorker:

388 # Set of requests that have been part of a batch, regardless of status.

389 self._reqs_to_process: set[ReqId] = set()

390

391 - # invalid blocks from failed NIXL operations

392 - self._invalid_block_ids: set[int] = set()

391 + # Invalid blocks from failed NIXL operations (thread-safe queue of block ids)

392 + self._invalid_block_ids: queue.Queue[set[int]] = queue.Queue()

393 # requests that skipped transfer (handshake or transfer failures)

394 - self._failed_recv_reqs: set[ReqId] = set()

394 + # Uses Queue for thread-safe cross-thread coordination with the

395 + # background handshake thread, matching the _ready_requests pattern.

396 + self._failed_recv_reqs: queue.Queue[ReqId] = queue.Queue()

397

398 # Handshake metadata of this worker for NIXL transfers.

399 self.xfer_handshake_metadata: NixlHandshakePayload | None = None

@@ -768,11 +770,7 @@ class NixlConnectorWorker:

770 error=e,

771 meta=meta,

772 )

771 - if (

772 - req_meta := self._recving_metadata.get(req_id)

773 - ) and not self._is_hma_required:

774 - self._invalid_block_ids.update(req_meta.local_block_ids[0])

775 - self._failed_recv_reqs.add(req_id)

773 + self._handle_failed_transfer(req_id, None)

774

775 fut.add_done_callback(request_ready)

776

@@ -1696,17 +1694,26 @@ class NixlConnectorWorker:

1694 done_sending = self._get_new_notifs()

1695 done_recving = self._pop_done_transfers(self._recving_transfers)

1696

1699 - # add requests that skipped transfer to done_recving

1700 - done_recving.update(self._failed_recv_reqs)

1701 - self._failed_recv_reqs.clear()

1697 + # Drain queue of requests where handshake or transfer setup failed.

1698 + failed_recv_reqs = set[ReqId]()

1699 + while not self._failed_recv_reqs.empty():

1700 + try:

1701 + failed_recv_reqs.add(self._failed_recv_reqs.get_nowait())

1702 + except queue.Empty:

1703 + break

1704 +

1705 + # Add failed requests to done_recving for scheduler tracking

1706 + # (blocks are already marked invalid, scheduler will handle recompute)

1707 + done_recving.update(failed_recv_reqs)

1708

1709 if len(done_sending) > 0 or len(done_recving) > 0:

1710 logger.debug(

1711 "Rank %s, get_finished: %s requests done sending "

1706 - "and %s requests done recving",

1712 + "and %s requests done recving (%s failed)",

1713 self.tp_rank,

1714 len(done_sending),

1715 len(done_recving),

1716 + len(failed_recv_reqs),

1717 )

1718

1719 block_ids_for_blocksize_post_process = defaultdict(list)

@@ -1715,6 +1722,15 @@ class NixlConnectorWorker:

1722 # clean up metadata for completed requests

1723 meta = self._recving_metadata.pop(req_id, None)

1724 assert meta is not None, f"{req_id} not found in recving_metadata list"

1725 +

1726 + # Skip KV sync and post-processing for failed requests

1727 + if req_id in failed_recv_reqs:

1728 + logger.warning(

1729 + "Skipping KV post-processing for failed request %s",

1730 + req_id,

1731 + )

1732 + continue

1733 +

1734 assert meta.remote is not None

1735 if self.use_host_buffer:

1736 self.sync_recved_kv_to_device(req_id, meta)

@@ -1889,7 +1905,7 @@ class NixlConnectorWorker:

1905 transfers[req_id] = in_progress

1906 return done_req_ids

1907

1892 - def _handle_failed_transfer(self, req_id: str, handle: int):

1908 + def _handle_failed_transfer(self, req_id: str, handle: int | None):

1909 """

1910 Handle a failed transfer by marking all (logical) blocks as invalid and

1911 recording the failure.

@@ -1901,8 +1917,10 @@ class NixlConnectorWorker:

1917 # Use .get() here as the metadata cleanup is handled by get_finished()

1918 # TODO (NickLucche) handle failed transfer for HMA.

1919 if (meta := self._recving_metadata.get(req_id)) and not self._is_hma_required:

1904 - self._invalid_block_ids.update(meta.local_block_ids[0])

1905 - self.nixl_wrapper.release_xfer_handle(handle)

1920 + self._invalid_block_ids.put(set(meta.local_block_ids[0]))

1921 + self._failed_recv_reqs.put(req_id)

1922 + if handle is not None:

1923 + self.nixl_wrapper.release_xfer_handle(handle)

1924 self.xfer_stats.record_failed_transfer()

1925

1926 def start_load_kv(self, metadata: NixlConnectorMetadata):

@@ -2221,14 +2239,7 @@ class NixlConnectorWorker:

2239 dst_engine_id=dst_engine_id,

2240 remote_rank=remote_rank,

2241 )

2224 - if (

2225 - meta := self._recving_metadata.get(request_id)

2226 - ) and not self._is_hma_required:

2227 - self._invalid_block_ids.update(meta.local_block_ids[0])

2228 - self.xfer_stats.record_failed_transfer()

2229 - if handle is not None:

2230 - self.nixl_wrapper.release_xfer_handle(handle)

2231 - self._failed_recv_reqs.add(request_id)

2242 + self._handle_failed_transfer(request_id, handle)

2243

2244 def get_mapped_blocks(

2245 self, block_ids: np.ndarray, block_size_ratio: int

@@ -2366,8 +2377,13 @@ class NixlConnectorWorker:

2377 This is called by the scheduler to identify blocks that need

2378 to be retried after a NIXL transfer failure.

2379 """

2369 - result = self._invalid_block_ids

2370 - self._invalid_block_ids = set()

2380 + # Drain the queue (thread-safe, no lock needed).

2381 + result: set[int] = set()

2382 + while not self._invalid_block_ids.empty():

2383 + try:

2384 + result.update(self._invalid_block_ids.get_nowait())

2385 + except queue.Empty:

2386 + break

2387 return result

2388

2389 def __del__(self):