[P/D] NixlConnector use cache device index for memory registration (vllm-project#18969)

Signed-off-by: Piotr Tarasiewicz <[email protected]>
Signed-off-by: amit <[email protected]>

Author: ptarasiewiczNV

vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py

@@ -488,7 +488,8 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
             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)