diff --git a/iris/ccl/triton/all_gather.py b/iris/ccl/triton/all_gather.py index 16dd586a3..69a538041 100644 --- a/iris/ccl/triton/all_gather.py +++ b/iris/ccl/triton/all_gather.py @@ -64,6 +64,16 @@ 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) + + # 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 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): num_pid_in_group = GROUP_SIZE_M * num_pid_n group_id = tile_id // num_pid_in_group @@ -103,9 +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): - target_rank = rank_start + i * rank_stride - + # 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 @@ -122,21 +132,14 @@ 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: 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), - ) + # Remote destination: store through the pre-translated output base + 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) @triton.jit() diff --git a/iris/mem/triton/context.py b/iris/mem/triton/context.py index 63ee446e9..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 @@ -862,8 +863,28 @@ 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): + 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 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): + 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):