@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 699d62e6cf
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 14, 05:30 UTC

[NIXL][BUGFIX] delay done_recving queue cleanup to bottom of get_finished (#27297)

Signed-off-by: Chendi Xue <chendi.xue@intel.com>

Chendi.Xue committed Oct 24, 2025 at 12:01 UTC 699d62e6cf4d9a570e2441141c65f687b3b1eef2
1 file changed +8 -11
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+8 -11
@@ -1452,11 +1452,17 @@ class NixlConnectorWorker:
1452 len(done_recving),
1453 )
1454
1455 - # clean up metadata for completed requests
1455 + block_ids_to_permute = []
1456 for req_id in done_recving:
1457 + # clean up metadata for completed requests
1458 meta = self._recving_metadata.pop(req_id, None)
1458 - if self.use_host_buffer and meta:
1459 + assert meta is not None, f"{req_id} not found in recving_metadata list"
1460 + if self.use_host_buffer:
1461 self.sync_recved_kv_to_device(req_id, meta)
1462 + if self.enable_permute_local_kv:
1463 + block_ids_to_permute += meta.local_block_ids
1464 + if len(block_ids_to_permute) > 0:
1465 + self.permute_device_kv(block_ids_to_permute)
1466
1467 # Handle timeout to avoid stranding blocks on remote.
1468 now = time.perf_counter()
@@ -1477,15 +1483,6 @@ class NixlConnectorWorker:
1483 del self._reqs_to_send[req_id]
1484 done_sending.add(req_id)
1485
1480 - if self.enable_permute_local_kv and len(done_recving) > 0:
1481 - block_ids = []
1482 - for req_id in done_recving:
1483 - meta = self._recving_metadata.pop(req_id)
1484 - assert meta, f"{req_id} not found in recving_metadata list"
1485 - block_ids += meta.local_block_ids
1486 -
1487 - self.permute_device_kv(block_ids)
1488 -
1486 return done_sending, done_recving
1487
1488 def _get_new_notifs(self) -> set[str]: