diff --git a/benchmark/ccl/bench_all_gather.py b/benchmark/ccl/bench_all_gather.py index e491ea934..dfff8cc4b 100644 --- a/benchmark/ccl/bench_all_gather.py +++ b/benchmark/ccl/bench_all_gather.py @@ -16,7 +16,7 @@ @bench.axis("M", M_VALUES) @bench.axis("N", N_VALUES) @bench.axis("dtype", DTYPES) -@bench.axis("variant", ["persistent", "tdm"]) +@bench.axis("variant", ["persistent", "gluon", "tdm"]) @bench.axis("comm_sms", [64, 128, 256]) @bench.axis("backend", ["iris", "rccl"]) def all_gather(state, ctx): @@ -55,6 +55,8 @@ def all_gather(state, ctx): num_warps=16, comm_sms=comm_sms, ) + elif variant == "gluon": + config = Config(use_gluon=True, block_size_m=8, block_size_n=256, comm_sms=comm_sms) else: config = Config(comm_sms=comm_sms) diff --git a/iris/ccl/gluon/all_gather.py b/iris/ccl/gluon/all_gather.py index 5ae3fa129..37ee79663 100644 --- a/iris/ccl/gluon/all_gather.py +++ b/iris/ccl/gluon/all_gather.py @@ -13,7 +13,6 @@ """ try: - import triton.language as tl from triton.experimental import gluon from triton.experimental.gluon import language as gl from triton.language.extra.hip import utils as hip_utils @@ -114,9 +113,15 @@ def persistent_all_gather_gluon( ELEMS_PER_THREAD: gl.constexpr = TOTAL_ELEMS // THREADS flat_layout: gl.constexpr = gl.BlockedLayout([ELEMS_PER_THREAD], [THREADS // WARPS], [WARPS], [0]) - # Hoist local heap base outside the tile loop: eliminates redundant - # gl.load(heap_bases) calls in the inner store loop. + # Heap bases do not change between tiles. local_base = gl.load(ctx.heap_bases + iris_rank) + output_bases = (output_ptr,) + for rank_idx in gl.static_range(1, world_size): + dest_idx = (group_rank + rank_idx) % world_size + target_iris_rank = rank_start + dest_idx * rank_stride + ptr_delta = gl.load(ctx.heap_bases + target_iris_rank) - local_base + remote_base = gl.cast(gl.cast(output_ptr, gl.uint64) + ptr_delta, output_ptr.dtype) + output_bases += (remote_base,) for tile_id in range(pid, total_tiles, COMM_SMS): # Swizzled tile index computation for better L2 locality @@ -146,26 +151,16 @@ def persistent_all_gather_gluon( # Output: this rank's data goes to output[group_rank * M + row, col] output_row = group_rank * M + row output_offsets = output_row * stride_out_m + col * stride_out_n + output_ptrs = output_ptr + output_offsets # Traffic-shaped stores to all ranks: stagger write order per rank # so each rank writes to a different target at any given moment, # avoiding memory controller contention on the receiver side. - for rank_idx in range(world_size): - dest_idx = (group_rank + rank_idx) % world_size - target_iris_rank = rank_start + dest_idx * rank_stride - output_ptrs = output_ptr + output_offsets - - if dest_idx == group_rank: + for rank_idx in gl.static_range(world_size): + if rank_idx == 0: gl.store(output_ptrs, data, mask=mask, cache_modifier=".wt") else: - # Hoisted translation: compute ptr_delta from pre-loaded - # local_base rather than calling ctx.store() which would - # do 2x gl.load(heap_bases) per call. - target_base = gl.load(ctx.heap_bases + target_iris_rank) - ptr_delta = target_base - local_base - output_ptrs_int = tl.cast(output_ptrs, gl.uint64) - remote_ptrs_int = output_ptrs_int + ptr_delta - remote_ptrs = tl.cast(remote_ptrs_int, output_ptrs.dtype) + remote_ptrs = output_bases[rank_idx] + output_offsets gl.store(remote_ptrs, data, mask=mask) diff --git a/tests/ccl/test_all_gather_gluon.py b/tests/ccl/test_all_gather_gluon.py index 438f01cfe..614a694e9 100644 --- a/tests/ccl/test_all_gather_gluon.py +++ b/tests/ccl/test_all_gather_gluon.py @@ -32,18 +32,20 @@ ], ) @pytest.mark.parametrize( - "M, N, block_size_m, block_size_n", + "M, N, block_size_m, block_size_n, comm_sms", [ # block_size_n must be a multiple of (threads_per_warp * num_warps). # With num_warps=4 the minimum is 256 on a 64-wide wavefront, 128 on a 32-wide one. # elems_per_thread = block_size_n / (threads_per_warp * num_warps): higher = wider # vector loads. The per-case counts below are for a 64-wide wavefront; double them at 32. - (256, 256, 32, 256), # Small: elems_per_thread=1 (scalar loads) - (1024, 512, 32, 512), # Medium: elems_per_thread=2 (dword loads) - (8192, 8192, 32, 1024), # Large: elems_per_thread=4 (dwordx4, optimal) + (256, 256, 32, 256, 64), # Small: elems_per_thread=1 (scalar loads) + (1024, 512, 32, 512, 64), # Medium: elems_per_thread=2 (dword loads) + (8192, 8192, 32, 1024, 64), # Large: elems_per_thread=4 (dwordx4, optimal) + (13, 300, 8, 256, 1), # Reuse peer addresses across tiles with partial rows and columns. + (7, 64, 8, 256, 1), ], ) -def test_all_gather_gluon(dtype, M, N, block_size_m, block_size_n): +def test_all_gather_gluon(dtype, M, N, block_size_m, block_size_n, comm_sms): """Test all-gather functionality using Gluon by comparing against PyTorch's implementation.""" # Ensure torch.distributed is initialized (should be done by test runner) if not dist.is_initialized(): @@ -61,8 +63,6 @@ def test_all_gather_gluon(dtype, M, N, block_size_m, block_size_n): # Each rank has an M x N input tensor # Output is (world_size * M, N) - concatenated along dimension 0 pytorch_input_tensor = torch.randn(M, N, dtype=dtype, device=f"cuda:{rank}") - # Fill with deterministic values for easier debugging - pytorch_input_tensor.fill_(float(rank + 1)) # Create output tensor for PyTorch: (world_size * M, N) pytorch_output_tensor = torch.zeros(world_size * M, N, dtype=dtype, device=f"cuda:{rank}") @@ -80,19 +80,13 @@ def test_all_gather_gluon(dtype, M, N, block_size_m, block_size_n): # Run Iris Gluon all_gather shmem.barrier() - config = Config(use_gluon=True, block_size_m=block_size_m, block_size_n=block_size_n) + config = Config(use_gluon=True, block_size_m=block_size_m, block_size_n=block_size_n, comm_sms=comm_sms) shmem.ccl.all_gather(iris_output_tensor, iris_input_tensor, config=config) torch.cuda.synchronize() # Compare results - atol = 1e-3 if dtype == torch.float16 else 1e-5 - max_diff = torch.abs(iris_output_tensor - pytorch_output_tensor).max().item() - try: - assert torch.allclose(iris_output_tensor, pytorch_output_tensor, atol=atol), ( - f"Max difference: {max_diff}, expected < {atol}\n" - f"Rank {rank}: Iris Gluon output doesn't match PyTorch's all_gather_into_tensor" - ) + torch.testing.assert_close(iris_output_tensor, pytorch_output_tensor, rtol=0, atol=0) finally: # Final barrier to ensure all ranks complete before test cleanup # This helps with test isolation when running multiple tests diff --git a/tests/ccl/test_process_groups.py b/tests/ccl/test_process_groups.py index 1ddd1116d..94ba7d894 100644 --- a/tests/ccl/test_process_groups.py +++ b/tests/ccl/test_process_groups.py @@ -178,7 +178,7 @@ def test_all_reduce_with_groups(variant, group_type, dtype=torch.float32, M=256, # ============================================================================= -@pytest.mark.parametrize("variant", ["persistent", "tdm"]) +@pytest.mark.parametrize("variant", ["persistent", "gluon", "tdm"]) @pytest.mark.parametrize("group_type", ["consecutive", "strided"]) def test_all_gather_with_groups(group_type, variant, dtype=torch.float32, M=128, N=64): """Test all-gather with ProcessGroups.""" @@ -222,7 +222,7 @@ def test_all_gather_with_groups(group_type, variant, dtype=torch.float32, M=128, # N=64 fp32 is a 256-byte row: one 32 x 64 tile per 32 rows. config = Config(use_gluon=True, all_gather_variant="tdm", block_size_m=32, block_size_n=64) else: - config = Config() + config = Config(use_gluon=variant == "gluon") shmem.ccl.all_gather(iris_output_tensor, iris_input_tensor, group=my_group, config=config) torch.cuda.synchronize()