rccl-tests/gin-sdma: cap A2A GIN put to 128 MiB to fix >1 GiB SDMA hang

projects/rccl-tests/src/CMakeLists.txt

@@ -73,6 +73,7 @@ set(COMMON_FILES
   collector.cu
   common.h
   common.cu
+  gin_sdma_collective_policy.h
   nccl1_compat.h
   rccl_compat.h
   rccl_float8.h

projects/rccl-tests/src/alltoall.cu

@@ -248,10 +248,10 @@ __global__ void GinAlltoAllKernel(ncclWindow_t sendwin, size_t sendoffset, ncclW
   int tid = threadIdx.x + blockIdx.x * blockDim.x;
   int nthreads = blockDim.x * gridDim.x;
 
-  /* send to all peers via GIN */
+  /* send to all peers via GIN (chunked to <=128 MiB/put; see common.h::ginPutChunked) */
   const size_t size = count * sizeof(T);
   for (int r=tid; r<devComm.nRanks; r+=nthreads) {
-    gin.put(ncclTeamWorld(devComm), r,
+    ginPutChunked(gin, ncclTeamWorld(devComm), r,
         recvwin, recvoffset + devComm.rank * size,
         sendwin, sendoffset + r * size,
         size, ncclGin_SignalInc{signalIndex});
@@ -298,13 +298,13 @@ __global__ void HybridAlltoAllKernel(ncclWindow_t sendwin, size_t sendoffset, nc
     int nthreads = blockDim.x;
 
     for (int r = tid; r < startLsa; r += nthreads) {
-      gin.put(world, r,
+      ginPutChunked(gin, world, r,
           recvwin, recvoffset + world.rank * size,
           sendwin, sendoffset + r * size,
           size, ncclGin_SignalInc{signalIndex});
     }
     for (int r = startLsa + lsaSize + tid; r < world.nRanks; r += nthreads) {
-      gin.put(world, r,
+      ginPutChunked(gin, world, r,
           recvwin, recvoffset + world.rank * size,
           sendwin, sendoffset + r * size,
           size, ncclGin_SignalInc{signalIndex});

projects/rccl-tests/src/common.h

@@ -11,6 +11,12 @@
 #define NCCL_TESTS_VERSION "2.18.3"
 
 #include "rccl/rccl.h"
+// GIN-SDMA shared policy constants (single source of truth for the put-chunk
+// cap that ginPutChunked below clamps to). Cherry-picked as a standalone,
+// namespaced (gin_sdma::), no-GPU/no-RCCL header from
+// users/dondai/gin-stage2j-sdma-A1457-a2av-wip to fix A2A >1 GiB WITHOUT
+// pulling in the other GIN-SDMA collective kernel designs.
+#include "gin_sdma_collective_policy.h"
 // nccl_device.h provides the device-API public types referenced below
 // (ncclCommProperties_t, full ncclDevCommRequirements, NCCL_GIN_TYPE_NONE, ...).
 // As of NCCL 2.29 these types are referenced unconditionally by the test engine
@@ -412,6 +418,64 @@ extern int is_main_proc;
 extern thread_local int is_main_thread;
 
 #if defined(ENABLE_DEVICE_API) && NCCL_VERSION_CODE >= NCCL_VERSION(2,28,0)
+// Overflow-safe AND hang-safe gin.put for the Anvil-SDMA backend. Two distinct
+// single-descriptor limits apply, and we clamp to the smaller of the two:
+//   1) Correctness: the SDMA linear-copy count field is 30 bits and 1-based
+//      (count = bytes-1), so a put of >1 GiB (kGinPutMaxBytes) silently
+//      truncates and corrupts data (a 2 GiB transfer copies only 1 GiB).
+//   2) Reliability: on MI355X + ROCm 7.13 a single copy descriptor at/above
+//      256 MiB stalls the SDMA engine on the fused copy+signal packet, so the
+//      copy never lands and the SignalInc never fires -> waitSignal hangs. A
+//      128 MiB cap (gin_sdma::kGinSdmaSafeCopyBytes) is measured hang-free with
+//      no bandwidth loss; see that constant for the evidence.
+// Split the transfer into <=kGinSdmaSafeCopyBytes segments and carry the
+// caller's remote action (e.g. SignalInc) ONLY on the final segment. A
+// <=128 MiB message is a single put with no extra overhead.
+//
+// ORDERING PRECONDITION (why the trailing signal is safe): the Anvil-SDMA
+// backend selects the hardware SDMA queue per *issuer*, as
+// effectiveChannel = (blockId*warpsPerBlock + warpId)*stride % numChannels
+// (gin_anvil_sdma.h) -- NOT per peer. A single hardware queue is a strict
+// in-order FIFO; distinct channels are independent completion domains with NO
+// cross-queue ordering (only ncclGinApi_Flush drains all of a peer's channels).
+// Because this loop runs entirely within ONE thread, every segment of this
+// message -- data segments AND the final signal-carrying put -- resolves to the
+// same (blockId, warpId) and therefore the same in-order queue, even when
+// NCCL_GIN_ANVIL_SDMA_NUM_CHANNELS > 1. So the final SignalInc is FIFO-ordered
+// behind all prior segments of this message and "signal fired" still means "the
+// whole message landed"; per-message waitSignal accounting is unchanged.
+//
+// CALLER CONTRACT: issue an entire per-peer message (all segments + the
+// signal-carrying final put) from a SINGLE thread/warp, as the AllToAll GIN
+// sites do (one thread owns each disjoint (peer, offset) tuple). Do NOT split a
+// single peer's message across multiple warps/CTAs while carrying the signal on
+// only one of them: under NUM_CHANNELS > 1 the segments could land on different
+// queues and the SignalInc could fire before an earlier segment completes
+// (silent data corruption). If you ever need that pattern, drain with a Flush
+// (or force NUM_CHANNELS=1 / SPREAD_CHANNELS=0) before signaling.
+template <typename RemoteAction>
+__device__ __forceinline__ void ginPutChunked(
+    ncclGin& gin, ncclTeam team, int peer,
+    ncclWindow_t dstWin, size_t dstOff,
+    ncclWindow_t srcWin, size_t srcOff,
+    size_t bytes, RemoteAction finalAction) {
+  const size_t kMax = gin_sdma::kGinSdmaSafeCopyBytes < gin_sdma::kGinPutMaxBytes
+                          ? gin_sdma::kGinSdmaSafeCopyBytes
+                          : gin_sdma::kGinPutMaxBytes;
+  size_t off = 0;
+  do {
+    const size_t rem = bytes - off;
+    const size_t seg = rem > kMax ? kMax : rem;
+    if (off + seg >= bytes) {
+      // Final (or only) segment carries the signal / remote action.
+      gin.put(team, peer, dstWin, dstOff + off, srcWin, srcOff + off, seg, finalAction);
+    } else {
+      gin.put(team, peer, dstWin, dstOff + off, srcWin, srcOff + off, seg);
+    }
+    off += seg;
+  } while (off < bytes);
+}
+
 template <typename F>
 testResult_t testLaunchDeviceKernel(F kernel, void* sendbuff, size_t sendoffset, void* recvbuff, size_t recvoffset, size_t count, ncclDataType_t type, ncclRedOp_t op, int root, ncclComm_t comm, cudaStream_t stream) {
   if (kernel == nullptr) return testNotImplemented;