[Bugfix][Store] Open CUDA IPC via exporter PCI bus ID

mooncake-store/include/device/cuda_ipc_buffer_handle.h

@@ -15,7 +15,12 @@ struct CudaIpcBufferHandle {
     std::array<uint8_t, kCudaIpcHandleSize> handle{};
     uint64_t offset = 0;
     uint64_t size = 0;
+    // Exporter-local ordinal under its CUDA_VISIBLE_DEVICES; keep for logging
+    // / backward compatibility. Prefer pci_bus_id when opening.
     int32_t device_id = -1;
+    // Stable GPU identity across CUDA_VISIBLE_DEVICES remapping
+    // (cudaDeviceGetPCIBusId). Empty means a legacy exporter.
+    std::string pci_bus_id;
 };
 
 struct CudaIpcShmBufferRef {
@@ -38,7 +43,8 @@ struct CudaIpcReadRequest {
 
 }  // namespace mooncake
 
-YLT_REFL(mooncake::CudaIpcBufferHandle, handle, offset, size, device_id);
+YLT_REFL(mooncake::CudaIpcBufferHandle, handle, offset, size, device_id,
+         pci_bus_id);
 YLT_REFL(mooncake::CudaIpcShmBufferRef, ptr, size);
 YLT_REFL(mooncake::CudaIpcWriteRequest, key, metadata, payload);
 YLT_REFL(mooncake::CudaIpcReadRequest, key, destination, source_offset, size);

mooncake-store/src/device/cuda_ipc_buffer.cpp

@@ -103,6 +103,13 @@ tl::expected<CudaIpcBufferHandle, ErrorCode> ExportCudaIpcBuffer(
 
     cudaIpcMemHandle_t ipc_handle{};
     cudaError_t ret = cudaIpcGetMemHandle(&ipc_handle, (void *)base_ptr);
+
+    // PCI bus ID survives CUDA_VISIBLE_DEVICES remapping across processes
+    // (#4980). Prefer this over attr.device when the peer opens the handle.
+    char pci_bus_id[32] = {};
+    const cudaError_t pci_ret =
+        cudaDeviceGetPCIBusId(pci_bus_id, sizeof(pci_bus_id), attr.device);
+
     if (current_device >= 0) cudaSetDevice(current_device);
     if (ret != cudaSuccess) {
         ClearCudaError();
@@ -114,6 +121,11 @@ tl::expected<CudaIpcBufferHandle, ErrorCode> ExportCudaIpcBuffer(
     exported.offset = offset;
     exported.size = payload_size;
     exported.device_id = attr.device;
+    if (pci_ret == cudaSuccess && pci_bus_id[0] != '\0') {
+        exported.pci_bus_id = pci_bus_id;
+    } else {
+        ClearCudaError();
+    }
     return exported;
 #else
     (void)ptr;
@@ -178,18 +190,40 @@ tl::expected<CudaIpcBufferMapping, ErrorCode> CudaIpcBufferMapping::Open(
     static_assert(sizeof(cudaIpcMemHandle_t) == kCudaIpcHandleSize,
                   "Unexpected CUDA IPC handle size");
 
-    if (handle.size == 0 || handle.device_id < 0 ||
-        AddOverflows(handle.offset, handle.size)) {
+    // Never cudaSetDevice(handle.device_id): that ordinal is relative to the
+    // peer's CUDA_VISIBLE_DEVICES and may map to a different physical GPU
+    // (or an invalid index) in this process (#4980). Prefer pci_bus_id.
+    if (handle.size == 0 || AddOverflows(handle.offset, handle.size)) {
         return tl::unexpected(ErrorCode::INVALID_PARAMS);
     }
 
-    int current_device = -1;
-    cudaGetDevice(&current_device);
-    if (cudaSetDevice(handle.device_id) != cudaSuccess) {
+    int previous_device = -1;
+    if (cudaGetDevice(&previous_device) != cudaSuccess || previous_device < 0) {
         ClearCudaError();
         return tl::unexpected(ErrorCode::INTERNAL_ERROR);
     }
 
+    int open_device = previous_device;
+    bool switched_device = false;
+    if (!handle.pci_bus_id.empty()) {
+        if (cudaDeviceGetByPCIBusId(&open_device, handle.pci_bus_id.c_str()) !=
+            cudaSuccess) {
+            ClearCudaError();
+            // Exporter GPU is not visible under this process's
+            // CUDA_VISIBLE_DEVICES.
+            return tl::unexpected(ErrorCode::INVALID_PARAMS);
+        }
+        if (open_device != previous_device) {
+            if (cudaSetDevice(open_device) != cudaSuccess) {
+                ClearCudaError();
+                return tl::unexpected(ErrorCode::INTERNAL_ERROR);
+            }
+            switched_device = true;
+        }
+    }
+    // Legacy handles without pci_bus_id: open on the caller's current device
+    // (same as TENT NVLinkTransport). Do not trust handle.device_id.
+
     cudaIpcMemHandle_t ipc_handle{};
     std::memcpy(&ipc_handle, handle.handle.data(), sizeof(ipc_handle));
 
@@ -209,14 +243,16 @@ tl::expected<CudaIpcBufferMapping, ErrorCode> CudaIpcBufferMapping::Open(
             ret = cudaErrorInvalidValue;
         }
     }
-    if (current_device >= 0) cudaSetDevice(current_device);
+    if (switched_device) {
+        if (cudaSetDevice(previous_device) != cudaSuccess) ClearCudaError();
+    }
     if (ret != cudaSuccess || base == nullptr) {
         ClearCudaError();
         return tl::unexpected(ErrorCode::INVALID_PARAMS);
     }
 
     auto *ptr = static_cast<char *>(base) + handle.offset;
-    return CudaIpcBufferMapping(base, ptr, handle.device_id);
+    return CudaIpcBufferMapping(base, ptr, open_device);
 #else
     (void)handle;
     return tl::unexpected(ErrorCode::INVALID_PARAMS);