Fix MNNVL warmup hang: skip warmup when fabric mem is available (#1644)

* Fix MNNVL warmup hang: skip warmup when fabric mem is available

On GB200 MNNVL clusters, the warmup handshake in ConnectionContext
allocates send/recv buffers from the CPU heap. The NVLink transport
can only access cuMemCreate(CU_MEM_HANDLE_TYPE_FABRIC) memory
cross-node, so remote writes to these heap buffers silently fail
and the state machine retries forever.

Since MNNVL fabric guarantees connectivity between all peers in a
ComputeDomain, we can safely skip the warmup write entirely when
supportFabricMem() is true. The store key exchange alone is
sufficient proof of peer reachability.

Changes:
- Add supportFabricMem() to connection_poller.cpp (same check as
  nvlink_transport.cpp and mooncake_ep_buffer.cpp)
- Skip warmup buffer allocation and warmup write when on MNNVL,
  transition directly to CONNECTED after opening the segment
- Guard destructor against null warmup buffers
- Link CUDA driver library in setup.py for cuDeviceGetAttribute

Fixes #1639

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_;
 

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"),

mooncake-pg/src/connection_poller.cpp

@@ -1,6 +1,7 @@
 #include <c10/util/Exception.h>
 #include <connection_poller.h>
 #include <ATen/cuda/CUDAContext.h>
+#include <cuda.h>
 #include <cuda_runtime.h>
 #include <torch/torch.h>
 #include <atomic>
@@ -9,10 +10,34 @@
 #include <thread>
 #include <torch/csrc/distributed/c10d/Backend.hpp>
 #include <algorithm>
+#include <cstring>
 #include <limits>
 #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<std::mutex> 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(