diff --git a/mooncake-pg/include/connection_poller.h b/mooncake-pg/include/connection_poller.h index 4319aef8..f9f364cd 100644 --- a/mooncake-pg/include/connection_poller.h +++ b/mooncake-pg/include/connection_poller.h @@ -68,8 +68,12 @@ class ConnectionContext { PeerConnection peerStates_[kMaxNumRanks]; + // On MNNVL, warmup is skipped because CPU heap buffers aren't + // fabric-accessible for cross-node NVLink writes. + bool skip_warmup_; + // warmup_send_region_ and warmup_recv_region_ are managed by - // ConnectionContext. + // ConnectionContext. nullptr when skip_warmup_ is true. int32_t* warmup_send_region_; int32_t* warmup_recv_region_; diff --git a/mooncake-pg/setup.py b/mooncake-pg/setup.py index dbfc1a33..8a393773 100644 --- a/mooncake-pg/setup.py +++ b/mooncake-pg/setup.py @@ -3,7 +3,7 @@ import re from setuptools import setup import torch -from torch.utils.cpp_extension import BuildExtension, CUDAExtension +from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CUDA_HOME torch_version = re.match(r"\d+(?:\.\d+)*", torch.__version__).group() @@ -13,6 +13,18 @@ module_name = "mooncake.pg" + version_suffix abi_flag = int(torch._C._GLIBCXX_USE_CXX11_ABI) current_dir = os.path.abspath(os.path.dirname(__file__)) +# Link against the CUDA driver stub library if available. +# Same approach as mooncake-ep/setup.py. +cuda_libraries = ["ibverbs", "mlx5"] +cuda_library_dirs = [] + +if CUDA_HOME is not None: + cuda_stub_dir = os.path.join(CUDA_HOME, "lib64", "stubs") + cuda_stub_lib = os.path.join(cuda_stub_dir, "libcuda.so") + if os.path.exists(cuda_stub_lib): + cuda_libraries.insert(0, "cuda") + cuda_library_dirs.append(cuda_stub_dir) + setup( name=module_name, @@ -47,7 +59,8 @@ setup( "-g0", ], }, - libraries=["ibverbs", "mlx5"], + libraries=cuda_libraries, + library_dirs=cuda_library_dirs, extra_link_args=[ "-Wl,-rpath,$ORIGIN", "-L" + os.path.join(current_dir, "../mooncake-wheel/mooncake"), diff --git a/mooncake-pg/src/connection_poller.cpp b/mooncake-pg/src/connection_poller.cpp index 41fd55a7..9b1563d3 100644 --- a/mooncake-pg/src/connection_poller.cpp +++ b/mooncake-pg/src/connection_poller.cpp @@ -1,6 +1,7 @@ #include #include #include +#include #include #include #include @@ -9,10 +10,34 @@ #include #include #include +#include #include #include "mooncake_worker.cuh" namespace mooncake { + +// Same check as nvlink_transport.cpp and mooncake_ep_buffer.cpp. +// On MNNVL clusters all GPUs support fabric mem handles, meaning +// NVLink transport can only access cuMemCreate(FABRIC) memory +// cross-node -- CPU heap buffers are invisible to remote peers. +static bool supportFabricMem() { + const char* nvlink_ipc = getenv("MC_USE_NVLINK_IPC"); + + bool fabric_enabled = nvlink_ipc && strcmp(nvlink_ipc, "0") == 0; + if (!fabric_enabled) return false; + + int num_devices = 0; + cudaError_t err = cudaGetDeviceCount(&num_devices); + if (err != cudaSuccess || num_devices == 0) return false; + + for (int dev = 0; dev < num_devices; ++dev) { + int supported = 0; + cuDeviceGetAttribute( + &supported, CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED, dev); + if (!supported) return false; + } + return true; +} ConnectionContext::ConnectionContext(int backendIndex, int rank, int size, uint64_t* local2global_rank_map, c10::intrusive_ptr<::c10d::Store> store, @@ -26,7 +51,17 @@ ConnectionContext::ConnectionContext(int backendIndex, int rank, int size, store_(std::move(store)), meta_(std::move(meta)), p2p_proxy_(std::move(p2p_proxy)), - engine_(engine) { + engine_(engine), + skip_warmup_(supportFabricMem()) { + if (skip_warmup_) { + // On MNNVL clusters, CPU heap buffers aren't fabric-accessible so + // remote NVLink writes to them will fail. The fabric topology already + // guarantees connectivity, so we skip the warmup handshake entirely. + warmup_send_region_ = nullptr; + warmup_recv_region_ = nullptr; + return; + } + warmup_send_region_ = new int32_t[kMaxNumRanks]; warmup_send_region_[0] = 1; int rc = engine_->registerLocalMemory( @@ -46,10 +81,14 @@ ConnectionContext::~ConnectionContext() { } } - engine_->unregisterLocalMemory(warmup_send_region_); - engine_->unregisterLocalMemory(warmup_recv_region_); - delete[] warmup_send_region_; - delete[] warmup_recv_region_; + if (warmup_send_region_) { + engine_->unregisterLocalMemory(warmup_send_region_); + delete[] warmup_send_region_; + } + if (warmup_recv_region_) { + engine_->unregisterLocalMemory(warmup_recv_region_); + delete[] warmup_recv_region_; + } } void ConnectionContext::waitUntilAllConnected() { @@ -138,7 +177,19 @@ bool ConnectionContext::pollPeer(int pollingRank) { memcpy(&meta_->segmentInfos[pollingRank], buffer_data.data(), sizeof(SegmentInfo)); - if (pollingRank <= rank_) { + if (skip_warmup_) { + // MNNVL: fabric guarantees connectivity, skip warmup write + // since CPU heap buffers aren't fabric-accessible anyway. + meta_->peerConnected[pollingRank] = true; + global_peerConnected_[globalPollingRank] = true; + peerState.state = PeerConnectionState::CONNECTED; + { + std::lock_guard lock(backend_wakeup_mutex_); + totalConnectedPeers_.fetch_add(1, + std::memory_order_release); + if (isAllPeerConnected()) backend_wakeup_cv_.notify_all(); + } + } else if (pollingRank <= rank_) { // Send a warmup request to establish connections auto batchID = engine_->allocateBatchID(1); engine_->submitTransfer(