[CCF Archive] Store object type eviction policy submission #3
|
|
@ -172,7 +172,7 @@ jobs:
|
|||
export LD_LIBRARY_PATH=$GITHUB_WORKSPACE/build/mooncake-common:$GITHUB_WORKSPACE/build/mooncake-store/src:$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src:$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src/common/base:$GITHUB_WORKSPACE/build/mooncake-common/etcd
|
||||
export CGO_ENABLED=1
|
||||
export CGO_CFLAGS="-I$GITHUB_WORKSPACE/mooncake-store/include -I$GITHUB_WORKSPACE/mooncake-transfer-engine/include"
|
||||
export CGO_LDFLAGS="-L$GITHUB_WORKSPACE/build/mooncake-store/src -L$GITHUB_WORKSPACE/build/mooncake-store/src/cachelib_memory_allocator -L$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src -L$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src/common/base -L$GITHUB_WORKSPACE/build/mooncake-common -L$GITHUB_WORKSPACE/build/mooncake-common/etcd -lmooncake_store -lcachelib_memory_allocator -ltransfer_engine -lbase -lasio -letcd_wrapper -lstdc++ -lnuma -lglog -lgflags -libverbs -ljsoncpp -lzstd -lcurl -luring -lasan -lm -lgcov -lxxhash"
|
||||
export CGO_LDFLAGS="-L$GITHUB_WORKSPACE/build/mooncake-store/src -L$GITHUB_WORKSPACE/build/mooncake-store/src/cachelib_memory_allocator -L$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src -L$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src/common/base -L$GITHUB_WORKSPACE/build/mooncake-common -L$GITHUB_WORKSPACE/build/mooncake-common/etcd -lmooncake_store -lcachelib_memory_allocator -ltransfer_engine -lbase -lasio -letcd_wrapper -lstdc++ -lnuma -lglog -lgflags -libverbs -lmlx5 -ljsoncpp -lzstd -lcurl -luring -lasan -lm -lgcov -lxxhash"
|
||||
# Link cudart if CUDA is available (needed for D2H staging in mooncake_store)
|
||||
if [ -d /usr/local/cuda/lib64 ]; then export CGO_LDFLAGS="$CGO_LDFLAGS -L/usr/local/cuda/lib64 -lcudart"; fi
|
||||
ASAN_OPTIONS=detect_leaks=0:verify_asan_link_order=0 MC_METADATA_SERVER=http://127.0.0.1:8080/metadata go test -v ./tests/...
|
||||
|
|
|
|||
|
|
@ -145,6 +145,7 @@ if (WITH_EP)
|
|||
"-DTORCH_CUDA_ARCH_LIST=${_torch_cuda_arch_list_pipe}"
|
||||
"-DSTAGING_DIR=${EP_PG_STAGING_DIR}"
|
||||
"-DENGINE_SO_PATH=$<TARGET_FILE:engine>"
|
||||
"-DEP_USE_MUSA=$<IF:$<BOOL:${USE_MUSA}>,1,0>"
|
||||
-P "${CMAKE_CURRENT_SOURCE_DIR}/mooncake-ep/BuildEpExt.cmake"
|
||||
COMMENT "Building Mooncake EP Python extension(s)"
|
||||
DEPENDS engine
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
set(CMAKE_C_STANDARD 99)
|
||||
set(CMAKE_CXX_STANDARD 20)
|
||||
set(CMAKE_CUDA_STANDARD 20)
|
||||
|
||||
option(ENABLE_DEBUG_SYMBOLS "Include debug symbols (-g) in compilation" ON)
|
||||
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@
|
|||
# TORCH_CUDA_ARCH_LIST - pipe-separated CUDA arch list forwarded to torch
|
||||
# STAGING_DIR - destination directory for the built .so files
|
||||
# ENGINE_SO_PATH - absolute path to the built engine.cpython-XYZ.so
|
||||
# EP_USE_MUSA - set to "1" when building for MUSA (MTLink path)
|
||||
|
||||
cmake_minimum_required(VERSION 3.16)
|
||||
|
||||
|
|
@ -34,6 +35,11 @@ endif()
|
|||
set(ENV{MAKEFLAGS} "")
|
||||
set(ENV{MFLAGS} "")
|
||||
set(ENV{TORCH_CUDA_ARCH_LIST} "${TORCH_CUDA_ARCH_LIST}")
|
||||
if(EP_USE_MUSA)
|
||||
set(ENV{MOONCAKE_EP_USE_MUSA} "1")
|
||||
else()
|
||||
unset(ENV{MOONCAKE_EP_USE_MUSA})
|
||||
endif()
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 2. Ensure engine.so exists in mooncake-wheel/mooncake/ for setup.py linking.
|
||||
|
|
|
|||
|
|
@ -6,8 +6,8 @@
|
|||
#include <cuda.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <fstream>
|
||||
#include <mooncake_ibgda/memheap.h>
|
||||
#include <mooncake_ibgda/mlx5gda.h>
|
||||
#include <transport/device/ibgda/memheap.h>
|
||||
#include <transport/device/ibgda/mlx5gda.h>
|
||||
#include <mooncake_ep_api.cuh>
|
||||
#include <mooncake_ep_configs.cuh>
|
||||
#include <mooncake_ep_event.h>
|
||||
|
|
|
|||
|
|
@ -39,7 +39,6 @@ setup(
|
|||
"src/ep_py.cpp",
|
||||
"src/mooncake_ep_buffer.cpp",
|
||||
"src/mooncake_ep_kernel.cu",
|
||||
"src/mooncake_ibgda/mlx5gda.cpp",
|
||||
],
|
||||
extra_compile_args={
|
||||
"cxx": [f"-D_GLIBCXX_USE_CXX11_ABI={abi_flag}", "-std=c++20", "-O3", "-g0"],
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
add_library(mooncake_ep ep_py.cpp mooncake_ep_buffer.cpp mooncake_ep_kernel.cu mooncake_ibgda/mlx5gda.cpp)
|
||||
add_library(mooncake_ep ep_py.cpp mooncake_ep_buffer.cpp mooncake_ep_kernel.cu)
|
||||
|
||||
set_target_properties(mooncake_ep PROPERTIES POSITION_INDEPENDENT_CODE ON)
|
||||
target_link_libraries(mooncake_ep PUBLIC ${TORCH_LIBRARIES} transfer_engine ibverbs mlx5)
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@
|
|||
#include <mooncake_ep_configs.cuh>
|
||||
#include <mooncake_ep_exception.cuh>
|
||||
#include <mooncake_ep_launch.cuh>
|
||||
#include <mooncake_ibgda/mlx5gda.h>
|
||||
#include <transport/device/ibgda/mlx5gda.h>
|
||||
#include <mooncake_ep_utils.cuh>
|
||||
|
||||
namespace mooncake {
|
||||
|
|
|
|||
|
|
@ -35,7 +35,7 @@ EXT_LDFLAGS="-L$BUILD_DIR/mooncake-transfer-engine/src"
|
|||
EXT_LDFLAGS+=" -L$BUILD_DIR/mooncake-transfer-engine/src/common/base"
|
||||
EXT_LDFLAGS+=" -L$BUILD_DIR/mooncake-common"
|
||||
EXT_LDFLAGS+=" -L$BUILD_DIR/mooncake-common/src"
|
||||
EXT_LDFLAGS+=" -ltransfer_engine -lbase -lasio -lstdc++ -lnuma -lglog -libverbs -ljsoncpp -lmooncake_common"
|
||||
EXT_LDFLAGS+=" -ltransfer_engine -lbase -lasio -lstdc++ -lnuma -lglog -libverbs -lmlx5 -ljsoncpp -lmooncake_common -lm"
|
||||
|
||||
if [ -d "/usr/local/cuda/lib64/stubs" ]; then
|
||||
EXT_LDFLAGS+=" -L/usr/local/cuda/lib64/stubs"
|
||||
|
|
|
|||
|
|
@ -42,7 +42,7 @@ CGO_LDFLAGS+=" -L${BUILD_DIR}/mooncake-transfer-engine/src/common/base"
|
|||
CGO_LDFLAGS+=" -L${BUILD_DIR}/mooncake-common"
|
||||
CGO_LDFLAGS+=" -L${BUILD_DIR}/mooncake-common/src"
|
||||
CGO_LDFLAGS+=" -lmooncake_store -lcachelib_memory_allocator -ltransfer_engine -lbase -lasio -lmooncake_common -lxxhash"
|
||||
CGO_LDFLAGS+=" -lstdc++ -lnuma -lglog -lgflags -libverbs -ljsoncpp -lzstd -lcurl"
|
||||
CGO_LDFLAGS+=" -lstdc++ -lnuma -lglog -lgflags -libverbs -lmlx5 -ljsoncpp -lzstd -lcurl -lm"
|
||||
|
||||
if [ -d "/usr/local/cuda/lib64" ]; then
|
||||
CGO_LDFLAGS+=" -L/usr/local/cuda/lib64 -lcudart"
|
||||
|
|
|
|||
|
|
@ -331,6 +331,7 @@ fn main() {
|
|||
("curl", &["curl"]),
|
||||
("cuda", &["cuda"]),
|
||||
("cudart", &["cudart"]),
|
||||
("mlx5", &["mlx5"]), // IBGDA device transport (mlx5 DevX) pulled into transfer_engine, CUDA-only
|
||||
("uring", &["uring"]),
|
||||
] {
|
||||
if has_library(&search_dirs, candidates) {
|
||||
|
|
|
|||
|
|
@ -61,3 +61,48 @@ endif()
|
|||
if(USE_UBSHMEM)
|
||||
target_link_libraries(transfer_engine_bench PUBLIC transfer_engine)
|
||||
endif()
|
||||
|
||||
# Device Transport Example — two-rank P2P write + signal via Device API.
|
||||
# Requires 2 GPUs with P2P access. Built but not registered with CTest
|
||||
# (manual execution only, like rdma_transport_test).
|
||||
#
|
||||
# Off by default: this .cu target pulls in transfer_engine.h, whose common.h
|
||||
# uses C++20 std::string APIs (e.g. starts_with), so it needs the CUDA20
|
||||
# dialect. Older CMake (such as the CI Docker image) cannot enable CUDA20 and
|
||||
# fails at generate time. Opt in explicitly with -DBUILD_DEVICE_TRANSPORT_EXAMPLE=ON
|
||||
# when building manually on a host with a recent CMake and 2 GPUs.
|
||||
option(BUILD_DEVICE_TRANSPORT_EXAMPLE
|
||||
"Build the two-rank Device API P2P example (requires CUDA20-capable CMake and 2 GPUs)"
|
||||
OFF)
|
||||
if(USE_CUDA AND BUILD_DEVICE_TRANSPORT_EXAMPLE)
|
||||
enable_language(CUDA)
|
||||
add_executable(device_transport_example
|
||||
${WORKSPACE}/device_transport_example.cu)
|
||||
set_source_files_properties(
|
||||
${WORKSPACE}/device_transport_example.cu
|
||||
PROPERTIES LANGUAGE CUDA)
|
||||
target_include_directories(device_transport_example PRIVATE
|
||||
${CMAKE_SOURCE_DIR}/mooncake-transfer-engine/include)
|
||||
target_link_libraries(device_transport_example PUBLIC
|
||||
transfer_engine gflags::gflags glog::glog)
|
||||
# common.h (pulled in via transfer_engine.h) uses C++20 std::string APIs
|
||||
# such as starts_with; nvcc does not inherit CMAKE_CXX_STANDARD, so set the
|
||||
# CUDA standard explicitly on this target. CUDA_EXTENSIONS OFF requests
|
||||
# plain -std=c++20 (nvcc has no gnu++20 dialect mapping).
|
||||
set_target_properties(device_transport_example PROPERTIES
|
||||
CUDA_STANDARD 20 CUDA_STANDARD_REQUIRED ON CUDA_EXTENSIONS OFF)
|
||||
# Convert torch-style arch list (e.g. "8.0;9.0") to CMake CUDA format
|
||||
# (e.g. "80;90") by stripping the dot.
|
||||
if(TORCH_CUDA_ARCH_LIST)
|
||||
set(_cuda_arch_list "")
|
||||
foreach(_arch IN LISTS TORCH_CUDA_ARCH_LIST)
|
||||
string(REPLACE "." "" _arch_clean "${_arch}")
|
||||
list(APPEND _cuda_arch_list "${_arch_clean}")
|
||||
endforeach()
|
||||
set_target_properties(device_transport_example PROPERTIES
|
||||
CUDA_ARCHITECTURES "${_cuda_arch_list}")
|
||||
else()
|
||||
set_target_properties(device_transport_example PROPERTIES
|
||||
CUDA_ARCHITECTURES "80;90")
|
||||
endif()
|
||||
endif()
|
||||
|
|
|
|||
|
|
@ -0,0 +1,305 @@
|
|||
// Copyright 2024 KVCache.AI
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
// Device Transport Example — two-rank P2P write + signal via Device API.
|
||||
//
|
||||
// This example demonstrates the full Device API lifecycle:
|
||||
// 1. Host side: P2pTransport for IPC handle exchange and peer mapping.
|
||||
// 2. Device side: CommCtx + mc_route_put + mc_signal for GPU-initiated
|
||||
// P2P data transfer and notification.
|
||||
//
|
||||
// Usage (two terminals on the same node):
|
||||
// $ ./device_transport_example --rank=0
|
||||
// $ ./device_transport_example --rank=1
|
||||
//
|
||||
// The example uses file-based IPC handle exchange (no external dependencies).
|
||||
// Requires 2 GPUs with P2P access (NVLink or PCIe).
|
||||
|
||||
#include <gflags/gflags.h>
|
||||
#include <glog/logging.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include <atomic>
|
||||
#include <chrono>
|
||||
#include <cstring>
|
||||
#include <fstream>
|
||||
#include <thread>
|
||||
#include <vector>
|
||||
|
||||
#include "transfer_engine.h"
|
||||
#include "transport/device/comm_device.cuh"
|
||||
#include "transport/device/device_transport.h"
|
||||
|
||||
DEFINE_int32(rank, -1, "Rank of this process (0 or 1)");
|
||||
DEFINE_int32(world_size, 2, "Total number of ranks");
|
||||
DEFINE_int32(gpu_id, -1, "GPU ID (defaults to rank)");
|
||||
DEFINE_string(metadata_server, "P2PHANDSHAKE",
|
||||
"Metadata server (P2PHANDSHAKE for no external deps)");
|
||||
DEFINE_string(local_server_name, "",
|
||||
"Local server name (default: 127.0.0.1:<port>)");
|
||||
DEFINE_string(ipc_dir, "/tmp", "Directory for IPC handle exchange files");
|
||||
DEFINE_int32(kDataBytes, 4096, "Bytes to transfer in the P2P write");
|
||||
DEFINE_int32(kSignalOffset, 0, "Offset within buffer for the signal word");
|
||||
|
||||
namespace {
|
||||
|
||||
static void checkCuda(cudaError_t err, const char* msg) {
|
||||
if (err != cudaSuccess) {
|
||||
LOG(FATAL) << msg << ": " << cudaGetErrorString(err);
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// File-based IPC handle exchange
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
static std::string ipcFilePath(int rank) {
|
||||
return FLAGS_ipc_dir + "/device_api_ex_rank_" + std::to_string(rank) +
|
||||
".bin";
|
||||
}
|
||||
|
||||
static void writeIpcHandle(int rank, const std::vector<int32_t>& handle) {
|
||||
std::string path = ipcFilePath(rank);
|
||||
std::ofstream ofs(path, std::ios::binary | std::ios::trunc);
|
||||
CHECK(ofs) << "Failed to open " << path << " for writing";
|
||||
uint32_t size = static_cast<uint32_t>(handle.size());
|
||||
ofs.write(reinterpret_cast<const char*>(&size), sizeof(size));
|
||||
ofs.write(reinterpret_cast<const char*>(handle.data()),
|
||||
handle.size() * sizeof(int32_t));
|
||||
ofs.close();
|
||||
LOG(INFO) << "Rank " << rank << " wrote IPC handle (" << handle.size()
|
||||
<< " int32s) to " << path;
|
||||
}
|
||||
|
||||
static std::vector<int32_t> readIpcHandle(int rank) {
|
||||
std::string path = ipcFilePath(rank);
|
||||
// Poll until the file appears (peer may not have written it yet).
|
||||
for (int attempt = 0; attempt < 300; ++attempt) {
|
||||
std::ifstream ifs(path, std::ios::binary);
|
||||
if (ifs) {
|
||||
uint32_t size = 0;
|
||||
ifs.read(reinterpret_cast<char*>(&size), sizeof(size));
|
||||
CHECK(size > 0 && size <= 256)
|
||||
<< "Invalid IPC handle size: " << size;
|
||||
std::vector<int32_t> handle(size);
|
||||
ifs.read(reinterpret_cast<char*>(handle.data()),
|
||||
size * sizeof(int32_t));
|
||||
ifs.close();
|
||||
LOG(INFO) << "Rank " << FLAGS_rank << " read IPC handle from rank "
|
||||
<< rank << " (" << size << " int32s)";
|
||||
return handle;
|
||||
}
|
||||
std::this_thread::sleep_for(std::chrono::milliseconds(100));
|
||||
}
|
||||
LOG(FATAL) << "Timeout waiting for IPC handle from rank " << rank;
|
||||
return {};
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Device-side kernel
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// GDR buffer layout:
|
||||
// [0 .. kDataBytes) — data region
|
||||
// [kDataBytes .. kDataBytes+4) — signal word (int32_t)
|
||||
static constexpr int kSignalWordOffset = 4096; // must match kDataBytes default
|
||||
|
||||
// Rank 0: P2P-write data to rank 1, then signal.
|
||||
__global__ void senderKernel(mooncake::device::CommCtx ctx, int dst_rank,
|
||||
int data_bytes) {
|
||||
using namespace mooncake::device;
|
||||
if (threadIdx.x != 0 || blockIdx.x != 0) return;
|
||||
|
||||
// Write a known pattern to the peer's data region via P2P.
|
||||
char* local_data = reinterpret_cast<char*>(ctx.p2p.local_base);
|
||||
void* peer_data = mc_route_put(ctx, dst_rank, local_data);
|
||||
if (peer_data == nullptr) {
|
||||
printf("[Rank 0] FAIL: P2P not available to rank %d\n", dst_rank);
|
||||
asm("trap;");
|
||||
return;
|
||||
}
|
||||
|
||||
// Fill local buffer with pattern, then copy to peer via P2P store.
|
||||
for (int i = 0; i < data_bytes; ++i) {
|
||||
local_data[i] = static_cast<char>(i & 0xFF);
|
||||
}
|
||||
// Use 16-byte non-temporal stores for the bulk copy.
|
||||
for (int i = 0; i < data_bytes; i += 16) {
|
||||
int4 val = mc_ld_nc(reinterpret_cast<const int4*>(local_data + i));
|
||||
mc_st_na(
|
||||
reinterpret_cast<int4*>(reinterpret_cast<char*>(peer_data) + i),
|
||||
val);
|
||||
}
|
||||
|
||||
// Signal: write 1 to the peer's signal word.
|
||||
int* local_sig = reinterpret_cast<int*>(local_data + kSignalWordOffset);
|
||||
*local_sig = 1;
|
||||
mc_signal(ctx, dst_rank, 0 /*channel*/, 1 /*qps_per_rank*/, local_sig, 1);
|
||||
}
|
||||
|
||||
// Rank 1: wait for signal from rank 0, then verify data.
|
||||
__global__ void receiverKernel(mooncake::device::CommCtx ctx, int src_rank,
|
||||
int data_bytes) {
|
||||
using namespace mooncake::device;
|
||||
if (threadIdx.x != 0 || blockIdx.x != 0) return;
|
||||
|
||||
char* local_data = reinterpret_cast<char*>(ctx.p2p.local_base);
|
||||
int* local_sig = reinterpret_cast<int*>(local_data + kSignalWordOffset);
|
||||
|
||||
// Spin-wait for the signal.
|
||||
// The sender writes the signal via P2P store (mc_p2p_signal →
|
||||
// mc_st_release). We use ld_acquire to observe it.
|
||||
while (mc_ld_acquire(local_sig) == 0) {
|
||||
__threadfence_block();
|
||||
}
|
||||
|
||||
// Verify the data written by rank 0.
|
||||
int mismatches = 0;
|
||||
for (int i = 0; i < data_bytes; ++i) {
|
||||
if (local_data[i] != static_cast<char>(i & 0xFF)) {
|
||||
++mismatches;
|
||||
}
|
||||
}
|
||||
if (mismatches > 0) {
|
||||
printf("[Rank 1] FAIL: %d byte mismatches in received data\n",
|
||||
mismatches);
|
||||
} else {
|
||||
printf("[Rank 1] PASS: all %d bytes match expected pattern\n",
|
||||
data_bytes);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
int main(int argc, char** argv) {
|
||||
gflags::ParseCommandLineFlags(&argc, &argv, true);
|
||||
google::InitGoogleLogging(argv[0]);
|
||||
|
||||
int rank = FLAGS_rank;
|
||||
int world_size = FLAGS_world_size;
|
||||
CHECK(rank == 0 || rank == 1) << "--rank must be 0 or 1";
|
||||
CHECK_EQ(world_size, 2) << "This example requires exactly 2 ranks";
|
||||
CHECK_EQ(FLAGS_kDataBytes % 16, 0) << "kDataBytes must be a multiple of 16";
|
||||
CHECK_LE(FLAGS_kDataBytes, kSignalWordOffset)
|
||||
<< "kDataBytes cannot exceed " << kSignalWordOffset;
|
||||
|
||||
int gpu_id = FLAGS_gpu_id >= 0 ? FLAGS_gpu_id : rank;
|
||||
checkCuda(cudaSetDevice(gpu_id), "cudaSetDevice");
|
||||
LOG(INFO) << "Rank " << rank << " using GPU " << gpu_id;
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// 1. Create TransferEngine and get P2pTransport.
|
||||
// -----------------------------------------------------------------------
|
||||
auto engine = std::make_unique<mooncake::TransferEngine>(false);
|
||||
|
||||
std::string local_name = FLAGS_local_server_name;
|
||||
if (local_name.empty()) {
|
||||
local_name = "127.0.0.1:" + std::to_string(12345 + rank);
|
||||
}
|
||||
engine->init(FLAGS_metadata_server, local_name);
|
||||
|
||||
auto* p2p = engine->getOrCreateP2pTransport(world_size);
|
||||
CHECK_NOTNULL(p2p);
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// 2. Allocate GDR buffer and fill with initial pattern.
|
||||
// -----------------------------------------------------------------------
|
||||
constexpr size_t kBufSize = kSignalWordOffset + sizeof(int32_t);
|
||||
void* gdr_buffer = p2p->allocateBuffer(kBufSize);
|
||||
CHECK_NOTNULL(gdr_buffer);
|
||||
LOG(INFO) << "Rank " << rank << " allocated GDR buffer: " << gdr_buffer;
|
||||
|
||||
// Zero-initialize the signal word.
|
||||
checkCuda(cudaMemset(gdr_buffer, 0, kBufSize), "cudaMemset");
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// 3. Export IPC handle and exchange with peer.
|
||||
// -----------------------------------------------------------------------
|
||||
auto my_handle = p2p->exportIpcHandle(gdr_buffer);
|
||||
CHECK(!my_handle.empty()) << "exportIpcHandle returned empty vector";
|
||||
writeIpcHandle(rank, my_handle);
|
||||
|
||||
// Read peer's handle.
|
||||
int peer_rank = 1 - rank;
|
||||
auto peer_handle = readIpcHandle(peer_rank);
|
||||
|
||||
// Build remote_handles vector: index by rank.
|
||||
std::vector<std::vector<int32_t>> remote_handles(world_size);
|
||||
remote_handles[rank] = my_handle;
|
||||
remote_handles[peer_rank] = peer_handle;
|
||||
|
||||
std::vector<int> active_ranks_mask(world_size, 1);
|
||||
|
||||
p2p->importPeerHandles(gdr_buffer, rank, world_size, remote_handles,
|
||||
active_ranks_mask);
|
||||
|
||||
CHECK(p2p->allPeersAccessible())
|
||||
<< "P2P not accessible between rank " << rank << " and peer";
|
||||
LOG(INFO) << "Rank " << rank << " P2P peer mapping complete";
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// 4. Build CommCtx and launch kernel.
|
||||
// -----------------------------------------------------------------------
|
||||
mooncake::device::CommCtx ctx{};
|
||||
ctx.rank = rank;
|
||||
ctx.p2p.available = p2p->availableTablePtr();
|
||||
ctx.p2p.peer_ptrs = p2p->peerPtrsTablePtr();
|
||||
ctx.p2p.local_base = gdr_buffer;
|
||||
// IBGDA fields left as nullptr (not used for P2P-only example).
|
||||
|
||||
// Pass CommCtx by value — CUDA copies it to kernel parameter space.
|
||||
if (rank == 0) {
|
||||
LOG(INFO) << "Rank 0 launching sender kernel...";
|
||||
senderKernel<<<1, 1>>>(ctx, peer_rank, FLAGS_kDataBytes);
|
||||
} else {
|
||||
LOG(INFO) << "Rank 1 launching receiver kernel...";
|
||||
receiverKernel<<<1, 1>>>(ctx, peer_rank, FLAGS_kDataBytes);
|
||||
}
|
||||
checkCuda(cudaDeviceSynchronize(), "cudaDeviceSynchronize");
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// 5. Barrier: rank 1 signals completion, rank 0 waits before cleanup.
|
||||
// The IPC handle is only valid while the original allocation exists,
|
||||
// so rank 0 must not free its buffer until rank 1 has opened it.
|
||||
// -----------------------------------------------------------------------
|
||||
if (rank == 1) {
|
||||
// Signal rank 0 that we're done.
|
||||
std::string done_path = FLAGS_ipc_dir + "/device_api_ex_done.bin";
|
||||
std::ofstream ofs(done_path);
|
||||
ofs << "1";
|
||||
ofs.close();
|
||||
} else {
|
||||
// Wait for rank 1 to finish.
|
||||
std::string done_path = FLAGS_ipc_dir + "/device_api_ex_done.bin";
|
||||
for (int attempt = 0; attempt < 300; ++attempt) {
|
||||
std::ifstream ifs(done_path);
|
||||
if (ifs) break;
|
||||
std::this_thread::sleep_for(std::chrono::milliseconds(100));
|
||||
}
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------
|
||||
// 6. Cleanup.
|
||||
// -----------------------------------------------------------------------
|
||||
p2p->freeBuffer(gdr_buffer);
|
||||
|
||||
// Remove IPC handle files.
|
||||
std::remove(ipcFilePath(rank).c_str());
|
||||
if (rank == 0) {
|
||||
std::remove((FLAGS_ipc_dir + "/device_api_ex_done.bin").c_str());
|
||||
}
|
||||
|
||||
LOG(INFO) << "Rank " << rank << " done.";
|
||||
return 0;
|
||||
}
|
||||
|
|
@ -10,3 +10,16 @@ install(FILES transfer_metadata.h DESTINATION include)
|
|||
install(FILES ub_allocator.h DESTINATION include)
|
||||
install(FILES common/base/status.h DESTINATION include/common/base)
|
||||
install(FILES transport/transport.h DESTINATION include/transport)
|
||||
|
||||
# Device API headers (header-only, consumed by EP kernel)
|
||||
install(FILES transport/device/device_transport.h DESTINATION include/transport/device)
|
||||
install(FILES transport/device/device_ops.cuh DESTINATION include/transport/device)
|
||||
install(FILES transport/device/comm_device.cuh DESTINATION include/transport/device)
|
||||
install(FILES transport/device/p2p_device.cuh DESTINATION include/transport/device)
|
||||
install(FILES transport/device/ibgda_device.cuh DESTINATION include/transport/device)
|
||||
install(FILES transport/device/cuda/cuda_ops.cuh DESTINATION include/transport/device/cuda)
|
||||
install(FILES transport/device/musa/musa_ops.cuh DESTINATION include/transport/device/musa)
|
||||
|
||||
# IBGDA library headers
|
||||
install(DIRECTORY transport/device/ibgda/ DESTINATION include/transport/device/ibgda
|
||||
FILES_MATCHING PATTERN "*.h")
|
||||
|
|
|
|||
|
|
@ -1,3 +1,5 @@
|
|||
#pragma once
|
||||
|
||||
#include <hip/hip_runtime.h>
|
||||
#include <string>
|
||||
|
||||
|
|
|
|||
|
|
@ -1,3 +1,5 @@
|
|||
#pragma once
|
||||
|
||||
#include <string>
|
||||
#include <musa.h>
|
||||
#include <musa_runtime.h>
|
||||
|
|
@ -20,6 +22,11 @@ const static std::string GPU_PREFIX = "musa:";
|
|||
#define CU_POINTER_ATTRIBUTE_RANGE_SIZE MU_POINTER_ATTRIBUTE_RANGE_SIZE
|
||||
#define CU_DEVICE_ATTRIBUTE_MEMORY_POOLS_SUPPORTED \
|
||||
MU_DEVICE_ATTRIBUTE_MEMORY_POOLS_SUPPORTED
|
||||
#define CU_DEVICE_ATTRIBUTE_DMA_BUF_SUPPORTED \
|
||||
MU_DEVICE_ATTRIBUTE_DMA_BUF_SUPPORTED
|
||||
#define CU_POINTER_ATTRIBUTE_DEVICE_ORDINAL MU_POINTER_ATTRIBUTE_DEVICE_ORDINAL
|
||||
|
||||
#define CUcontext MUcontext
|
||||
|
||||
#define CUdevice MUdevice
|
||||
#define CUdeviceptr MUdeviceptr
|
||||
|
|
@ -93,3 +100,22 @@ const static std::string GPU_PREFIX = "musa:";
|
|||
#define cudaEvent_t musaEvent_t
|
||||
#define cudaHostGetDevicePointer musaHostGetDevicePointer
|
||||
#define cudaSuccess musaSuccess
|
||||
#define cudaDeviceGetAttribute musaDeviceGetAttribute
|
||||
#define cudaEvent_t musaEvent_t
|
||||
#define cudaMemcpyDeviceToDevice musaMemcpyDeviceToDevice
|
||||
#define cudaDevAttrClockRate musaDevAttrClockRate
|
||||
#define cudaLaunchConfig_t musaLaunchConfig_t
|
||||
#define cudaLaunchAttribute musaLaunchAttribute
|
||||
#define cudaLaunchAttributeCooperative musaLaunchAttributeCooperative
|
||||
#define cudaLaunchKernelEx musaLaunchKernelEx
|
||||
#define CUDA_R_16BF MUSA_R_16BF
|
||||
#define CUDA_R_32F MUSA_R_32F
|
||||
|
||||
// IBGDA-specific mappings
|
||||
#define cuInit muInit
|
||||
#define cuDevicePrimaryCtxRetain muDevicePrimaryCtxRetain
|
||||
#define cuDevicePrimaryCtxRelease muDevicePrimaryCtxRelease
|
||||
#define cuCtxSetCurrent muCtxSetCurrent
|
||||
#define cudaHostRegisterMapped musaHostRegisterMapped
|
||||
#define cudaHostRegisterIoMemory musaHostRegisterIoMemory
|
||||
#define cudaHostGetDevicePointer musaHostGetDevicePointer
|
||||
|
|
|
|||
|
|
@ -1,3 +1,5 @@
|
|||
#pragma once
|
||||
|
||||
#include <string>
|
||||
#include <acl/acl.h>
|
||||
|
||||
|
|
|
|||
|
|
@ -25,6 +25,12 @@ class TransferEngineImpl;
|
|||
namespace tent {
|
||||
class TransferEngine;
|
||||
};
|
||||
#if defined(USE_CUDA) || defined(USE_MUSA)
|
||||
namespace device {
|
||||
class P2pTransport;
|
||||
class RdmaTransport;
|
||||
} // namespace device
|
||||
#endif
|
||||
using TransferRequest = Transport::TransferRequest;
|
||||
using TransferStatus = Transport::TransferStatus;
|
||||
using TransferStatusEnum = Transport::TransferStatusEnum;
|
||||
|
|
@ -150,6 +156,16 @@ class TransferEngine {
|
|||
|
||||
Transport* getTransport(const std::string& proto);
|
||||
|
||||
#if defined(USE_CUDA) || defined(USE_MUSA)
|
||||
// Device transport accessors (P2P + IBGDA). Lazily created on first
|
||||
// call and owned by the TransferEngine. These allow EP (and future
|
||||
// CPU-proxy paths) to obtain device transports from an engine instance
|
||||
// instead of calling the global factory functions directly.
|
||||
device::P2pTransport* getOrCreateP2pTransport(int num_ranks);
|
||||
device::RdmaTransport* getOrCreateRdmaTransport(
|
||||
const std::vector<std::string>& device_filter = {});
|
||||
#endif
|
||||
|
||||
/**
|
||||
* @brief Check if TCP is the only installed transport.
|
||||
*
|
||||
|
|
|
|||
|
|
@ -33,6 +33,9 @@
|
|||
#include "transfer_metadata.h"
|
||||
#include "transfer_engine.h"
|
||||
#include "transport/transport.h"
|
||||
#if defined(USE_CUDA) || defined(USE_MUSA)
|
||||
#include "transport/device/device_transport.h"
|
||||
#endif
|
||||
#ifdef WITH_METRICS
|
||||
#include "ylt/metric/counter.hpp"
|
||||
#include "ylt/metric/histogram.hpp"
|
||||
|
|
@ -335,6 +338,13 @@ class TransferEngineImpl {
|
|||
return multi_transports_->getTransport(proto);
|
||||
}
|
||||
|
||||
#if defined(USE_CUDA) || defined(USE_MUSA)
|
||||
// Device transport accessors — lazily created, owned by this impl.
|
||||
device::P2pTransport* getOrCreateP2pTransport(int num_ranks);
|
||||
device::RdmaTransport* getOrCreateRdmaTransport(
|
||||
const std::vector<std::string>& device_filter = {});
|
||||
#endif
|
||||
|
||||
bool isTcpOnly() const { return multi_transports_->isTcpOnly(); }
|
||||
|
||||
int syncSegmentCache(const std::string& segment_name = "") {
|
||||
|
|
@ -411,6 +421,13 @@ class TransferEngineImpl {
|
|||
std::vector<std::string> filter_;
|
||||
bool use_barex_ = false;
|
||||
|
||||
#if defined(USE_CUDA) || defined(USE_MUSA)
|
||||
// Device transports (P2P + IBGDA) — lazily created, owned by this impl.
|
||||
// Referenced by EP and future CPU-proxy paths.
|
||||
std::unique_ptr<device::P2pTransport> p2p_transport_;
|
||||
std::unique_ptr<device::RdmaTransport> rdma_transport_;
|
||||
#endif
|
||||
|
||||
#ifdef WITH_METRICS
|
||||
// Latency bucket in microseconds
|
||||
inline static const std::vector<double> kTaskLatencyBuckets = {
|
||||
|
|
|
|||
|
|
@ -0,0 +1,136 @@
|
|||
// Communication device API — top-level context and routing.
|
||||
//
|
||||
// CommCtx bundles all transport state the kernel needs. The kernel
|
||||
// constructs one from the raw pointers passed in and calls mc_route_put /
|
||||
// mc_signal / mc_red_add instead of touching transport internals directly.
|
||||
#pragma once
|
||||
|
||||
#include "transport/device/device_ops.cuh"
|
||||
#include "transport/device/p2p_device.cuh"
|
||||
#include "transport/device/ibgda_device.cuh"
|
||||
|
||||
namespace mooncake {
|
||||
namespace device {
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// CommCtx
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
struct CommCtx {
|
||||
P2PContext p2p;
|
||||
IbgdaContext ibgda;
|
||||
int rank;
|
||||
};
|
||||
|
||||
// Construct CommCtx from the raw kernel arguments.
|
||||
// raddrs/rkeys/qp_devctxs may be nullptr on MUSA (ignored).
|
||||
__device__ __forceinline__ CommCtx make_comm_ctx(
|
||||
void* gdr_buffer, const int32_t* nvlink_available,
|
||||
void* const* ipc_peer_ptrs, void* raddrs, void* rkeys, void* qp_devctxs,
|
||||
const void* rdma_send_signal_buffer, const void* rdma_recv_signal_buffer,
|
||||
int rank, int num_ranks, int num_qps) {
|
||||
CommCtx ctx;
|
||||
ctx.rank = rank;
|
||||
|
||||
ctx.p2p.available = nvlink_available;
|
||||
ctx.p2p.peer_ptrs = ipc_peer_ptrs;
|
||||
ctx.p2p.local_base = gdr_buffer;
|
||||
|
||||
ctx.ibgda.qp_devctxs = reinterpret_cast<mlx5gda_qp_devctx*>(qp_devctxs);
|
||||
ctx.ibgda.raddrs = reinterpret_cast<const uint64_t*>(raddrs);
|
||||
ctx.ibgda.rkeys = reinterpret_cast<const uint32_t*>(rkeys);
|
||||
ctx.ibgda.local_atomic_base = rdma_send_signal_buffer;
|
||||
ctx.ibgda.remote_atomic_base = rdma_recv_signal_buffer;
|
||||
|
||||
return ctx;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Routing helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
__device__ __forceinline__ bool mc_comm_p2p_available(const CommCtx& ctx,
|
||||
int dst_rank) {
|
||||
return mc_p2p_available(ctx.p2p, dst_rank);
|
||||
}
|
||||
|
||||
// Translate a local GDR pointer to the peer's mapped VA.
|
||||
__device__ __forceinline__ void* mc_comm_peer_ptr(const CommCtx& ctx,
|
||||
int dst_rank,
|
||||
const void* local_ptr) {
|
||||
return mc_p2p_peer_ptr(ctx.p2p, dst_rank, local_ptr);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// mc_route_put
|
||||
//
|
||||
// Returns the destination pointer for a warp-cooperative copy:
|
||||
// - local rank: recv_ptr itself (caller does UNROLLED_WARP_COPY)
|
||||
// - P2P rank: peer-mapped recv_ptr (caller does UNROLLED_WARP_COPY)
|
||||
// - IBGDA rank: nullptr (caller must stage data then call mc_rdma_put)
|
||||
// ---------------------------------------------------------------------------
|
||||
__device__ __forceinline__ void* mc_route_put(const CommCtx& ctx, int dst_rank,
|
||||
void* recv_ptr) {
|
||||
if (dst_rank == ctx.rank) return recv_ptr;
|
||||
if (mc_comm_p2p_available(ctx, dst_rank))
|
||||
return mc_comm_peer_ptr(ctx, dst_rank, recv_ptr);
|
||||
return nullptr; // IBGDA path
|
||||
}
|
||||
|
||||
// Issue an IBGDA RDMA WRITE. Call only when mc_route_put returned nullptr.
|
||||
// lane_id: only lane 0 issues the WQE.
|
||||
__device__ __forceinline__ void mc_rdma_put(
|
||||
const CommCtx& ctx, int channel, int dst_rank, int qps_per_rank,
|
||||
const void* send_ptr,
|
||||
void* recv_ptr, // local VA of the recv slot (for raddr computation)
|
||||
uint32_t nbytes, int lane_id) {
|
||||
if (lane_id == 0) {
|
||||
uint64_t recv_raddr =
|
||||
ctx.ibgda.raddrs[dst_rank] +
|
||||
(reinterpret_cast<const char*>(recv_ptr) -
|
||||
reinterpret_cast<const char*>(ctx.p2p.local_base));
|
||||
mc_ibgda_put(ctx.ibgda, channel, dst_rank, ctx.rank, qps_per_rank,
|
||||
send_ptr, recv_raddr, nbytes);
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// mc_signal / mc_red_add
|
||||
//
|
||||
// Route a signal (store) or reduction (atomic add) to dst_rank.
|
||||
// sig_ptr is a local VA within the GDR buffer.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
__device__ __forceinline__ void mc_signal(const CommCtx& ctx, int dst_rank,
|
||||
int channel, int qps_per_rank,
|
||||
int* sig_ptr, int32_t val) {
|
||||
if (dst_rank == ctx.rank) {
|
||||
mc_st_release(sig_ptr, val);
|
||||
return;
|
||||
}
|
||||
if (mc_comm_p2p_available(ctx, dst_rank)) {
|
||||
mc_p2p_signal(ctx.p2p, dst_rank, sig_ptr, val);
|
||||
} else {
|
||||
uint64_t recv_raddr =
|
||||
ctx.ibgda.raddrs[dst_rank] +
|
||||
(reinterpret_cast<const char*>(sig_ptr) -
|
||||
reinterpret_cast<const char*>(ctx.p2p.local_base));
|
||||
uint64_t laddr =
|
||||
ctx.ibgda.raddrs[ctx.rank] +
|
||||
(reinterpret_cast<const char*>(sig_ptr) -
|
||||
reinterpret_cast<const char*>(ctx.ibgda.remote_atomic_base)) +
|
||||
(reinterpret_cast<const char*>(ctx.ibgda.local_atomic_base) -
|
||||
reinterpret_cast<const char*>(ctx.p2p.local_base));
|
||||
mc_ibgda_red_add(ctx.ibgda, channel, dst_rank, ctx.rank, qps_per_rank,
|
||||
laddr, recv_raddr, val);
|
||||
}
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void mc_red_add(const CommCtx& ctx, int dst_rank,
|
||||
int channel, int qps_per_rank,
|
||||
int* sig_ptr, int32_t val) {
|
||||
mc_signal(ctx, dst_rank, channel, qps_per_rank, sig_ptr, val);
|
||||
}
|
||||
|
||||
} // namespace device
|
||||
} // namespace mooncake
|
||||
|
|
@ -0,0 +1,144 @@
|
|||
// CUDA PTX implementations of device-side memory ordering primitives.
|
||||
// Included by device_ops.cuh when MOONCAKE_EP_USE_MUSA is not defined.
|
||||
#pragma once
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include <cooperative_groups.h>
|
||||
|
||||
namespace mooncake {
|
||||
namespace device {
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Acquire loads — cross-GPU visibility (sys scope)
|
||||
// ---------------------------------------------------------------------------
|
||||
__device__ __forceinline__ int mc_ld_acquire(const int* ptr) {
|
||||
int ret;
|
||||
asm volatile("ld.acquire.sys.global.s32 %0, [%1];" : "=r"(ret) : "l"(ptr));
|
||||
return ret;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ uint64_t mc_ld_acquire_u64(const uint64_t* ptr) {
|
||||
uint64_t ret;
|
||||
asm volatile("ld.acquire.sys.global.u64 %0, [%1];" : "=l"(ret) : "l"(ptr));
|
||||
return ret;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Release stores — cross-GPU visibility (sys scope), non-temporal (no alloc)
|
||||
// ---------------------------------------------------------------------------
|
||||
__device__ __forceinline__ void mc_st_release(const int* ptr, int val) {
|
||||
asm volatile("st.release.sys.global.L1::no_allocate.s32 [%0], %1;"
|
||||
:
|
||||
: "l"(ptr), "r"(val));
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void mc_st_release_u32(const uint32_t* ptr,
|
||||
uint32_t val) {
|
||||
asm volatile("st.release.sys.global.L1::no_allocate.b32 [%0], %1;"
|
||||
:
|
||||
: "l"(ptr), "r"(val));
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void mc_st_release_u64(const uint64_t* ptr,
|
||||
uint64_t val) {
|
||||
asm volatile("st.release.sys.global.L1::no_allocate.b64 [%0], %1;"
|
||||
:
|
||||
: "l"(ptr), "l"(val));
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Atomic add — release semantics, sys scope
|
||||
// ---------------------------------------------------------------------------
|
||||
__device__ __forceinline__ int mc_atomic_add_release(const int* ptr, int val) {
|
||||
int ret;
|
||||
asm volatile("atom.add.release.sys.global.s32 %0, [%1], %2;"
|
||||
: "=r"(ret)
|
||||
: "l"(ptr), "r"(val));
|
||||
return ret;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Non-coherent loads (read-only cache, no L1 alloc) — for bulk data reads
|
||||
// ---------------------------------------------------------------------------
|
||||
__device__ __forceinline__ int4 mc_ld_nc(const int4* ptr) {
|
||||
int4 ret;
|
||||
asm volatile(
|
||||
"ld.global.nc.L1::no_allocate.L2::256B.v4.s32 {%0,%1,%2,%3}, [%4];"
|
||||
: "=r"(ret.x), "=r"(ret.y), "=r"(ret.z), "=r"(ret.w)
|
||||
: "l"(ptr));
|
||||
return ret;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ int mc_ld_nc_s32(const int* ptr) {
|
||||
int ret;
|
||||
asm volatile("ld.global.nc.L1::no_allocate.s32 %0, [%1];"
|
||||
: "=r"(ret)
|
||||
: "l"(ptr));
|
||||
return ret;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ float mc_ld_nc_f32(const float* ptr) {
|
||||
float ret;
|
||||
asm volatile("ld.global.nc.L1::no_allocate.f32 %0, [%1];"
|
||||
: "=f"(ret)
|
||||
: "l"(ptr));
|
||||
return ret;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ int64_t mc_ld_nc_s64(const int64_t* ptr) {
|
||||
int64_t ret;
|
||||
asm volatile("ld.global.nc.L1::no_allocate.s64 %0, [%1];"
|
||||
: "=l"(ret)
|
||||
: "l"(ptr));
|
||||
return ret;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Non-temporal stores (no L1 alloc) — for bulk data writes
|
||||
// ---------------------------------------------------------------------------
|
||||
__device__ __forceinline__ void mc_st_na(const int4* ptr, const int4& val) {
|
||||
asm volatile("st.global.L1::no_allocate.v4.s32 [%0], {%1,%2,%3,%4};"
|
||||
:
|
||||
: "l"(ptr), "r"(val.x), "r"(val.y), "r"(val.z), "r"(val.w));
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Named barrier init — no-op on CUDA (hardware named barriers need no setup).
|
||||
// ---------------------------------------------------------------------------
|
||||
__device__ __forceinline__ void mc_bar_init() {}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Named barrier (warp-group scope) — CUDA PTX bar.sync
|
||||
// On MUSA this maps to __syncthreads(); see musa_ops.cuh for details.
|
||||
// ---------------------------------------------------------------------------
|
||||
__device__ __forceinline__ void mc_bar_sync(int bar_id, int num_threads) {
|
||||
asm volatile("bar.sync %0, %1;" : : "r"(bar_id), "r"(num_threads));
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Grid-level sync — cooperative_groups::this_grid().sync()
|
||||
// On MUSA this is a no-op because the host always uses separate kernel
|
||||
// launches (return_recv_hook=true), so SEND and RECV never run in the same
|
||||
// kernel invocation.
|
||||
// ---------------------------------------------------------------------------
|
||||
__device__ __forceinline__ void mc_grid_sync() {
|
||||
cooperative_groups::this_grid().sync();
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Byte-swap helpers (for mlx5 big-endian WQE fields)
|
||||
// ---------------------------------------------------------------------------
|
||||
__device__ __forceinline__ uint16_t mc_bswap16(uint16_t x) {
|
||||
return __byte_perm(x, x, 0x2301);
|
||||
}
|
||||
__device__ __forceinline__ uint32_t mc_bswap32(uint32_t x) {
|
||||
return __byte_perm(x, x, 0x0123);
|
||||
}
|
||||
__device__ __forceinline__ uint64_t mc_bswap64(uint64_t x) {
|
||||
uint32_t hi = __byte_perm((uint32_t)(x >> 32), 0, 0x0123);
|
||||
uint32_t lo = __byte_perm((uint32_t)(x), 0, 0x0123);
|
||||
return ((uint64_t)lo << 32) | hi;
|
||||
}
|
||||
|
||||
} // namespace device
|
||||
} // namespace mooncake
|
||||
|
|
@ -0,0 +1,12 @@
|
|||
// Device-side memory ordering primitives — platform-portable selector.
|
||||
//
|
||||
// This is the ONLY file in the device API that contains a platform #ifdef.
|
||||
// All other device API headers include this file and use the mc_* functions.
|
||||
// The kernel itself has zero platform #ifdef branches.
|
||||
#pragma once
|
||||
|
||||
#ifdef MOONCAKE_EP_USE_MUSA
|
||||
#include "transport/device/musa/musa_ops.cuh"
|
||||
#else
|
||||
#include "transport/device/cuda/cuda_ops.cuh"
|
||||
#endif
|
||||
|
|
@ -0,0 +1,168 @@
|
|||
// Copyright 2024 KVCache.AI
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#pragma once
|
||||
|
||||
// EP DeviceTransport — platform-agnostic host-side abstraction for the two
|
||||
// communication paths used by the EP kernel:
|
||||
//
|
||||
// P2pTransport — intra-node GPU-initiated P2P (NVLink on CUDA, MTLink on
|
||||
// MUSA). Manages IPC handle exchange and peer pointer
|
||||
// table.
|
||||
// RdmaTransport — inter-node GPU-initiated RDMA (IBGDA / mlx5gda).
|
||||
// Manages QP lifecycle, MR, and device context table.
|
||||
//
|
||||
// EP code includes only this header and calls the abstract interface.
|
||||
// Platform-specific implementations live in device/ and are
|
||||
// selected at build time by the factory functions at the bottom of this file.
|
||||
//
|
||||
// The header intentionally avoids including cuda_alike.h so it can be included
|
||||
// from pure C++ translation units. Implementations include cuda_alike.h.
|
||||
|
||||
#include <cstdint>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
namespace mooncake {
|
||||
namespace device {
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// P2pTransport
|
||||
//
|
||||
// Manages intra-node GPU P2P IPC handles and the device-visible peer pointer
|
||||
// table used by the EP kernel.
|
||||
// ---------------------------------------------------------------------------
|
||||
class P2pTransport {
|
||||
public:
|
||||
virtual ~P2pTransport() = default;
|
||||
|
||||
// Allocate the GDR buffer that will be shared via IPC.
|
||||
// Returns a device pointer; size must be > 0.
|
||||
virtual void* allocateBuffer(size_t bytes) = 0;
|
||||
|
||||
// Free a buffer previously returned by allocateBuffer.
|
||||
virtual void freeBuffer(void* ptr) = 0;
|
||||
|
||||
// Export an IPC handle for the buffer allocated by this rank.
|
||||
// Returns a byte blob (serialised as int32_t array for Python compat).
|
||||
// Returns empty vector if IPC is not needed (e.g. fabric memory).
|
||||
virtual std::vector<int32_t> exportIpcHandle(void* ptr) = 0;
|
||||
|
||||
// Import peer IPC handles and populate the device-visible tables.
|
||||
// remote_handles[i] is the handle exported by rank i (may be empty for
|
||||
// ranks that use fabric memory or are on a different node).
|
||||
// active_ranks_mask[i] == 1 means rank i is participating.
|
||||
// After this call, availableTablePtr() and peerPtrsTablePtr() are valid.
|
||||
virtual void importPeerHandles(
|
||||
void* local_ptr, int rank, int num_ranks,
|
||||
const std::vector<std::vector<int32_t>>& remote_handles,
|
||||
const std::vector<int>& active_ranks_mask) = 0;
|
||||
|
||||
// Device pointer to int32_t[num_ranks]: 1 if P2P is available to rank i.
|
||||
virtual int32_t* availableTablePtr() = 0;
|
||||
|
||||
// Device pointer to void*[num_ranks]: peer GDR buffer base pointers.
|
||||
virtual void** peerPtrsTablePtr() = 0;
|
||||
|
||||
// True if all active ranks have P2P access (fast-path condition).
|
||||
virtual bool allPeersAccessible() const = 0;
|
||||
|
||||
// Verify that peer-mapped memory is writable (via memcpy).
|
||||
// Returns true if a small test write to each peer's mapped buffer
|
||||
// succeeds. On failure, sets all_peers_accessible to false.
|
||||
virtual bool verifyPeerAccess() = 0;
|
||||
};
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// RdmaLocalMetadata — exchanged between ranks during IBGDA bootstrap.
|
||||
// ---------------------------------------------------------------------------
|
||||
struct RdmaLocalMetadata {
|
||||
int64_t raddr;
|
||||
int32_t rkey;
|
||||
int64_t subnet_prefix;
|
||||
int64_t interface_id;
|
||||
std::vector<int32_t> qpns;
|
||||
std::vector<int32_t> lids;
|
||||
};
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// RdmaTransport
|
||||
//
|
||||
// Manages IBGDA QP lifecycle, MR registration, and the device-visible RDMA
|
||||
// context tables (raddrs, rkeys, qp_devctxs) used by the EP kernel.
|
||||
// ---------------------------------------------------------------------------
|
||||
class RdmaTransport {
|
||||
public:
|
||||
virtual ~RdmaTransport() = default;
|
||||
|
||||
// Initialise the RDMA transport for the given NIC.
|
||||
// device_name: e.g. "mlx5_1". Pass empty string for auto-detect.
|
||||
// Returns 0 on success, non-zero on failure (IBGDA disabled).
|
||||
virtual int initialize(const std::string& device_name, int num_ranks,
|
||||
int num_qps) = 0;
|
||||
|
||||
// Register the GDR buffer for RDMA access.
|
||||
virtual int registerMemory(void* ptr, size_t bytes) = 0;
|
||||
|
||||
// Allocate the GPU-side control buffer (QP/CQ structures).
|
||||
virtual int allocateControlBuffer() = 0;
|
||||
|
||||
// Create QPs in RST→INIT state. Call after allocateControlBuffer.
|
||||
// stream is a cudaStream_t / musaStream_t cast to void*.
|
||||
virtual int createQueuePairs(void* stream) = 0;
|
||||
|
||||
// Destroy and recreate QPs (used when active_ranks changes).
|
||||
virtual int recreateQueuePairs(void* stream) = 0;
|
||||
|
||||
// Connect QPs to peers using exchanged metadata.
|
||||
// is_roce: true for RoCE, false for IB.
|
||||
virtual int connectPeers(int local_rank, bool is_roce,
|
||||
const std::vector<int64_t>& remote_addrs,
|
||||
const std::vector<int32_t>& remote_keys,
|
||||
const std::vector<int32_t>& remote_qpns,
|
||||
const std::vector<int32_t>& remote_lids,
|
||||
const std::vector<int64_t>& subnet_prefixes,
|
||||
const std::vector<int64_t>& interface_ids,
|
||||
const std::vector<int>& active_ranks_mask) = 0;
|
||||
|
||||
// Metadata for this rank, to be exchanged with peers.
|
||||
virtual RdmaLocalMetadata localMetadata() const = 0;
|
||||
|
||||
// Device pointers to the tables consumed by the EP kernel.
|
||||
virtual void* raddrsPtr() = 0; // uint64_t[num_ranks]
|
||||
virtual void* rkeysPtr() = 0; // uint32_t[num_ranks]
|
||||
virtual void* qpDevCtxsPtr() = 0; // mlx5gda_qp_devctx[num_qps]
|
||||
|
||||
virtual bool isRoce() const = 0;
|
||||
virtual int gidIndex() const = 0;
|
||||
};
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Factory functions — implemented in device_transport.cpp.
|
||||
// Returns nullptr if the transport is not available on this platform.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// Create the platform-native P2P transport (NVLink on CUDA, MTLink on MUSA).
|
||||
std::unique_ptr<P2pTransport> createP2pDeviceTransport(int num_ranks);
|
||||
|
||||
// Create the IBGDA RDMA transport backed by TE's RdmaContext.
|
||||
// device_filter: optional whitelist of NIC names (e.g. {"mlx5_1", "mlx5_2"}).
|
||||
// Empty vector = auto-detect via TE's Topology::discover() with no filter.
|
||||
// Non-empty = restrict discovery to these NICs, then pick the closest one.
|
||||
std::unique_ptr<RdmaTransport> createIbgdaDeviceTransport(
|
||||
const std::vector<std::string>& device_filter = {});
|
||||
|
||||
} // namespace device
|
||||
} // namespace mooncake
|
||||
|
|
@ -3,7 +3,12 @@
|
|||
|
||||
#include <stddef.h>
|
||||
#include <stdint.h>
|
||||
#ifdef USE_MUSA
|
||||
#include <musa_runtime.h>
|
||||
#define cudaStream_t musaStream_t
|
||||
#else
|
||||
#include <cuda_runtime.h>
|
||||
#endif
|
||||
#include <infiniband/verbs.h>
|
||||
#include <infiniband/mlx5dv.h>
|
||||
|
||||
|
|
@ -0,0 +1,206 @@
|
|||
// IBGDA device context and operations.
|
||||
//
|
||||
// Wraps mlx5gda_qp_devctx and issues RDMA writes/atomics via
|
||||
// device-side WQE construction.
|
||||
//
|
||||
// On MUSA: BF (Blue Flame) doorbell is not available (musaHostRegisterIoMemory
|
||||
// doesn't support MMIO), so qp->bf is NULL and the kernel uses DBR-only mode.
|
||||
// All other IBGDA logic (WQE construction, CQ polling, DBR write) is shared.
|
||||
#pragma once
|
||||
|
||||
#include <cstdint>
|
||||
#include "transport/device/device_ops.cuh"
|
||||
|
||||
#ifndef MOONCAKE_EP_USE_MUSA
|
||||
#include <cuda/atomic>
|
||||
#endif
|
||||
#include <transport/device/ibgda/mlx5gda.h>
|
||||
|
||||
// mlx5 32-bit atomic-add WQE segment (not in mlx5gda.h; defined here for use
|
||||
// in mc_ibgda_write_rdma_atomic_add_wqe).
|
||||
struct mlx5_wqe_atomic_add_32_seg {
|
||||
__be32 add_data;
|
||||
__be32 field_boundary;
|
||||
__be64 compare;
|
||||
};
|
||||
|
||||
namespace mooncake {
|
||||
namespace device {
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// IbgdaContext
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
struct IbgdaContext {
|
||||
mlx5gda_qp_devctx* qp_devctxs; // device ptr: [num_qps]
|
||||
const uint64_t* raddrs; // device ptr: [num_ranks] remote GDR base
|
||||
const uint32_t* rkeys; // device ptr: [num_ranks] remote rkey
|
||||
const void* local_atomic_base; // local scratch base for atomic responses
|
||||
const void* remote_atomic_base; // symmetric remote signal base
|
||||
};
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
__device__ __forceinline__ mlx5gda_qp_devctx* mc_ibgda_channel(
|
||||
const IbgdaContext& ctx, int channel, int dst_rank, int qps_per_rank) {
|
||||
int qp_idx = dst_rank * qps_per_rank + (channel % qps_per_rank);
|
||||
return ctx.qp_devctxs + qp_idx;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void mc_ibgda_lock(mlx5gda_qp_devctx* qp) {
|
||||
#ifdef MOONCAKE_EP_USE_MUSA
|
||||
uint32_t old;
|
||||
do {
|
||||
old = atomicCAS(&qp->mutex, 0u, 1u);
|
||||
} while (old != 0);
|
||||
#else
|
||||
cuda::atomic_ref<uint32_t, cuda::thread_scope_system> lock(qp->mutex);
|
||||
while (lock.exchange(1u, cuda::memory_order_acquire) != 0u);
|
||||
#endif
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void mc_ibgda_unlock(mlx5gda_qp_devctx* qp) {
|
||||
#ifdef MOONCAKE_EP_USE_MUSA
|
||||
mc_st_release_u32(&qp->mutex, 0u);
|
||||
#else
|
||||
cuda::atomic_ref<uint32_t, cuda::thread_scope_system> lock(qp->mutex);
|
||||
lock.store(0u, cuda::memory_order_release);
|
||||
#endif
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void mc_ibgda_poll_cq(mlx5gda_qp_devctx* qp,
|
||||
uint16_t expect) {
|
||||
uint16_t wq_tail = qp->wq_tail;
|
||||
while (static_cast<int16_t>(wq_tail - expect) <= 0) {
|
||||
uint16_t cq_be =
|
||||
*reinterpret_cast<volatile uint16_t*>(&qp->cq->wqe_counter);
|
||||
uint8_t opcode = qp->cq->op_own >> 4;
|
||||
if (opcode == 0xD)
|
||||
printf("[EP IBGDA] Requester error: syndrome=0x%lx\n",
|
||||
qp->cq->timestamp >> 56);
|
||||
if (!(opcode == 0x0 || opcode == 0xF)) {
|
||||
printf("[EP IBGDA] Unexpected CQE opcode=0x%x, trapping\n", opcode);
|
||||
__trap();
|
||||
}
|
||||
wq_tail = mc_bswap16(cq_be) + 1;
|
||||
}
|
||||
if (wq_tail != qp->wq_tail) qp->wq_tail = wq_tail;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void mc_ibgda_post_send_db(mlx5gda_qp_devctx* qp) {
|
||||
uint32_t num_posted = static_cast<uint32_t>(qp->wq_head);
|
||||
// DBR write — always done (NIC polls doorbell record in GPU memory)
|
||||
mc_st_release_u32(reinterpret_cast<uint32_t*>(&qp->dbr->send_counter),
|
||||
mc_bswap32(num_posted));
|
||||
// BF (Blue Flame) doorbell — only if BF register is mapped into GPU VA.
|
||||
// On MUSA, musaHostRegisterIoMemory fails for MMIO addresses, so bf is
|
||||
// NULL and we rely on DBR-only mode (slightly higher latency).
|
||||
if (qp->bf != nullptr) {
|
||||
auto* last_wqe = qp->wq + ((num_posted - 1) & qp->wqeid_mask);
|
||||
mc_st_release_u64(reinterpret_cast<uint64_t*>(qp->bf + qp->bf_offset),
|
||||
*reinterpret_cast<uint64_t*>(last_wqe));
|
||||
qp->bf_offset ^= MLX5GDA_BF_SIZE;
|
||||
}
|
||||
}
|
||||
|
||||
// Issue an RDMA WRITE WQE. laddr/raddr are device VAs; keys are big-endian.
|
||||
__device__ __forceinline__ void mc_ibgda_write_rdma_write_wqe(
|
||||
mlx5gda_qp_devctx* qp, uint64_t laddr, __be32 lkey, uint64_t raddr,
|
||||
__be32 rkey, uint32_t bytes) {
|
||||
auto* wqe = reinterpret_cast<mlx5gda_rdma_write_wqe*>(
|
||||
qp->wq + (qp->wq_head & qp->wqeid_mask));
|
||||
|
||||
wqe->ctrl = {};
|
||||
wqe->ctrl.qpn_ds = mc_bswap32((qp->qpn << 8) | 3);
|
||||
wqe->ctrl.fm_ce_se = MLX5_WQE_CTRL_CQ_UPDATE;
|
||||
wqe->ctrl.opmod_idx_opcode = mc_bswap32(
|
||||
(static_cast<uint32_t>(qp->wq_head) << 8) | MLX5_OPCODE_RDMA_WRITE);
|
||||
|
||||
wqe->raddr.raddr = mc_bswap64(raddr);
|
||||
wqe->raddr.rkey = rkey;
|
||||
wqe->raddr.reserved = 0;
|
||||
|
||||
wqe->data.byte_count = mc_bswap32(bytes);
|
||||
wqe->data.lkey = lkey;
|
||||
wqe->data.addr = mc_bswap64(laddr);
|
||||
|
||||
++qp->wq_head;
|
||||
}
|
||||
|
||||
// Issue an RDMA ATOMIC MASKED FETCH-AND-ADD WQE (32-bit add_data).
|
||||
// This matches the original CUDA IBGDA EP kernel. A regular 64-bit
|
||||
// MLX5_OPCODE_ATOMIC_FA does not implement the 32-bit signal-buffer add used
|
||||
// by dispatch/combine.
|
||||
__device__ __forceinline__ void mc_ibgda_write_rdma_atomic_add_wqe(
|
||||
mlx5gda_qp_devctx* qp, int32_t value, uint64_t laddr, __be32 lkey,
|
||||
uint64_t raddr, __be32 rkey) {
|
||||
auto* wqe = reinterpret_cast<mlx5gda_rdma_atomic_wqe*>(
|
||||
qp->wq + (qp->wq_head & qp->wqeid_mask));
|
||||
|
||||
wqe->ctrl = {};
|
||||
wqe->ctrl.qpn_ds = mc_bswap32((qp->qpn << 8) | 4);
|
||||
wqe->ctrl.fm_ce_se = MLX5_WQE_CTRL_CQ_UPDATE;
|
||||
wqe->ctrl.opmod_idx_opcode =
|
||||
mc_bswap32(MLX5_OPCODE_ATOMIC_MASKED_FA |
|
||||
(static_cast<uint32_t>(qp->wq_head) << 8) | 0x08000000);
|
||||
|
||||
wqe->raddr.raddr = mc_bswap64(raddr);
|
||||
wqe->raddr.rkey = rkey;
|
||||
wqe->raddr.reserved = 0;
|
||||
|
||||
// atomic_seg: add_data (32-bit, big-endian), field_boundary=0, compare=0
|
||||
auto* atomic_seg =
|
||||
reinterpret_cast<mlx5_wqe_atomic_add_32_seg*>(&wqe->atomic);
|
||||
atomic_seg->add_data = mc_bswap32(static_cast<uint32_t>(value));
|
||||
atomic_seg->field_boundary = 0;
|
||||
atomic_seg->compare = 0;
|
||||
|
||||
wqe->data.byte_count = mc_bswap32(static_cast<uint32_t>(4));
|
||||
wqe->data.lkey = lkey;
|
||||
wqe->data.addr = mc_bswap64(laddr);
|
||||
|
||||
++qp->wq_head;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// High-level IBGDA operations
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// RDMA WRITE: send `nbytes` from `send_ptr` to `recv_ptr` on `dst_rank`.
|
||||
// Must be called by lane 0 only.
|
||||
__device__ __forceinline__ void mc_ibgda_put(const IbgdaContext& ctx,
|
||||
int channel, int dst_rank,
|
||||
int src_rank, int qps_per_rank,
|
||||
const void* send_ptr,
|
||||
uint64_t recv_raddr,
|
||||
uint32_t nbytes) {
|
||||
auto* qp = mc_ibgda_channel(ctx, channel, dst_rank, qps_per_rank);
|
||||
mc_ibgda_lock(qp);
|
||||
mc_ibgda_write_rdma_write_wqe(qp, reinterpret_cast<uint64_t>(send_ptr),
|
||||
mc_bswap32(ctx.rkeys[src_rank]), recv_raddr,
|
||||
mc_bswap32(ctx.rkeys[dst_rank]), nbytes);
|
||||
mc_ibgda_post_send_db(qp);
|
||||
mc_ibgda_unlock(qp);
|
||||
}
|
||||
|
||||
// RDMA ATOMIC ADD: add `value` to the 32-bit word at `recv_raddr` on
|
||||
// `dst_rank`. Must be called by lane 0 only.
|
||||
__device__ __forceinline__ void mc_ibgda_red_add(
|
||||
const IbgdaContext& ctx, int channel, int dst_rank, int src_rank,
|
||||
int qps_per_rank,
|
||||
uint64_t laddr, // local scratch VA for the atomic result
|
||||
uint64_t recv_raddr, // remote VA of the signal word
|
||||
int32_t value) {
|
||||
auto* qp = mc_ibgda_channel(ctx, channel, dst_rank, qps_per_rank);
|
||||
mc_ibgda_lock(qp);
|
||||
mc_ibgda_write_rdma_atomic_add_wqe(
|
||||
qp, value, laddr, mc_bswap32(ctx.rkeys[src_rank]), recv_raddr,
|
||||
mc_bswap32(ctx.rkeys[dst_rank]));
|
||||
mc_ibgda_post_send_db(qp);
|
||||
mc_ibgda_unlock(qp);
|
||||
}
|
||||
|
||||
} // namespace device
|
||||
} // namespace mooncake
|
||||
|
|
@ -0,0 +1,137 @@
|
|||
// MUSA implementations of device-side memory ordering primitives.
|
||||
// Included by device_ops.cuh when MOONCAKE_EP_USE_MUSA is defined.
|
||||
//
|
||||
// MUSA has no PTX-style acquire/release instructions. We emulate them with
|
||||
// plain loads/stores + __threadfence_system() for cross-GPU (MTLink)
|
||||
// visibility.
|
||||
//
|
||||
// Known MUSA SDK 4.3.3 compiler bugs to avoid:
|
||||
// - atomicAdd_system / atomicCAS_system → infinite SelectionDAG loop.
|
||||
// Use block-scope atomicAdd + __threadfence_system() instead.
|
||||
// - Named barriers (bar.sync) → not available; use __syncthreads().
|
||||
// - cooperative_groups::this_grid().sync() → not available; host uses
|
||||
// separate kernel launches (return_recv_hook=true) so grid sync is a no-op.
|
||||
#pragma once
|
||||
|
||||
#include <musa_runtime.h>
|
||||
|
||||
namespace mooncake {
|
||||
namespace device {
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Acquire loads
|
||||
// ---------------------------------------------------------------------------
|
||||
__device__ __forceinline__ int mc_ld_acquire(const int* ptr) {
|
||||
int ret = *const_cast<volatile const int*>(ptr);
|
||||
__threadfence_system();
|
||||
return ret;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ uint64_t mc_ld_acquire_u64(const uint64_t* ptr) {
|
||||
uint64_t ret = *const_cast<volatile const uint64_t*>(ptr);
|
||||
__threadfence_system();
|
||||
return ret;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Release stores
|
||||
// ---------------------------------------------------------------------------
|
||||
__device__ __forceinline__ void mc_st_release(const int* ptr, int val) {
|
||||
__threadfence_system();
|
||||
*const_cast<volatile int*>(ptr) = val;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void mc_st_release_u32(const uint32_t* ptr,
|
||||
uint32_t val) {
|
||||
__threadfence_system();
|
||||
*const_cast<volatile uint32_t*>(ptr) = val;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void mc_st_release_u64(const uint64_t* ptr,
|
||||
uint64_t val) {
|
||||
__threadfence_system();
|
||||
*const_cast<volatile uint64_t*>(ptr) = val;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Atomic add — block-scope atomicAdd + system fence (avoids SDK bug)
|
||||
// ---------------------------------------------------------------------------
|
||||
__device__ __forceinline__ int mc_atomic_add_release(const int* ptr, int val) {
|
||||
__threadfence_system();
|
||||
return atomicAdd(const_cast<int*>(ptr), val);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Non-coherent loads — MUSA has no nc/no_allocate cache hints; use volatile.
|
||||
// int4 volatile copy not supported by MUSA compiler; copy field-by-field.
|
||||
// ---------------------------------------------------------------------------
|
||||
__device__ __forceinline__ int4 mc_ld_nc(const int4* ptr) {
|
||||
const volatile int* vp = reinterpret_cast<const volatile int*>(ptr);
|
||||
int4 ret;
|
||||
ret.x = vp[0];
|
||||
ret.y = vp[1];
|
||||
ret.z = vp[2];
|
||||
ret.w = vp[3];
|
||||
return ret;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ int mc_ld_nc_s32(const int* ptr) {
|
||||
return *const_cast<volatile const int*>(ptr);
|
||||
}
|
||||
|
||||
__device__ __forceinline__ float mc_ld_nc_f32(const float* ptr) {
|
||||
return *const_cast<volatile const float*>(ptr);
|
||||
}
|
||||
|
||||
__device__ __forceinline__ int64_t mc_ld_nc_s64(const int64_t* ptr) {
|
||||
return *const_cast<volatile const int64_t*>(ptr);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Non-temporal stores — MUSA has no nt/no_allocate hints; plain store.
|
||||
// int4 volatile assignment not supported by MUSA compiler; copy field-by-field.
|
||||
// ---------------------------------------------------------------------------
|
||||
__device__ __forceinline__ void mc_st_na(const int4* ptr, const int4& val) {
|
||||
volatile int* vp = reinterpret_cast<volatile int*>(const_cast<int4*>(ptr));
|
||||
vp[0] = val.x;
|
||||
vp[1] = val.y;
|
||||
vp[2] = val.z;
|
||||
vp[3] = val.w;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Named barrier — MUSA has no bar.sync.
|
||||
// On MUSA, mc_bar_sync is implemented as __syncthreads() (full CTA barrier).
|
||||
// Kernels that call mc_bar_sync from a subset of threads must ensure the
|
||||
// remaining threads call it the same number of times from another code path.
|
||||
// See mooncake_ep_kernel.cu for the warp-31 dispatch-send workaround.
|
||||
// ---------------------------------------------------------------------------
|
||||
__device__ __forceinline__ void mc_bar_sync(int /*bar_id*/,
|
||||
int /*num_threads*/) {
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Grid sync — not available on MUSA. Host always uses separate kernel
|
||||
// launches, so SEND and RECV never share a kernel invocation. No-op.
|
||||
// ---------------------------------------------------------------------------
|
||||
__device__ __forceinline__ void mc_grid_sync() {}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Byte-swap helpers — MUSA has no __byte_perm; implement manually.
|
||||
// ---------------------------------------------------------------------------
|
||||
__device__ __forceinline__ uint16_t mc_bswap16(uint16_t x) {
|
||||
return (uint16_t)(((x & 0x00FFu) << 8) | ((x & 0xFF00u) >> 8));
|
||||
}
|
||||
__device__ __forceinline__ uint32_t mc_bswap32(uint32_t x) {
|
||||
return ((x & 0x000000FFu) << 24) | ((x & 0x0000FF00u) << 8) |
|
||||
((x & 0x00FF0000u) >> 8) | ((x & 0xFF000000u) >> 24);
|
||||
}
|
||||
__device__ __forceinline__ uint64_t mc_bswap64(uint64_t x) {
|
||||
uint32_t hi = mc_bswap32((uint32_t)(x >> 32));
|
||||
uint32_t lo = mc_bswap32((uint32_t)(x));
|
||||
return ((uint64_t)lo << 32) | hi;
|
||||
}
|
||||
|
||||
} // namespace device
|
||||
} // namespace mooncake
|
||||
|
|
@ -0,0 +1,46 @@
|
|||
// P2P device context and operations — unified NVLink (CUDA) / MTLink (MUSA).
|
||||
//
|
||||
// P2PContext holds the device-visible tables populated by P2pTransport on
|
||||
// the host side. The kernel constructs one from the raw pointers passed in
|
||||
// and calls mc_p2p_* helpers instead of accessing the tables directly.
|
||||
#pragma once
|
||||
|
||||
#include <cstdint>
|
||||
#include "transport/device/device_ops.cuh"
|
||||
|
||||
namespace mooncake {
|
||||
namespace device {
|
||||
|
||||
struct P2PContext {
|
||||
const int32_t* available; // device ptr: [num_ranks], 1 = P2P reachable
|
||||
void* const* peer_ptrs; // device ptr: [num_ranks], peer GDR base ptrs
|
||||
void* local_base; // this rank's GDR buffer base (for offset math)
|
||||
};
|
||||
|
||||
__device__ __forceinline__ bool mc_p2p_available(const P2PContext& ctx,
|
||||
int dst_rank) {
|
||||
return ctx.available[dst_rank] != 0 && ctx.peer_ptrs[dst_rank] != nullptr;
|
||||
}
|
||||
|
||||
// Translate a local pointer (within the GDR buffer) to the peer's mapped VA.
|
||||
__device__ __forceinline__ void* mc_p2p_peer_ptr(const P2PContext& ctx,
|
||||
int dst_rank,
|
||||
const void* local_ptr) {
|
||||
const auto offset = reinterpret_cast<const char*>(local_ptr) -
|
||||
reinterpret_cast<const char*>(ctx.local_base);
|
||||
return reinterpret_cast<char*>(ctx.peer_ptrs[dst_rank]) + offset;
|
||||
}
|
||||
|
||||
// Write a 32-bit signal value to the peer's signal slot via P2P store.
|
||||
// Single-writer assumption: uses release store, not atomic add.
|
||||
__device__ __forceinline__ void mc_p2p_signal(const P2PContext& ctx,
|
||||
int dst_rank,
|
||||
const int* local_sig_ptr,
|
||||
int32_t val) {
|
||||
auto* peer_sig =
|
||||
reinterpret_cast<int*>(mc_p2p_peer_ptr(ctx, dst_rank, local_sig_ptr));
|
||||
mc_st_release(peer_sig, val);
|
||||
}
|
||||
|
||||
} // namespace device
|
||||
} // namespace mooncake
|
||||
|
|
@ -64,7 +64,7 @@ endif()
|
|||
|
||||
if(USE_CUDA)
|
||||
target_include_directories(transfer_engine PRIVATE /usr/local/cuda/include)
|
||||
target_link_libraries(transfer_engine PUBLIC cuda cudart rt)
|
||||
target_link_libraries(transfer_engine PUBLIC cuda cudart rt mlx5)
|
||||
if(USE_NVMEOF)
|
||||
target_link_libraries(transfer_engine PUBLIC nvmeof_transport cufile)
|
||||
endif()
|
||||
|
|
|
|||
|
|
@ -179,6 +179,17 @@ Transport* TransferEngine::getTransport(const std::string& proto) {
|
|||
return impl_->getTransport(proto);
|
||||
}
|
||||
|
||||
#if defined(USE_CUDA) || defined(USE_MUSA)
|
||||
device::P2pTransport* TransferEngine::getOrCreateP2pTransport(int num_ranks) {
|
||||
return impl_->getOrCreateP2pTransport(num_ranks);
|
||||
}
|
||||
|
||||
device::RdmaTransport* TransferEngine::getOrCreateRdmaTransport(
|
||||
const std::vector<std::string>& device_filter) {
|
||||
return impl_->getOrCreateRdmaTransport(device_filter);
|
||||
}
|
||||
#endif
|
||||
|
||||
bool TransferEngine::isTcpOnly() const { return impl_->isTcpOnly(); }
|
||||
|
||||
int TransferEngine::syncSegmentCache(const std::string& segment_name) {
|
||||
|
|
@ -580,6 +591,19 @@ Transport* TransferEngine::getTransport(const std::string& proto) {
|
|||
return impl_->getTransport(proto);
|
||||
}
|
||||
|
||||
#if defined(USE_CUDA) || defined(USE_MUSA)
|
||||
device::P2pTransport* TransferEngine::getOrCreateP2pTransport(int num_ranks) {
|
||||
if (use_tent_) return nullptr;
|
||||
return impl_->getOrCreateP2pTransport(num_ranks);
|
||||
}
|
||||
|
||||
device::RdmaTransport* TransferEngine::getOrCreateRdmaTransport(
|
||||
const std::vector<std::string>& device_filter) {
|
||||
if (use_tent_) return nullptr;
|
||||
return impl_->getOrCreateRdmaTransport(device_filter);
|
||||
}
|
||||
#endif
|
||||
|
||||
bool TransferEngine::isTcpOnly() const {
|
||||
if (use_tent_)
|
||||
// TENT already rejects TCP loopback transfers when MC_STORE_MEMCPY
|
||||
|
|
|
|||
|
|
@ -413,6 +413,24 @@ int TransferEngineImpl::uninstallTransport(const std::string& proto) {
|
|||
return 0;
|
||||
}
|
||||
|
||||
#if defined(USE_CUDA) || defined(USE_MUSA)
|
||||
device::P2pTransport* TransferEngineImpl::getOrCreateP2pTransport(
|
||||
int num_ranks) {
|
||||
if (!p2p_transport_) {
|
||||
p2p_transport_ = device::createP2pDeviceTransport(num_ranks);
|
||||
}
|
||||
return p2p_transport_.get();
|
||||
}
|
||||
|
||||
device::RdmaTransport* TransferEngineImpl::getOrCreateRdmaTransport(
|
||||
const std::vector<std::string>& device_filter) {
|
||||
if (!rdma_transport_) {
|
||||
rdma_transport_ = device::createIbgdaDeviceTransport(device_filter);
|
||||
}
|
||||
return rdma_transport_.get();
|
||||
}
|
||||
#endif
|
||||
|
||||
int TransferEngineImpl::getRpcPort() {
|
||||
return metadata_->localRpcMeta().rpc_port;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -72,3 +72,11 @@ if (USE_EFA)
|
|||
target_sources(transport PUBLIC $<TARGET_OBJECTS:efa_transport>)
|
||||
target_link_libraries(transport PRIVATE fabric)
|
||||
endif()
|
||||
|
||||
if(USE_CUDA OR USE_MUSA)
|
||||
add_subdirectory(device)
|
||||
target_sources(transport PUBLIC $<TARGET_OBJECTS:device_transport>)
|
||||
# device_transport (ibgda_device_transport.cpp / mlx5gda.cpp) calls libmlx5
|
||||
# DevX symbols (mlx5dv_devx_*, mlx5dv_init_obj) directly.
|
||||
target_link_libraries(transport PUBLIC mlx5)
|
||||
endif()
|
||||
|
|
|
|||
|
|
@ -0,0 +1,22 @@
|
|||
# Device transport sources — conditional on GPU vendor.
|
||||
# mlx5gda.cpp (IBGDA QP lifecycle) is compiled directly into the
|
||||
# device_transport OBJECT library, like every other transport module, so its
|
||||
# mlx5gda_* symbols flow into transfer_engine without a separate static lib.
|
||||
# This keeps them visible to consumers that link transfer_engine via hand-
|
||||
# written ldflags (Go p2p-store / mooncake-store), which bypass CMake's
|
||||
# target_link_libraries propagation.
|
||||
set(DEVICE_TRANSPORT_SOURCES p2p_device_transport.cpp)
|
||||
if(USE_CUDA OR USE_MUSA)
|
||||
list(APPEND DEVICE_TRANSPORT_SOURCES ibgda_device_transport.cpp mlx5gda.cpp)
|
||||
endif()
|
||||
|
||||
add_library(device_transport OBJECT ${DEVICE_TRANSPORT_SOURCES})
|
||||
target_include_directories(device_transport PRIVATE ${CMAKE_SOURCE_DIR}/include)
|
||||
|
||||
if(USE_CUDA)
|
||||
target_include_directories(device_transport PRIVATE /usr/local/cuda/include)
|
||||
endif()
|
||||
if(USE_MUSA)
|
||||
target_include_directories(device_transport PRIVATE /usr/local/musa/include)
|
||||
target_compile_definitions(device_transport PRIVATE USE_MUSA)
|
||||
endif()
|
||||
|
|
@ -0,0 +1,430 @@
|
|||
// Copyright 2024 KVCache.AI
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
// IBGDA device transport — wraps mlx5gda QP lifecycle behind RdmaTransport.
|
||||
//
|
||||
// Uses TE's Topology for NIC auto-detection when device_name is empty.
|
||||
// Uses ibv_open_device / ibv_alloc_pd directly (same as the original
|
||||
// MooncakeEpBuffer::init_ibgda()), but encapsulated in a reusable class.
|
||||
|
||||
#include "transport/device/device_transport.h"
|
||||
|
||||
#include <arpa/inet.h>
|
||||
#include <glog/logging.h>
|
||||
#include <infiniband/mlx5dv.h>
|
||||
#include <infiniband/verbs.h>
|
||||
|
||||
#include <cstring>
|
||||
#include <stdexcept>
|
||||
|
||||
#include "cuda_alike.h"
|
||||
#include "transport/device/ibgda/memheap.h"
|
||||
#include "transport/device/ibgda/mlx5gda.h"
|
||||
#include "topology.h"
|
||||
|
||||
namespace mooncake {
|
||||
namespace device {
|
||||
|
||||
static constexpr size_t kCtrlBufSize = 1024ULL * 1024 * 1024; // 1 GiB
|
||||
|
||||
// Check if IPv6 address is IPv4-mapped (::ffff:x.x.x.x)
|
||||
static bool isIpv4Mapped(const struct in6_addr* a) {
|
||||
return ((a->s6_addr32[0] | a->s6_addr32[1]) == 0 &&
|
||||
a->s6_addr32[2] == htonl(0x0000ffff));
|
||||
}
|
||||
|
||||
static int findBestGidIndex(ibv_context* ctx, uint8_t port,
|
||||
const ibv_port_attr& port_attr) {
|
||||
for (int i = 0; i < port_attr.gid_tbl_len; ++i) {
|
||||
ibv_gid_entry entry;
|
||||
if (ibv_query_gid_ex(ctx, port, i, &entry, 0)) continue;
|
||||
|
||||
if (entry.gid_type == IBV_GID_TYPE_ROCE_V2) {
|
||||
bool v4mapped = isIpv4Mapped(
|
||||
reinterpret_cast<const struct in6_addr*>(entry.gid.raw));
|
||||
if (v4mapped) return i;
|
||||
} else if (entry.gid_type == IBV_GID_TYPE_IB) {
|
||||
return i;
|
||||
}
|
||||
}
|
||||
|
||||
return -1;
|
||||
}
|
||||
|
||||
// Auto-detect the best NIC for the current GPU using TE's Topology.
|
||||
// filter: if non-empty, only consider NICs in this list.
|
||||
static std::string autoDetectNic(const std::vector<std::string>& filter) {
|
||||
Topology topo;
|
||||
if (topo.discover(filter) != 0) return "";
|
||||
const auto& hca_list = topo.getHcaList();
|
||||
if (hca_list.empty()) return "";
|
||||
|
||||
// Build a location string for the current GPU so Topology picks the
|
||||
// topologically closest NIC. Fall back to wildcard if cudaGetDevice fails.
|
||||
int device_id = 0;
|
||||
cudaGetDevice(&device_id);
|
||||
std::string location = "cuda:" + std::to_string(device_id);
|
||||
|
||||
int idx = topo.selectDevice(location);
|
||||
if (idx < 0) idx = topo.selectDevice("*"); // wildcard fallback
|
||||
if (idx < 0 || idx >= static_cast<int>(hca_list.size())) return "";
|
||||
return hca_list[idx];
|
||||
}
|
||||
|
||||
class IbgdaDeviceTransportImpl : public RdmaTransport {
|
||||
public:
|
||||
explicit IbgdaDeviceTransportImpl(std::vector<std::string> filter)
|
||||
: device_filter_(std::move(filter)) {}
|
||||
|
||||
~IbgdaDeviceTransportImpl() override { teardown(); }
|
||||
|
||||
int initialize(const std::string& device_name, int num_ranks,
|
||||
int num_qps) override {
|
||||
if (num_qps < num_ranks) {
|
||||
LOG(ERROR) << "[EP IBGDA] num_qps (" << num_qps
|
||||
<< ") must be >= num_ranks (" << num_ranks << ")";
|
||||
return -1;
|
||||
}
|
||||
num_ranks_ = num_ranks;
|
||||
num_qps_ = num_qps;
|
||||
|
||||
std::string nic =
|
||||
device_name.empty() ? autoDetectNic(device_filter_) : device_name;
|
||||
if (nic.empty()) {
|
||||
LOG(WARNING) << "[EP IBGDA] No RDMA NIC found";
|
||||
return -1;
|
||||
}
|
||||
|
||||
int num_devices = 0;
|
||||
ibv_device** dev_list = ibv_get_device_list(&num_devices);
|
||||
if (!dev_list) {
|
||||
LOG(ERROR) << "[EP IBGDA] ibv_get_device_list failed";
|
||||
return -1;
|
||||
}
|
||||
|
||||
ibv_device* dev = nullptr;
|
||||
for (int i = 0; i < num_devices; ++i) {
|
||||
if (nic == ibv_get_device_name(dev_list[i])) {
|
||||
dev = dev_list[i];
|
||||
break;
|
||||
}
|
||||
}
|
||||
if (!dev) {
|
||||
LOG(ERROR) << "[EP IBGDA] NIC '" << nic << "' not found";
|
||||
ibv_free_device_list(dev_list);
|
||||
return -1;
|
||||
}
|
||||
|
||||
ctx_ = ibv_open_device(dev);
|
||||
ibv_free_device_list(dev_list);
|
||||
if (!ctx_) {
|
||||
LOG(ERROR) << "[EP IBGDA] ibv_open_device failed for " << nic;
|
||||
return -1;
|
||||
}
|
||||
|
||||
const uint8_t port = 1;
|
||||
ibv_port_attr port_attr{};
|
||||
if (ibv_query_port(ctx_, port, &port_attr)) {
|
||||
LOG(ERROR) << "[EP IBGDA] ibv_query_port failed";
|
||||
return -1;
|
||||
}
|
||||
|
||||
gid_index_ = findBestGidIndex(ctx_, port, port_attr);
|
||||
if (gid_index_ < 0) {
|
||||
LOG(ERROR) << "[EP IBGDA] No suitable GID on " << nic;
|
||||
return -1;
|
||||
}
|
||||
|
||||
if (ibv_query_gid(ctx_, port, gid_index_, &gid_)) {
|
||||
LOG(ERROR) << "[EP IBGDA] ibv_query_gid failed";
|
||||
return -1;
|
||||
}
|
||||
|
||||
is_roce_ = (port_attr.link_layer == IBV_LINK_LAYER_ETHERNET);
|
||||
lid_ = port_attr.lid;
|
||||
device_name_ = nic;
|
||||
|
||||
pd_ = ibv_alloc_pd(ctx_);
|
||||
if (!pd_) {
|
||||
LOG(ERROR) << "[EP IBGDA] ibv_alloc_pd failed";
|
||||
return -1;
|
||||
}
|
||||
|
||||
mlx5dv_obj dv_obj{};
|
||||
dv_obj.pd.in = pd_;
|
||||
dv_obj.pd.out = &mpd_;
|
||||
if (mlx5dv_init_obj(&dv_obj, MLX5DV_OBJ_PD)) {
|
||||
LOG(ERROR) << "[EP IBGDA] mlx5dv_init_obj failed";
|
||||
return -1;
|
||||
}
|
||||
|
||||
// Allocate device-visible tables
|
||||
if (cudaMalloc(&raddrs_, num_ranks_ * sizeof(uint64_t)) !=
|
||||
cudaSuccess ||
|
||||
cudaMalloc(&rkeys_, num_ranks_ * sizeof(uint32_t)) != cudaSuccess ||
|
||||
cudaMalloc(&qp_devctxs_, num_qps_ * sizeof(mlx5gda_qp_devctx)) !=
|
||||
cudaSuccess) {
|
||||
LOG(ERROR) << "[EP IBGDA] Failed to allocate device-visible tables";
|
||||
return -1;
|
||||
}
|
||||
|
||||
LOG(INFO) << "[EP IBGDA] Initialized on " << nic
|
||||
<< " (gid_index=" << gid_index_ << ", roce=" << is_roce_
|
||||
<< ")";
|
||||
return 0;
|
||||
}
|
||||
|
||||
int registerMemory(void* ptr, size_t bytes) override {
|
||||
mr_ =
|
||||
ibv_reg_mr(pd_, ptr, bytes,
|
||||
IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_READ |
|
||||
IBV_ACCESS_REMOTE_WRITE | IBV_ACCESS_REMOTE_ATOMIC);
|
||||
if (!mr_) {
|
||||
LOG(ERROR) << "[EP IBGDA] ibv_reg_mr failed";
|
||||
return -1;
|
||||
}
|
||||
mr_ptr_ = ptr;
|
||||
return 0;
|
||||
}
|
||||
|
||||
int allocateControlBuffer() override {
|
||||
cudaError_t err = cudaMalloc(&ctrl_buf_, kCtrlBufSize);
|
||||
if (err != cudaSuccess) {
|
||||
LOG(ERROR) << "[EP IBGDA] cudaMalloc ctrl_buf failed: "
|
||||
<< cudaGetErrorString(err);
|
||||
return -1;
|
||||
}
|
||||
|
||||
ctrl_buf_umem_ = mlx5dv_devx_umem_reg(ctx_, ctrl_buf_, kCtrlBufSize,
|
||||
IBV_ACCESS_LOCAL_WRITE);
|
||||
if (!ctrl_buf_umem_) {
|
||||
LOG(ERROR) << "[EP IBGDA] mlx5dv_devx_umem_reg failed (errno="
|
||||
<< errno << ")";
|
||||
return -1;
|
||||
}
|
||||
LOG(INFO) << "[EP IBGDA] ctrl_buf UMEM registered via VA path";
|
||||
|
||||
ctrl_buf_heap_ = memheap_create(kCtrlBufSize);
|
||||
if (!ctrl_buf_heap_) {
|
||||
LOG(ERROR) << "[EP IBGDA] memheap_create failed";
|
||||
return -1;
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
int createQueuePairs(void* stream_ptr) override {
|
||||
auto stream = static_cast<cudaStream_t>(stream_ptr);
|
||||
for (int i = 0; i < num_qps_; ++i) {
|
||||
mlx5gda_qp* qp =
|
||||
mlx5gda_create_rc_qp(mpd_, ctrl_buf_, ctrl_buf_umem_,
|
||||
ctrl_buf_heap_, pd_, 16384, 1, stream);
|
||||
if (!qp) {
|
||||
LOG(ERROR) << "[EP IBGDA] mlx5gda_create_rc_qp failed at " << i;
|
||||
return -1;
|
||||
}
|
||||
if (mlx5gda_modify_rc_qp_rst2init(qp, 0)) {
|
||||
LOG(ERROR) << "[EP IBGDA] rst2init failed at " << i;
|
||||
mlx5gda_destroy_qp(ctrl_buf_heap_, qp);
|
||||
return -1;
|
||||
}
|
||||
cudaStreamSynchronize(stream);
|
||||
mlx5gda_qp_devctx devctx{
|
||||
.qpn = qp->qpn,
|
||||
.wqeid_mask = qp->num_wqebb - 1,
|
||||
.wq = reinterpret_cast<mlx5gda_wqebb*>(
|
||||
static_cast<char*>(ctrl_buf_) + qp->wq_offset),
|
||||
.cq = reinterpret_cast<mlx5_cqe64*>(
|
||||
static_cast<char*>(ctrl_buf_) + qp->send_cq->cq_offset),
|
||||
.dbr = reinterpret_cast<mlx5gda_wq_dbr*>(
|
||||
static_cast<char*>(ctrl_buf_) + qp->dbr_offset),
|
||||
.bf = static_cast<char*>(qp->uar->reg_addr),
|
||||
};
|
||||
cudaMemcpy(
|
||||
static_cast<char*>(qp_devctxs_) + i * sizeof(mlx5gda_qp_devctx),
|
||||
&devctx, sizeof(mlx5gda_qp_devctx), cudaMemcpyHostToDevice);
|
||||
qps_.push_back(qp);
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
int recreateQueuePairs(void* stream_ptr) override {
|
||||
auto stream = static_cast<cudaStream_t>(stream_ptr);
|
||||
for (auto* qp : qps_) {
|
||||
if (qp) mlx5gda_destroy_qp(ctrl_buf_heap_, qp);
|
||||
}
|
||||
qps_.clear();
|
||||
return createQueuePairs(stream_ptr);
|
||||
}
|
||||
|
||||
int connectPeers(int local_rank, bool is_roce,
|
||||
const std::vector<int64_t>& remote_addrs,
|
||||
const std::vector<int32_t>& remote_keys,
|
||||
const std::vector<int32_t>& remote_qpns,
|
||||
const std::vector<int32_t>& remote_lids,
|
||||
const std::vector<int64_t>& subnet_prefixes,
|
||||
const std::vector<int64_t>& interface_ids,
|
||||
const std::vector<int>& active_ranks_mask) override {
|
||||
for (int i = 0; i < num_qps_; ++i) {
|
||||
int peer_rank = i * num_ranks_ / num_qps_;
|
||||
if (active_ranks_mask[peer_rank] == 0) continue;
|
||||
|
||||
ibv_ah_attr ah_attr{};
|
||||
if (is_roce) {
|
||||
ibv_gid remote_gid{};
|
||||
remote_gid.global.subnet_prefix = subnet_prefixes[peer_rank];
|
||||
remote_gid.global.interface_id = interface_ids[peer_rank];
|
||||
ah_attr.is_global = 1;
|
||||
ah_attr.grh.dgid = remote_gid;
|
||||
ah_attr.grh.sgid_index = gid_index_;
|
||||
// Match the legacy IBGDA path. mlx5gda previously hard-coded
|
||||
// QPC hop_limit=255 and ignored this field; after moving QP
|
||||
// setup into the device transport this value reaches hardware.
|
||||
ah_attr.grh.hop_limit = 255;
|
||||
ah_attr.port_num = 1;
|
||||
ah_attr.dlid = qps_[i]->port_attr.lid | 0xC000;
|
||||
} else {
|
||||
ah_attr.dlid = static_cast<uint16_t>(remote_lids[i]);
|
||||
ah_attr.port_num = 0;
|
||||
}
|
||||
|
||||
if (mlx5gda_modify_rc_qp_init2rtr(qps_[i], ah_attr, remote_qpns[i],
|
||||
IBV_MTU_4096)) {
|
||||
LOG(ERROR) << "[EP IBGDA] init2rtr failed for QP " << i
|
||||
<< " (roce=" << is_roce << " gid_idx=" << gid_index_
|
||||
<< " remote_qpn=" << remote_qpns[i]
|
||||
<< " udp_sport=" << ah_attr.dlid
|
||||
<< " hop_limit=" << (int)ah_attr.grh.hop_limit
|
||||
<< ")";
|
||||
return -1;
|
||||
}
|
||||
if (mlx5gda_modify_rc_qp_rtr2rts(qps_[i])) {
|
||||
LOG(ERROR) << "[EP IBGDA] rtr2rts failed for QP " << i;
|
||||
return -1;
|
||||
}
|
||||
}
|
||||
|
||||
// Populate device-visible raddrs/rkeys tables
|
||||
for (int i = 0; i < num_ranks_; ++i) {
|
||||
if (active_ranks_mask[i] == 0) continue;
|
||||
uint64_t raddr = static_cast<uint64_t>(remote_addrs[i]);
|
||||
uint32_t rkey = (i == local_rank)
|
||||
? static_cast<uint32_t>(mr_->lkey)
|
||||
: static_cast<uint32_t>(remote_keys[i]);
|
||||
cudaMemcpy(static_cast<char*>(raddrs_) + i * sizeof(uint64_t),
|
||||
&raddr, sizeof(uint64_t), cudaMemcpyHostToDevice);
|
||||
cudaMemcpy(static_cast<char*>(rkeys_) + i * sizeof(uint32_t), &rkey,
|
||||
sizeof(uint32_t), cudaMemcpyHostToDevice);
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
RdmaLocalMetadata localMetadata() const override {
|
||||
RdmaLocalMetadata meta;
|
||||
meta.raddr = mr_ ? reinterpret_cast<int64_t>(mr_->addr) : 0;
|
||||
meta.rkey = mr_ ? static_cast<int32_t>(mr_->rkey) : 0;
|
||||
meta.subnet_prefix = static_cast<int64_t>(gid_.global.subnet_prefix);
|
||||
meta.interface_id = static_cast<int64_t>(gid_.global.interface_id);
|
||||
for (auto* qp : qps_) {
|
||||
meta.qpns.push_back(static_cast<int32_t>(qp->qpn));
|
||||
meta.lids.push_back(static_cast<int32_t>(lid_));
|
||||
}
|
||||
return meta;
|
||||
}
|
||||
|
||||
void* raddrsPtr() override { return raddrs_; }
|
||||
void* rkeysPtr() override { return rkeys_; }
|
||||
void* qpDevCtxsPtr() override { return qp_devctxs_; }
|
||||
bool isRoce() const override { return is_roce_; }
|
||||
int gidIndex() const override { return gid_index_; }
|
||||
|
||||
private:
|
||||
void teardown() {
|
||||
for (auto* qp : qps_) {
|
||||
if (qp) mlx5gda_destroy_qp(ctrl_buf_heap_, qp);
|
||||
}
|
||||
qps_.clear();
|
||||
if (ctrl_buf_heap_) {
|
||||
memheap_destroy(ctrl_buf_heap_);
|
||||
ctrl_buf_heap_ = nullptr;
|
||||
}
|
||||
if (ctrl_buf_umem_) {
|
||||
mlx5dv_devx_umem_dereg(ctrl_buf_umem_);
|
||||
ctrl_buf_umem_ = nullptr;
|
||||
}
|
||||
if (ctrl_buf_) {
|
||||
cudaFree(ctrl_buf_);
|
||||
ctrl_buf_ = nullptr;
|
||||
}
|
||||
if (mr_) {
|
||||
ibv_dereg_mr(mr_);
|
||||
mr_ = nullptr;
|
||||
}
|
||||
if (raddrs_) {
|
||||
cudaFree(raddrs_);
|
||||
raddrs_ = nullptr;
|
||||
}
|
||||
if (rkeys_) {
|
||||
cudaFree(rkeys_);
|
||||
rkeys_ = nullptr;
|
||||
}
|
||||
if (qp_devctxs_) {
|
||||
cudaFree(qp_devctxs_);
|
||||
qp_devctxs_ = nullptr;
|
||||
}
|
||||
if (pd_) {
|
||||
ibv_dealloc_pd(pd_);
|
||||
pd_ = nullptr;
|
||||
}
|
||||
if (ctx_) {
|
||||
ibv_close_device(ctx_);
|
||||
ctx_ = nullptr;
|
||||
}
|
||||
}
|
||||
|
||||
// IB resources
|
||||
ibv_context* ctx_ = nullptr;
|
||||
ibv_pd* pd_ = nullptr;
|
||||
mlx5dv_pd mpd_{};
|
||||
ibv_mr* mr_ = nullptr;
|
||||
void* mr_ptr_ = nullptr;
|
||||
ibv_gid gid_{};
|
||||
int gid_index_ = -1;
|
||||
uint16_t lid_ = 0;
|
||||
bool is_roce_ = false;
|
||||
std::string device_name_;
|
||||
std::vector<std::string> device_filter_;
|
||||
|
||||
// Control buffer
|
||||
void* ctrl_buf_ = nullptr; // GPU VA
|
||||
mlx5dv_devx_umem* ctrl_buf_umem_ = nullptr;
|
||||
memheap* ctrl_buf_heap_ = nullptr;
|
||||
|
||||
// QPs
|
||||
std::vector<mlx5gda_qp*> qps_;
|
||||
int num_ranks_ = 0;
|
||||
int num_qps_ = 0;
|
||||
|
||||
// Device-visible tables
|
||||
void* raddrs_ = nullptr;
|
||||
void* rkeys_ = nullptr;
|
||||
void* qp_devctxs_ = nullptr;
|
||||
};
|
||||
|
||||
std::unique_ptr<RdmaTransport> createIbgdaDeviceTransport(
|
||||
const std::vector<std::string>& device_filter) {
|
||||
return std::make_unique<IbgdaDeviceTransportImpl>(device_filter);
|
||||
}
|
||||
|
||||
} // namespace device
|
||||
} // namespace mooncake
|
||||
|
|
@ -1,14 +1,14 @@
|
|||
#include <cmath>
|
||||
|
||||
#include <cuda_runtime_api.h>
|
||||
#include "cuda_alike.h"
|
||||
|
||||
#include <infiniband/verbs.h>
|
||||
#include <infiniband/mlx5dv.h>
|
||||
|
||||
#include <mooncake_ibgda/memheap.h>
|
||||
#include <mooncake_ibgda/mlx5gda.h>
|
||||
#include <mooncake_ibgda/mlx5_ifc.h>
|
||||
#include <mooncake_ibgda/mlx5_prm.h>
|
||||
#include <transport/device/ibgda/memheap.h>
|
||||
#include <transport/device/ibgda/mlx5gda.h>
|
||||
#include <transport/device/ibgda/mlx5_ifc.h>
|
||||
#include <transport/device/ibgda/mlx5_prm.h>
|
||||
|
||||
template <typename T>
|
||||
inline T IBGDA_ILOG2(T _n) {
|
||||
|
|
@ -30,10 +30,22 @@ constexpr T round_up_pow2(T n) {
|
|||
#define IBGDA_ROUND_UP_POW2_OR_0(_n) (((_n) == 0) ? 0 : round_up_pow2(_n))
|
||||
|
||||
static void print_cuda_error(const char* msg) {
|
||||
#ifdef USE_MUSA
|
||||
const char* err_str = musaGetErrorString(musaGetLastError());
|
||||
#else
|
||||
const char* err_str = cudaGetErrorString(cudaGetLastError());
|
||||
#endif
|
||||
fprintf(stderr, "%s: %s\n", msg, err_str);
|
||||
}
|
||||
|
||||
// Create UAR for BF (Blue Flame) doorbell ringing.
|
||||
// On CUDA: registers the BF MMIO region into GPU address space so the
|
||||
// GPU kernel can directly write the doorbell (lowest latency).
|
||||
// On MUSA: musaHostRegisterIoMemory is not supported for MMIO addresses,
|
||||
// so we skip BF registration and return a UAR with reg_addr=NULL.
|
||||
// The GPU kernel will use DBR-only mode (write to memory-mapped
|
||||
// doorbell record, NIC polls it) — slightly higher latency but
|
||||
// functionally correct.
|
||||
static struct mlx5dv_devx_uar* create_uar(struct ibv_context* ctx) {
|
||||
struct mlx5dv_devx_uar* uar =
|
||||
mlx5dv_devx_alloc_uar(ctx, MLX5DV_UAR_ALLOC_TYPE_BF);
|
||||
|
|
@ -41,6 +53,15 @@ static struct mlx5dv_devx_uar* create_uar(struct ibv_context* ctx) {
|
|||
errno = EIO;
|
||||
return NULL;
|
||||
}
|
||||
#ifdef USE_MUSA
|
||||
// MUSA cannot map MMIO addresses into GPU VA. Skip the
|
||||
// musaHostRegister(IoMemory) call entirely — attempting it
|
||||
// corrupts the MUSA runtime, causing all subsequent device-side
|
||||
// fill operations to fail with "illegal memory access".
|
||||
// Use DBR-only mode: the kernel writes to the doorbell record
|
||||
// in GPU memory instead of the BF MMIO register.
|
||||
uar->reg_addr = NULL;
|
||||
#else
|
||||
if (cudaHostRegister(uar->reg_addr, MLX5GDA_BF_SIZE * 2,
|
||||
cudaHostRegisterPortable | cudaHostRegisterMapped |
|
||||
cudaHostRegisterIoMemory) != cudaSuccess) {
|
||||
|
|
@ -49,13 +70,16 @@ static struct mlx5dv_devx_uar* create_uar(struct ibv_context* ctx) {
|
|||
mlx5dv_devx_free_uar(uar);
|
||||
return NULL;
|
||||
}
|
||||
#endif
|
||||
return uar;
|
||||
}
|
||||
|
||||
static void destroy_uar(struct mlx5dv_devx_uar* uar) {
|
||||
if (!uar) return;
|
||||
if (cudaHostUnregister(uar->reg_addr) != cudaSuccess) {
|
||||
print_cuda_error("Failed to unregister MMIO memory");
|
||||
if (uar->reg_addr) {
|
||||
if (cudaHostUnregister(uar->reg_addr) != cudaSuccess) {
|
||||
print_cuda_error("Failed to unregister MMIO memory");
|
||||
}
|
||||
}
|
||||
mlx5dv_devx_free_uar(uar);
|
||||
}
|
||||
|
|
@ -444,7 +468,8 @@ int mlx5gda_modify_rc_qp_init2rtr(struct mlx5gda_qp* qp,
|
|||
|
||||
memcpy(DEVX_ADDR_OF(qpc, qpc, primary_address_path.rmac_47_32),
|
||||
&dah.av->rmac, sizeof(dah.av->rmac));
|
||||
DEVX_SET(qpc, qpc, primary_address_path.hop_limit, 255);
|
||||
DEVX_SET(qpc, qpc, primary_address_path.hop_limit,
|
||||
ah_attr.grh.hop_limit);
|
||||
DEVX_SET(qpc, qpc, primary_address_path.src_addr_index,
|
||||
ah_attr.grh.sgid_index);
|
||||
DEVX_SET(qpc, qpc, primary_address_path.udp_sport, ah_attr.dlid);
|
||||
|
|
@ -0,0 +1,268 @@
|
|||
// Copyright 2024 KVCache.AI
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
// P2P device transport — unified NVLink (CUDA) + MTLink (MUSA) implementation.
|
||||
//
|
||||
// Uses cuda_alike.h so all cuda* APIs map to musa* when USE_MUSA is defined.
|
||||
// No #ifdef USE_MUSA / MOONCAKE_EP_USE_MUSA in this file.
|
||||
|
||||
#include "transport/device/device_transport.h"
|
||||
|
||||
#include <glog/logging.h>
|
||||
#include <algorithm>
|
||||
#include <cstring>
|
||||
|
||||
#include "cuda_alike.h"
|
||||
|
||||
namespace mooncake {
|
||||
namespace device {
|
||||
|
||||
class P2pDeviceTransportImpl : public P2pTransport {
|
||||
public:
|
||||
explicit P2pDeviceTransportImpl(int num_ranks) : num_ranks_(num_ranks) {
|
||||
cudaMalloc(&available_table_, num_ranks_ * sizeof(int32_t));
|
||||
cudaMemset(available_table_, 0, num_ranks_ * sizeof(int32_t));
|
||||
cudaMallocHost(&peer_ptrs_host_, num_ranks_ * sizeof(void*));
|
||||
cudaMalloc(&peer_ptrs_dev_, num_ranks_ * sizeof(void*));
|
||||
for (int i = 0; i < num_ranks_; ++i) peer_ptrs_host_[i] = nullptr;
|
||||
cudaMemset(peer_ptrs_dev_, 0, num_ranks_ * sizeof(void*));
|
||||
}
|
||||
|
||||
~P2pDeviceTransportImpl() override {
|
||||
if (available_table_) cudaFree(available_table_);
|
||||
if (peer_ptrs_dev_) cudaFree(peer_ptrs_dev_);
|
||||
if (peer_ptrs_host_) {
|
||||
for (int i = 0; i < num_ranks_; ++i) {
|
||||
if (peer_ptrs_host_[i] && peer_ptrs_host_[i] != local_ptr_) {
|
||||
cudaIpcCloseMemHandle(peer_ptrs_host_[i]);
|
||||
}
|
||||
}
|
||||
cudaFreeHost(peer_ptrs_host_);
|
||||
}
|
||||
}
|
||||
|
||||
void* allocateBuffer(size_t bytes) override {
|
||||
void* ptr = nullptr;
|
||||
cudaError_t err = cudaMalloc(&ptr, bytes);
|
||||
if (err != cudaSuccess) {
|
||||
LOG(ERROR) << "[EP P2P] cudaMalloc(" << bytes
|
||||
<< ") failed: " << cudaGetErrorString(err);
|
||||
return nullptr;
|
||||
}
|
||||
return ptr;
|
||||
}
|
||||
|
||||
void freeBuffer(void* ptr) override { cudaFree(ptr); }
|
||||
|
||||
std::vector<int32_t> exportIpcHandle(void* ptr) override {
|
||||
cudaIpcMemHandle_t handle;
|
||||
cudaError_t err = cudaIpcGetMemHandle(&handle, ptr);
|
||||
if (err != cudaSuccess) {
|
||||
LOG(ERROR) << "[EP P2P] cudaIpcGetMemHandle failed: "
|
||||
<< cudaGetErrorString(err);
|
||||
return {};
|
||||
}
|
||||
constexpr size_t kHandleBytes = sizeof(cudaIpcMemHandle_t);
|
||||
constexpr size_t kNumInt32s =
|
||||
(kHandleBytes + sizeof(int32_t) - 1) / sizeof(int32_t);
|
||||
std::vector<int32_t> result(kNumInt32s);
|
||||
memcpy(result.data(), &handle, kHandleBytes);
|
||||
return result;
|
||||
}
|
||||
|
||||
void importPeerHandles(
|
||||
void* local_ptr, int rank, int num_ranks,
|
||||
const std::vector<std::vector<int32_t>>& remote_handles,
|
||||
const std::vector<int>& active_ranks_mask) override {
|
||||
local_ptr_ = local_ptr;
|
||||
int device_id = 0;
|
||||
cudaGetDevice(&device_id);
|
||||
int device_count = 0;
|
||||
cudaGetDeviceCount(&device_count);
|
||||
CHECK_GT(device_count, 0) << "No CUDA/MUSA devices found";
|
||||
|
||||
std::vector<int32_t> available(num_ranks_, 0);
|
||||
available[rank] = 1;
|
||||
peer_ptrs_host_[rank] = local_ptr;
|
||||
|
||||
int node_id = rank / device_count;
|
||||
int group_start = node_id * device_count;
|
||||
int group_end = std::min(group_start + device_count, num_ranks_);
|
||||
|
||||
for (int dst = group_start; dst < group_end; ++dst) {
|
||||
if (active_ranks_mask[dst] == 0) continue;
|
||||
if (dst == rank) continue;
|
||||
|
||||
int dst_device = dst % device_count;
|
||||
int can_access = 0;
|
||||
cudaDeviceCanAccessPeer(&can_access, device_id, dst_device);
|
||||
LOG(INFO) << "[EP P2P] rank " << rank << " (device " << device_id
|
||||
<< ") -> rank " << dst << " (device " << dst_device
|
||||
<< "): canAccessPeer=" << can_access;
|
||||
if (!can_access) continue;
|
||||
|
||||
cudaError_t err = cudaDeviceEnablePeerAccess(dst_device, 0);
|
||||
if (err != cudaSuccess &&
|
||||
err != cudaErrorPeerAccessAlreadyEnabled) {
|
||||
LOG(WARNING) << "[EP P2P] rank " << rank
|
||||
<< " failed to enable peer access to device "
|
||||
<< dst_device << ": " << cudaGetErrorString(err);
|
||||
continue;
|
||||
}
|
||||
if (err == cudaErrorPeerAccessAlreadyEnabled) cudaGetLastError();
|
||||
|
||||
if (dst >= static_cast<int>(remote_handles.size())) continue;
|
||||
const auto& h = remote_handles[dst];
|
||||
if (h.empty()) continue;
|
||||
|
||||
constexpr size_t kHandleBytes = sizeof(cudaIpcMemHandle_t);
|
||||
constexpr size_t kNumInt32s =
|
||||
(kHandleBytes + sizeof(int32_t) - 1) / sizeof(int32_t);
|
||||
if (h.size() < kNumInt32s) continue;
|
||||
|
||||
cudaIpcMemHandle_t handle;
|
||||
memcpy(&handle, h.data(), kHandleBytes);
|
||||
|
||||
void* peer_ptr = nullptr;
|
||||
err = cudaIpcOpenMemHandle(&peer_ptr, handle,
|
||||
cudaIpcMemLazyEnablePeerAccess);
|
||||
if (err != cudaSuccess) {
|
||||
LOG(WARNING) << "[EP P2P] rank " << rank
|
||||
<< " failed to open IPC handle for rank " << dst
|
||||
<< ": " << cudaGetErrorString(err);
|
||||
continue;
|
||||
}
|
||||
LOG(INFO) << "[EP P2P] rank " << rank
|
||||
<< " opened IPC handle for rank " << dst
|
||||
<< ": peer_ptr=" << peer_ptr;
|
||||
available[dst] = 1;
|
||||
peer_ptrs_host_[dst] = peer_ptr;
|
||||
}
|
||||
|
||||
// Determine if all active ranks have P2P
|
||||
all_peers_accessible_ = true;
|
||||
for (int i = 0; i < num_ranks_; ++i) {
|
||||
if (active_ranks_mask[i] == 0) continue;
|
||||
if (!available[i] || !peer_ptrs_host_[i]) {
|
||||
all_peers_accessible_ = false;
|
||||
break;
|
||||
}
|
||||
}
|
||||
// Multi-node: P2P only within a node
|
||||
if (all_peers_accessible_ && num_ranks_ > 1) {
|
||||
int first_node = 0 / device_count;
|
||||
int last_node = (num_ranks_ - 1) / device_count;
|
||||
if (first_node != last_node) all_peers_accessible_ = false;
|
||||
}
|
||||
|
||||
cudaMemcpy(available_table_, available.data(),
|
||||
num_ranks_ * sizeof(int32_t), cudaMemcpyHostToDevice);
|
||||
cudaMemcpy(peer_ptrs_dev_, peer_ptrs_host_, num_ranks_ * sizeof(void*),
|
||||
cudaMemcpyHostToDevice);
|
||||
}
|
||||
|
||||
int32_t* availableTablePtr() override { return available_table_; }
|
||||
void** peerPtrsTablePtr() override { return peer_ptrs_dev_; }
|
||||
bool allPeersAccessible() const override { return all_peers_accessible_; }
|
||||
|
||||
bool verifyPeerAccess() override {
|
||||
if (!all_peers_accessible_) return false;
|
||||
|
||||
int device_id = 0;
|
||||
cudaGetDevice(&device_id);
|
||||
|
||||
bool all_ok = true;
|
||||
for (int i = 0; i < num_ranks_; ++i) {
|
||||
if (i == device_id) continue;
|
||||
if (!peer_ptrs_host_[i] || peer_ptrs_host_[i] == local_ptr_)
|
||||
continue;
|
||||
|
||||
// Test: write a pattern to the peer buffer via cudaMemcpy, then
|
||||
// read it back. This verifies the IPC mapping is writable.
|
||||
constexpr int kTestBytes = 256;
|
||||
std::vector<uint8_t> pattern(kTestBytes);
|
||||
for (int j = 0; j < kTestBytes; ++j)
|
||||
pattern[j] = (uint8_t)(j ^ 0xA5);
|
||||
|
||||
cudaError_t err = cudaMemcpy(peer_ptrs_host_[i], pattern.data(),
|
||||
kTestBytes, cudaMemcpyHostToDevice);
|
||||
if (err != cudaSuccess) {
|
||||
LOG(WARNING) << "[EP P2P] verifyPeerAccess: rank " << device_id
|
||||
<< " cannot write to peer " << i
|
||||
<< " mapped buffer: " << cudaGetErrorString(err);
|
||||
all_ok = false;
|
||||
continue;
|
||||
}
|
||||
|
||||
// Read back and verify
|
||||
std::vector<uint8_t> readback(kTestBytes, 0);
|
||||
err = cudaMemcpy(readback.data(), peer_ptrs_host_[i], kTestBytes,
|
||||
cudaMemcpyDeviceToHost);
|
||||
if (err != cudaSuccess) {
|
||||
LOG(WARNING) << "[EP P2P] verifyPeerAccess: rank " << device_id
|
||||
<< " cannot read back from peer " << i
|
||||
<< " mapped buffer: " << cudaGetErrorString(err);
|
||||
all_ok = false;
|
||||
continue;
|
||||
}
|
||||
|
||||
bool match =
|
||||
(memcmp(readback.data(), pattern.data(), kTestBytes) == 0);
|
||||
if (!match) {
|
||||
LOG(WARNING) << "[EP P2P] verifyPeerAccess: rank " << device_id
|
||||
<< " readback mismatch from peer " << i;
|
||||
all_ok = false;
|
||||
} else {
|
||||
LOG(INFO) << "[EP P2P] verifyPeerAccess: rank " << device_id
|
||||
<< " peer " << i << " OK (memcpy write/read)";
|
||||
}
|
||||
}
|
||||
|
||||
if (!all_ok) {
|
||||
all_peers_accessible_ = false;
|
||||
// Update device table to reflect failure
|
||||
std::vector<int32_t> avail_h(num_ranks_, 0);
|
||||
cudaMemcpy(avail_h.data(), available_table_,
|
||||
num_ranks_ * sizeof(int32_t), cudaMemcpyDeviceToHost);
|
||||
for (int i = 0; i < num_ranks_; ++i) {
|
||||
if (i == device_id) continue;
|
||||
if (!peer_ptrs_host_[i] || peer_ptrs_host_[i] == local_ptr_)
|
||||
continue;
|
||||
// Leave self as available; only clear failed peers
|
||||
}
|
||||
// Re-upload with all peers marked unavailable except self
|
||||
std::vector<int32_t> cleared(num_ranks_, 0);
|
||||
cleared[device_id] = 1;
|
||||
cudaMemcpy(available_table_, cleared.data(),
|
||||
num_ranks_ * sizeof(int32_t), cudaMemcpyHostToDevice);
|
||||
}
|
||||
|
||||
return all_ok;
|
||||
}
|
||||
|
||||
private:
|
||||
int num_ranks_;
|
||||
void* local_ptr_ = nullptr;
|
||||
int32_t* available_table_ = nullptr;
|
||||
void** peer_ptrs_host_ = nullptr;
|
||||
void** peer_ptrs_dev_ = nullptr;
|
||||
bool all_peers_accessible_ = false;
|
||||
};
|
||||
|
||||
std::unique_ptr<P2pTransport> createP2pDeviceTransport(int num_ranks) {
|
||||
return std::make_unique<P2pDeviceTransportImpl>(num_ranks);
|
||||
}
|
||||
|
||||
} // namespace device
|
||||
} // namespace mooncake
|
||||
Loading…
Reference in New Issue