feat(rccl): TDM async-to-LDS path for LL128 user buffers on gfx1250

projects/rccl/CHANGELOG.md

@@ -106,6 +106,8 @@ Full documentation for RCCL is available at [https://rccl.readthedocs.io](https:
 * Added `ncclGinFenceLevel` barrier semantics on the Anvil SDMA GIN backend (`NCCL_GIN_TYPE=7`). `gin.get`, `flushAsync` and `wait` are implemented instead of trapping, and a barrier's signal is now ordered behind the payload on the peer's SDMA queue, so `Put`, `Get` and the default `Put | Get` fences drain what they promise. `Put` covers a rank's puts to itself, and `ncclGinAllContexts` fences every GIN context. Single-node validated on gfx950.
 * Added gfx1250 (MI450) Tensor Data Mover staging to the symmetric AllGather, AllReduce and ReduceScatter kernels, where NCCL uses NVIDIA's TMA. Their deep loops move each tile through the warp's shared-memory scratch with the DMA engine instead of per-lane vector loads; on ROCm that path previously compiled out entirely, so the `Tma` kernels were never generated. Launch tuning was also corrected so these kernels can use up to a quarter of a GPU's WGPs, where only 14 were reached before, which is where most of the large-message gain comes from. Opt-in via `NCCL_SYM_TMA_ENABLE=1`, and reached only when the application opts into symmetric memory. `NCCL_SYM_TMA_ENABLE=2` additionally forces these kernels wherever the collective has one, skipping both the tuner's cost comparison and the deep-loop size bar, which is intended for A/B measurement rather than production. The new kernels are generated only when gfx1250 is in `GPU_TARGETS`, and every other architecture keeps the vector path unchanged.
 * Removed GIN rocSHMEM GDA dependency on rocSHMEM GDA bitcode: GIN rocSHMEM GDA is now header-only. The CMake `roc::rccl` target provides the required include directories in its `INSTALL_INTERFACE` property.
+* Added `RCCL_TDM_LL128_ENABLE` (default `0`): an experimental gfx1250 (MI450) Tensor Data Mover path for the LL128 protocol's user destination buffer, staging each warp slice in LDS and pushing it out as one bulk transfer instead of per-lane vector stores. Excluded from the default build: it requires `--enable-tdm-prim-ll128` at build time and `RCCL_TDM_LL128_ENABLE=1` at runtime, and compiles only on gfx1250 with an SDK providing the tensor data mover builtins. The runtime flag is captured once in the primitive's constructor rather than read from `ncclShmem` per slice, so the disabled case costs one register test per slice. Only the destination leg is staged: an async source load through the same per-warp window corrupted one slice in roughly 2% of 2-node ReduceScatter runs at 4 GiB and measured no faster, so it was removed rather than fixed. The receive and send FIFO legs are not eligible either, because both interleave protocol flag words into the payload. Measured on gfx1250 under forced RING/LL128 (out-of-place bus bandwidth, median of 3, 4 and 8 ranks): AllGather gains 5-10% above 256 MiB, AllReduce 3-5%, ReduceScatter 1-2%, while messages below roughly 32 MiB lose 2-5% to the staging overhead. With the tuner left to choose algorithm and protocol the path is not selected and throughput is unchanged, because LL128 is not the winning protocol in the size range where the staging pays off. The accumulator load is deliberately kept off this path because it has no spin to overlap; the accumulator-reduced store still goes through it.
+* Added host-side and device-side tests for the LL128 Tensor Data Mover store leg (`Ll128TdmStoreTest`): the build flag reaching the device pass, the slot indices tiling the staged window exactly once, and the staged store compared byte-for-byte against the scalar path it replaces. These require one GPU rather than the eight the existing LL128 suites need.
 
 ### Changed
 

projects/rccl/CMakeLists.txt

@@ -114,6 +114,7 @@ option(BUILD_NCCL4PY                           "Build the nccl4py Python binding
 option(RCCL_HIP_FATBIN_AT_TAIL                 "Place .hip_fatbin at the end of librccl.so"    ON)
 option(GENERATE_SYM_KERNELS                    "Generate symmetric memory kernels"             ON)
 option(ENABLE_TDM_SIMPLE                       "Build the experimental gfx1250 TDM SIMPLE copy path" OFF)
+option(ENABLE_TDM_PRIM_LL128                   "Build the experimental gfx1250 TDM LL128 user-buffer path" OFF)
 option(INSTALL_DEPENDENCIES                    "Force install dependencies"                    OFF)
 option(REPORT_KERNEL_RESOURCE_USE              "Append -Rpass-analysis=kernel to CXX flags"    OFF)
 option(ROCTX                                   "Enable ROCTX"                                  ${RCCL_DEFAULT_ROCTX})
@@ -128,6 +129,10 @@ option(DWORDX4_INTRINSICS                      "Turn off compilation with dwordx
 if(ENABLE_TDM_SIMPLE)
   add_compile_definitions(ENABLE_TDM_SIMPLE=1)
 endif()
+
+if(ENABLE_TDM_PRIM_LL128)
+  add_compile_definitions(ENABLE_TDM_PRIM_LL128=1)
+endif()
 # rocSHMEM integration (mutually exclusive — enabling both causes duplicate
 # rocshmem state between librccl.so and the executable at runtime):
 #   ENABLE_ROCSHMEM: alltoall_wg offload.  Links librocshmem.a into librccl.so

projects/rccl/docs/api-reference/env-variables.rst

@@ -1105,6 +1105,18 @@ threshold table introduced for gfx1250.
       - | ``-1``: Use the compile-time default of 256 MiB (default).
         | ``N`` (bytes): Set the per-slot payload capacity to ``N``; ``ceARTmpBuf`` is ``NCCL_CE_NUM_SLOTS`` (2) times that.
 
+    * - | ``RCCL_TDM_LL128_ENABLE``
+        | Routes the LL128 user destination buffer through the gfx1250 tensor data
+          mover: each warp's slice is staged in LDS and pushed out as one bulk
+          transfer instead of per-lane vector stores. Gains are concentrated in
+          large messages under LL128 (AllGather is the best case); below roughly
+          32 MiB the staging overhead makes it a small loss. The library must also
+          be built with ``--enable-tdm-prim-ll128``, and the path only compiles on
+          gfx1250 with an SDK providing the tensor data mover builtins. On any
+          other architecture, or without the build flag, this variable has no
+          effect.
+      - | ``0``: Disabled (default). The LL128 store uses the per-lane vector path.
+        | ``1``: Enabled, where the build and architecture support it.
     * - | ``RCCL_CE_AR_2SHOT_MIN_BYTES``
         | Overrides the minimum message size for the CE 2-shot AllReduce path.
           When ``-1`` (default), the value is resolved from ``ceNonRegMin[AllReduce]``

projects/rccl/docs/install/building-installing.rst

@@ -88,6 +88,7 @@ The RCCL build and installation helper script options are as follows:
           --enable_backtrace      Build with custom backtrace support
           --enable-mpi-tests      Enable MPI-based tests (requires --debug and MPI installation; set MPI_PATH if not in /opt/ompi)
           --enable-tdm-simple     Build the experimental gfx1250 TDM SIMPLE copy path
+          --enable-tdm-prim-ll128 Build the experimental gfx1250 TDM LL128 user-buffer path
        -f|--fast                  Quick-build RCCL (local gpu arch only, no backtrace)
           --force-reduce-pipeline Force reduce_copy sw pipeline to be used for every reduce-based collectives and datatypes
        -h|--help                  Prints this help message

projects/rccl/install.sh

@@ -40,6 +40,7 @@ time_trace=false
 use_ninja=false
 force_reduce_pipeline=false
 enable_tdm_simple=false
+enable_tdm_prim_ll128=false
 generate_sym_kernels=true
 device_linker=true
 warp_speed_enabled=true # note that this flag will be overridden to false for non MI350/MI300 platforms
@@ -92,6 +93,7 @@ function display_help()
     echo "       --enable-mpi-tests      Enable MPI-based tests (requires --debug and MPI installation; set MPI_PATH if not in /opt/ompi)"
     echo "       --enable-rccl-ep-tests  Build the rccl_ep multi-GPU tests (requires --enable-mpi-tests and a runtime with working cuMem symmetric memory)"
     echo "       --enable-tdm-simple     Build the experimental gfx1250 TDM SIMPLE copy path"
+    echo "       --enable-tdm-prim-ll128 Build the experimental gfx1250 TDM LL128 user-buffer path"
     echo "    -f|--fast                  Quick-build RCCL (local gpu arch only, no backtrace)"
     echo "       --force-reduce-pipeline Force reduce_copy sw pipeline to be used for every reduce-based collectives and datatypes"
     echo "    -h|--help                  Prints this help message"
@@ -149,7 +151,7 @@ function display_help()
 # check if we have a modern version of getopt that can handle whitespace and long parameters
 getopt -T
 if [[ "$?" -eq 4 ]]; then
-    GETOPT_PARSE=$(getopt --name "${0}" --options cdfhij:lprtq --longoptions address-sanitizer,all_unrolls,amdgpu_targets:,cmake-options:,debug,debug-fast,dependencies,device-linker,disable-colltrace,disable-kernarg-preload,disable-roctx,disable-sym-kernels,disable-warp-speed,dump-asm,enable-code-coverage,enable-full-coverage,enable_backtrace,enable-mpi-tests,enable-rccl-ep-tests,enable-tdm-simple,fast,force-reduce-pipeline,generate-sym-kernels,help,install,jobs:,kernel-resource-use,local_gpu_only,log-trace,ninja,no_clean,no-device-linker,npkit-enable,openmp-test-enable,package_build,prefix:,quiet-warnings,rm-legacy-include-dir,rocshmem,rocshmem-gin,roctx-enable,sqtt-enable,run_tests_all,run_tests_quick,static,tests_build,time-trace,verbose -- "$@")
+    GETOPT_PARSE=$(getopt --name "${0}" --options cdfhij:lprtq --longoptions address-sanitizer,all_unrolls,amdgpu_targets:,cmake-options:,debug,debug-fast,dependencies,device-linker,disable-colltrace,disable-kernarg-preload,disable-roctx,disable-sym-kernels,disable-warp-speed,dump-asm,enable-code-coverage,enable-full-coverage,enable_backtrace,enable-mpi-tests,enable-rccl-ep-tests,enable-tdm-simple,enable-tdm-prim-ll128,fast,force-reduce-pipeline,generate-sym-kernels,help,install,jobs:,kernel-resource-use,local_gpu_only,log-trace,ninja,no_clean,no-device-linker,npkit-enable,openmp-test-enable,package_build,prefix:,quiet-warnings,rm-legacy-include-dir,rocshmem,rocshmem-gin,roctx-enable,sqtt-enable,run_tests_all,run_tests_quick,static,tests_build,time-trace,verbose -- "$@")
 else
     echo "Need a new version of getopt"
     exit 1
@@ -182,6 +184,7 @@ while true; do
          --enable_backtrace)         build_bfd=true;                                                                                   shift ;;
          --enable-mpi-tests)         enable_mpi_tests=true;                                                                            shift ;;
          --enable-rccl-ep-tests)     enable_rccl_ep_tests=true;                                                                        shift ;;
+         --enable-tdm-prim-ll128)    enable_tdm_prim_ll128=true;                                                                       shift ;;
          --enable-tdm-simple)        enable_tdm_simple=true;                                                                           shift ;;
     -f | --fast)                     build_local_gpu_only=true;                                                                        shift ;;
          --force-reduce-pipeline)    force_reduce_pipeline=true;                                                                       shift ;;
@@ -482,6 +485,10 @@ if [[ "${enable_tdm_simple}" == true ]]; then
     cmake_common_options="${cmake_common_options} -DENABLE_TDM_SIMPLE=ON"
 fi
 
+if [[ "${enable_tdm_prim_ll128}" == true ]]; then
+    cmake_common_options="${cmake_common_options} -DENABLE_TDM_PRIM_LL128=ON"
+fi
+
 # Disable symmetric memory kernels
 if [[ "${generate_sym_kernels}" == false ]]; then
     cmake_common_options="${cmake_common_options} -DGENERATE_SYM_KERNELS=OFF"

projects/rccl/src/CMakeLists.txt

@@ -63,6 +63,7 @@ set(SRC_FILES
   device/hierarchical_shuffle.h
   device/network/unpack/unpack.h
   device/network/unpack/unpack_defs.h
+  device/ll128Layout.h
   device/onerank.cu
   device/op128.h
   device/primitives.h
@@ -87,6 +88,7 @@ set(SRC_FILES
   device/symmetric/reduce_scatter_gin.cuh
   device/tdm/asyncCopy.h
   device/tdm/cachePolicy.h
+  device/tdm/ll128Tdm.h
   device/tdm/syncPolicy.h
   device/tdm/tdmCopy.h
   algorithms/dda/device/device_buffer.cc

projects/rccl/src/device/ll128Layout.h

@@ -0,0 +1,27 @@
+/*************************************************************************
+ * Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * See LICENSE.txt for license information
+ ************************************************************************/
+
+#ifndef LL128_LAYOUT_H_
+#define LL128_LAYOUT_H_
+
+#include <hip/hip_runtime.h>
+
+// Dense 16-byte-chunk index for register pair `g` of lane `wid`. The LL128 wire
+// format interleaves a flag word into every line, so the payload slots are not
+// contiguous in register order; this compensates for those holes to keep the
+// packed user buffer gap-free and coalesced.
+//
+// Templated on the line geometry rather than reading NCCL_LL128_LINEELEMS /
+// WARP_SIZE directly so it stays usable from host code: the unit tests compute
+// the expected layout on the CPU, and the TDM store leg's correctness is
+// precisely the claim that these indices densely pack [0, DataBytes).
+template <int LineElems, int WarpSz>
+__host__ __device__ __forceinline__ constexpr int ll128SlotIx(int g, int wid) {
+  constexpr int LineSkip = 2 * WarpSz / LineElems;
+  return g * WarpSz - LineSkip * (g / 2) + wid - (g % 2) * (wid / (LineElems / 2));
+}
+
+#endif  // LL128_LAYOUT_H_

projects/rccl/src/device/prims_ll128.h

@@ -6,8 +6,19 @@
  * See LICENSE.txt for license information
  ************************************************************************/
 
+#include "tdm/tdmCopy.h"
+#include "ll128Layout.h"
+
 #define NCCL_LL128_FLAGTHREAD (NCCL_LL128_LINEELEMS - 1)
 
+// gfx1250 tensor-data-mover path for the LL128 user buffers; see tdm/ll128Tdm.h.
+// Built in with --enable-tdm-prim-ll128, then selected per comm with RCCL_TDM_LL128_ENABLE=1.
+#ifndef ENABLE_TDM_PRIM_LL128
+#define ENABLE_TDM_PRIM_LL128 0
+#endif
+
+#define TDM_LL128_ON (TDM_SUPPORTED && ENABLE_TDM_PRIM_LL128)
+
 #ifndef RCCL_USE_WBINVL1_VOL
 #if defined(__GFX8__) || defined(__gfx906__) || defined(__gfx908__) || defined(__gfx90a__)
 #define RCCL_USE_WBINVL1_VOL 1
@@ -127,6 +138,10 @@ class Primitives<T, RedOp, Fan, Direct, ProtoLL128, P2p, isNetOffload, Metadata,
   uint64_t* barriers;
   uint64_t barrier_next = 0;
   bool skip_fence = false;
+#if TDM_LL128_ON
+  // Read from shmem once here, not per slice: the hot path tests a register.
+  bool tdmEnable = false;
+#endif
 
   inline __device__ void barrier() {
 #if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__)
@@ -236,6 +251,10 @@ class Primitives<T, RedOp, Fan, Direct, ProtoLL128, P2p, isNetOffload, Metadata,
     }
   }
 
+#if TDM_LL128_ON
+#include "tdm/ll128Tdm.h"  // member functions; included inside the class body on purpose
+#endif
+
   template <int WordPerThread>
   __device__ __forceinline__ void loadRegsBegin(uint64_t (&regs)[WordPerThread], T const* src, int eltN) {
     constexpr int EltPer16B = 16 / sizeof(T);
@@ -244,12 +263,10 @@ class Primitives<T, RedOp, Fan, Direct, ProtoLL128, P2p, isNetOffload, Metadata,
     // LINEELEMS/2 for the 64-byte-line case; gfx1250 doubles LINEELEMS
     // (128-byte non-tearing line) so the stride and flag-subgroup width
     // both double accordingly.
-    constexpr int LineElems = NCCL_LL128_LINEELEMS;
-    constexpr int LineSkip = 2 * WARP_SIZE / LineElems;
     int ix[WordPerThread / 2];
 #pragma unroll
     for (int g = 0; g < WordPerThread / 2; g++) {
-      ix[g] = g * WARP_SIZE - LineSkip * (g / 2) + wid - (g % 2) * (wid / (LineElems / 2));
+      ix[g] = ll128SlotIx<NCCL_LL128_LINEELEMS, WARP_SIZE>(g, wid);
     }
     if (reinterpret_cast<uintptr_t>(src) % 16 == 0) {
       /* We are aligned to 16 bytes, so load directly to registers no shmem.
@@ -266,6 +283,11 @@ class Primitives<T, RedOp, Fan, Direct, ProtoLL128, P2p, isNetOffload, Metadata,
         }
       }
     } else {
+#if TDM_LL128_ON
+      // This fallback stages through the same per-warp scratch the TDM store leg
+      // hands the mover, so drain the previous slice's store before clobbering it.
+      tdmDrain();
+#endif
       // Not aligned. Stage the smallest 16 byte aligned region subsuming the
       // buffer into shmem.
       int misalignment = reinterpret_cast<uintptr_t>(src) % 16;
@@ -312,15 +334,16 @@ class Primitives<T, RedOp, Fan, Direct, ProtoLL128, P2p, isNetOffload, Metadata,
     for (int g = 1; g < WordPerThread / 2; g += 2) {
       if (flagThread) regs[2 * g - 1] = regs[2 * g];
     }
+#if TDM_LL128_ON
+    if (tdmStoreRegs<WordPerThread>(dst, regs, eltN)) return;
+#endif
 
     // Write to dst if 4-byte aligned, shmem otherwise.
     int misalignment = reinterpret_cast<uintptr_t>(dst) % 16;
     uint64_t* shm8 = shmemCvtPtr((uint64_t*)ncclScratchForWarp(warpInBlock));
-    constexpr int LineElems = NCCL_LL128_LINEELEMS;
-    constexpr int LineSkip = 2 * WARP_SIZE / LineElems;
 #pragma unroll
     for (int g = 0; g < WordPerThread / 2; g++) {
-      int ix = g * WARP_SIZE - LineSkip * (g / 2) + wid - (g % 2) * (wid / (LineElems / 2));
+      int ix = ll128SlotIx<NCCL_LL128_LINEELEMS, WARP_SIZE>(g, wid);
       if (!flagThread || g % 2 == 0) {
         if (misalignment == 0 && (ix + 1) * EltPer16B <= eltN) {
           storeUser128((uint64_t*)(dst + ix * EltPer16B), regs[2 * g + 0], regs[2 * g + 1]);
@@ -497,6 +520,9 @@ class Primitives<T, RedOp, Fan, Direct, ProtoLL128, P2p, isNetOffload, Metadata,
       dstPtr += DataEltPerSlice * nwarps;
       nelem -= DataEltPerSlice * nwarps;
     }
+#if TDM_LL128_ON
+    tdmDrain();
+#endif
 
     barrier();
 
@@ -576,6 +602,9 @@ class Primitives<T, RedOp, Fan, Direct, ProtoLL128, P2p, isNetOffload, Metadata,
     loadSendSync();
     userRegUsed = (e != nullptr) && (e->regUsed || e->netRegUsed);
     setDataPtrs(inputBuf, outputBuf, e != nullptr ? e->acc : nullptr);
+#if TDM_LL128_ON
+    tdmEnable = ncclShmem.comm.tdmLl128Enable;
+#endif
 #if RCCL_HAVE_GLOBAL_DWORDX4_BUILTINS
     skip_fence = !ncclShmem.comm.cheapPostSendFenceOff;
 #else