Skip to content

feat(rccl): stage symmetric kernel tiles with gfx1250 TDM instead of TMA #11850

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
alex-breslow-amd merged 12 commits into from
Sep 26, 2026
Merged

feat(rccl): stage symmetric kernel tiles with gfx1250 TDM instead of TMA #11850

Show file tree
Hide file tree
Changes from 8 commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
82efe18
feat(rccl): stage symmetric kernel tiles with gfx1250 TDM instead of TMA
alex-breslow-amd Sep 18, 2026
0e908e5
Use device scope where feasible instead of system scope
alex-breslow-amd Sep 21, 2026
7f01960
Add tuning model changes
alex-breslow-amd Sep 21, 2026
a13ac61
Improve tuning logic and performance. ~+35% uptick for ReduceScatter
alex-breslow-amd Sep 21, 2026
fc830ce
Remove overly verbose comment
alex-breslow-amd Sep 21, 2026
fe74859
Revert device-scope TDM staging; keep a single system-scope policy
alex-breslow-amd Sep 23, 2026
a7988c2
Address Nusrat's feedback
alex-breslow-amd Sep 24, 2026
6bf6555
Add more tests. Address feedback
alex-breslow-amd Sep 24, 2026
1b0acc7
perf(rccl): skip TDM's head peel on the 256 B-aligned AllGather tile
alex-breslow-amd Sep 25, 2026
f54d703
Fix wording in documentation in header for TDM library because it's t…
alex-breslow-amd Sep 25, 2026
d1bc477
perf(rccl): issue the one-way TDM bulk as a 1-D tile
alex-breslow-amd Sep 25, 2026
97c6cc5
Merge branch 'develop' into users/abreslow/tmaSymmetricPort
alex-breslow-amd Sep 25, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion projects/rccl/CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ Full documentation for RCCL is available at [https://rccl.readthedocs.io](https:
* Multiple GIN proxy progress threads via `NCCL_GIN_PROXY_NTHREADS`.
* Enabled hierarchical Copy Engine `ncclAllGather` and `ncclAlltoAll` on multi-node communicators: the intra-node phase runs over SDMA/CE across the LSA team and the inter-node phase over one-sided RMA through the CPU (GIN) proxy. The path is off by default and takes two opt-ins: `NCCL_CTA_POLICY=2` (equivalently `config.CTAPolicy = NCCL_CTA_POLICY_ZERO`), and symmetric registration of both the send and the receive buffer through `ncclCommWindowRegister`. On architectures other than gfx1250 (MI450) a third is needed, `NCCL_CUMEM_ENABLE=1`, because symmetric memory depends on cuMem and cuMem auto-detection enables itself only on gfx1250. `NCCL_GIN_ENABLE`, `NCCL_NUM_RMA_CTX`, `NCCL_WIN_ENABLE` and `NCCL_DMABUF_ENABLE` need no change, since their defaults already qualify; the inter-node rail draws from the internal RMA context pool (`NCCL_NUM_RMA_INT_CTX`, default 4) rather than the user-addressable contexts, so `numRmaCtx` only has to stay above zero. Selection also requires an available GIN backend and the same number of ranks on every node; communicators that do not qualify keep using the kernel path. Validated for correctness and path selection from 1 to 32 nodes.
* 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.

### Changed
* **Breaking: `NCCL_GIN_TYPE` values for AMD backends are not compatible with 2.30.7.** NCCL 2.31 inserted EFA GDA at value 5, so rocSHMEM GDA moved 5→6 and Anvil SDMA moved 6→7. The IB proxy remains `2`. Jobs that still set `NCCL_GIN_TYPE=6` now select rocSHMEM GDA, not Anvil SDMA. See `src/gin/README.md`.
Expand All @@ -28,7 +29,7 @@ Full documentation for RCCL is available at [https://rccl.readthedocs.io](https:
* Fixed hierarchical Copy Engine collectives dereferencing a NULL RMA proxy context when `NCCL_RMA_DISABLE=1` was combined with the hierarchical opt-ins. `NCCL_RMA_DISABLE` suppresses the window registration that creates those contexts, but the hierarchical availability check did not consult it, so a multi-node communicator with `NCCL_CTA_POLICY=2` and symmetric send and receive windows admitted the path and then crashed at launch. The internal RMA context predicate now defers to `ncclRmaProxyEnabled`, so such a communicator declines the hierarchical path and uses the kernel path instead.

### Known issues
* The following NCCL 2.31 features are NVIDIA-specific and are not available in RCCL: Compute Fabric Transport, the GDAKI and EFA GDA GIN backends, PAT combined with NVLS, NVLink-multicast AllGather, TMA-based symmetric kernels, and the CuTeDSL and `nccl4rust` bindings.
* The following NCCL 2.31 features are NVIDIA-specific and are not available in RCCL: Compute Fabric Transport, the GDAKI and EFA GDA GIN backends, PAT combined with NVLS, NVLink-multicast AllGather, and the CuTeDSL and `nccl4rust` bindings.

## RCCL 2.30.7 for ROCm 10.1.0 (Unreleased)

Expand Down
6 changes: 4 additions & 2 deletions projects/rccl/src/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -849,9 +849,11 @@ if (gen_py_result)
endif()

if (GENERATE_SYM_KERNELS)
# Execute the python script to generate required symmetric memory kernels
# Execute the python script to generate required symmetric memory kernels.
# GPU_TARGETS selects the arch-specific algos (the Tma* kernels need gfx1250's
# Tensor Data Mover).
execute_process(
COMMAND ${Python3_EXECUTABLE} ${RCCL_SOURCE_DIR}/src/device/symmetric/generate.py ${GEN_SYM_DIR}
COMMAND ${Python3_EXECUTABLE} ${RCCL_SOURCE_DIR}/src/device/symmetric/generate.py ${GEN_SYM_DIR} "${GPU_TARGETS}"
WORKING_DIRECTORY ${RCCL_SOURCE_DIR}
RESULT_VARIABLE gen_sym_py_result
ERROR_VARIABLE gen_sym_py_error
Expand Down
59 changes: 22 additions & 37 deletions projects/rccl/src/device/symmetric/all_gather.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,52 +27,44 @@ static __device__ void bcastDeep(ncclSymkArgsHandler const& handler, int tn, int

Pack* inpPacks = (Pack*)input.localPtr() + intptr_t(w) * UnrollPacks * WARP_SIZE +
(
#if __CUDA_ARCH__ >= 1000
#if NCCL_SYMK_ASYNC_TILE
EnableTma ? 0 :
#endif
lane);

ncclSymPtr<Pack> outPacks = (ncclSymPtr<Pack>)output + intptr_t(w) * UnrollPacks * WARP_SIZE +
(
#if __CUDA_ARCH__ >= 1000
#if NCCL_SYMK_ASYNC_TILE
EnableTma ? 0 :
#endif
lane);

Pack tmp[UnrollPacks];

#if __CUDA_ARCH__ >= 1000
#if NCCL_SYMK_ASYNC_TILE
int lw = threadIdx.x / WARP_SIZE;
extern __shared__ char smemScratch[];
using tmaSmemStruct_t = tmaSmemStruct<Pack, UnrollPacks>;
constexpr int smemSizePerWarp = ncclTmaShmemScratchWarpSize();
tmaSmemStruct_t* tmaSmem = reinterpret_cast<tmaSmemStruct_t*>(smemScratch + lw * smemSizePerWarp);
tmaSmemStruct_t* tmaSmem = ncclSymkTileSmem<tmaSmemStruct_t>(lw);
constexpr size_t tileSize = UnrollPacks * WARP_SIZE * BytePerPack;
#endif
bool skip = false; // all lanes issue loads/stores

#if __CUDA_ARCH__ >= 1000
#if NCCL_SYMK_ASYNC_TILE
size_t pending = 0;
if NCCL_IF_CONSTEXPR (EnableTma) {
if (lane == 0) {
// lane0 issues async.cp.bulk commands
init(&tmaSmem->bar, 1);
} else {
// other lanes can skip the loop
skip = true;
}
ncclSymkTileBarInit(&tmaSmem->bar, /*arrivers=*/1, lane);
// TMA issues from lane 0, so the rest of the warp has nothing left to do.
// TDM needs the whole warp to build the descriptor, so every lane stays in.
skip = NCCL_SYMK_TILE_TMA && lane != 0;
}
#endif

nIters -= w;
if (0 < nIters) {
#if __CUDA_ARCH__ >= 1000
#if NCCL_SYMK_ASYNC_TILE
if NCCL_IF_CONSTEXPR (EnableTma) {
if (lane == 0) {
cuda::device::memcpy_async_tx(tmaSmem->buff[0], inpPacks, cuda::aligned_size_t<16>(tileSize), tmaSmem->bar);
cuda::barrier<cuda::thread_scope_block>::arrival_token token =
cuda::device::barrier_arrive_tx(tmaSmem->bar, 1, tileSize);
tmaSmem->bar.wait(std::move(token));
}
ncclSymkTileLoad(tmaSmem->buff[0], inpPacks, tileSize, tmaSmem->bar, pending, lane);
ncclSymkTileLoadWait</*Arrivers=*/1>(tmaSmem->bar, pending, lane);
} else
#endif
{
Expand All @@ -97,9 +89,9 @@ static __device__ void bcastDeep(ncclSymkArgsHandler const& handler, int tn, int
NVCC_PRAGMA_UNROLL_AUTO
for (int ur = 0; ur < UnrollPeers - partial; ur++) {
if (partial && dr == nRanks) break;
#if __CUDA_ARCH__ >= 1000
#if NCCL_SYMK_ASYNC_TILE
if NCCL_IF_CONSTEXPR (EnableTma) {
ptx::cp_async_bulk(ptx::space_global, ptx::space_shared, outPacks.lsaPtr(r), tmaSmem->buff[0], tileSize);
ncclSymkTileStore(outPacks.lsaPtr(r), tmaSmem->buff[0], tileSize, lane);
} else
#endif
{
Expand All @@ -110,12 +102,9 @@ static __device__ void bcastDeep(ncclSymkArgsHandler const& handler, int tn, int
}
if (++r == nRanks) r = 0;
}
#if __CUDA_ARCH__ >= 1000
#if NCCL_SYMK_ASYNC_TILE
if NCCL_IF_CONSTEXPR (EnableTma) {
if (lane == 0) {
ptx::cp_async_bulk_commit_group();
ptx::cp_async_bulk_wait_group_read(ptx::n32_t<0>());
}
ncclSymkTileStoreWait(lane);
Comment thread
alex-breslow-amd marked this conversation as resolved.
Resolved
}
#endif
}
Expand All @@ -124,14 +113,10 @@ static __device__ void bcastDeep(ncclSymkArgsHandler const& handler, int tn, int
outPacks += intptr_t(wn) * UnrollPacks * WARP_SIZE;
nIters -= wn;
if (nIters <= 0) break;
#if __CUDA_ARCH__ >= 1000
#if NCCL_SYMK_ASYNC_TILE
if NCCL_IF_CONSTEXPR (EnableTma) {
if (lane == 0) {
cuda::device::memcpy_async_tx(tmaSmem->buff[0], inpPacks, cuda::aligned_size_t<16>(tileSize), tmaSmem->bar);
cuda::barrier<cuda::thread_scope_block>::arrival_token token =
cuda::device::barrier_arrive_tx(tmaSmem->bar, 1, tileSize);
tmaSmem->bar.wait(std::move(token));
}
ncclSymkTileLoad(tmaSmem->buff[0], inpPacks, tileSize, tmaSmem->bar, pending, lane);
ncclSymkTileLoadWait</*Arrivers=*/1>(tmaSmem->bar, pending, lane);
} else
#endif
{
Expand Down Expand Up @@ -185,15 +170,15 @@ static __device__ void bcast(ncclSymkArgsHandler const& handler, int tn, int t,

uint32_t alignment = uint32_t(input.offset - output.offset);
uint32_t nPreBytes =
#if __CUDA_ARCH__ >= 1000
#if NCCL_SYMK_ASYNC_TILE
(EnableTma && alignment % 256 == 0) ? (256 - input.offset) % 256 :
#endif
(16 - input.offset) % 16;

nPreBytes = min((size_t)nPreBytes, nBytes);
uintptr_t cursor = nPreBytes;

#if __CUDA_ARCH__ >= 1000
#if NCCL_SYMK_ASYNC_TILE
if NCCL_IF_CONSTEXPR (EnableTma) {
if (alignment % 256 == 0) {
constexpr int BytePerPack = ncclSymkBytePerPack, UnrollPacks = ncclSymkAlign256BDeepUnrollPacks, UnrollPeers = 2;
Expand Down
104 changes: 37 additions & 67 deletions projects/rccl/src/device/symmetric/all_reduce.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,46 +33,38 @@ static __device__ __forceinline__ void allreduceDeep(ncclSymkArgsHandler const&

ncclSymPtr<Pack> inpPacks = (ncclSymPtr<Pack>)input + intptr_t(w) * UnrollPacks * WARP_SIZE +
(
#if __CUDA_ARCH__ >= 1000
#if NCCL_SYMK_ASYNC_TILE
EnableTma ? 0 :
#endif
lane);

ncclSymPtr<Pack> outPacks = (ncclSymPtr<Pack>)output + intptr_t(w) * UnrollPacks * WARP_SIZE +
(
#if __CUDA_ARCH__ >= 1000
#if NCCL_SYMK_ASYNC_TILE
EnableTma ? 0 :
#endif
lane);

Pack acc0[UnrollPacks];

#if __CUDA_ARCH__ >= 1000
#if NCCL_SYMK_ASYNC_TILE
int lw = threadIdx.x / WARP_SIZE;
extern __shared__ char smemScratch[];
using tmaSmemStruct_t = tmaSmemStruct<Pack, UnrollPacks, UnrollPeers>;
constexpr int smemSizePerWarp = ncclTmaShmemScratchWarpSize();
tmaSmemStruct_t* tmaSmem = reinterpret_cast<tmaSmemStruct_t*>(smemScratch + lw * smemSizePerWarp);
tmaSmemStruct_t* tmaSmem = ncclSymkTileSmem<tmaSmemStruct_t>(lw);
constexpr size_t tilePack = UnrollPacks * WARP_SIZE;
constexpr size_t tileSize = tilePack * BytePerPack;
size_t tmaSize = 0;

if NCCL_IF_CONSTEXPR (EnableTma) {
if (lane == 0) {
init(&tmaSmem->bar, WARP_SIZE);
}
ncclSymkTileBarInit(&tmaSmem->bar, /*arrivers=*/WARP_SIZE, lane);
}
#endif

nIters -= w;
if (0 < nIters) {
#if __CUDA_ARCH__ >= 1000
#if NCCL_SYMK_ASYNC_TILE
if NCCL_IF_CONSTEXPR (EnableTma) {
if (lane == 0) {
cuda::device::memcpy_async_tx(tmaSmem->buff[0], inpPacks.peerPtr(world, rank),
cuda::aligned_size_t<16>(tileSize), tmaSmem->bar);
tmaSize += tileSize;
}
ncclSymkTileLoad(tmaSmem->buff[0], inpPacks.peerPtr(world, rank), tileSize, tmaSmem->bar, tmaSize, lane);
} else
#endif
{
Expand All @@ -92,17 +84,10 @@ static __device__ __forceinline__ void allreduceDeep(ncclSymkArgsHandler const&
if (++r == nRanks) r = 0;
{
Pack tmp1[UnrollPacks];
#if __CUDA_ARCH__ >= 1000
#if NCCL_SYMK_ASYNC_TILE
if NCCL_IF_CONSTEXPR (EnableTma) {
if (lane == 0) {
cuda::device::memcpy_async_tx(tmaSmem->buff[1], inpPacks.peerPtr(world, r),
cuda::aligned_size_t<16>(tileSize), tmaSmem->bar);
tmaSize += tileSize;
}
cuda::barrier<cuda::thread_scope_block>::arrival_token token =
cuda::device::barrier_arrive_tx(tmaSmem->bar, 1, tmaSize);
tmaSmem->bar.wait(std::move(token));
tmaSize = 0;
ncclSymkTileLoad(tmaSmem->buff[1], inpPacks.peerPtr(world, r), tileSize, tmaSmem->bar, tmaSize, lane);
ncclSymkTileLoadWait</*Arrivers=*/WARP_SIZE>(tmaSmem->bar, tmaSize, lane);
} else
#endif
{
Expand All @@ -113,7 +98,7 @@ static __device__ __forceinline__ void allreduceDeep(ncclSymkArgsHandler const&
}
NVCC_PRAGMA_UNROLL_AUTO
for (int u = 0; u < UnrollPacks; u++) {
#if __CUDA_ARCH__ >= 1000
#if NCCL_SYMK_ASYNC_TILE
if NCCL_IF_CONSTEXPR (EnableTma) {
acc0[u] = tmaSmem->buff[0][lane + WARP_SIZE * u];
tmp1[u] = tmaSmem->buff[1][lane + WARP_SIZE * u];
Expand All @@ -133,23 +118,20 @@ static __device__ __forceinline__ void allreduceDeep(ncclSymkArgsHandler const&
if (partial && dr == nRanks) break;

Pack tmp1[UnrollPeers][UnrollPacks];
#if __CUDA_ARCH__ >= 1000
#if NCCL_SYMK_ASYNC_TILE
if NCCL_IF_CONSTEXPR (EnableTma) {
// lane 0 waits for all threads to reduce tmp1 before next batch of TMA loads
// The next batch reloads these same tiles, so hold until every lane has
// finished reducing out of them.
__syncwarp();
}
#endif

NVCC_PRAGMA_UNROLL_AUTO
for (int ur = 0; ur < UnrollPeers - partial; ur++) {
if (partial && ur != 0 && dr + ur == nRanks) break;
#if __CUDA_ARCH__ >= 1000
#if NCCL_SYMK_ASYNC_TILE
if NCCL_IF_CONSTEXPR (EnableTma) {
if (lane == 0) {
cuda::device::memcpy_async_tx(tmaSmem->buff[ur], inpPacks.peerPtr(world, r),
cuda::aligned_size_t<16>(tileSize), tmaSmem->bar);
tmaSize += tileSize;
}
ncclSymkTileLoad(tmaSmem->buff[ur], inpPacks.peerPtr(world, r), tileSize, tmaSmem->bar, tmaSize, lane);
} else
#endif
{
Expand All @@ -160,20 +142,17 @@ static __device__ __forceinline__ void allreduceDeep(ncclSymkArgsHandler const&
}
if (++r == nRanks) r = 0;
}
#if __CUDA_ARCH__ >= 1000
#if NCCL_SYMK_ASYNC_TILE
if NCCL_IF_CONSTEXPR (EnableTma) {
cuda::barrier<cuda::thread_scope_block>::arrival_token token =
cuda::device::barrier_arrive_tx(tmaSmem->bar, 1, tmaSize);
tmaSmem->bar.wait(std::move(token));
tmaSize = 0;
ncclSymkTileLoadWait</*Arrivers=*/WARP_SIZE>(tmaSmem->bar, tmaSize, lane);
}
#endif
NVCC_PRAGMA_UNROLL_AUTO
for (int ur = 0; ur < UnrollPeers - partial; ur++) {
if (partial && ur != 0 && dr + ur == nRanks) break;
NVCC_PRAGMA_UNROLL(UnrollPacks)
for (int u = 0; u < UnrollPacks; u++) {
#if __CUDA_ARCH__ >= 1000
#if NCCL_SYMK_ASYNC_TILE
if NCCL_IF_CONSTEXPR (EnableTma) {
tmp1[ur][u] = tmaSmem->buff[ur][lane + WARP_SIZE * u];
}
Expand All @@ -186,7 +165,7 @@ static __device__ __forceinline__ void allreduceDeep(ncclSymkArgsHandler const&

NVCC_PRAGMA_UNROLL_AUTO
for (int u = 0; u < UnrollPacks; u++) {
#if __CUDA_ARCH__ >= 1000
#if NCCL_SYMK_ASYNC_TILE
if NCCL_IF_CONSTEXPR (EnableTma) {
tmaSmem->buff[0][lane + WARP_SIZE * u] = applyCast<Acc, T>(acc1[u]);
} else
Expand All @@ -196,11 +175,10 @@ static __device__ __forceinline__ void allreduceDeep(ncclSymkArgsHandler const&
}
}

#if __CUDA_ARCH__ >= 1000
#if NCCL_SYMK_ASYNC_TILE
if NCCL_IF_CONSTEXPR (EnableTma) {
// threads flush data to point of consistency for async proxy
ptx::fence_proxy_async(ptx::space_shared);
__syncwarp();
// Publish the lanes' reduced tile to the engine that is about to read it.
ncclSymkTileFenceSmem();
}
#endif

Expand All @@ -213,12 +191,9 @@ static __device__ __forceinline__ void allreduceDeep(ncclSymkArgsHandler const&
NVCC_PRAGMA_UNROLL_AUTO
for (int ur = 0; ur < UnrollPeers - partial; ur++) {
if (partial && dr == nRanks) break;
#if __CUDA_ARCH__ >= 1000
#if NCCL_SYMK_ASYNC_TILE
if NCCL_IF_CONSTEXPR (EnableTma) {
if (lane == 0) {
ptx::cp_async_bulk(ptx::space_global, ptx::space_shared, outPacks.peerPtr(world, r), tmaSmem->buff[0],
tileSize);
}
ncclSymkTileStore(outPacks.peerPtr(world, r), tmaSmem->buff[0], tileSize, lane);
} else
#endif
{
Expand All @@ -231,12 +206,10 @@ static __device__ __forceinline__ void allreduceDeep(ncclSymkArgsHandler const&
}
}
}
#if __CUDA_ARCH__ >= 1000
#if NCCL_SYMK_ASYNC_TILE
if NCCL_IF_CONSTEXPR (EnableTma) {
if (lane == 0) {
ptx::cp_async_bulk_commit_group();
ptx::cp_async_bulk_wait_group_read(ptx::n32_t<0>());
}
// Drain every peer's store before the next iteration refills the tile.
ncclSymkTileStoreWait(lane);
__syncwarp();
}
#endif
Expand All @@ -246,14 +219,10 @@ static __device__ __forceinline__ void allreduceDeep(ncclSymkArgsHandler const&
nIters -= wn;
if (nIters <= 0) break;

#if __CUDA_ARCH__ >= 1000
#if NCCL_SYMK_ASYNC_TILE
// Load data for next iteration.
if NCCL_IF_CONSTEXPR (EnableTma) {
if (lane == 0) {
cuda::device::memcpy_async_tx(tmaSmem->buff[0], inpPacks.peerPtr(world, rank),
cuda::aligned_size_t<16>(tileSize), tmaSmem->bar);
tmaSize += tileSize;
}
ncclSymkTileLoad(tmaSmem->buff[0], inpPacks.peerPtr(world, rank), tileSize, tmaSmem->bar, tmaSize, lane);
} else
#endif
{
Expand Down Expand Up @@ -344,20 +313,21 @@ static __device__ void allreduce(ncclSymkArgsHandler const& handler, int tn, int
uint32_t nBlocks_rcp32 = nccl::utility::idivRcp32_upto64(nBlocks);
uint32_t nRanks_nBlocks_rcp32 = nccl::utility::imulRcp32(nRanks, nRanks_rcp32, nBlocks, nBlocks_rcp32);

// True only where the deep loop really stages through a TDM engine.
constexpr bool AsyncTile = ncclSymkAsyncTile && EnableTma;

uint32_t nPreBytes = (16u - input.offset) % 16u;
nPreBytes = min((size_t)nPreBytes, nBytes);
uintptr_t cursor = nPreBytes;

if ((input.offset - output.offset) % 16 == 0) {
constexpr int BytePerPack = ncclSymkBytePerPack,
UnrollPacks =
#if __CUDA_ARCH__ >= 1000
EnableTma ? ncclSymkDeepUnrollPacks :
#endif
ncclSymkUnrollPacks,
UnrollPacks = AsyncTile ? ncclSymkDeepUnrollPacks(sizeof(T)) : ncclSymkUnrollPacks,
UnrollPeers = 2;

constexpr int BytePerChunk = EnableTma ? ncclSymkDeepBytePerChunk : ncclSymkBytePerChunk;
// Derived from UnrollPacks so the two cannot disagree: a chunk wider than what allreduceDeep()
// reduces would leave the difference unreduced.
constexpr int BytePerChunk = ncclSymkGetBytesPerChunk(ncclSymkMinWarpsPerBlock, UnrollPacks);
uint32_t chunks = (nBytes - cursor) / BytePerChunk;
chunks -= imodFast32(chunks, nRanks * nBlocks, nRanks_nBlocks_rcp32);
if (chunks != 0) {
Expand Down
Loading
Loading