@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 381c61e008

[Bugfix][NIXL] Report a full prefix cache hit as finished (#56640)

Signed-off-by: Yizheng Jiao <jyizheng@gmail.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>

Yizheng Jiao committed Sep 14, 2026 at 20:15 UTC 381c61e008682945d43d6c8808ed633477364b4b

7 files changed +142 -6

tests/v1/kv_connector/unit/test_nixl_connector.py

+72

@@ -341,6 +341,10 @@ def test_abort_immediately_remote_prefill_enqueues_empty_recv():

341 assert req_meta.remote.request_id == f"prefill-{42}"

342 # do_remote_prefill is consumed by request_finished to prevent re-issuing.

343 assert request.kv_transfer_params["do_remote_prefill"] is False

344 + # The scheduler is not waiting on this recv -- the request is already gone

345 + # from self.requests, so reporting it would trip `assert req_id in

346 + # self.requests` in _update_from_kv_xfer_finished.

347 + assert req_meta.awaiting_kvs is False

348

349

350 @patch(

@@ -2695,6 +2699,74 @@ class FailingNixlWrapper(FakeNixlWrapper):

2699 return super().check_xfer_state(handle)

2700

2701

2702 +@patch(

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

2704 + FailingNixlWrapper,

2705 +)

2706 +@pytest.mark.parametrize("awaiting_kvs", [True, False])

2707 +@pytest.mark.parametrize("notification_fails", [False, True])

2708 +def test_empty_recv_is_reported_only_when_awaited(

2709 + default_vllm_config, dist_init, awaiting_kvs, notification_fails

2710 +):

2711 + """An empty recv is reported iff the scheduler parked the request on it.

2712 +

2713 + _read_blocks returns early when the local block list is empty -- D already

2714 + holds the KV and only the producer needs telling. Two kinds of caller reach

2715 + that path, and they need opposite handling:

2716 +

2717 + Parked (awaiting_kvs=True, a full prefix cache hit). Without a report the

2718 + request is named in neither _recving_transfers nor _failed_recv_reqs and

2719 + never reaches finished_recving. The scheduler has no other way to release a

2720 + WAITING_FOR_REMOTE_KVS request, and there is no timeout, so it sits in

2721 + skipped_waiting holding its blocks for the life of the process.

2722 +

2723 + Notify-only (awaiting_kvs=False): request_finished seeding an empty recv to

2724 + free P's blocks for a request aborted before it was scheduled, or a readback

2725 + on a request that stays RUNNING. Reporting either one crashes the engine

2726 + core -- _update_from_kv_xfer_finished asserts the id is still in

2727 + self.requests and that the request is parked or finished.

2728 +

2729 + A failed notification changes neither: the KV is local either way, and the

2730 + producer frees its own blocks on a timeout.

2731 + """

2732 + vllm_config = create_vllm_config()

2733 + connector = NixlConnector(

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

2735 + )

2736 + connector.connector_worker = FakeNixlConnectorWorker(

2737 + vllm_config,

2738 + connector.engine_id,

2739 + hand_shake_latency=0.0,

2740 + kv_cache_config=connector._kv_cache_config,

2741 + )

2742 + connector.connector_worker.nixl_wrapper.fail_send_notif = notification_fails

2743 +

2744 + request_id = f"test_empty_recv_awaited_{awaiting_kvs}"

2745 + metadata = NixlConnectorMetadata()

2746 + metadata.add_new_req_to_recv(

2747 + request_id=request_id,

2748 + local_block_ids=(), # empty: the whole prompt is already cached locally

2749 + kv_transfer_params={

2750 + "remote_block_ids": [[20, 21, 22]],

2751 + "remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID,

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

2753 + "remote_host": "localhost",

2754 + "remote_port": 1234,

2755 + "remote_tp_size": 1,

2756 + },

2757 + awaiting_kvs=awaiting_kvs,

2758 + )

2759 + connector.bind_connector_metadata(metadata)

2760 + dummy_ctx = ForwardContext(no_compile_layers={}, attn_metadata={}, slot_mapping={})

2761 + connector.start_load_kv(dummy_ctx)

2762 + connector.bind_connector_metadata(NixlConnectorMetadata())

2763 + time.sleep(0.1)

2764 + connector.start_load_kv(dummy_ctx)

2765 +

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

2767 + assert (request_id in done_recving) is awaiting_kvs

2768 +

2769 +

2770 @patch(

2771 "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",

2772 FailingNixlWrapper,

tests/v1/kv_connector/unit/test_nixl_connector_hma.py

+42 -1

@@ -149,8 +149,49 @@ def test_update_state_after_alloc_tracks_cached_blocks_per_group():

149

150 scheduler.update_state_after_alloc(request, blocks, num_external_tokens=2)

151

152 - _, _, local_num_computed_blocks = scheduler._reqs_need_recv[request.request_id]

152 + _, _, local_num_computed_blocks, awaiting_kvs = scheduler._reqs_need_recv[

153 + request.request_id

154 + ]

155 assert local_num_computed_blocks == (2, 1)

156 + # There are external tokens to pull, so the scheduler parks the request and

157 + # is waiting to see it in finished_recving.

158 + assert awaiting_kvs is True

159 +

160 +

161 +@pytest.mark.cpu_test

162 +def test_full_local_hit_is_not_awaited_by_the_scheduler():

163 + """A recv with num_external_tokens == 0 must not be flagged as awaited.

164 +

165 + The request keeps running -- there is nothing to pull, only a notification

166 + owed to the producer -- so the worker must not report it in

167 + finished_recving. Doing so trips `assert RequestStatus.is_finished` in

168 + _update_from_kv_xfer_finished, which only tolerates a finished recv for a

169 + request that is parked or already done."""

170 + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.pull_scheduler import (

171 + NixlPullConnectorScheduler,

172 + )

173 + from vllm.v1.core.kv_cache_manager import KVCacheBlocks

174 + from vllm.v1.core.kv_cache_utils import KVCacheBlock

175 +

176 + scheduler = object.__new__(NixlPullConnectorScheduler)

177 + scheduler._reqs_in_batch = set()

178 + scheduler._reqs_need_save = {}

179 + scheduler._reqs_need_recv = {}

180 + scheduler.use_host_buffer = False

181 + scheduler.is_bidirectional_kv_xfer_enabled = False

182 + scheduler._is_hma_required = False

183 + scheduler.kv_cache_config = MagicMock(

184 + select_transfer_block_ids=lambda block_ids: block_ids

185 + )

186 +

187 + blocks = KVCacheBlocks(blocks=([KVCacheBlock(block_id=0, _block_hash=object())],))

188 + request = create_request(do_remote_prefill=True)

189 +

190 + scheduler.update_state_after_alloc(request, blocks, num_external_tokens=0)

191 +

192 + _, local_block_ids, _, awaiting_kvs = scheduler._reqs_need_recv[request.request_id]

193 + assert not local_block_ids # full local hit: nothing to pull

194 + assert awaiting_kvs is False

195

196

197 @pytest.mark.cpu_test

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

+8 -2

@@ -113,7 +113,7 @@ class NixlBaseConnectorScheduler:

113 # New requests are added by update_state_after_alloc in

114 # the scheduler. Used to make metadata passed to Worker.

115 self._reqs_need_recv: dict[

116 - ReqId, tuple[Request, BlockIds, tuple[int, ...]]

116 + ReqId, tuple[Request, BlockIds, tuple[int, ...], bool]

117 ] = {}

118 self._reqs_need_save: dict[ReqId, Request] = {}

119 # Reqs to send and their expiration time

@@ -459,13 +459,19 @@ class NixlBaseConnectorScheduler:

459 meta = NixlConnectorMetadata()

460

461 # Loop through scheduled reqs and convert to ReqMeta.

462 - for req_id, (req, block_ids, cached) in self._reqs_need_recv.items():

462 + for req_id, (

463 + req,

464 + block_ids,

465 + cached,

466 + awaiting_kvs,

467 + ) in self._reqs_need_recv.items():

468 assert req.kv_transfer_params is not None

469 meta.add_new_req_to_recv(

470 request_id=req_id,

471 local_block_ids=block_ids,

472 kv_transfer_params=req.kv_transfer_params,

473 local_num_computed_blocks=cached,

474 + awaiting_kvs=awaiting_kvs,

475 )

476

477 if self.use_host_buffer:

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

@@ -246,6 +246,9 @@ class ReqMeta:

246 remote_block_size: int | None = None

247 # Remote producer pipeline-parallel size (push mode, D side).

248 pp_size: int = 1

249 + # True only when the scheduler parked the request in WAITING_FOR_REMOTE_KVS

250 + # and expects it in finished_recving; notify-only recvs must not be reported.

251 + awaiting_kvs: bool = False

252

253

254 class NixlConnectorMetadata(KVConnectorMetadata):

@@ -275,6 +278,7 @@ class NixlConnectorMetadata(KVConnectorMetadata):

278 local_block_ids: BlockIds,

279 kv_transfer_params: dict[str, Any],

280 local_num_computed_blocks: tuple[int, ...] = (),

281 + awaiting_kvs: bool = False,

282 ) -> ReqMeta:

283 return ReqMeta(

284 local_block_ids=local_block_ids,

@@ -285,6 +289,7 @@ class NixlConnectorMetadata(KVConnectorMetadata):

289 remote_block_size=kv_transfer_params.get("remote_block_size"),

290 pp_size=kv_transfer_params.get("pp_size", 1),

291 local_num_computed_blocks=local_num_computed_blocks,

292 + awaiting_kvs=awaiting_kvs,

293 )

294

295 def add_new_req_to_save(

@@ -303,11 +308,13 @@ class NixlConnectorMetadata(KVConnectorMetadata):

308 local_block_ids: BlockIds,

309 kv_transfer_params: dict[str, Any],

310 local_num_computed_blocks: tuple[int, ...] = (),

311 + awaiting_kvs: bool = False,

312 ):

313 req = self._add_new_req(

314 local_block_ids,

315 kv_transfer_params,

316 local_num_computed_blocks,

317 + awaiting_kvs,

318 )

319 req.remote = RemoteMeta(

320 block_ids=kv_transfer_params["remote_block_ids"],

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

+4 -1

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

172 request,

173 local_block_ids,

174 local_num_computed_blocks,

175 + # Parked in WAITING_FOR_REMOTE_KVS only when there is

176 + # something to pull; a full local hit stays RUNNING.

177 + num_external_tokens > 0,

178 )

179

180 else:

@@ -224,7 +227,7 @@ class NixlPullConnectorScheduler(NixlBaseConnectorScheduler):

227 # To avoid stranding the prefill blocks in the prefill instance,

228 # we must add empty block_ids to _reqs_need_recv so that our

229 # worker side will notify and free blocks in the prefill instance.

227 - self._reqs_need_recv[request.request_id] = (request, [], ())

230 + self._reqs_need_recv[request.request_id] = (request, [], (), False)

231 params["do_remote_prefill"] = False

232 return False, None

233

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

@@ -303,6 +303,7 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):

303 local_dram_handle=local_dram_handle,

304 remote_xfer_side_handle=remote_xfer_side_handle,

305 expected_consumers=plan.local_consumers,

306 + awaiting_kvs=meta.awaiting_kvs,

307 ):

308 return

309

@@ -327,6 +328,7 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):

328 local_dram_handle: int | None,

329 remote_xfer_side_handle: int,

330 expected_consumers: int,

331 + awaiting_kvs: bool,

332 ) -> bool:

333 """

334 Post a READ point-to-point xfer request from a single local worker to

@@ -387,6 +389,11 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):

389 remote_agent_name=agent_name,

390 )

391 self.xfer_stats.record_failed_notification()

392 + # Report even on notification failure: the KV is already local, and

393 + # an unreported parked request would hold its blocks forever.

394 + # Notify-only recvs must stay unreported (scheduler asserts).

395 + if awaiting_kvs:

396 + self._recving_transfers.setdefault(request_id, [])

397 return True

398

399 if read_spec.block_ids_by_region:

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

+2 -2

@@ -198,7 +198,7 @@ class NixlPushConnectorScheduler(NixlBaseConnectorScheduler):

198 # ReqMeta without a KeyError — the actual remote block IDs are

199 # learned by P over the NIXL handshake at WRITE time.

200 params["remote_block_ids"] = ()

201 - self._reqs_need_recv[request.request_id] = (request, local_block_ids, ())

201 + self._reqs_need_recv[request.request_id] = (request, local_block_ids, (), False)

202

203 # Mark as processed so a re-entry (e.g. preemption + reschedule)

204 # doesn't re-stage the registration.

@@ -240,7 +240,7 @@ class NixlPushConnectorScheduler(NixlBaseConnectorScheduler):

240 # recv so the worker emits a notif that lets P free them.

241 # Seed remote_block_ids so add_new_req_to_recv won't KeyError.

242 params["remote_block_ids"] = ()

243 - self._reqs_need_recv[request.request_id] = (request, [], ())

243 + self._reqs_need_recv[request.request_id] = (request, [], (), False)

244 params["do_remote_prefill"] = False

245 return False, None

246