From 708d5c943f182246096e3355e6b29fc69c5e2f8e Mon Sep 17 00:00:00 2001 From: David Sidler Date: Wed, 4 Nov 2026 18:00:59 -0600 Subject: [PATCH 1/6] Hoist heap-base loads out of the Triton ctx.all_gather store loop ctx.all_gather called ctx.store once per destination, and every call loaded heap_bases twice (source and destination base). A load issued after a store cannot be hoisted above it, and waiting for that load also waits for every earlier store, so the compiler emitted s_waitcnt vmcnt(0) before each of the 128 stores of a 256x256 fp16 tile, serializing the remote stores. Load the local and all peer heap bases once before the first store, keep each peer as a scalar pointer delta, and write to peers in a per-rank rotated order starting with the local copy. The assembly now has a single vmcnt(0) in the epilogue and no loads between stores. Co-Authored-By: Claude --- iris/mem/triton/context.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/iris/mem/triton/context.py b/iris/mem/triton/context.py index 63ee446e9..bc8a6f192 100644 --- a/iris/mem/triton/context.py +++ b/iris/mem/triton/context.py @@ -862,8 +862,22 @@ def all_gather(self, tile: Tile, dst_view: TensorView, dim: tl.constexpr): else: dst_ptr, combined_mask = dst_view.offset_tile_ptr(tile, offset_n=self.rank * N_local, src_mask=None) - for dest_rank in range(self.world_size): - self.store(dst_ptr, tile.data, to_rank=dest_rank, mask=combined_mask, hint=(1, tile.block_n)) + # Load every peer's heap base before the first store: a load issued after a store cannot be + # hoisted above it, and waiting for that load also waits for all earlier stores to complete. + local_base = tl.load(self.heap_bases + self.rank) + ptr_int = tl.cast(dst_ptr, tl.uint64) + ptr_deltas = () + for step in tl.static_range(1, self.world_size): + peer = (self.rank + step) % self.world_size + ptr_deltas += (tl.load(self.heap_bases + peer) - local_base,) + + # Start with the local copy, then write to peers in a per-rank rotated order so that each + # destination sees a different writer at any given moment. + tl.store(dst_ptr, tile.data, mask=combined_mask) + for step in tl.static_range(1, self.world_size): + remote_ptr = tl.cast(ptr_int + ptr_deltas[step - 1], dst_ptr.dtype) + remote_ptr = tl.max_contiguous(tl.multiple_of(remote_ptr, (1, tile.block_n)), (1, tile.block_n)) + tl.store(remote_ptr, tile.data, mask=combined_mask) @triton.jit def gather(self, tile: TileView, src_view: TensorView, source_rank: tl.constexpr, hint: tl.constexpr = None): From 49b01d737654b19abcbf6656fe2e9d866d06f91f Mon Sep 17 00:00:00 2001 From: David Sidler Date: Thu, 5 Nov 2026 00:45:21 -0600 Subject: [PATCH 2/6] Hoist heap-base loads out of the Triton CCL all-gather kernels persistent_all_gather called iris.store once per remote destination inside the tile loop, so every store re-loaded heap_bases twice and was preceded by an s_waitcnt vmcnt(0). persistent_all_gather_partitioned did the same translation after every tile's store even though its destination never changes. Translate the output base once per peer before the tile loop and store through the pre-translated pointers. The persistent kernel now writes the local copy first and peers in a per-rank rotated order, matching the Gluon all-gather. Co-Authored-By: Claude --- iris/ccl/triton/all_gather.py | 92 +++++++++++++++++------------------ 1 file changed, 45 insertions(+), 47 deletions(-) diff --git a/iris/ccl/triton/all_gather.py b/iris/ccl/triton/all_gather.py index 16dd586a3..f92f9d5e4 100644 --- a/iris/ccl/triton/all_gather.py +++ b/iris/ccl/triton/all_gather.py @@ -7,7 +7,6 @@ import triton import triton.language as tl -import iris from iris.host.tracing.kernel_artifacts import iris_launch @@ -64,6 +63,18 @@ def persistent_all_gather( num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) total_tiles = num_pid_m * num_pid_n tl.assume(total_tiles > 0) + + # Heap bases do not change between tiles, so translate the output base once per peer up front. + # A heap_bases load issued after a store cannot be hoisted above it, and waiting for that load + # also waits for every earlier store, which would serialize the remote stores. + local_base = tl.load(heap_bases + iris_rank) + output_ptr_int = tl.cast(output_ptr, tl.uint64) + remote_output_ptrs = () + for step in tl.static_range(1, world_size): + dest_idx = (group_rank + step) % world_size + ptr_delta = tl.load(heap_bases + rank_start + dest_idx * rank_stride) - local_base + remote_output_ptrs += (tl.cast(output_ptr_int + ptr_delta, output_ptr.dtype),) + for tile_id in range(pid, total_tiles, COMM_SMS): num_pid_in_group = GROUP_SIZE_M * num_pid_n group_id = tile_id // num_pid_in_group @@ -101,42 +112,30 @@ def persistent_all_gather( # Load local input data once for this tile data = tl.load(input_ptr_source, mask=input_mask, other=0.0) - # Send local shard data to all destination ranks - # Each rank's input goes to output[group_rank * M : (group_rank + 1) * M, :] on all ranks - for i in tl.static_range(world_size): - target_rank = rank_start + i * rank_stride - - # Compute global output row indices: offset by group_rank * M - rm_output = rm_input + group_rank * M - - # Output mask: only write where input was valid - output_mask = (rm_output[:, None] < (group_rank + 1) * M) & (rn[None, :] < N) - - # Combine masks: must be valid in both input and output - combined_mask = input_mask & output_mask - - # Compute output offset - output_base_m = rm_output[:, None] * stride_out_m - output_base_n = rn[None, :] * stride_out_n - output_offset = output_base_m + output_base_n - output_ptr_target = output_ptr + output_offset - output_ptr_target = tl.multiple_of(output_ptr_target, (BLOCK_SIZE_M, BLOCK_SIZE_N)) - - if i == group_rank: - # Local destination (i == group_rank): use direct store - tl.store(output_ptr_target, data, mask=combined_mask, cache_modifier=".wt") - else: - # Remote destination: use iris.store to send data to remote destination - # Use iris_rank for iris RMA operations (heap_bases indexing) - iris.store( - output_ptr_target, - data, - iris_rank, - target_rank, - heap_bases, - mask=combined_mask, - hint=(1, BLOCK_SIZE_N), - ) + # Compute global output row indices: offset by group_rank * M + rm_output = rm_input + group_rank * M + + # Output mask: only write where input was valid + output_mask = (rm_output[:, None] < (group_rank + 1) * M) & (rn[None, :] < N) + + # Combine masks: must be valid in both input and output + combined_mask = input_mask & output_mask + + # Compute output offset + output_base_m = rm_output[:, None] * stride_out_m + output_base_n = rn[None, :] * stride_out_n + output_offset = output_base_m + output_base_n + + # Each rank's input goes to output[group_rank * M : (group_rank + 1) * M, :] on all ranks. + # Write the local copy first, then the peers in a per-rank rotated order so that each + # destination sees a different writer at any given moment. + output_ptr_target = output_ptr + output_offset + output_ptr_target = tl.multiple_of(output_ptr_target, (BLOCK_SIZE_M, BLOCK_SIZE_N)) + tl.store(output_ptr_target, data, mask=combined_mask, cache_modifier=".wt") + for step in tl.static_range(1, world_size): + remote_ptr = remote_output_ptrs[step - 1] + output_offset + remote_ptr = tl.max_contiguous(tl.multiple_of(remote_ptr, (1, BLOCK_SIZE_N)), (1, BLOCK_SIZE_N)) + tl.store(remote_ptr, data, mask=combined_mask) @triton.jit() @@ -206,6 +205,11 @@ def persistent_all_gather_partitioned( total_tiles = num_pid_m * num_pid_n tl.assume(total_tiles > 0) + # This program writes to a single destination for all of its tiles, so translate the output base once + # instead of loading heap_bases after every tile's store. For the local destination the delta is zero. + ptr_delta = tl.load(heap_bases + target_rank) - tl.load(heap_bases + iris_rank) + remote_output_ptr = tl.cast(tl.cast(output_ptr, tl.uint64) + ptr_delta, output_ptr.dtype) + # Iterate over tiles with this PID's offset and stride within the rank group for tile_id in range(pid_in_rank_group, total_tiles, pids_per_rank): num_pid_in_group = GROUP_SIZE_M * num_pid_n @@ -265,16 +269,10 @@ def persistent_all_gather_partitioned( # Local destination: use direct store tl.store(output_ptr_target, data, mask=combined_mask, cache_modifier=".wt") else: - # Remote destination: use iris.store to send data to remote destination - iris.store( - output_ptr_target, - data, - iris_rank, - target_rank, - heap_bases, - mask=combined_mask, - hint=(1, BLOCK_SIZE_N), - ) + # Remote destination: store through the pre-translated base + remote_ptr = remote_output_ptr + output_offset + remote_ptr = tl.max_contiguous(tl.multiple_of(remote_ptr, (1, BLOCK_SIZE_N)), (1, BLOCK_SIZE_N)) + tl.store(remote_ptr, data, mask=combined_mask) def launch( From 67fd05e37ad3ecbc929c5d0951e5d1320203579b Mon Sep 17 00:00:00 2001 From: David Sidler Date: Thu, 5 Nov 2026 01:35:57 -0600 Subject: [PATCH 3/6] Leave the partitioned CCL all-gather kernel as it was Translating the destination base once before the tile loop made persistent_all_gather_partitioned 5-12% slower at 64 or fewer programs (for example 4.2 ms -> 4.6 ms at M=16384, N=4096, 64 programs), so restore the original per-tile iris.store there. Only the persistent kernel keeps the hoisting. Co-Authored-By: Claude --- iris/ccl/triton/all_gather.py | 20 +++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/iris/ccl/triton/all_gather.py b/iris/ccl/triton/all_gather.py index f92f9d5e4..6ae8311df 100644 --- a/iris/ccl/triton/all_gather.py +++ b/iris/ccl/triton/all_gather.py @@ -7,6 +7,7 @@ import triton import triton.language as tl +import iris from iris.host.tracing.kernel_artifacts import iris_launch @@ -205,11 +206,6 @@ def persistent_all_gather_partitioned( total_tiles = num_pid_m * num_pid_n tl.assume(total_tiles > 0) - # This program writes to a single destination for all of its tiles, so translate the output base once - # instead of loading heap_bases after every tile's store. For the local destination the delta is zero. - ptr_delta = tl.load(heap_bases + target_rank) - tl.load(heap_bases + iris_rank) - remote_output_ptr = tl.cast(tl.cast(output_ptr, tl.uint64) + ptr_delta, output_ptr.dtype) - # Iterate over tiles with this PID's offset and stride within the rank group for tile_id in range(pid_in_rank_group, total_tiles, pids_per_rank): num_pid_in_group = GROUP_SIZE_M * num_pid_n @@ -269,10 +265,16 @@ def persistent_all_gather_partitioned( # Local destination: use direct store tl.store(output_ptr_target, data, mask=combined_mask, cache_modifier=".wt") else: - # Remote destination: store through the pre-translated base - remote_ptr = remote_output_ptr + output_offset - remote_ptr = tl.max_contiguous(tl.multiple_of(remote_ptr, (1, BLOCK_SIZE_N)), (1, BLOCK_SIZE_N)) - tl.store(remote_ptr, data, mask=combined_mask) + # Remote destination: use iris.store to send data to remote destination + iris.store( + output_ptr_target, + data, + iris_rank, + target_rank, + heap_bases, + mask=combined_mask, + hint=(1, BLOCK_SIZE_N), + ) def launch( From 809551aea869cf8864bc318cb60b86f7cea6fec9 Mon Sep 17 00:00:00 2001 From: David Sidler Date: Thu, 5 Nov 2026 01:58:43 -0600 Subject: [PATCH 4/6] Shrink the CCL all-gather change to the hoist and the store call Keep the original destination loop, its local-store branch and its ascending order. Only translate the output base for every destination before the tile loop and replace the iris.store call with a tl.store through the pre-translated pointer. Co-Authored-By: Claude --- iris/ccl/triton/all_gather.py | 61 ++++++++++++++++++----------------- 1 file changed, 31 insertions(+), 30 deletions(-) diff --git a/iris/ccl/triton/all_gather.py b/iris/ccl/triton/all_gather.py index 6ae8311df..7fb2fa145 100644 --- a/iris/ccl/triton/all_gather.py +++ b/iris/ccl/triton/all_gather.py @@ -65,15 +65,13 @@ def persistent_all_gather( total_tiles = num_pid_m * num_pid_n tl.assume(total_tiles > 0) - # Heap bases do not change between tiles, so translate the output base once per peer up front. - # A heap_bases load issued after a store cannot be hoisted above it, and waiting for that load - # also waits for every earlier store, which would serialize the remote stores. + # Translate the output base once per destination up front. A heap_bases load issued after a store + # cannot be hoisted above it, and waiting for that load also waits for every earlier store. local_base = tl.load(heap_bases + iris_rank) output_ptr_int = tl.cast(output_ptr, tl.uint64) remote_output_ptrs = () - for step in tl.static_range(1, world_size): - dest_idx = (group_rank + step) % world_size - ptr_delta = tl.load(heap_bases + rank_start + dest_idx * rank_stride) - local_base + for i in tl.static_range(world_size): + ptr_delta = tl.load(heap_bases + rank_start + i * rank_stride) - local_base remote_output_ptrs += (tl.cast(output_ptr_int + ptr_delta, output_ptr.dtype),) for tile_id in range(pid, total_tiles, COMM_SMS): @@ -113,30 +111,33 @@ def persistent_all_gather( # Load local input data once for this tile data = tl.load(input_ptr_source, mask=input_mask, other=0.0) - # Compute global output row indices: offset by group_rank * M - rm_output = rm_input + group_rank * M - - # Output mask: only write where input was valid - output_mask = (rm_output[:, None] < (group_rank + 1) * M) & (rn[None, :] < N) - - # Combine masks: must be valid in both input and output - combined_mask = input_mask & output_mask - - # Compute output offset - output_base_m = rm_output[:, None] * stride_out_m - output_base_n = rn[None, :] * stride_out_n - output_offset = output_base_m + output_base_n - - # Each rank's input goes to output[group_rank * M : (group_rank + 1) * M, :] on all ranks. - # Write the local copy first, then the peers in a per-rank rotated order so that each - # destination sees a different writer at any given moment. - output_ptr_target = output_ptr + output_offset - output_ptr_target = tl.multiple_of(output_ptr_target, (BLOCK_SIZE_M, BLOCK_SIZE_N)) - tl.store(output_ptr_target, data, mask=combined_mask, cache_modifier=".wt") - for step in tl.static_range(1, world_size): - remote_ptr = remote_output_ptrs[step - 1] + output_offset - remote_ptr = tl.max_contiguous(tl.multiple_of(remote_ptr, (1, BLOCK_SIZE_N)), (1, BLOCK_SIZE_N)) - tl.store(remote_ptr, data, mask=combined_mask) + # Send local shard data to all destination ranks + # Each rank's input goes to output[group_rank * M : (group_rank + 1) * M, :] on all ranks + for i in tl.static_range(world_size): + # Compute global output row indices: offset by group_rank * M + rm_output = rm_input + group_rank * M + + # Output mask: only write where input was valid + output_mask = (rm_output[:, None] < (group_rank + 1) * M) & (rn[None, :] < N) + + # Combine masks: must be valid in both input and output + combined_mask = input_mask & output_mask + + # Compute output offset + output_base_m = rm_output[:, None] * stride_out_m + output_base_n = rn[None, :] * stride_out_n + output_offset = output_base_m + output_base_n + output_ptr_target = output_ptr + output_offset + output_ptr_target = tl.multiple_of(output_ptr_target, (BLOCK_SIZE_M, BLOCK_SIZE_N)) + + if i == group_rank: + # Local destination (i == group_rank): use direct store + tl.store(output_ptr_target, data, mask=combined_mask, cache_modifier=".wt") + else: + # Remote destination: store through the pre-translated output base + remote_ptr = remote_output_ptrs[i] + output_offset + remote_ptr = tl.max_contiguous(tl.multiple_of(remote_ptr, (1, BLOCK_SIZE_N)), (1, BLOCK_SIZE_N)) + tl.store(remote_ptr, data, mask=combined_mask) @triton.jit() From e8fc0cfe7143548d674849befad248daa8c419c4 Mon Sep 17 00:00:00 2001 From: David Sidler Date: Thu, 5 Nov 2026 17:47:21 -0600 Subject: [PATCH 5/6] Rotate the CCL all-gather destinations and hint the local store Write the local copy first and the peers in a per-rank rotated order in persistent_all_gather, matching ctx.all_gather and the Gluon all-gather, so each destination sees one writer at a time. Also apply the same max_contiguous/multiple_of hint to the local store in ctx.all_gather that the remote stores and the old self.store path had. Co-Authored-By: Claude --- iris/ccl/triton/all_gather.py | 10 ++++++---- iris/mem/triton/context.py | 3 ++- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/iris/ccl/triton/all_gather.py b/iris/ccl/triton/all_gather.py index 7fb2fa145..69a538041 100644 --- a/iris/ccl/triton/all_gather.py +++ b/iris/ccl/triton/all_gather.py @@ -113,7 +113,9 @@ def persistent_all_gather( # Send local shard data to all destination ranks # Each rank's input goes to output[group_rank * M : (group_rank + 1) * M, :] on all ranks - for i in tl.static_range(world_size): + # Write the local copy first, then the peers in a per-rank rotated order so that each destination + # sees a different writer at any given moment. + for step in tl.static_range(world_size): # Compute global output row indices: offset by group_rank * M rm_output = rm_input + group_rank * M @@ -130,12 +132,12 @@ def persistent_all_gather( output_ptr_target = output_ptr + output_offset output_ptr_target = tl.multiple_of(output_ptr_target, (BLOCK_SIZE_M, BLOCK_SIZE_N)) - if i == group_rank: - # Local destination (i == group_rank): use direct store + if step == 0: + # Local destination (step 0, i.e. group_rank): use direct store tl.store(output_ptr_target, data, mask=combined_mask, cache_modifier=".wt") else: # Remote destination: store through the pre-translated output base - remote_ptr = remote_output_ptrs[i] + output_offset + remote_ptr = remote_output_ptrs[(group_rank + step) % world_size] + output_offset remote_ptr = tl.max_contiguous(tl.multiple_of(remote_ptr, (1, BLOCK_SIZE_N)), (1, BLOCK_SIZE_N)) tl.store(remote_ptr, data, mask=combined_mask) diff --git a/iris/mem/triton/context.py b/iris/mem/triton/context.py index bc8a6f192..40dc422d3 100644 --- a/iris/mem/triton/context.py +++ b/iris/mem/triton/context.py @@ -873,7 +873,8 @@ def all_gather(self, tile: Tile, dst_view: TensorView, dim: tl.constexpr): # Start with the local copy, then write to peers in a per-rank rotated order so that each # destination sees a different writer at any given moment. - tl.store(dst_ptr, tile.data, mask=combined_mask) + local_ptr = tl.max_contiguous(tl.multiple_of(dst_ptr, (1, tile.block_n)), (1, tile.block_n)) + tl.store(local_ptr, tile.data, mask=combined_mask) for step in tl.static_range(1, self.world_size): remote_ptr = tl.cast(ptr_int + ptr_deltas[step - 1], dst_ptr.dtype) remote_ptr = tl.max_contiguous(tl.multiple_of(remote_ptr, (1, tile.block_n)), (1, tile.block_n)) From 401f0b3c87f17fe204e36eed6db2fa5c19763acd Mon Sep 17 00:00:00 2001 From: David Sidler Date: Fri, 6 Nov 2026 18:28:19 -0600 Subject: [PATCH 6/6] Start each program at a different peer in ctx.all_gather on CDNA3/4 All programs on a rank reach the all-gather epilogue at about the same time and, with a per-rank rotation alone, all start writing to the same peer, so they share one link. On gfx942/gfx950 offset the peer order by program id so concurrent programs use different outgoing links. Other targets keep the per-rank rotation. Co-Authored-By: Claude --- iris/mem/triton/context.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/iris/mem/triton/context.py b/iris/mem/triton/context.py index 40dc422d3..888f4b36a 100644 --- a/iris/mem/triton/context.py +++ b/iris/mem/triton/context.py @@ -8,6 +8,7 @@ import triton import triton.language as tl from triton.language.core import _aggregate as aggregate +from triton.language.target_info import is_hip_cdna3, is_hip_cdna4 from iris.mem.utils import get_xcc_id, get_cu_id, read_realtime # noqa: F401 — used by Tracing from iris.mem.triton.tracing import Tracing from iris.mem.triton.types import Tile, TileView, TensorView @@ -868,11 +869,16 @@ def all_gather(self, tile: Tile, dst_view: TensorView, dim: tl.constexpr): ptr_int = tl.cast(dst_ptr, tl.uint64) ptr_deltas = () for step in tl.static_range(1, self.world_size): - peer = (self.rank + step) % self.world_size + if is_hip_cdna3() or is_hip_cdna4(): + # Programs that start at the same peer queue behind each other on one link, so start each + # program at a different peer. + peer = (self.rank + 1 + (step - 1 + tl.program_id(0)) % (self.world_size - 1)) % self.world_size + else: + # Rotate by rank so that each destination sees a different writer at any given moment. + peer = (self.rank + step) % self.world_size ptr_deltas += (tl.load(self.heap_bases + peer) - local_base,) - # Start with the local copy, then write to peers in a per-rank rotated order so that each - # destination sees a different writer at any given moment. + # Start with the local copy, then write to the peers in the order chosen above. local_ptr = tl.max_contiguous(tl.multiple_of(dst_ptr, (1, tile.block_n)), (1, tile.block_n)) tl.store(local_ptr, tile.data, mask=combined_mask) for step in tl.static_range(1, self.world_size):