From 8bf507d766894a8104946e026a2eeb84cbf34242 Mon Sep 17 00:00:00 2001 From: ptarasiewiczNV <104908264+ptarasiewiczNV@users.noreply.github.com> Date: Sat, 31 May 2025 17:19:18 +0200 Subject: [PATCH] [P/D] NixlConnector use cache device index for memory registration (#18969) Signed-off-by: Piotr Tarasiewicz --- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py index f02434aeb5ca8..6a34721574685 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py @@ -488,7 +488,8 @@ class NixlConnectorWorker: for cache in cache_list: base_addr = cache.data_ptr() region_len = self.num_blocks * self.block_len - caches_data.append((base_addr, region_len, self.rank, "")) + caches_data.append( + (base_addr, region_len, cache.device.index, "")) kv_caches_base_addr.append(base_addr) self.kv_caches_base_addr[self.engine_id] = kv_caches_base_addr self.num_regions = len(caches_data)