GIN proxy on AINIC (Pollara): IB-CAST RMA backend, DMA-BUF for cuMem, spread over UDMA sets (#4480)

comms/rcclx/develop/projects/rccl/src/algorithms/rccl_ep/python/rccl_ep_capi.hip

@@ -1051,8 +1051,11 @@ int ep_configure_v2(void* handle, int num_experts, int num_topk, int grouped_onl
   // ranks write y and cw -- each its own slot(rank, t) in a local owner's window, while rows for a
   // remote owner go to the forwarder's staging -- and the reduce epilogue reads only those. So
   // y and cw keep this node's sources alone, (num_ranks - node_size) * T rows smaller.
+  // Such a window always gets one y row per slot: the hybrid combine is grouped only (see
+  // ep_combine_v4), and the y base moves down by node_base * T layout rows, which with
+  // num_topk rows per slot would not line up with the combine's slot rows on any node but 0.
   const bool y_local = h->hybrid && h->hyb_multi && env_int("RCCL_EP_HYB_Y_LOCAL") == 1;
-  h->lay = EpWindowLayout::make(h->cfg, sizeof(bf16_t), h->grouped_only ? 1 : 0,
+  h->lay = EpWindowLayout::make(h->cfg, sizeof(bf16_t), (h->grouped_only || y_local) ? 1 : 0,
                                 y_local ? h->node_size : 0, y_local ? h->node_base : 0);
   h->win_bytes = h->lay.total_bytes;
   if (h->hybrid && h->hyb_multi) {

comms/rcclx/develop/projects/rccl/src/gin/gin_host_proxy.cc

@@ -21,6 +21,12 @@ NCCL_PARAM(GinProxyQueueSize, "GIN_PROXY_QUEUE_SIZE", -1);
 extern int64_t ncclParamIbDataDirect();
 extern int64_t ncclParamDmaBufEnable();
 
+// [RCCL] GIN registers cuMem (VMM) GPU memory, which peermem cannot pin: with
+// NCCL_DMABUF_ENABLE=0 the fallback ibv_reg_mr fails ("Cannot allocate memory") and the
+// devComm cannot be created. rocmwrap.cc already turns DMA-BUF on for hipMem when cuMem is on
+// and NCCL_DMABUF_ENABLE=0; honor that here too.
+static bool ginProxyDmaBufEnabled() { return ncclParamDmaBufEnable() != 0 || ncclCuMemEnable(); }
+
 struct ginProxyGfdState {
   ncclGinProxyOp_t op;
   uint16_t counterId;
@@ -84,7 +90,7 @@ struct ginProxyCtx {
 static ncclRma_t* rmaBackend;
 
 static ncclResult_t getDmaBufFd(void* addr, size_t length, int* fd, bool forceNonDataDirect = false) {
-  if (ncclParamDmaBufEnable() == 0) return ncclInvalidUsage;
+  if (!ginProxyDmaBufEnabled()) return ncclInvalidUsage;
 
   // GIN's symmetric windows are cuMem/VMM allocations registered with the NIC via ibv_reg_dmabuf_mr.
   // the cuMem/hipMemGetHandleForAddressRange that exports the DMA-BUF FD requires this HIP version.
@@ -426,7 +432,7 @@ static ncclResult_t ncclGinProxyRegMrSym(void* ginCtx, void* addr, size_t size,
     NCCLCHECK(rmaBackend->regMrSym(cComm->collComm, addr, size, type, mrFlags, mhandle));
   } else if (type == NCCL_PTR_CUDA) {
     ncclResult_t dmabufResult = ncclInvalidUsage;
-    if (ncclParamDmaBufEnable() && (cComm->props.ptrSupport & NCCL_PTR_DMABUF)) {
+    if (ginProxyDmaBufEnabled() && (cComm->props.ptrSupport & NCCL_PTR_DMABUF)) {
       ncclResult_t registrationResult = ncclSuccess;
       int dmabufFd = -1;
       dmabufResult = getDmaBufFd(addr, size, &dmabufFd);

comms/rcclx/develop/projects/rccl/src/plugin/rma.cc

@@ -20,6 +20,7 @@ typedef ncclRma_t* getNcclRma_t(void* rmaPluginLib);
 
 extern getNcclRma_t getNcclRma_v14;
 extern getNcclRma_t getNcclRma_v13;
+extern bool rcclUseAinic();
 NCCL_PARAM(RmaPluginRefCount, "RMA_PLUGIN_REF_COUNT", 0);
 #define NCCL_RMA_VERSION_COUNT 2
 int ncclRmaVersion[NCCL_RMA_VERSION_COUNT] = {14, 13};
@@ -212,13 +213,19 @@ static void initPluginLibsOnceFunc() {
   // Add internal ib plugin. IB-CAST uses its own RMA backend layered on the CAST
   // transport (also backs GIN over IB-CAST via the generic ncclGinProxy); all other
   // cases use the default IB RMA proxy.
+  // Mirror the NET selection in plugin/net.cc: NCCL_NET=ROCM-IB is an alias of
+  // IB-CAST, and on AINIC autodetect (NCCL_NET unset) the data path is IB-CAST
+  // too. Without this, GIN/RMA on AINIC fall back to the generic net_ib backend
+  // (a second, non-ionic-aware device enumeration), which faults on Pollara
+  // ("async fatal event on QP ... local access violation work queue error").
+  // Same rule ROCM-27070 (ROCm/rocm-systems#9085) added to plugin/gin.cc before
+  // the v14 split.
   {
     const char* envNet = ncclGetEnv("NCCL_NET");
-    if (envNet && strcasecmp(envNet, "IB-CAST") == 0) {
-      pluginLibs[pluginCounter].ncclRma = &IbCastRmaIbProxy;
-    } else {
-      pluginLibs[pluginCounter].ncclRma = &ncclRmaIbProxy;
-    }
+    if (envNet && strcasecmp(envNet, "ROCM-IB") == 0) envNet = "IB-CAST";
+    bool useCast = (envNet && strcasecmp(envNet, "IB-CAST") == 0) || (!envNet && rcclUseAinic());
+    pluginLibs[pluginCounter].ncclRma = useCast ? &IbCastRmaIbProxy : &ncclRmaIbProxy;
+    INFO(NCCL_INIT | NCCL_NET, "RMA/Plugin: internal RMA backend %s", useCast ? "IB-CAST" : "IB");
   }
   pluginLibs[pluginCounter].state = ncclRmaPluginStateInitReady;
   pluginLibs[pluginCounter].version = ncclRmaVersion[0];

comms/rcclx/develop/projects/rccl/src/transport/net_ib_cast/gin.cc

@@ -418,18 +418,26 @@ ncclResult_t IbCastRmaIbProxyCreateContext(void* collComm, ncclRmaConfig_t* conf
     NCCLCHECKGOTO(ncclIbMalloc((void**)&gc->fullSendComm, sizeof(void*) * nranks), ret, end);
     NCCLCHECKGOTO(ncclIbMalloc((void**)&gc->fullRecvComm, sizeof(void*) * nranks), ret, end);
     gc->rank = cComm->rank;
+    // [RCCL] On AINIC a QP's channel id picks its UDMA engine set (connect.cc): the first QP of
+    // each (NIC, channel id) takes the set after the last one handed out, so new ids alternate.
+    // Without an id every context's QPs got channel 0 and all of GIN shared one set: ~30 GB/s per
+    // NIC, against 40-45 GB/s for RCCL send/recv over 8 channels. Ids from the top of the range,
+    // which the NET transport's channels never reach, make consecutive contexts alternate.
+    ncclNet_ctxt_t chCtxt = {};
+    chCtxt.chId = (MAXCHANNELS - 1) - (c % (MAXCHANNELS / 2));
 
     for (int i = 0; i < nranks; i += config->rankStride) {
       int connectPeer = (cComm->rank + i) % nranks;
       int acceptPeer = (cComm->rank - i + nranks) % nranks;
       do {
         if (gc->fullSendComm[connectPeer] == NULL)
           NCCLCHECKGOTO(IbCastConnectImpl(cComm->ctx, cComm->dev, handles + NCCL_NET_HANDLE_MAXSIZE * connectPeer,
-                                          &gc->fullSendComm[connectPeer], NULL,
+                                          &gc->fullSendComm[connectPeer], (ncclNetDeviceHandle_t**)&chCtxt,
                                           ncclParamCastGinIbTc() != -1 ? ncclParamCastGinIbTc() : ncclParamIbCastTc()),
                         ret, end);
         if (gc->fullRecvComm[acceptPeer] == NULL)
-          NCCLCHECKGOTO(netIbCast.accept(lComm, &gc->fullRecvComm[acceptPeer], NULL), ret, end);
+          NCCLCHECKGOTO(netIbCast.accept(lComm, &gc->fullRecvComm[acceptPeer], (ncclNetDeviceHandle_t**)&chCtxt),
+                        ret, end);
       } while ((gc->fullSendComm[connectPeer] == NULL) || (gc->fullRecvComm[acceptPeer] == NULL));
       NCCLCHECKGOTO(IbCastGinIbP2PBarrier(cComm), ret, end);
     }