Revert "[TransferEngine] Fix minor bugs in NVLink transport and benchmark (#468)" (#469)

alogfans · web-flow · commit 38c435fcc6dc · 2025-06-11T16:50:29.000+08:00
This reverts commit ffaad6a.
diff --git a/mooncake-transfer-engine/example/transfer_engine_bench.cpp b/mooncake-transfer-engine/example/transfer_engine_bench.cpp
@@ -32,8 +32,8 @@
 
 #ifdef USE_CUDA
 #include <bits/stdint-uintn.h>
-#include <cuda.h>
 #include <cuda_runtime.h>
+#include <cuda.h>
 
 #ifdef USE_NVMEOF
 #include <cufile.h>
@@ -100,12 +100,12 @@ static void *allocateMemoryPool(size_t size, int socket_id,
         int gpu_id = FLAGS_gpu_id;
         void *d_buf;
         checkCudaError(cudaSetDevice(gpu_id), "Failed to set device");
-        if (FLAGS_protocol == "nvlink") {
-            d_buf = mooncake::NvlinkTransport::allocatePinnedLocalMemory(size);
-        } else {
-            checkCudaError(cudaMalloc(&d_buf, size),
-                           "Failed to allocate device memory");
-        }
+#ifdef USE_NVLINK
+        d_buf = mooncake::NvlinkTransport::allocatePinnedLocalMemory(size);
+#else
+        checkCudaError(cudaMalloc(&d_buf, size),
+                       "Failed to allocate device memory");
+#endif
         return d_buf;
     }
 #endif
@@ -114,14 +114,14 @@ static void *allocateMemoryPool(size_t size, int socket_id,
 
 static void freeMemoryPool(void *addr, size_t size) {
 #ifdef USE_CUDA
-    if (FLAGS_protocol == "nvlink") {
-        CUmemGenericAllocationHandle handle;
-        auto result = cuMemRetainAllocationHandle(&handle, addr);
-        if (result == CUDA_SUCCESS) {
-            mooncake::NvlinkTransport::freePinnedLocalMemory(addr);
-            return;
-        }
+#ifdef USE_NVLINK
+    CUmemGenericAllocationHandle handle;
+    auto result = cuMemRetainAllocationHandle(&handle, addr);
+    if (result == CUDA_SUCCESS) {
+        mooncake::NvlinkTransport::freePinnedLocalMemory(addr);
+        return;
     }
+#endif
     // check pointer on GPU
     cudaPointerAttributes attributes;
     checkCudaError(cudaPointerGetAttributes(&attributes, addr),
@@ -406,7 +406,12 @@ int target() {
     buffer_num = FLAGS_use_vram ? 1 : NR_SOCKETS;
     if (FLAGS_use_vram) LOG(INFO) << "VRAM is used";
     for (int i = 0; i < buffer_num; ++i) {
+#ifdef USE_NVLINK
+        addr[i] = mooncake::NvlinkTransport::allocatePinnedLocalMemory(
+            FLAGS_buffer_size);
+#else
         addr[i] = allocateMemoryPool(FLAGS_buffer_size, i, FLAGS_use_vram);
+#endif
         std::string name_prefix = FLAGS_use_vram ? "cuda:" : "cpu:";
         int rc = engine->registerLocalMemory(addr[i], FLAGS_buffer_size,
                                              name_prefix + std::to_string(i));
@@ -426,7 +431,11 @@ int target() {
     while (target_running) sleep(1);
     for (int i = 0; i < buffer_num; ++i) {
         engine->unregisterLocalMemory(addr[i]);
+#ifdef USE_NVLINK
+        mooncake::NvlinkTransport::freePinnedLocalMemory(addr[i]);
+#else
         freeMemoryPool(addr[i], FLAGS_buffer_size);
+#endif
     }
 
     return 0;
diff --git a/mooncake-transfer-engine/src/transport/nvlink_transport/nvlink_transport.cpp b/mooncake-transfer-engine/src/transport/nvlink_transport/nvlink_transport.cpp
@@ -281,12 +281,6 @@ int NvlinkTransport::registerLocalMemory(void *addr, size_t length,
             return 0;
         }
 
-        cudaError_t err = cudaSetDevice(0);
-        if (err != cudaSuccess) {
-            LOG(ERROR) << "NvlinkTransport: cudaSetDevice failed";
-            return -1;
-        }
-
         // Find whole physical page for memory registration
         void *real_addr;
         size_t real_size;
