[ CRL] Kernel proxy bypass RMA Proxy with direct netAdaptor posting (flagos-ai#514)

Author: MC952-arch

flagcx/adaptor/device/cuda_adaptor.cc

@@ -5,7 +5,12 @@
 #include "adaptor.h"
 #include "alloc.h"
 #include "param.h"
+#include <mutex>
 #include <unistd.h>
+#include <unordered_map>
+
+static std::mutex gVmmHandleMapMtx;
+static std::unordered_map<void *, CUmemGenericAllocationHandle> gVmmHandleMap;
 
 std::map<flagcxMemcpyType_t, cudaMemcpyKind> memcpy_type_map = {
     {flagcxMemcpyHostToDevice, cudaMemcpyHostToDevice},
@@ -145,22 +150,31 @@ flagcxResult_t cudaAdaptorGdrMemAlloc(void **ptr, size_t size,
   DEVCHECK(cuDeviceGetAttribute(
       &flag, CU_DEVICE_ATTRIBUTE_GPU_DIRECT_RDMA_WITH_CUDA_VMM_SUPPORTED,
       currentDev));
+  INFO(FLAGCX_INIT,
+       "[gdrMemAlloc] dev=%d GPU_DIRECT_RDMA_WITH_CUDA_VMM_SUPPORTED=%d "
+       "size=%zu",
+       cudaDev, flag, size);
   if (flag)
     memprop.allocFlags.gpuDirectRDMACapable = 1;
   DEVCHECK(cuMemGetAllocationGranularity(&memGran, &memprop,
                                          CU_MEM_ALLOC_GRANULARITY_RECOMMENDED));
   ALIGN_SIZE(handleSize, memGran);
+  INFO(FLAGCX_INIT,
+       "[gdrMemAlloc] memGran=%zu handleSize=%zu gpuDirectRDMACapable=%d",
+       memGran, handleSize, (int)memprop.allocFlags.gpuDirectRDMACapable);
   /* Allocate the physical memory on the device */
   DEVCHECK(cuMemCreate(&handle, handleSize, &memprop, 0));
   /* Reserve a virtual address range */
   cuRes = cuMemAddressReserve((CUdeviceptr *)ptr, handleSize, memGran, 0, 0);
   if (cuRes != CUDA_SUCCESS) {
+    WARN("[gdrMemAlloc] cuMemAddressReserve FAILED: %d", (int)cuRes);
     cuMemRelease(handle);
     return flagcxUnhandledDeviceError;
   }
   /* Map the virtual address range to the physical allocation */
   cuRes = cuMemMap((CUdeviceptr)*ptr, handleSize, 0, handle, 0);
   if (cuRes != CUDA_SUCCESS) {
+    WARN("[gdrMemAlloc] cuMemMap FAILED: %d", (int)cuRes);
     cuMemAddressFree((CUdeviceptr)*ptr, handleSize);
     cuMemRelease(handle);
     *ptr = NULL;
@@ -173,14 +187,21 @@ flagcxResult_t cudaAdaptorGdrMemAlloc(void **ptr, size_t size,
   accessDesc.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
   cuRes = cuMemSetAccess((CUdeviceptr)*ptr, handleSize, &accessDesc, 1);
   if (cuRes != CUDA_SUCCESS) {
+    WARN("[gdrMemAlloc] cuMemSetAccess FAILED: %d", (int)cuRes);
     cuMemUnmap((CUdeviceptr)*ptr, handleSize);
     cuMemAddressFree((CUdeviceptr)*ptr, handleSize);
     cuMemRelease(handle);
     *ptr = NULL;
     return flagcxUnhandledDeviceError;
   }
-  /* Release the create-time handle reference; the mapping holds its own. */
-  cuMemRelease(handle);
+  INFO(FLAGCX_INIT, "[gdrMemAlloc] VMM alloc OK: ptr=%p size=%zu", *ptr,
+       handleSize);
+  /* Retain the handle so cuMemGetHandleForAddressRange can export DMA-BUF fds.
+     Released in cudaAdaptorGdrMemFree. */
+  {
+    std::lock_guard<std::mutex> lk(gVmmHandleMapMtx);
+    gVmmHandleMap[*ptr] = handle;
+  }
 #else
   DEVCHECK(cudaMalloc(ptr, size));
   cudaPointerAttributes attrs;
@@ -206,6 +227,16 @@ flagcxResult_t cudaAdaptorGdrMemFree(void *ptr, void *memHandle) {
   DEVCHECK(cuMemGetAddressRange(NULL, &size, (CUdeviceptr)ptr));
   DEVCHECK(cuMemUnmap((CUdeviceptr)ptr, size));
   DEVCHECK(cuMemAddressFree((CUdeviceptr)ptr, size));
+
+  // Release the VMM handle we retained at alloc time
+  {
+    std::lock_guard<std::mutex> lk(gVmmHandleMapMtx);
+    auto it = gVmmHandleMap.find(ptr);
+    if (it != gVmmHandleMap.end()) {
+      cuMemRelease(it->second);
+      gVmmHandleMap.erase(it);
+    }
+  }
 #else
   DEVCHECK(cudaFree(ptr));
 #endif
@@ -509,8 +540,11 @@ flagcxResult_t
 cudaAdaptorMemGetHandleForAddressRange(void *handleOut, void *buffer,
                                        size_t size, unsigned long long flags) {
   CUdeviceptr dptr = (CUdeviceptr)buffer;
-  DEVCHECK(cuMemGetHandleForAddressRange(
-      handleOut, dptr, size, CU_MEM_RANGE_HANDLE_TYPE_DMA_BUF_FD, flags));
+  CUresult err = cuMemGetHandleForAddressRange(
+      handleOut, dptr, size, CU_MEM_RANGE_HANDLE_TYPE_DMA_BUF_FD, flags);
+  if (err != CUDA_SUCCESS) {
+    return flagcxUnhandledDeviceError;
+  }
   return flagcxSuccess;
 }
 
@@ -544,7 +578,13 @@ flagcxResult_t cudaAdaptorSymPhysAlloc(void *ptr, size_t size,
     return flagcxSystemError;
 
   // Retain the physical allocation handle from the VMM-backed pointer
-  DEVCHECK(cuMemRetainAllocationHandle(cuHandle, ptr));
+  CUresult retainRes = cuMemRetainAllocationHandle(cuHandle, ptr);
+  if (retainRes != CUDA_SUCCESS) {
+    WARN("[symPhysAlloc] cuMemRetainAllocationHandle FAILED: %d ptr=%p",
+         (int)retainRes, ptr);
+    free(cuHandle);
+    return flagcxUnhandledDeviceError;
+  }
 
   // Discover actual physical allocation size (already granularity-aligned)
   size_t actualAllocSize = 0;
@@ -556,8 +596,16 @@ flagcxResult_t cudaAdaptorSymPhysAlloc(void *ptr, size_t size,
     free(cuHandle);
     return flagcxInvalidArgument;
   }
-  DEVCHECK(cuMemExportToShareableHandle(
-      shareableHandle, *cuHandle, CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR, 0));
+  CUresult exportRes = cuMemExportToShareableHandle(
+      shareableHandle, *cuHandle, CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR, 0);
+  if (exportRes != CUDA_SUCCESS) {
+    WARN("[symPhysAlloc] cuMemExportToShareableHandle FAILED: %d",
+         (int)exportRes);
+    free(cuHandle);
+    return flagcxUnhandledDeviceError;
+  }
+  INFO(FLAGCX_INIT, "[symPhysAlloc] ptr=%p allocSize=%zu fd=%d", ptr,
+       actualAllocSize, *(int *)shareableHandle);
   *handleSize = sizeof(int); // POSIX fd is an int
   *physHandle = cuHandle;
   return flagcxSuccess;

flagcx/adaptor/flagcx_device.cc

@@ -1008,7 +1008,7 @@ extern "C" flagcxResult_t flagcxDevCommDestroy(flagcxComm_t comm,
   if (comm != nullptr && comm->heteroComm != nullptr &&
       comm->heteroComm->devCommHandle == devComm) {
     if (devComm->signalBuffer) {
-      flagcxOneSideSignalDeregister(comm->heteroComm);
+      flagcxOneSideSignalDeregister(comm);
     }
     comm->heteroComm->devCommHandle = nullptr;
   }
@@ -1244,11 +1244,10 @@ extern "C" flagcxResult_t flagcxDevMemDestroy(flagcxComm_t comm,
     return flagcxSuccess;
   }
 
-  // Mark IPC table entry as no longer in use (actual cleanup deferred to
-  // flagcxCommDestroy.
-  if (comm != nullptr && devMem->ipcIndex >= 0 &&
-      devMem->ipcIndex < FLAGCX_MAX_IPC_ENTRIES) {
-    comm->ipcTable[devMem->ipcIndex].inUse = false;
+  // Release IPC table slot (resources moved to deferred queue for cleanup at
+  // comm destroy).
+  if (devMem->ipcIndex >= 0) {
+    releaseIpcTableSlot(comm, devMem->ipcIndex);
   }
 
   // Free window allocation if present
@@ -1464,6 +1463,65 @@ flagcxResult_t flagcxCommCleanupIpcTable(flagcxComm_t comm) {
   return flagcxSuccess;
 }
 
+// ==========================================================================
+// Deferred IPC table slot release.
+// ==========================================================================
+void releaseIpcTableSlot(flagcxComm_t comm, int slot) {
+  if (comm == nullptr || slot < 0 || slot >= FLAGCX_MAX_IPC_ENTRIES)
+    return;
+  struct flagcxIpcTableEntry *e = &comm->ipcTable[slot];
+  if (e->hostPeerPtrs == nullptr && e->devPeerPtrs == nullptr) {
+    e->inUse = false;
+    return;
+  }
+
+  // Move resources to deferred linked list for cleanup at comm destroy
+  struct flagcxDeferredIpcEntry *d =
+      (struct flagcxDeferredIpcEntry *)malloc(sizeof(*d));
+  if (d == nullptr) {
+    // OOM: leave slot occupied so flagcxCommCleanupIpcTable handles it at
+    // destroy. The slot won't be reusable, but resources are still safe.
+    WARN(
+        "releaseIpcTableSlot: OOM, keeping slot %d occupied until comm destroy",
+        slot);
+    e->inUse = false;
+    return;
+  }
+  d->hostPeerPtrs = e->hostPeerPtrs;
+  d->devPeerPtrs = e->devPeerPtrs;
+  d->nPeers = e->nPeers;
+  d->basePtr = e->basePtr;
+  d->next = nullptr;
+  flagcxIntruQueueEnqueue(&comm->deferredIpcQueue, d);
+
+  // Clear slot — now reusable by buildIpcPeerPointers
+  e->hostPeerPtrs = nullptr;
+  e->devPeerPtrs = nullptr;
+  e->nPeers = 0;
+  e->basePtr = nullptr;
+  e->inUse = false;
+}
+
+flagcxResult_t flagcxCommDrainDeferredIpc(flagcxComm_t comm) {
+  if (comm == nullptr)
+    return flagcxSuccess;
+  while (!flagcxIntruQueueEmpty(&comm->deferredIpcQueue)) {
+    struct flagcxDeferredIpcEntry *d =
+        flagcxIntruQueueDequeue(&comm->deferredIpcQueue);
+    if (d->hostPeerPtrs) {
+      for (int j = 0; j < d->nPeers; j++) {
+        if (d->hostPeerPtrs[j] && d->hostPeerPtrs[j] != d->basePtr)
+          deviceAdaptor->ipcMemHandleClose(d->hostPeerPtrs[j]);
+      }
+      free(d->hostPeerPtrs);
+    }
+    if (d->devPeerPtrs)
+      deviceAdaptor->deviceFree(d->devPeerPtrs, flagcxMemDevice, NULL);
+    free(d);
+  }
+  return flagcxSuccess;
+}
+
 // ==========================================================================
 // Deferred device/host-pinned memory free.
 // ==========================================================================

flagcx/adaptor/include/device_api/default_comm_traits.h

@@ -494,8 +494,10 @@ struct CommTraits<Default<PlatformTag>> {
     // ---- MR offset helper ----
     FLAGCX_DEVICE_INLINE_DECORATOR
     static size_t toDataOffset(const Window &win, size_t off) {
-      void *ptr = win.getLocalPointer(off);
-      return (uintptr_t)ptr - win.mrBase;
+      // Use rawPtr (the original buffer VA used for MR registration) rather
+      // than getLocalPointer() which may return a VMM flat-mapped VA that
+      // differs from the MR-registered VA.
+      return (uintptr_t)win.getRawPtr() + off - win.mrBase;
     }
 
     // ---- Action decomposition helpers ----

flagcx/adaptor/net/ibrc_adaptor.cc

@@ -2681,7 +2681,6 @@ flagcxResult_t flagcxIbIputSignal(void *sendComm, uint64_t srcOff,
     void *dstPtr = (void *)(dstInfo->baseVas[dstRank] + dstOff);
     uint32_t lkey = srcInfo->lkeys[srcRank];
     uint32_t rkey = dstInfo->rkeys[dstRank];
-
     wr[0].opcode = IBV_WR_RDMA_WRITE;
     wr[0].send_flags = 0; // No CQE — only signal gets CQE
     wr[0].wr_id = req - comm->base.reqs;

flagcx/core/flagcx_hetero.cc

@@ -789,6 +789,10 @@ flagcxResult_t flagcxHeteroFlushRma(flagcxHeteroComm_t comm, int peer,
       return flagcxRemoteError;
     usleep(100);
   }
+  // Final rmaError check: kernel proxy or network failures set rmaError;
+  // catch errors that occurred after doneSeqs reached the target.
+  if (__atomic_load_n(&proxy->rmaError, __ATOMIC_ACQUIRE))
+    return flagcxRemoteError;
   return flagcxSuccess;
 }
 
@@ -827,6 +831,10 @@ flagcxResult_t flagcxHeteroFlushAllRma(flagcxHeteroComm_t comm) {
       usleep(100);
     }
   }
+  // Final rmaError check: kernel proxy or network failures set rmaError;
+  // catch errors that occurred after doneSeqs reached the target.
+  if (__atomic_load_n(&proxy->rmaError, __ATOMIC_ACQUIRE))
+    return flagcxRemoteError;
   return flagcxSuccess;
 }
 

flagcx/core/include/global_comm.h

@@ -35,6 +35,16 @@ struct flagcxDeferredFree {
   int memType; // flagcxMemDevice, flagcxMemHost, etc.
 };
 
+// Deferred IPC entry — moved from ipcTable when a slot is released at runtime.
+// Actual ipcMemHandleClose + deviceFree happens at comm destroy.
+struct flagcxDeferredIpcEntry {
+  void **hostPeerPtrs;
+  void **devPeerPtrs;
+  int nPeers;
+  void *basePtr;
+  struct flagcxDeferredIpcEntry *next;
+};
+
 // Deferred DevComm buffer handle — buffers that cannot be freed immediately
 // in flagcxDevCommDestroy because peers may still hold IPC mappings to them.
 // Drained at flagcxCommDestroy time.
@@ -110,6 +120,11 @@ struct flagcxComm {
   // IPC peer pointer table — deferred cleanup
   struct flagcxIpcTableEntry ipcTable[FLAGCX_MAX_IPC_ENTRIES];
 
+  // Deferred IPC entry queue — IPC resources moved here when slots are released
+  // at runtime; actual ipcMemHandleClose + deviceFree deferred to comm destroy.
+  flagcxIntruQueue<struct flagcxDeferredIpcEntry, &flagcxDeferredIpcEntry::next>
+      deferredIpcQueue;
+
   // Deferred DevComm buffer queue — buffers stashed here during
   // flagcxDevCommDestroy, drained at flagcxCommDestroy.
   flagcxIntruQueue<struct flagcxDevCommBufferHandle,

flagcx/core/include/onesided.h

@@ -23,14 +23,27 @@ struct flagcxOneSideHandleInfo {
   void *localMrHandle; // local rank's MR handle for deregMr
   void *localRecvComm; // recvComm used for MR registration (PD match)
   // Full-mesh IB connections (including self loopback, aligned with NCCL GIN)
-  void **fullSendComms; // [nRanks] per-peer sendComm (NULL if not owner)
-  void **fullRecvComms; // [nRanks] per-peer recvComm (NULL if not owner)
+  void **fullSendComms; // [nRanks] per-peer sendComm — alias for
+                        // contextSendComms[0]
+  void **fullRecvComms; // [nRanks] per-peer recvComm — alias for
+                        // contextRecvComms[0]
   int nRanks;           // number of ranks (for cleanup iteration)
 
+  // Per-context QP arrays for thread isolation (NCCL GIN pattern).
+  // Context 0 = RMA proxy; contexts 1..N = kernel proxy threads.
+  // Each context has its own full-mesh of RC QPs so no QP is shared
+  // across threads. All contexts share the same MR handles/rkeys (same PD).
+  void ***contextSendComms; // [nContexts][nRanks]
+  void ***contextRecvComms; // [nContexts][nRanks]
+  int nContexts;            // 1 + nKernelProxies
+
   // Symmetric memory window for intra-node D2D bypass (CE path).
   // NULL if VMM not available or window not registered with
   // FLAGCX_WIN_COLL_SYMMETRIC.
   struct flagcxSymWindow *symWin;
+
+  // IPC table slot index for signal buffer D2D bypass (-1 if none).
+  int signalIpcSlot;
 };
 
 // Internal implementation used by sym_heap and flagcxCommRegister

flagcx/core/include/proxy.h

@@ -30,13 +30,19 @@ enum flagcxProxyOpState {
 struct flagcxProxyKernelState {
   pthread_t threads[FLAGCX_DEVICE_CTA_COUNT];
   flagcxFifo_t fifos[FLAGCX_DEVICE_CTA_COUNT];
-  int contextCount = 1;
+  int contextCount = 0;
   flagcxStream_t stream;
   int stop = 0;
   // Synchronization for initialization
   pthread_mutex_t initMutex;
   pthread_cond_t initCond;
   int ready = 0;
+  int initFailed = 0;
+  // Shared per-peer spinlocks for PUT_VALUE staging slot protection.
+  // All kernel proxy threads lock pvLocks[peer] before writing the staging
+  // buffer and posting iput, preventing cross-thread corruption.
+  pthread_spinlock_t *pvLocks = nullptr;
+  int pvLocksCount = 0;
 };
 
 struct flagcxProxyArgs;