zbtrs2
/
Mooncake
forked from mooncake-track/Mooncake
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

* [TransferEngine] Fix compilation bug in NVLink xport

* [TransferEngine] Fix minor bugs in nvlink benchmark
This commit is contained in:
Feng Ren 2025-06-11 16:02:41 +08:00 committed by GitHub
parent f09c501b2a
commit ffaad6aa18
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 20 additions and 23 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

37
mooncake-transfer-engine/example/transfer_engine_bench.cpp
View File

@ -32,8 +32,8 @@
#ifdef USE_CUDA
#include <bits/stdint-uintn.h>
#include <cuda_runtime.h>
#include <cuda.h>
#include <cuda_runtime.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");
#ifdef USE_NVLINK
d_buf = mooncake::NvlinkTransport::allocatePinnedLocalMemory(size);
#else
checkCudaError(cudaMalloc(&d_buf, size),
"Failed to allocate device memory");
#endif
if (FLAGS_protocol == "nvlink") {
d_buf = mooncake::NvlinkTransport::allocatePinnedLocalMemory(size);
} else {
checkCudaError(cudaMalloc(&d_buf, size),
"Failed to allocate device memory");
}
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
#ifdef USE_NVLINK
CUmemGenericAllocationHandle handle;
auto result = cuMemRetainAllocationHandle(&handle, addr);
if (result == CUDA_SUCCESS) {
mooncake::NvlinkTransport::freePinnedLocalMemory(addr);
return;
if (FLAGS_protocol == "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,12 +406,7 @@ 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));
@ -431,11 +426,7 @@ 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;

6
mooncake-transfer-engine/src/transport/nvlink_transport/nvlink_transport.cpp
View File

@ -281,6 +281,12 @@ 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;