diff --git a/benchmark/ops/bench_all_gather_matmul.py b/benchmark/ops/bench_all_gather_matmul.py index daf98492e..0c9a87529 100644 --- a/benchmark/ops/bench_all_gather_matmul.py +++ b/benchmark/ops/bench_all_gather_matmul.py @@ -14,6 +14,7 @@ import torch import torch.distributed as dist +import tritonblas import iris.bench as bench from iris.ops.all_gather_matmul_hbm_buffer import ( all_gather_matmul_hbm_buffer as _hbm_buffer, @@ -57,6 +58,40 @@ def _run(): state.exec(_run) +@bench.register +@bench.axis("num_ranks", [2, 4, 8]) +@bench.axis("M", [1024, 4096, 16384]) +@bench.axis("N", [3584]) +@bench.axis("K", [8192]) +@bench.axis("dtype", [torch.float16]) +def tritonblas_rccl_all_gather_matmul(state, ctx): + """RCCL all_gather + tritonBLAS matmul baseline (same GEMM as the iris ops).""" + M, N, K = state["M"], state["N"], state["K"] + dtype = state["dtype"] + world_size = dist.get_world_size() + rank = dist.get_rank() + K_local = K // world_size + + # Per-rank seed for A (each rank holds different shards); shared seed for B + A_sharded = ctx.randn((M, K_local), dtype=dtype, generator=torch.Generator("cuda").manual_seed(42 + rank)) + B = ctx.randn((K, N), dtype=dtype, generator=torch.Generator("cuda").manual_seed(123)) + A_gathered_list = [ctx.zeros((M, K_local), dtype=dtype) for _ in range(world_size)] + C = ctx.zeros((M, N), dtype=dtype) + + selector = tritonblas.OrigamiMatmulSelector(M, N, K, dtype, dtype, dtype, A_sharded.device) + config = tritonblas.matmul_preamble(selector) + + state.set_flops(2 * M * N * K) + state.set_bytes((world_size - 1) * M * K_local * A_sharded.element_size()) + + def _run(): + dist.all_gather(A_gathered_list, A_sharded) + A_gathered = torch.cat(A_gathered_list, dim=1) + tritonblas.matmul_lt(A_gathered, B, C, selector, config) + + state.exec(_run) + + @bench.register @bench.axis("num_ranks", [2, 4, 8]) @bench.axis("M", [1024, 4096, 16384]) diff --git a/benchmark/ops/bench_all_gather_matmul_copy_engine.py b/benchmark/ops/bench_all_gather_matmul_copy_engine.py new file mode 100644 index 000000000..3d862c272 --- /dev/null +++ b/benchmark/ops/bench_all_gather_matmul_copy_engine.py @@ -0,0 +1,112 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. + +"""Benchmarks for all-gather + GEMM copy-engine variants.""" + +import torch +import iris.bench as bench + +from iris.ops.all_gather_matmul_copy_engine import ( + all_gather_matmul_copy_engine as _copy_engine, + all_gather_matmul_copy_engine_preamble, +) +from tritonblas.matmul import _make_matmul_selector + + +def _make_selector(M: int, N: int, K: int, dtype: torch.dtype, device: torch.device): + return _make_matmul_selector( + M, + N, + K, + dtype, + dtype, + dtype, + device, + streamk=False, + ) + + +def _register_copy_engine(state, ctx, *, device_initiated: bool) -> None: + if ctx.get_copy_engine_ctx() is None: + state.skip("SDMA engine is not enabled; run with --engines sdma") + M, N, K = state["M"], state["N"], state["K"] + dtype = state["dtype"] + rank = ctx.get_rank() + world_size = ctx.get_num_ranks() + if K % world_size != 0: + state.skip(f"K={K} must be divisible by world_size={world_size}") + + K_local = K // world_size + device = torch.device(f"cuda:{torch.cuda.current_device()}") + selector = _make_selector(M, N, K, dtype, device) + + if M % selector.block_m != 0: + state.skip(f"M={M} must be divisible by block_size_m={selector.block_m}") + if K % selector.block_k != 0: + state.skip(f"K={K} must be divisible by block_size_k={selector.block_k}") + if K_local % selector.block_k != 0: + state.skip(f"K_local={K_local} must be divisible by block_size_k={selector.block_k}") + + A_sharded = ctx.zeros((M, K_local), dtype=dtype) + torch.manual_seed(123 + rank) + A_sharded_data = torch.randn((M, K_local), device="cuda", dtype=dtype) + A_sharded.copy_(A_sharded_data) + torch.manual_seed(456) + B = torch.randn((K, N), device="cuda", dtype=dtype) + C = ctx.zeros((M, N), dtype=dtype) + + workspace = all_gather_matmul_copy_engine_preamble( + ctx, + A_sharded, + B, + selector=selector, + ) + m_tiles_per_batch = workspace.m_tiles_per_batch + + flag_iteration = [0] + + def _run(): + _copy_engine( + ctx, + C, + A_sharded, + B, + async_op=False, + workspace=workspace, + flag_iteration=flag_iteration[0], + device_initiated=device_initiated, + ) + flag_iteration[0] += 1 + + state.set_flops(2 * M * N * K) + state.set_bytes((world_size - 1) * M * K_local * A_sharded.element_size()) + state.add_counter("group_size_m", float(selector.group_m)) + state.add_counter("m_tiles_per_batch", float(m_tiles_per_batch)) + state.add_counter("device_initiated", 1.0 if device_initiated else 0.0) + + state.exec(_run) + + +@bench.register +@bench.axis("num_ranks", [8]) +@bench.axis("M", [1024, 4096, 16384]) +@bench.axis("N", [3584]) +@bench.axis("K", [8192]) +@bench.axis("dtype", [torch.float16]) +def all_gather_matmul_copy_engine_host(state, ctx): + _register_copy_engine(state, ctx, device_initiated=False) + + +@bench.register +@bench.axis("num_ranks", [8]) +@bench.axis("M", [1024, 4096, 16384]) +@bench.axis("N", [3584]) +@bench.axis("K", [8192]) +@bench.axis("dtype", [torch.float16]) +def all_gather_matmul_copy_engine_device(state, ctx): + _register_copy_engine(state, ctx, device_initiated=True) + + +if __name__ == "__main__": + bench.main() diff --git a/benchmark/ops/bench_matmul_all_gather.py b/benchmark/ops/bench_matmul_all_gather.py index fd7dff480..98aae662e 100644 --- a/benchmark/ops/bench_matmul_all_gather.py +++ b/benchmark/ops/bench_matmul_all_gather.py @@ -2,11 +2,23 @@ # SPDX-License-Identifier: MIT # Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. -"""Benchmark for fused GEMM + all-gather (iris.ops).""" +"""Benchmarks for fused GEMM + all-gather (iris.ops) and its RCCL baselines.""" import torch +import torch.distributed as dist +import tritonblas import iris.bench as bench from iris.ops import FusedConfig +from tritonblas.matmul import persistent_matmul_lt + + +def _make_inputs(ctx, M_local, N, K, dtype): + # Per-rank seed for A (each rank holds a different shard); shared seed for B. + # Same seeding as the all-gather + matmul benchmarks. + rank = ctx.get_rank() + A = ctx.randn((M_local, K), dtype=dtype, generator=torch.Generator("cuda").manual_seed(42 + rank)) + B = ctx.randn((K, N), dtype=dtype, generator=torch.Generator("cuda").manual_seed(123)) + return A, B @bench.register @@ -21,9 +33,7 @@ def matmul_all_gather(state, ctx): world_size = ctx.get_num_ranks() M = M_local * world_size - A = ctx.zeros((M_local, K), dtype=dtype) - A.fill_(1.0) - B = torch.randn((K, N), device="cuda", dtype=dtype) + A, B = _make_inputs(ctx, M_local, N, K, dtype) C = ctx.zeros((M, N), dtype=dtype) config = FusedConfig() @@ -37,5 +47,62 @@ def matmul_all_gather(state, ctx): ) +@bench.register +@bench.axis("num_ranks", [2, 4, 8]) +@bench.axis("M_local", [1024, 4096, 16384]) +@bench.axis("N", [3584]) +@bench.axis("K", [8192]) +@bench.axis("dtype", [torch.float16]) +def pytorch_matmul_all_gather(state, ctx): + """PyTorch/RCCL baseline: torch.mm + all_gather_into_tensor.""" + M_local, N, K = state["M_local"], state["N"], state["K"] + dtype = state["dtype"] + world_size = ctx.get_num_ranks() + M = M_local * world_size + + A, B = _make_inputs(ctx, M_local, N, K, dtype) + C_local = ctx.zeros((M_local, N), dtype=dtype) + C = ctx.zeros((M, N), dtype=dtype) + + state.set_flops(2 * M_local * N * K) + state.set_bytes((world_size - 1) * M_local * N * A.element_size()) + + state.exec( + lambda: ( + torch.mm(A, B, out=C_local), + dist.all_gather_into_tensor(C, C_local), + ), + ) + + +@bench.register +@bench.axis("num_ranks", [2, 4, 8]) +@bench.axis("M_local", [1024, 4096, 16384]) +@bench.axis("N", [3584]) +@bench.axis("K", [8192]) +@bench.axis("dtype", [torch.float16]) +def tritonblas_matmul_all_gather(state, ctx): + """tritonBLAS GEMM + RCCL all_gather_into_tensor baseline (same GEMM as the iris ops).""" + M_local, N, K = state["M_local"], state["N"], state["K"] + dtype = state["dtype"] + world_size = ctx.get_num_ranks() + M = M_local * world_size + + A, B = _make_inputs(ctx, M_local, N, K, dtype) + C_local = ctx.zeros((M_local, N), dtype=dtype) + C = ctx.zeros((M, N), dtype=dtype) + selector = tritonblas.OrigamiMatmulSelector(M_local, N, K, A.dtype, B.dtype, C_local.dtype, A.device) + + state.set_flops(2 * M_local * N * K) + state.set_bytes((world_size - 1) * M_local * N * A.element_size()) + + state.exec( + lambda: ( + persistent_matmul_lt(A, B, C_local, selector, config=None, work_stealing=False), + dist.all_gather_into_tensor(C, C_local), + ), + ) + + if __name__ == "__main__": bench.main() diff --git a/benchmark/ops/bench_matmul_all_gather_copy_engine.py b/benchmark/ops/bench_matmul_all_gather_copy_engine.py new file mode 100644 index 000000000..626f7712a --- /dev/null +++ b/benchmark/ops/bench_matmul_all_gather_copy_engine.py @@ -0,0 +1,131 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. + +"""Benchmarks for fused GEMM + all-gather copy-engine variants.""" + +import torch +import iris.bench as bench + +from iris.ops.matmul_all_gather_copy_engine import ( + matmul_all_gather_copy_engine as _device_copy_engine, + matmul_all_gather_copy_engine_preamble as _device_preamble, +) +from iris.ops.matmul_all_gather_host_copy_engine import ( + matmul_all_gather_host_copy_engine as _host_copy_engine, + matmul_all_gather_host_copy_engine_preamble as _host_preamble, +) +from tritonblas.matmul import _make_matmul_selector + + +def _make_selector(M_local: int, N: int, K: int, dtype: torch.dtype, device: torch.device): + return _make_matmul_selector( + M_local, + N, + K, + dtype, + dtype, + dtype, + device, + streamk=False, + ) + + +def _register_copy_engine(state, ctx, *, device_initiated: bool) -> None: + if ctx.get_copy_engine_ctx() is None: + state.skip("SDMA engine is not enabled; run with --engines sdma") + M_local, N, K = state["M_local"], state["N"], state["K"] + dtype = state["dtype"] + rank = ctx.get_rank() + world_size = ctx.get_num_ranks() + M = M_local * world_size + + device = torch.device(f"cuda:{torch.cuda.current_device()}") + selector = _make_selector(M_local, N, K, dtype, device) + + if M_local % selector.block_m != 0: + state.skip(f"M_local={M_local} must be divisible by block_size_m={selector.block_m}") + if K % selector.block_k != 0: + state.skip(f"K={K} must be divisible by block_size_k={selector.block_k}") + + torch.manual_seed(123 + rank) + A = ctx.randn((M_local, K), dtype=dtype) + torch.manual_seed(456) + B = ctx.randn((K, N), dtype=dtype) + C = ctx.zeros((M, N), dtype=dtype) + + flag_iteration = [0] + + if device_initiated: + workspace = _device_preamble( + ctx, + A, + B, + selector=selector, + ) + + def _run(): + _device_copy_engine( + ctx, + C, + A, + B, + async_op=False, + workspace=workspace, + flag_iteration=flag_iteration[0], + ) + flag_iteration[0] += 1 + + else: + workspace = _host_preamble( + ctx, + A, + B, + trace=False, + selector=selector, + ) + + def _run(): + _host_copy_engine( + ctx, + C, + A, + B, + async_op=False, + workspace=workspace, + flag_iteration=flag_iteration[0], + trace=False, + ) + flag_iteration[0] += 1 + + state.set_flops(2 * M_local * N * K) + state.set_bytes((world_size - 1) * M_local * N * A.element_size()) + state.add_counter("group_size_m", float(selector.group_m)) + state.add_counter("m_tiles_per_batch", float(workspace.m_tiles_per_batch)) + state.add_counter("device_initiated", 1.0 if device_initiated else 0.0) + + state.exec(_run) + + +@bench.register +@bench.axis("num_ranks", [8]) +@bench.axis("M_local", [1024, 4096, 16384]) +@bench.axis("N", [3584]) +@bench.axis("K", [8192]) +@bench.axis("dtype", [torch.float16]) +def matmul_all_gather_copy_engine_host(state, ctx): + _register_copy_engine(state, ctx, device_initiated=False) + + +@bench.register +@bench.axis("num_ranks", [8]) +@bench.axis("M_local", [1024, 4096, 16384]) +@bench.axis("N", [3584]) +@bench.axis("K", [8192]) +@bench.axis("dtype", [torch.float16]) +def matmul_all_gather_copy_engine_device(state, ctx): + _register_copy_engine(state, ctx, device_initiated=True) + + +if __name__ == "__main__": + bench.main() diff --git a/iris/__init__.py b/iris/__init__.py index 38125e1ed..2b1f89ee8 100644 --- a/iris/__init__.py +++ b/iris/__init__.py @@ -58,7 +58,13 @@ copy, get, put, + wait_then_put_rect, + wait_then_put_rects, + wait_then_put_signal_rect, + wait_then_put_signal_rects, quiet, + put_signal, + put_signal_rect, atomic_add, atomic_cas, atomic_xchg, @@ -112,7 +118,13 @@ "copy", "get", "put", + "wait_then_put_rect", + "wait_then_put_rects", + "wait_then_put_signal_rect", + "wait_then_put_signal_rects", "quiet", + "put_signal", + "put_signal_rect", "atomic_add", "atomic_cas", "atomic_xchg", diff --git a/iris/bench/_runner.py b/iris/bench/_runner.py index 0e11d9780..374777208 100644 --- a/iris/bench/_runner.py +++ b/iris/bench/_runner.py @@ -336,6 +336,7 @@ def _run_benchmarks_worker( skip_overrides: dict[str, list[Any]], heap_size: int, use_gluon: bool, + engines: tuple[str, ...], n_warmup: int, n_repeat: int, benchmark_filter: str | None, @@ -355,7 +356,7 @@ def _run_benchmarks_worker( dist.init_process_group(backend=backend) # Create iris context - ctx = _iris.iris(heap_size) + ctx = _iris.iris(heap_size, engines=engines) rank = ctx.get_rank() @@ -528,6 +529,12 @@ def _build_parser() -> argparse.ArgumentParser: action="store_true", help="Use Gluon backend", ) + parser.add_argument( + "--engines", + type=str, + default="", + help="Comma-separated Iris copy engines to set up, e.g. 'sdma' (needs `pip install iris[sdma]`)", + ) parser.add_argument( "--n_warmup", type=int, @@ -564,6 +571,7 @@ def main(argv: list[str] | None = None) -> None: """ parser = _build_parser() args, remaining = parser.parse_known_args(argv) + engines = tuple(e.strip() for e in args.engines.split(",") if e.strip()) # Parse --axis_= and --skip_= from remaining args axis_overrides: dict[str, list[Any]] = {} @@ -625,6 +633,7 @@ def main(argv: list[str] | None = None) -> None: skip_overrides, args.heap_size, args.use_gluon, + engines, args.n_warmup, args.n_repeat, args.benchmark_filter, @@ -663,6 +672,7 @@ def main(argv: list[str] | None = None) -> None: skip_overrides, args.heap_size, args.use_gluon, + engines, args.n_warmup, args.n_repeat, args.benchmark_filter, diff --git a/iris/host/iris.py b/iris/host/iris.py index a399c7e65..fb56083f1 100644 --- a/iris/host/iris.py +++ b/iris/host/iris.py @@ -1133,32 +1133,36 @@ def put( wait_val = int(wait_value if wait_value is not None else 0) signal_val = int(signal_value) sdma_ep.wait_flag_then_put( - from_rank, to_rank, channel, wait_ptr, wait_val, from_ptr, to_ptr, size, wait_bits + from_rank, to_rank, channel, wait_ptr, wait_val, to_ptr, from_ptr, size, wait_bits ) sdma_ep.signal(from_rank, to_rank, channel, signal_ptr, signal_val, signal_bits) elif has_wait: # Wait + copy wait_val = int(wait_value if wait_value is not None else 0) sdma_ep.wait_flag_then_put( - from_rank, to_rank, channel, wait_ptr, wait_val, from_ptr, to_ptr, size, wait_bits + from_rank, to_rank, channel, wait_ptr, wait_val, to_ptr, from_ptr, size, wait_bits ) elif has_signal: # Copy + signal signal_val = int(signal_value) - sdma_ep.put_signal(from_rank, to_rank, channel, from_ptr, to_ptr, size, signal_ptr, signal_val, signal_bits) + sdma_ep.put_signal(from_rank, to_rank, channel, to_ptr, from_ptr, size, signal_ptr, signal_val, signal_bits) else: # Simple copy - sdma_ep.put(from_rank, to_rank, channel, from_ptr, to_ptr, size) + sdma_ep.put(from_rank, to_rank, channel, to_ptr, from_ptr, size) if not async_op: sdma_ep.quiet(from_rank, to_rank, channel) def put_tile( self, - tile, + global_addr: int, to_rank: int, - to_ptr: int, - to_stride: int, + dst_addr: int, + data_size: int, + tile_height: int, + tile_width: int, + src_stride: int, + dst_stride: int, wait_flag: int = None, wait_value: int = None, signal_flag: int = None, @@ -1174,10 +1178,14 @@ def put_tile( Low-level API - caller provides pre-translated pointers for performance. Args: - tile: Pre-configured sdma_ep.Tile object with data pointer and dimensions set + global_addr: Source address at the first element of the tile to_rank: Destination rank - to_ptr: Destination pointer (already translated to remote address space) - to_stride: Destination row stride in bytes + dst_addr: Destination address (already translated to remote address space) + data_size: Element size in bytes + tile_height: Number of rows + tile_width: Number of elements per row + src_stride: Source row stride in elements + dst_stride: Destination row stride in elements wait_flag: Optional LOCAL flag pointer to poll before transfer wait_value: Expected value for wait_flag signal_flag: Optional REMOTE flag pointer to atomic-add after transfer (already translated) @@ -1185,22 +1193,6 @@ def put_tile( async_op: If True, don't wait for completion channel: SDMA channel to use - Examples: - >>> from xio import sdma_ep - >>> tile = sdma_ep.Tile() - >>> tile.pid_m = 0 - >>> tile.pid_n = 0 - >>> tile.block_m = 256 - >>> tile.block_n = 256 - >>> tile.elem_size = A.element_size() - >>> tile.src_stride = A.stride(0) * tile.elem_size - >>> tile.data = A.data_ptr() - >>> dst_ptr = shmem.translate(A.data_ptr(), src_rank, dst_rank) - >>> dst_stride = A.stride(0) * tile.elem_size - >>> wait_ptr = flag.data_ptr() - >>> signal_ptr = shmem.translate(flag.data_ptr(), src_rank, dst_rank) - >>> shmem.put_tile(tile, to_rank=1, to_ptr=dst_ptr, to_stride=dst_stride, - ... wait_flag=wait_ptr, wait_value=256, signal_flag=signal_ptr) """ sdma_ep = self._require_copy_engine() from_rank = self.get_rank() @@ -1216,34 +1208,85 @@ def put_tile( wait_val = int(wait_value if wait_value is not None else 0) signal_val = int(signal_value) sdma_ep.wait_flag_then_put_tile( - from_rank, to_rank, channel, wait_ptr, wait_val, tile, int(to_ptr), int(to_stride), wait_bits + from_rank, + to_rank, + channel, + wait_ptr, + wait_val, + int(dst_addr), + int(global_addr), + int(data_size), + int(tile_height), + int(tile_width), + int(dst_stride), + int(src_stride), + wait_bits, ) sdma_ep.signal(from_rank, to_rank, channel, signal_ptr, signal_val, signal_bits) elif has_wait: # Wait + tile copy wait_val = int(wait_value if wait_value is not None else 0) sdma_ep.wait_flag_then_put_tile( - from_rank, to_rank, channel, wait_ptr, wait_val, tile, int(to_ptr), int(to_stride), wait_bits + from_rank, + to_rank, + channel, + wait_ptr, + wait_val, + int(dst_addr), + int(global_addr), + int(data_size), + int(tile_height), + int(tile_width), + int(dst_stride), + int(src_stride), + wait_bits, ) elif has_signal: # Tile copy + signal signal_val = int(signal_value) sdma_ep.put_tile_signal( - from_rank, to_rank, channel, tile, int(to_ptr), int(to_stride), signal_ptr, signal_val, signal_bits + from_rank, + to_rank, + channel, + int(dst_addr), + int(global_addr), + int(data_size), + int(tile_height), + int(tile_width), + int(dst_stride), + int(src_stride), + signal_ptr, + signal_val, + signal_bits, ) else: # Simple tile copy - sdma_ep.put_tile(from_rank, to_rank, channel, tile, int(to_ptr), int(to_stride)) + sdma_ep.put_tile( + from_rank, + to_rank, + channel, + int(dst_addr), + int(global_addr), + int(data_size), + int(tile_height), + int(tile_width), + int(dst_stride), + int(src_stride), + ) if not async_op: sdma_ep.quiet(from_rank, to_rank, channel) def put_tiles( self, - tiles, + global_addrs, to_rank: int, - to_ptrs, - to_strides, + dst_addrs, + data_size: int, + tile_heights, + tile_widths, + src_strides, + dst_strides, wait_flag: int = None, wait_value: int = None, signal_flag: int = None, @@ -1257,10 +1300,14 @@ def put_tiles( Requires initialization with ``engines=("sdma",)``. Args: - tiles: Sequence of pre-configured sdma_ep.Tile objects + global_addrs: Source addresses at the first element of each tile to_rank: Destination rank - to_ptrs: Sequence of translated destination pointers - to_strides: Sequence of destination row strides in bytes + dst_addrs: Sequence of translated destination pointers + data_size: Element size in bytes + tile_heights: Number of rows for each tile + tile_widths: Number of elements per row for each tile + src_strides: Source row strides in elements + dst_strides: Destination row strides in elements wait_flag: Optional LOCAL flag pointer to poll before all transfers wait_value: Expected value for wait_flag signal_flag: Optional REMOTE flag pointer to atomic-add after all transfers @@ -1271,8 +1318,15 @@ def put_tiles( sdma_ep = self._require_copy_engine() from_rank = self.get_rank() - if len(tiles) != len(to_ptrs) or len(tiles) != len(to_strides): - raise ValueError("tiles, to_ptrs, and to_strides must have the same length") + if not ( + len(global_addrs) + == len(dst_addrs) + == len(tile_heights) + == len(tile_widths) + == len(src_strides) + == len(dst_strides) + ): + raise ValueError("tile transfer arrays must have the same length") wait_ptr, wait_bits = self._flag_pointer_and_bits(wait_flag, default_bits=32) signal_ptr, signal_bits = self._flag_pointer_and_bits(signal_flag, default_bits=32) @@ -1280,31 +1334,81 @@ def put_tiles( has_wait = wait_ptr != 0 has_signal = signal_ptr != 0 - to_ptr_list = [int(p) for p in to_ptrs] - to_stride_list = [int(s) for s in to_strides] + src_ptr_list = [int(p) for p in global_addrs] + to_ptr_list = [int(p) for p in dst_addrs] + tile_height_list = [int(v) for v in tile_heights] + tile_width_list = [int(v) for v in tile_widths] + src_stride_list = [int(v) for v in src_strides] + dst_stride_list = [int(v) for v in dst_strides] if has_wait and has_signal: # Wait + tiles copy + signal (two calls) wait_val = int(wait_value if wait_value is not None else 0) signal_val = int(signal_value) sdma_ep.wait_flag_then_put_tiles( - from_rank, to_rank, channel, wait_ptr, wait_val, list(tiles), to_ptr_list, to_stride_list, wait_bits + from_rank, + to_rank, + channel, + wait_ptr, + wait_val, + to_ptr_list, + src_ptr_list, + int(data_size), + tile_height_list, + tile_width_list, + dst_stride_list, + src_stride_list, + wait_bits, ) sdma_ep.signal(from_rank, to_rank, channel, signal_ptr, signal_val, signal_bits) elif has_wait: # Wait + tiles copy wait_val = int(wait_value if wait_value is not None else 0) sdma_ep.wait_flag_then_put_tiles( - from_rank, to_rank, channel, wait_ptr, wait_val, list(tiles), to_ptr_list, to_stride_list, wait_bits + from_rank, + to_rank, + channel, + wait_ptr, + wait_val, + to_ptr_list, + src_ptr_list, + int(data_size), + tile_height_list, + tile_width_list, + dst_stride_list, + src_stride_list, + wait_bits, ) elif has_signal: # Tiles copy + signal (loop + signal) signal_val = int(signal_value) - sdma_ep.put_tiles(from_rank, to_rank, channel, list(tiles), to_ptr_list, to_stride_list) + sdma_ep.put_tiles( + from_rank, + to_rank, + channel, + to_ptr_list, + src_ptr_list, + int(data_size), + tile_height_list, + tile_width_list, + dst_stride_list, + src_stride_list, + ) sdma_ep.signal(from_rank, to_rank, channel, signal_ptr, signal_val, signal_bits) else: # Simple tiles copy - sdma_ep.put_tiles(from_rank, to_rank, channel, list(tiles), to_ptr_list, to_stride_list) + sdma_ep.put_tiles( + from_rank, + to_rank, + channel, + to_ptr_list, + src_ptr_list, + int(data_size), + tile_height_list, + tile_width_list, + dst_stride_list, + src_stride_list, + ) if not async_op: sdma_ep.quiet(from_rank, to_rank, channel) diff --git a/iris/mem/triton/ops.py b/iris/mem/triton/ops.py index 7c355484f..06824a2f6 100644 --- a/iris/mem/triton/ops.py +++ b/iris/mem/triton/ops.py @@ -995,3 +995,494 @@ def quiet(copy_engine_ctx: tl.tensor, to_rank): # Poll until hardware read pointer catches up while tl.load(read_ptr, cache_modifier=".cv", volatile=True) < target: pass + + +@triton.jit +def put_signal( + from_ptr, + to_ptr, + from_rank, + to_rank, + heap_bases, + copy_engine_ctx: tl.tensor, + flag_ptr, + flag_value, + mask=None, + hint: tl.constexpr = None, +): + """ + 1D copy with atomic signal using SDMA. + + This is a simpler version of put_signal_rect for 1D (flat) copies. + """ + # Translate destination pointers + translated_to_ptr = __translate(to_ptr, from_rank, to_rank, heap_bases, hint) + translated_flag_ptr = __translate(flag_ptr, from_rank, to_rank, heap_bases, hint) + + # Get copy engine context for target rank + ctx = copy_engine_ctx + (6 * to_rank) + queue_ptr_u32 = tl.load(ctx + 0).to(tl.pointer_type(tl.uint32)) + read_ptr = tl.load(ctx + 1).to(tl.pointer_type(tl.uint64)) + write_ptr = tl.load(ctx + 2).to(tl.pointer_type(tl.uint64)) + doorbell_ptr = tl.load(ctx + 3).to(tl.pointer_type(tl.uint64)) + cached_write_ptr = tl.load(ctx + 4).to(tl.pointer_type(tl.uint64)) + committed_write_ptr = tl.load(ctx + 5).to(tl.pointer_type(tl.uint64)) + + # For 1D: LINEAR_COPY (48 bytes) + ATOMIC (32 bytes) = 80 bytes + command_in_bytes = 80 + + # Calculate size from mask + size = tl.num_elements(mask) * 4 # Assuming 4-byte elements for now + + base, offset = sdma_utils.acquire_fadd( + queue_ptr_u32, read_ptr, write_ptr, doorbell_ptr, cached_write_ptr, committed_write_ptr, command_in_bytes + ) + + sdma_utils.place_nop_packet(queue_ptr_u32, base, offset) + + packet_offset_bytes = base + offset + sdma_utils.place_copy_packet( + queue_ptr_u32, + packet_offset_bytes, + size, + from_ptr.to(tl.uint64), + translated_to_ptr.to(tl.uint64), + ) + sdma_utils.place_atomic_add_packet( + queue_ptr_u32, + packet_offset_bytes + 48, + translated_flag_ptr.to(tl.uint64), + flag_value, + ) + + pending_wptr = base + offset + command_in_bytes + sdma_utils.submit(write_ptr, doorbell_ptr, committed_write_ptr, base, pending_wptr) + + +@triton.jit +def put_signal_rect( + from_ptr, + to_ptr, + from_rank, + to_rank, + heap_bases, + copy_engine_ctx: tl.tensor, + flag_ptr, + flag_value, + width_bytes: tl.constexpr, + height: tl.constexpr, + src_pitch: tl.constexpr, + dst_pitch: tl.constexpr, + hint: tl.constexpr = None, +): + """ + Combines 2D rectangular copy (put) with atomic_add signal in one SDMA submission. + + Unlike put_signal(), this function takes explicit dimensions instead of a mask, + allowing arbitrarily large tiles without hitting Triton's tensor size limit. + + Args: + from_ptr: Source base pointer (scalar) in current rank's local memory + to_ptr: Destination base pointer (scalar) - will be translated to to_rank's address space + from_rank: Current rank ID + to_rank: Remote rank ID to write to + heap_bases: Array of heap base addresses for all ranks + copy_engine_ctx: Copy engine context for SDMA operations + flag_ptr: Pointer to flag location for signaling + flag_value: Value to atomically add to flag (typically 1) + width_bytes: Width of rectangle in bytes + height: Height of rectangle in rows + src_pitch: Source row stride in bytes + dst_pitch: Destination row stride in bytes + hint: Vectorization hint for translated pointers + + Returns: + None + + Example: + >>> @triton.jit + >>> def kernel(A_sharded, staged_a, flags, heap_bases, ctx): + >>> iris.put_signal_rect( + >>> src_ptr, dst_ptr, 0, 1, heap_bases, ctx, flags, 1, + >>> width_bytes=1024, height=256, + >>> src_pitch=stride_am * 2, dst_pitch=stride_sa_m * 2 + >>> ) + """ + # Translate destination pointers + translated_to_ptr = __translate(to_ptr, from_rank, to_rank, heap_bases, hint) + translated_flag_ptr = __translate(flag_ptr, from_rank, to_rank, heap_bases, hint) + + # Get copy engine context for target rank + ctx = copy_engine_ctx + (6 * to_rank) + queue_ptr_u32 = tl.load(ctx + 0).to(tl.pointer_type(tl.uint32)) + read_ptr = tl.load(ctx + 1).to(tl.pointer_type(tl.uint64)) + write_ptr = tl.load(ctx + 2).to(tl.pointer_type(tl.uint64)) + doorbell_ptr = tl.load(ctx + 3).to(tl.pointer_type(tl.uint64)) + cached_write_ptr = tl.load(ctx + 4).to(tl.pointer_type(tl.uint64)) + committed_write_ptr = tl.load(ctx + 5).to(tl.pointer_type(tl.uint64)) + + # Extract addresses (scalar pointers) + src_ptr_val = from_ptr.to(tl.uint64) + dst_ptr_val = translated_to_ptr.to(tl.uint64) + flag_dst_ptr_val = translated_flag_ptr.to(tl.uint64) + + # Reserve space for BOTH packets: SUB_WINDOW_COPY (80 bytes) + ATOMIC (32 bytes) = 112 bytes + command_in_bytes = 112 + + # Acquire space in queue + base, offset = sdma_utils.acquire_fadd( + queue_ptr_u32, read_ptr, write_ptr, doorbell_ptr, cached_write_ptr, committed_write_ptr, command_in_bytes + ) + + # Write padding NOPs if we wrapped around + sdma_utils.place_nop_packet(queue_ptr_u32, base, offset) + + packet_offset_bytes = base + offset + sdma_utils.place_sub_window_copy_packet( + queue_ptr_u32, + packet_offset_bytes, + src_ptr_val, + dst_ptr_val, + tile_width=width_bytes, + tile_height=height, + src_buffer_pitch=src_pitch, + dst_buffer_pitch=dst_pitch, + src_x=0, + src_y=0, + dst_x=0, + dst_y=0, + ) + sdma_utils.place_atomic_add_packet( + queue_ptr_u32, + packet_offset_bytes + 80, # After SUB_WINDOW_COPY packet + flag_dst_ptr_val, + flag_value, + ) + + pending_wptr = base + offset + command_in_bytes + sdma_utils.submit(write_ptr, doorbell_ptr, committed_write_ptr, base, pending_wptr) + + +@triton.jit +def wait_then_put_rect( + from_ptr, + to_ptr, + from_rank, + to_rank, + heap_bases, + copy_engine_ctx: tl.tensor, + wait_flag_ptr, + wait_value, + width_bytes: tl.constexpr, + height: tl.constexpr, + src_pitch: tl.constexpr, + dst_pitch: tl.constexpr, + hint: tl.constexpr = None, +): + """ + Enqueue a POLL_REGMEM followed by a 2D SUB_WINDOW_COPY in one SDMA submission. + + This is the device-side counterpart to the host wait-then-put-tile path. + The SDMA queue waits on a local flag and performs the copy autonomously + after the producer has completed the corresponding batch. + """ + translated_to_ptr = __translate(to_ptr, from_rank, to_rank, heap_bases, hint) + + ctx = copy_engine_ctx + (6 * to_rank) + queue_ptr_u32 = tl.load(ctx + 0).to(tl.pointer_type(tl.uint32)) + read_ptr = tl.load(ctx + 1).to(tl.pointer_type(tl.uint64)) + write_ptr = tl.load(ctx + 2).to(tl.pointer_type(tl.uint64)) + doorbell_ptr = tl.load(ctx + 3).to(tl.pointer_type(tl.uint64)) + cached_write_ptr = tl.load(ctx + 4).to(tl.pointer_type(tl.uint64)) + committed_write_ptr = tl.load(ctx + 5).to(tl.pointer_type(tl.uint64)) + + poll_packet_bytes = 24 + copy_packet_bytes = 80 + command_in_bytes = poll_packet_bytes + copy_packet_bytes + + base, offset = sdma_utils.acquire_fadd( + queue_ptr_u32, read_ptr, write_ptr, doorbell_ptr, cached_write_ptr, committed_write_ptr, command_in_bytes + ) + sdma_utils.place_nop_packet(queue_ptr_u32, base, offset) + + packet_offset_bytes = base + offset + sdma_utils.place_poll_regmem_packet( + queue_ptr_u32, + packet_offset_bytes, + wait_flag_ptr.to(tl.uint64), + wait_value, + ) + sdma_utils.place_sub_window_copy_packet( + queue_ptr_u32, + packet_offset_bytes + poll_packet_bytes, + from_ptr.to(tl.uint64), + translated_to_ptr.to(tl.uint64), + tile_width=width_bytes, + tile_height=height, + src_buffer_pitch=src_pitch, + dst_buffer_pitch=dst_pitch, + src_x=0, + src_y=0, + dst_x=0, + dst_y=0, + ) + + pending_wptr = base + offset + command_in_bytes + sdma_utils.submit(write_ptr, doorbell_ptr, committed_write_ptr, base, pending_wptr) + + +@triton.jit +def wait_then_put_rects( + from_base_ptr, + to_base_ptr, + from_rank, + to_rank, + heap_bases, + copy_engine_ctx: tl.tensor, + wait_flag_ptr, + wait_value, + transfer_row_offsets, + transfer_col_offsets, + transfer_width_bytes, + transfer_heights, + transfer_start, + transfer_count, + stride_n_bytes, + src_pitch: tl.constexpr, + dst_pitch: tl.constexpr, + MAX_RECTS: tl.constexpr, + hint: tl.constexpr = None, +): + """ + Enqueue one POLL_REGMEM followed by many 2D SUB_WINDOW_COPY packets. + + The copy list is provided as flattened metadata arrays plus a per-wave + start/count pair so the poster can submit an entire wave with one queue + reservation and one doorbell ring. + """ + translated_to_base_ptr = __translate(to_base_ptr, from_rank, to_rank, heap_bases, hint) + + ctx = copy_engine_ctx + (6 * to_rank) + queue_ptr_u32 = tl.load(ctx + 0).to(tl.pointer_type(tl.uint32)) + read_ptr = tl.load(ctx + 1).to(tl.pointer_type(tl.uint64)) + write_ptr = tl.load(ctx + 2).to(tl.pointer_type(tl.uint64)) + doorbell_ptr = tl.load(ctx + 3).to(tl.pointer_type(tl.uint64)) + cached_write_ptr = tl.load(ctx + 4).to(tl.pointer_type(tl.uint64)) + committed_write_ptr = tl.load(ctx + 5).to(tl.pointer_type(tl.uint64)) + + poll_packet_bytes = 24 + copy_packet_bytes = 80 + command_in_bytes = poll_packet_bytes + transfer_count * copy_packet_bytes + + base, offset = sdma_utils.acquire_fadd( + queue_ptr_u32, read_ptr, write_ptr, doorbell_ptr, cached_write_ptr, committed_write_ptr, command_in_bytes + ) + sdma_utils.place_nop_packet(queue_ptr_u32, base, offset) + + packet_offset_bytes = base + offset + sdma_utils.place_poll_regmem_packet( + queue_ptr_u32, + packet_offset_bytes, + wait_flag_ptr.to(tl.uint64), + wait_value, + ) + + from_base_val = from_base_ptr.to(tl.uint64) + to_base_val = translated_to_base_ptr.to(tl.uint64) + + for i in range(MAX_RECTS): + if i < transfer_count: + transfer_idx = transfer_start + i + row_offset = tl.load(transfer_row_offsets + transfer_idx) + col_offset = tl.load(transfer_col_offsets + transfer_idx) + width_bytes = tl.load(transfer_width_bytes + transfer_idx) + height = tl.load(transfer_heights + transfer_idx) + byte_offset = (row_offset.to(tl.uint64) * src_pitch) + (col_offset.to(tl.uint64) * stride_n_bytes) + copy_offset_bytes = packet_offset_bytes + poll_packet_bytes + i * copy_packet_bytes + sdma_utils.place_sub_window_copy_packet( + queue_ptr_u32, + copy_offset_bytes, + from_base_val + byte_offset, + to_base_val + byte_offset, + tile_width=width_bytes, + tile_height=height, + src_buffer_pitch=src_pitch, + dst_buffer_pitch=dst_pitch, + src_x=0, + src_y=0, + dst_x=0, + dst_y=0, + ) + + pending_wptr = base + offset + command_in_bytes + sdma_utils.submit(write_ptr, doorbell_ptr, committed_write_ptr, base, pending_wptr) + + +@triton.jit +def wait_then_put_signal_rect( + from_ptr, + to_ptr, + from_rank, + to_rank, + heap_bases, + copy_engine_ctx: tl.tensor, + wait_flag_ptr, + wait_value, + signal_flag_ptr, + signal_value, + width_bytes: tl.constexpr, + height: tl.constexpr, + src_pitch: tl.constexpr, + dst_pitch: tl.constexpr, + hint: tl.constexpr = None, +): + """ + Enqueue POLL_REGMEM + 2D SUB_WINDOW_COPY + ATOMIC in one SDMA submission. + + This is the device-side counterpart to host-side wait/copy/signal flows and is + useful for marking receiver-visible completion after the final copy in a queue. + """ + translated_to_ptr = __translate(to_ptr, from_rank, to_rank, heap_bases, hint) + translated_signal_ptr = __translate(signal_flag_ptr, from_rank, to_rank, heap_bases, hint) + + ctx = copy_engine_ctx + (6 * to_rank) + queue_ptr_u32 = tl.load(ctx + 0).to(tl.pointer_type(tl.uint32)) + read_ptr = tl.load(ctx + 1).to(tl.pointer_type(tl.uint64)) + write_ptr = tl.load(ctx + 2).to(tl.pointer_type(tl.uint64)) + doorbell_ptr = tl.load(ctx + 3).to(tl.pointer_type(tl.uint64)) + cached_write_ptr = tl.load(ctx + 4).to(tl.pointer_type(tl.uint64)) + committed_write_ptr = tl.load(ctx + 5).to(tl.pointer_type(tl.uint64)) + + poll_packet_bytes = 24 + copy_packet_bytes = 80 + atomic_packet_bytes = 32 + command_in_bytes = poll_packet_bytes + copy_packet_bytes + atomic_packet_bytes + + base, offset = sdma_utils.acquire_fadd( + queue_ptr_u32, read_ptr, write_ptr, doorbell_ptr, cached_write_ptr, committed_write_ptr, command_in_bytes + ) + sdma_utils.place_nop_packet(queue_ptr_u32, base, offset) + + packet_offset_bytes = base + offset + sdma_utils.place_poll_regmem_packet( + queue_ptr_u32, + packet_offset_bytes, + wait_flag_ptr.to(tl.uint64), + wait_value, + ) + sdma_utils.place_sub_window_copy_packet( + queue_ptr_u32, + packet_offset_bytes + poll_packet_bytes, + from_ptr.to(tl.uint64), + translated_to_ptr.to(tl.uint64), + tile_width=width_bytes, + tile_height=height, + src_buffer_pitch=src_pitch, + dst_buffer_pitch=dst_pitch, + src_x=0, + src_y=0, + dst_x=0, + dst_y=0, + ) + sdma_utils.place_atomic_add_packet( + queue_ptr_u32, + packet_offset_bytes + poll_packet_bytes + copy_packet_bytes, + translated_signal_ptr.to(tl.uint64), + signal_value, + ) + + pending_wptr = base + offset + command_in_bytes + sdma_utils.submit(write_ptr, doorbell_ptr, committed_write_ptr, base, pending_wptr) + + +@triton.jit +def wait_then_put_signal_rects( + from_base_ptr, + to_base_ptr, + from_rank, + to_rank, + heap_bases, + copy_engine_ctx: tl.tensor, + wait_flag_ptr, + wait_value, + signal_flag_ptr, + signal_value, + transfer_row_offsets, + transfer_col_offsets, + transfer_width_bytes, + transfer_heights, + transfer_start, + transfer_count, + stride_n_bytes, + src_pitch: tl.constexpr, + dst_pitch: tl.constexpr, + MAX_RECTS: tl.constexpr, + hint: tl.constexpr = None, +): + """ + Enqueue one POLL_REGMEM, many 2D SUB_WINDOW_COPY packets, and one ATOMIC. + """ + translated_to_base_ptr = __translate(to_base_ptr, from_rank, to_rank, heap_bases, hint) + translated_signal_ptr = __translate(signal_flag_ptr, from_rank, to_rank, heap_bases, hint) + + ctx = copy_engine_ctx + (6 * to_rank) + queue_ptr_u32 = tl.load(ctx + 0).to(tl.pointer_type(tl.uint32)) + read_ptr = tl.load(ctx + 1).to(tl.pointer_type(tl.uint64)) + write_ptr = tl.load(ctx + 2).to(tl.pointer_type(tl.uint64)) + doorbell_ptr = tl.load(ctx + 3).to(tl.pointer_type(tl.uint64)) + cached_write_ptr = tl.load(ctx + 4).to(tl.pointer_type(tl.uint64)) + committed_write_ptr = tl.load(ctx + 5).to(tl.pointer_type(tl.uint64)) + + poll_packet_bytes = 24 + copy_packet_bytes = 80 + atomic_packet_bytes = 32 + command_in_bytes = poll_packet_bytes + transfer_count * copy_packet_bytes + atomic_packet_bytes + + base, offset = sdma_utils.acquire_fadd( + queue_ptr_u32, read_ptr, write_ptr, doorbell_ptr, cached_write_ptr, committed_write_ptr, command_in_bytes + ) + sdma_utils.place_nop_packet(queue_ptr_u32, base, offset) + + packet_offset_bytes = base + offset + sdma_utils.place_poll_regmem_packet( + queue_ptr_u32, + packet_offset_bytes, + wait_flag_ptr.to(tl.uint64), + wait_value, + ) + + from_base_val = from_base_ptr.to(tl.uint64) + to_base_val = translated_to_base_ptr.to(tl.uint64) + for i in range(MAX_RECTS): + if i < transfer_count: + transfer_idx = transfer_start + i + row_offset = tl.load(transfer_row_offsets + transfer_idx) + col_offset = tl.load(transfer_col_offsets + transfer_idx) + width_bytes = tl.load(transfer_width_bytes + transfer_idx) + height = tl.load(transfer_heights + transfer_idx) + byte_offset = (row_offset.to(tl.uint64) * src_pitch) + (col_offset.to(tl.uint64) * stride_n_bytes) + copy_offset_bytes = packet_offset_bytes + poll_packet_bytes + i * copy_packet_bytes + sdma_utils.place_sub_window_copy_packet( + queue_ptr_u32, + copy_offset_bytes, + from_base_val + byte_offset, + to_base_val + byte_offset, + tile_width=width_bytes, + tile_height=height, + src_buffer_pitch=src_pitch, + dst_buffer_pitch=dst_pitch, + src_x=0, + src_y=0, + dst_x=0, + dst_y=0, + ) + + sdma_utils.place_atomic_add_packet( + queue_ptr_u32, + packet_offset_bytes + poll_packet_bytes + transfer_count * copy_packet_bytes, + translated_signal_ptr.to(tl.uint64), + signal_value, + ) + + pending_wptr = base + offset + command_in_bytes + sdma_utils.submit(write_ptr, doorbell_ptr, committed_write_ptr, base, pending_wptr) diff --git a/iris/ops/all_gather_matmul.py b/iris/ops/all_gather_matmul.py index c99e83157..53805f935 100644 --- a/iris/ops/all_gather_matmul.py +++ b/iris/ops/all_gather_matmul.py @@ -154,7 +154,7 @@ def _fused_all_gather_matmul_kernel( # Store result using tritonblas Tile rm, rn = out_tile.indices() - C_ptr = C + rm[:, None] * stride_cm + rn[None, :] * stride_cn + C_ptr = C + rm.to(tl.int64)[:, None] * stride_cm + rn.to(tl.int64)[None, :] * stride_cn mask = (rm[:, None] < M) & (rn[None, :] < N) tl.store(C_ptr, c, mask=mask) diff --git a/iris/ops/all_gather_matmul_copy_engine.py b/iris/ops/all_gather_matmul_copy_engine.py new file mode 100644 index 000000000..ff8be008e --- /dev/null +++ b/iris/ops/all_gather_matmul_copy_engine.py @@ -0,0 +1,512 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2025-2026 Advanced Micro Devices, Inc. All rights reserved. + +""" +Fused All-Gather + GEMM using copy engine (SDMA) for data movement. + +Key differences from HBM buffer variant: +- SMs only perform GEMM (no fetcher workgroups) +- Host orchestrates SDMA transfers of remote tiles to staged_a buffer +- GEMM processes local K-blocks first (from A_sharded), then remote K-blocks (from staged_a) +- Flags only track remote tiles, updated by copy engine via host_atomic_add_32 +""" + +from typing import Optional +import torch +import torch.distributed as dist +import triton +import triton.language as tl +import iris +from tritonblas.matmul import persistent_matmul_lt +from tritonblas.config import PrologueConfig + +from .workspace import FusedWorkspace + + +@triton.jit +def _batch_poster_kernel( + A_sharded, + staged_a, + flags_ptr, + flag_iteration, + M, + K_local, + stride_am, + stride_sa_m, + stride_sa_k, + context_tensor: tl.tensor, + heap_bases_ptr: tl.tensor, + copy_engine_ctx: tl.tensor, + cur_rank: tl.constexpr, + world_size: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + NUM_M_TILES: tl.constexpr, + M_TILES_PER_BATCH: tl.constexpr, +): + """Post one SDMA transfer per (batch, rank) including local copy.""" + ctx = iris.DeviceContext.initialize(context_tensor, cur_rank, world_size) + + pid = tl.program_id(0) + dst_rank = pid + + ptr_dtype = A_sharded.dtype.element_ty + if ptr_dtype == tl.float16 or ptr_dtype == tl.bfloat16: + elem_size = 2 + elif ptr_dtype == tl.float32 or ptr_dtype == tl.int32: + elem_size = 4 + elif ptr_dtype == tl.float64 or ptr_dtype == tl.int64: + elem_size = 8 + else: + elem_size = 4 + + num_batches = (NUM_M_TILES + M_TILES_PER_BATCH - 1) // M_TILES_PER_BATCH + rows_per_batch = M_TILES_PER_BATCH * BLOCK_SIZE_M + + for batch_id in range(num_batches): + src_m_offset = batch_id * rows_per_batch + remaining_rows = M - src_m_offset + tile_height = tl.minimum(remaining_rows, rows_per_batch) + + src_m_offset_i64 = (src_m_offset + 0 * stride_am).to(tl.int64) + stride_am_i64 = (stride_am + 0 * src_m_offset).to(tl.int64) + src_ptr = A_sharded + src_m_offset_i64 * stride_am_i64 + + dst_m_offset = src_m_offset + dst_k_offset = cur_rank * K_local + dst_m_offset_i64 = (dst_m_offset + 0 * stride_sa_m).to(tl.int64) + dst_k_offset_i64 = (dst_k_offset + 0 * stride_sa_k).to(tl.int64) + stride_sa_m_i64 = (stride_sa_m + 0 * dst_m_offset).to(tl.int64) + stride_sa_k_i64 = (stride_sa_k + 0 * dst_m_offset).to(tl.int64) + dst_ptr = staged_a + (dst_m_offset_i64 * stride_sa_m_i64 + dst_k_offset_i64 * stride_sa_k_i64) + + tile_width_bytes = K_local * elem_size + src_pitch_bytes = stride_am * elem_size + dst_pitch_bytes = stride_sa_m * elem_size + + iris.put_signal_rect( + src_ptr, + dst_ptr, + cur_rank, + dst_rank, + heap_bases_ptr, + copy_engine_ctx, + flags_ptr + batch_id, + 1, + width_bytes=tile_width_bytes, + height=tile_height, + src_pitch=src_pitch_bytes, + dst_pitch=dst_pitch_bytes, + ) + + +def all_gather_matmul_copy_engine_preamble( + shmem, + A_sharded: torch.Tensor, + B: torch.Tensor, + selector=None, + m_tiles_per_batch: Optional[int] = None, + staged_a_layout: str = "k_contiguous", +) -> FusedWorkspace: + """ + Allocate workspace for copy engine variant. + + Args: + staged_a_layout: "k_contiguous" (default, row-major (M,K)) or + "m_contiguous" (col-major, stored as (K,M) transposed). + """ + shmem._require_copy_engine() + + from tritonblas.matmul import _make_matmul_selector + + M, K_local = A_sharded.shape + K, N = B.shape + world_size = shmem.get_num_ranks() + + assert world_size * K_local == K + + if selector is None: + selector = _make_matmul_selector( + M, N, K, A_sharded.dtype, B.dtype, A_sharded.dtype, A_sharded.device, streamk=False + ) + + assert K_local % selector.block_k == 0 + assert K % selector.block_k == 0 + assert M % selector.block_m == 0 + + num_m_tiles = M // selector.block_m + num_tiles_n = (N + selector.block_n - 1) // selector.block_n + total_tiles = num_m_tiles * num_tiles_n + + if m_tiles_per_batch is None: + # Auto-calculate optimal m_tiles_per_batch + active_cus = getattr(selector, "_ACTIVE_CU", None) + if active_cus is None or active_cus <= 0: + active_cus = torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count + tiles_per_group = max(1, selector.group_m * num_tiles_n) + groups_per_wave = max(1, int(active_cus) // tiles_per_group) + m_tiles_per_batch = max(1, min(num_m_tiles, groups_per_wave * selector.group_m)) + + num_batches = (num_m_tiles + m_tiles_per_batch - 1) // m_tiles_per_batch + num_flags = num_batches + ws = FusedWorkspace( + operation="all_gather_matmul_copy_engine", + shape=(M, N, K), + dtype=A_sharded.dtype, + world_size=world_size, + variant=f"copy_engine_{staged_a_layout}_wave_aware_noninterleaved", + prepared=True, + ) + + # Allocate staged_a - full K dimension (NON-INTERLEAVED like HBM buffer) + # Each rank's K-blocks are stored contiguously for efficient bulk SDMA + if staged_a_layout == "m_contiguous": + storage = shmem.zeros((K, M), dtype=A_sharded.dtype) + ws.aux_buffer = storage.T # (M, K) with M-contiguous + else: + ws.aux_buffer = shmem.zeros((M, K), dtype=A_sharded.dtype) + + # Allocate per-batch flags + ws.locks = shmem.zeros((num_flags,), dtype=torch.int32) + ws.wait_expected = shmem.zeros((total_tiles,), dtype=torch.int32) + + # Store metadata + ws.selector = selector + ws.m_tiles_per_batch = m_tiles_per_batch + ws.num_m_tiles = num_m_tiles + ws.num_batches = num_batches + + shmem.info( + f"Allocated {num_flags} per-batch flags " + f"(tiles={num_m_tiles}, " + f"{m_tiles_per_batch} M-tiles per batch) " + f"flags buffer at 0x{ws.locks.data_ptr():x}" + ) + + # Share pointers across ranks for SDMA addressing + # Need: A_sharded (source), staged_a (destination), flags (signaling) + A_sharded_ptr_tensor = torch.tensor([A_sharded.data_ptr()], dtype=torch.int64, device="cuda") + A_sharded_ptrs = [torch.zeros(1, dtype=torch.int64, device="cuda") for _ in range(world_size)] + dist.all_gather(A_sharded_ptrs, A_sharded_ptr_tensor) + + staged_a_ptr_tensor = torch.tensor([ws.aux_buffer.data_ptr()], dtype=torch.int64, device="cuda") + staged_a_ptrs = [torch.zeros(1, dtype=torch.int64, device="cuda") for _ in range(world_size)] + dist.all_gather(staged_a_ptrs, staged_a_ptr_tensor) + + flags_ptr_tensor = torch.tensor([ws.locks.data_ptr()], dtype=torch.int64, device="cuda") + flags_ptrs = [torch.zeros(1, dtype=torch.int64, device="cuda") for _ in range(world_size)] + dist.all_gather(flags_ptrs, flags_ptr_tensor) + + # Store all remote pointers in workspace + ws.remote_pointers = { + "A_sharded": [ptr.item() for ptr in A_sharded_ptrs], + "staged_a": [ptr.item() for ptr in staged_a_ptrs], + "flags": [ptr.item() for ptr in flags_ptrs], + } + + # Note: heap_bases are already cached in shmem.heap_bases_cpu (done in iris.py __init__) + + buffer_mb = M * K * A_sharded.element_size() / (1024**2) + sa_stride_m, sa_stride_k = ws.aux_buffer.stride() + shmem.info( + f"Copy Engine: staged_a=({M},{K}) [{buffer_mb:.1f} MB] " + f"layout={staged_a_layout} strides=({sa_stride_m},{sa_stride_k}), " + f"NON-INTERLEAVED: each rank's K-blocks contiguous" + ) + + shmem.barrier() + return ws + + +def all_gather_matmul_copy_engine( + shmem, + output_tensor: torch.Tensor, + A_sharded: torch.Tensor, + B: torch.Tensor, + bias: Optional[torch.Tensor] = None, + async_op: bool = False, + workspace: Optional[FusedWorkspace] = None, + flag_iteration: int = 0, + staged_a_layout: str = "k_contiguous", + num_warps: Optional[int] = None, + num_stages: Optional[int] = None, + trace: bool = False, + verbose: bool = False, + device_initiated: bool = False, +) -> FusedWorkspace: + """ + All-gather + matmul with copy engine orchestrating remote tile transfers. + + Key differences from HBM buffer: + - No fetcher workgroups (only GEMM) + - Host uses SDMA to copy remote tiles (default) OR device WGs initiate SDMA (device_initiated=True) + - GEMM processes local tiles first, then remote tiles + + Args: + staged_a_layout: Buffer layout for gathered A. + "k_contiguous" — (M,K) row-major, K is fast dim. + "m_contiguous" — (M,K) with M as fast dim. + device_initiated: If True, use device-side WGs to initiate SDMA transfers instead of host. + trace: Not supported yet; must be False. + k_per_flag: Retained for call compatibility; ignored by the current per-batch design. + """ + # Fail fast: without the SDMA engine the GEMM launched below would wait forever on flags. + shmem._require_copy_engine() + if trace: + raise NotImplementedError("trace=True is not supported for all_gather_matmul_copy_engine yet") + + M, K_local = A_sharded.shape + K, N = B.shape + world_size = shmem.get_num_ranks() + rank = shmem.get_rank() + + assert world_size * K_local == K + assert output_tensor.shape == (M, N) + + if workspace is None: + workspace = all_gather_matmul_copy_engine_preamble(shmem, A_sharded, B, staged_a_layout=staged_a_layout) + + selector = workspace.selector + m_tiles_per_batch = workspace.m_tiles_per_batch + + assert M % selector.block_m == 0 + assert K % selector.block_k == 0 + assert K_local % selector.block_k == 0 + + num_k_blocks_local = K_local // selector.block_k + num_m_tiles = M // selector.block_m + num_tiles_n = (N + selector.block_n - 1) // selector.block_n + + # Local K-blocks will be copied via SDMA (host or device initiated) + stride_am, stride_ak = A_sharded.stride() + stride_bk, stride_bn = B.stride() + stride_cm, stride_cn = output_tensor.stride() + stride_sa_m, stride_sa_k = workspace.aux_buffer.stride() + + if bias is not None: + assert bias.shape[0] == M + bias_ptr = bias + stride_bias = bias.stride()[0] if bias.dim() > 0 else 1 + use_bias = True + else: + bias_ptr = output_tensor + stride_bias = 1 + use_bias = False + + # num_m_tiles, num_tiles_n already calculated above + total_tiles = num_m_tiles * num_tiles_n + + # Auto-detect num_sms from device if not specified + num_comm_wgs = selector._hardware.NUM_XCD if device_initiated else 0 + gemm_tiles = total_tiles + + # m_tiles_per_batch was already set above before calling preamble + # Calculate number of batches + num_batches = (num_m_tiles + m_tiles_per_batch - 1) // m_tiles_per_batch + + launch_kwargs = {"matrix_instr_nonkdim": 16} + if num_warps is not None: + launch_kwargs["num_warps"] = num_warps + if num_stages is not None: + launch_kwargs["num_stages"] = num_stages + + # ====================================================================== + # Host orchestration: SDMA copy setup + # ====================================================================== + torch.cuda.current_device() # Initialize CUDA context + + # SDMA queues already connected during iris init + if verbose and rank == 0: + shmem.info(f"[Rank {rank}] Copy engines connected, launching kernel...") + shmem.info( + f"Kernel params: num_m_tiles={num_m_tiles}, " + f"num_tiles_n={num_tiles_n}, num_k_blocks_local={num_k_blocks_local}, " + f"group_size_m={selector.group_m}, m_tiles_per_batch={m_tiles_per_batch}" + ) + shmem.info( + f"Pointers: A_sharded=0x{A_sharded.data_ptr():x}, " + f"B=0x{B.data_ptr():x}, C=0x{output_tensor.data_ptr():x}, " + f"bias_ptr=0x{bias_ptr.data_ptr():x}, " + f"staged_a=0x{workspace.aux_buffer.data_ptr():x}, " + f"flags=0x{workspace.locks.data_ptr():x} (n={workspace.locks.numel()})" + ) + + if verbose and rank == 0: + shmem.info( + f"Launching kernel: gemm_tiles={gemm_tiles}, device_initiated={device_initiated}, sdma_wgs={num_comm_wgs}" + ) + + tb_block_m = selector.block_m + prologue_config = PrologueConfig( + wait_buffer=workspace.locks, + expected_buffer=workspace.wait_expected, + expected_inc=world_size, + map_type="block", + block_group_m=m_tiles_per_batch, + block_group_n=num_tiles_n, + ) + + if use_bias: + import warnings + + warnings.warn( + "Bias is not yet supported in the persistent tritonBLAS path for all_gather_matmul_copy_engine. " + "Ignoring bias for this launch." + ) + + # ====================================================================== + # Launch kernel and orchestrate SDMA transfers + # ====================================================================== + + import time + + sdma_start_time = time.perf_counter() + + if device_initiated: + # Device-initiated: keep the known-good split path in production. + # The combined COMM_WGS+GEMM kernel remains in this file as an + # experimental option, but we do not use it by default. + poster_grid = world_size + _batch_poster_kernel[(poster_grid,)]( + A_sharded, + workspace.aux_buffer, + workspace.locks, + flag_iteration, + M, + K_local, + stride_am, + stride_sa_m, + stride_sa_k, + shmem.get_device_context(), + shmem.get_heap_bases(), + shmem.get_copy_engine_ctx(), + rank, + world_size, + selector.block_m, + num_m_tiles, + m_tiles_per_batch, + ) + persistent_matmul_lt( + workspace.aux_buffer, + B, + output_tensor, + selector, + bias=None, + prologue_config=prologue_config, + ) + else: + # Host-initiated: + # Launch GEMM first, then post batches to overlap SDMA. + elem_size = A_sharded.element_size() + staged_a_base_addr = workspace.aux_buffer.data_ptr() + flags_base_addr = workspace.locks.data_ptr() + tile_transfer_count = 0 + + def post_host_batch(batch_id: int, m_tile_start: int, num_m_tiles_in_batch: int) -> None: + nonlocal tile_transfer_count + + for dst_rank in range(world_size): + flag_idx = batch_id + flag_addr_local = flags_base_addr + flag_idx * 4 + flag_addr_remote = shmem.heap.translate(flag_addr_local, rank, dst_rank) + + # Source is the local shard, so batches only advance in M. + src_offset_bytes = (m_tile_start * selector.block_m * stride_am) * elem_size + src_ptr = A_sharded.data_ptr() + src_offset_bytes + + # Destination is this rank's global-K slot inside staged_a. + dst_offset_bytes = ( + m_tile_start * selector.block_m * stride_sa_m + rank * K_local * stride_sa_k + ) * elem_size + dst_ptr_local = staged_a_base_addr + dst_offset_bytes + dst_ptr_remote = shmem.heap.translate(dst_ptr_local, rank, dst_rank) + + # Use shmem.put_tile method with signal + shmem.put_tile( + global_addr=src_ptr, + to_rank=dst_rank, + dst_addr=dst_ptr_remote, + data_size=elem_size, + tile_height=num_m_tiles_in_batch * selector.block_m, + tile_width=K_local, + src_stride=stride_am, + dst_stride=stride_sa_m, + signal_flag=flag_addr_remote, + signal_value=1, + async_op=True, + channel=0, + ) + tile_transfer_count += 1 + + if verbose and batch_id == 0 and dst_rank == (rank + 1) % world_size: + shmem.info( + f"[Rank {rank}→{dst_rank}] Signaled batch={batch_id} flag_idx={flag_idx} " + f"({num_m_tiles_in_batch} rows × full local K shard)" + ) + + if verbose and rank == 0: + num_batches_calc = workspace.locks.numel() + shmem.info( + f"[Rank {rank}] Starting SDMA loop (batched M-tile transfers)... " + f"num_tiles_m={num_m_tiles}, num_k_blocks_local={num_k_blocks_local}, " + f"m_tiles_per_batch={m_tiles_per_batch}, tb_block_m={tb_block_m}" + ) + shmem.info(f"[Rank {rank}] Will transfer in {num_batches_calc} batches of {m_tiles_per_batch} M-tiles each") + + # TODO not always faster + # Prime batch 0 before GEMM launch so the first released tile-group can + # start immediately instead of stalling on an empty wait queue. + # TODO issue with time measurement + # first_batch_tiles = min(m_tiles_per_batch, num_m_tiles) + # post_host_batch(0, 0, first_batch_tiles) + + if verbose and rank == 0: + shmem.info(f"[Rank {rank}] Launching GEMM kernel after pre-posting batch 0...") + + persistent_matmul_lt( + workspace.aux_buffer, + B, + output_tensor, + selector, + bias=None, + prologue_config=prologue_config, + ) + + # Post the remaining batches while GEMM is already running. + batch_id = 0 + for m_tile_start in range(0, num_m_tiles, m_tiles_per_batch): + m_tile_end = min(m_tile_start + m_tiles_per_batch, num_m_tiles) + num_m_tiles_in_batch = m_tile_end - m_tile_start + post_host_batch(batch_id, m_tile_start, num_m_tiles_in_batch) + batch_id += 1 + + sdma_end_post_time = time.perf_counter() + + if verbose: + # Ensure all SDMA operations complete + for dst_rank in range(world_size): + shmem.quiet(to_rank=dst_rank, channel=0) + sdma_end_time = time.perf_counter() + + post_ms = (sdma_end_post_time - sdma_start_time) * 1000.0 + quiet_ms = (sdma_end_time - sdma_end_post_time) * 1000.0 + total_ms = (sdma_end_time - sdma_start_time) * 1000.0 + shmem.info( + f"[Rank {rank}] SDMA complete. " + f"Post: {post_ms:.2f}ms, Quiet: {quiet_ms:.2f}ms, Total: {total_ms:.2f}ms, " + f"transfers={tile_transfer_count}" + ) + sample_count = min(8, workspace.locks.numel()) + sample_flags = workspace.locks[:sample_count].cpu().tolist() + shmem.info( + f"[Rank {rank}] Flag sample after SDMA quiet: " + f"expected_inc={world_size}, flags[:{sample_count}]={sample_flags}" + ) + + # ====================================================================== + # Synchronize + # ====================================================================== + if not async_op: + torch.cuda.synchronize() # Wait for kernel completion + shmem.barrier() + + return workspace diff --git a/iris/ops/all_gather_matmul_hbm_buffer.py b/iris/ops/all_gather_matmul_hbm_buffer.py index 4e18a4ea1..abd9f053e 100644 --- a/iris/ops/all_gather_matmul_hbm_buffer.py +++ b/iris/ops/all_gather_matmul_hbm_buffer.py @@ -283,6 +283,7 @@ def _hbm_buffer_all_gather_matmul_kernel( for fg_idx in range(stage_pid, total_fg_stage, stage_fetch_sms): m_group = fg_idx // tiles_per_m_group within_group = fg_idx % tiles_per_m_group + # TODO there is a bug if #M-Tiles is not multiple of GROUP_SIZE_M k_flag_group = within_group // GROUP_SIZE_M m_in_group = within_group % GROUP_SIZE_M m_tile = stage_m_start + m_group * GROUP_SIZE_M + m_in_group @@ -304,7 +305,9 @@ def _hbm_buffer_all_gather_matmul_kernel( rk = k_block_global * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) rk = tl.max_contiguous(tl.multiple_of(rk, BLOCK_SIZE_K), BLOCK_SIZE_K) - staged_ptrs = staged_a + rm.to(tl.int64)[:, None] * stride_sa_m + rk[None, :] * stride_sa_k + staged_ptrs = ( + staged_a + rm.to(tl.int64)[:, None] * stride_sa_m + rk.to(tl.int64)[None, :] * stride_sa_k + ) for compile_rank in range(world_size): if src_rank_idx == compile_rank: @@ -373,7 +376,7 @@ def _hbm_buffer_all_gather_matmul_kernel( rk = k_block * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) rk = tl.max_contiguous(tl.multiple_of(rk, BLOCK_SIZE_K), BLOCK_SIZE_K) - a_ptrs = staged_a + rm.to(tl.int64)[:, None] * stride_sa_m + rk[None, :] * stride_sa_k + a_ptrs = staged_a + rm.to(tl.int64)[:, None] * stride_sa_m + rk.to(tl.int64)[None, :] * stride_sa_k a = tl.load(a_ptrs) B_ptrs = B + rk[:, None] * stride_bk + rn[None, :] * stride_bn @@ -389,7 +392,9 @@ def _hbm_buffer_all_gather_matmul_kernel( acc = acc + bias_val[:, None] c = acc.to(C.type.element_ty) - C_ptrs = C + rm[:, None] * stride_cm + rn[None, :] * stride_cn + stride_cm_i64 = tl.cast(stride_cm, tl.int64) + stride_cn_i64 = tl.cast(stride_cn, tl.int64) + C_ptrs = C + rm.to(tl.int64)[:, None] * stride_cm_i64 + rn.to(tl.int64)[None, :] * stride_cn_i64 c_mask = (rm[:, None] < M) & (rn[None, :] < N) tl.store(C_ptrs, c, mask=c_mask, cache_modifier=".wt") @@ -436,6 +441,10 @@ def all_gather_matmul_hbm_buffer_preamble( num_m_tiles = M // config.block_size_m num_k_blocks = K // config.block_size_k + # Keep the default valid for small and single-rank problems where fewer + # than eight K-blocks are available. + while k_per_flag > 1 and num_k_blocks % k_per_flag != 0: + k_per_flag //= 2 assert num_k_blocks % k_per_flag == 0 num_flag_groups_k = num_k_blocks // k_per_flag @@ -587,6 +596,8 @@ def all_gather_matmul_hbm_buffer( fetch_block_k = config.block_size_k num_k_blocks = K // config.block_size_k + while k_per_flag > 1 and num_k_blocks % k_per_flag != 0: + k_per_flag //= 2 assert num_k_blocks % k_per_flag == 0 if workspace is None: diff --git a/iris/ops/matmul_all_gather_copy_engine.py b/iris/ops/matmul_all_gather_copy_engine.py new file mode 100644 index 000000000..d9c67da16 --- /dev/null +++ b/iris/ops/matmul_all_gather_copy_engine.py @@ -0,0 +1,495 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved. + +""" +Fused GEMM + All-Gather operation using SDMA (copy engine) for scatter. + +Each rank has a row-sharded input A_local (M_local x K) and computes C_local = A_local @ B. +Then scatters C_local tiles to form the full C (M x N) where M = world_size * M_local. + +This variant uses SDMA hardware for data movement instead of compute shader scatter. +""" + +from typing import Optional +import time +import torch +import triton +import triton.language as tl +import iris + +from .workspace import FusedWorkspace +from tritonblas.matmul import persistent_matmul_lt +from tritonblas.config import EpilogueConfig +from .tritonblas_launch_wave_schedule import build_launch_wave_plan + + +@triton.jit() +def wait_cnt(): + tl.inline_asm_elementwise("s_waitcnt vmcnt(0)", "=r", [], dtype=tl.int32, is_pure=False, pack=1) + + +@triton.jit() +def _launch_wave_wait_poster_kernel( + C_local_base, + flags, + completion_signals, + flag_iteration, + wave_tile_counts, + wave_transfer_offsets, + wave_transfer_counts, + transfer_row_offsets, + transfer_col_offsets, + transfer_width_bytes, + transfer_heights, + heap_bases: tl.tensor, + copy_engine_ctx: tl.tensor, + cur_rank: tl.constexpr, + world_size: tl.constexpr, + NUM_WAVES: tl.constexpr, + MAX_RECTS_PER_WAVE: tl.constexpr, + SRC_PITCH: tl.constexpr, + DST_PITCH: tl.constexpr, + STRIDE_N: tl.constexpr, +): + dst_rank = tl.program_id(0) + if dst_rank >= world_size or dst_rank == cur_rank: + return + + ptr_dtype = C_local_base.dtype.element_ty + if ptr_dtype == tl.float16 or ptr_dtype == tl.bfloat16: + elem_size = 2 + elif ptr_dtype == tl.float32 or ptr_dtype == tl.int32: + elem_size = 4 + elif ptr_dtype == tl.float64 or ptr_dtype == tl.int64: + elem_size = 8 + else: + elem_size = 4 + + for wave_id in range(NUM_WAVES): + transfer_start = tl.load(wave_transfer_offsets + wave_id) + transfer_count = tl.load(wave_transfer_counts + wave_id) + if transfer_count != 0: + wait_value = (flag_iteration + 1) * tl.load(wave_tile_counts + wave_id) + is_last_wave = wave_id == (NUM_WAVES - 1) + + if is_last_wave: + iris.wait_then_put_signal_rects( + C_local_base, + C_local_base, + cur_rank, + dst_rank, + heap_bases, + copy_engine_ctx, + flags + wave_id, + wait_value, + completion_signals + cur_rank, + 1, + transfer_row_offsets, + transfer_col_offsets, + transfer_width_bytes, + transfer_heights, + transfer_start, + transfer_count, + STRIDE_N * elem_size, + SRC_PITCH * elem_size, + DST_PITCH * elem_size, + MAX_RECTS_PER_WAVE, + ) + else: + iris.wait_then_put_rects( + C_local_base, + C_local_base, + cur_rank, + dst_rank, + heap_bases, + copy_engine_ctx, + flags + wave_id, + wait_value, + transfer_row_offsets, + transfer_col_offsets, + transfer_width_bytes, + transfer_heights, + transfer_start, + transfer_count, + STRIDE_N * elem_size, + SRC_PITCH * elem_size, + DST_PITCH * elem_size, + MAX_RECTS_PER_WAVE, + ) + + +@triton.jit() +def _wait_completion_signals_kernel( + completion_signals, + expected_value, + cur_rank: tl.constexpr, + world_size: tl.constexpr, +): + src_rank = tl.program_id(0) + if src_rank >= world_size or src_rank == cur_rank: + return + while tl.load(completion_signals + src_rank, cache_modifier=".cv", volatile=True) < expected_value: + pass + + +def _selector_wave_size(selector, device: torch.device) -> int: + wave_size = getattr(selector, "_ACTIVE_CU", None) + if wave_size is None or wave_size <= 0: + wave_size = torch.cuda.get_device_properties(device).multi_processor_count + return int(wave_size) + + +def _ensure_tritonblas_launch_wave_workspace( + shmem, + workspace: FusedWorkspace, + selector, + device: torch.device, + m_local: int, + n: int, + world_size: int, + elem_size: int, +): + num_tiles_m = triton.cdiv(m_local, selector.block_m) + num_tiles_n = triton.cdiv(n, selector.block_n) + total_tiles = num_tiles_m * num_tiles_n + wave_size = _selector_wave_size(selector, device) + num_xcds = max(1, int(getattr(selector, "num_sms", 1))) + plan_key = (num_tiles_m, num_tiles_n, selector.group_m, total_tiles, wave_size, num_xcds) + + if getattr(workspace, "launch_wave_plan_key", None) != plan_key: + plan = build_launch_wave_plan( + num_tiles_m=num_tiles_m, + num_tiles_n=num_tiles_n, + group_size_m=selector.group_m, + launch_grid=total_tiles, + wave_size=wave_size, + num_xcds=num_xcds, + ) + workspace.launch_wave_plan = plan + workspace.launch_wave_plan_key = plan_key + workspace.locks = shmem.zeros((plan.num_waves,), dtype=torch.int32) + transfers_by_wave = [[] for _ in range(plan.num_waves)] + for transfer in plan.transfers: + transfers_by_wave[transfer.wave_id].append(transfer) + + wave_transfer_offsets = [] + wave_transfer_counts = [] + transfer_row_offsets = [] + transfer_col_offsets = [] + transfer_width_bytes = [] + transfer_heights = [] + max_rects_per_wave = 0 + running_offset = 0 + + for wave_transfers in transfers_by_wave: + wave_transfer_offsets.append(running_offset) + wave_transfer_counts.append(len(wave_transfers)) + max_rects_per_wave = max(max_rects_per_wave, len(wave_transfers)) + for transfer in wave_transfers: + row_offset = transfer.m_tile_start * selector.block_m + col_offset = transfer.n_tile_start * selector.block_n + batch_height = min(transfer.m_tile_count * selector.block_m, m_local - row_offset) + batch_width = min(transfer.n_tile_count * selector.block_n, n - col_offset) + transfer_row_offsets.append(row_offset) + transfer_col_offsets.append(col_offset) + transfer_width_bytes.append(batch_width * elem_size) + transfer_heights.append(batch_height) + running_offset += len(wave_transfers) + + workspace.wave_transfer_offsets = torch.tensor(wave_transfer_offsets, device=device, dtype=torch.int32) + workspace.wave_transfer_counts = torch.tensor(wave_transfer_counts, device=device, dtype=torch.int32) + workspace.transfer_row_offsets = torch.tensor(transfer_row_offsets, device=device, dtype=torch.int32) + workspace.transfer_col_offsets = torch.tensor(transfer_col_offsets, device=device, dtype=torch.int32) + workspace.transfer_width_bytes = torch.tensor(transfer_width_bytes, device=device, dtype=torch.int32) + workspace.transfer_heights = torch.tensor(transfer_heights, device=device, dtype=torch.int32) + workspace.wave_tile_counts = torch.tensor(plan.wave_tile_counts, device=device, dtype=torch.int32) + workspace.num_tiles_m = num_tiles_m + workspace.num_tiles_n = num_tiles_n + workspace.num_waves = plan.num_waves + workspace.num_batches = plan.num_waves + workspace.num_transfers = len(plan.transfers) + workspace.max_rects_per_wave = max_rects_per_wave + if getattr(workspace, "completion_signals", None) is None or workspace.completion_signals.numel() != world_size: + workspace.completion_signals = shmem.zeros((world_size,), dtype=torch.int32) + return workspace.launch_wave_plan + + +def _auto_m_tiles_per_batch(selector, M_local: int, N: int) -> int: + """Auto-calculate optimal m_tiles_per_batch based on selector and shape.""" + num_tiles_m = (M_local + selector.block_m - 1) // selector.block_m + num_tiles_n = (N + selector.block_n - 1) // selector.block_n + active_cus = getattr(selector, "_ACTIVE_CU", None) + if active_cus is None or active_cus <= 0: + active_cus = torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count + tiles_per_group = max(1, selector.group_m * num_tiles_n) + groups_per_wave = max(1, int(active_cus) // tiles_per_group) + return max(1, min(num_tiles_m, groups_per_wave * selector.group_m)) + + +def matmul_all_gather_copy_engine_preamble( + shmem, + A: torch.Tensor, + B: torch.Tensor, + m_tiles_per_batch: Optional[int] = None, + selector=None, +) -> FusedWorkspace: + """Allocate workspace for matmul_all_gather_copy_engine including per-batch flags.""" + shmem._require_copy_engine() + + from tritonblas.matmul import _make_matmul_selector + + M_local, K = A.shape + K2, N = B.shape + world_size = shmem.get_num_ranks() + + assert K == K2, f"Inner dimensions must match: A has K={K}, B has K={K2}" + + M = M_local * world_size + + # Create selector for block size configuration if not provided + if selector is None: + selector = _make_matmul_selector( + M_local, + N, + K, + A.dtype, + B.dtype, + A.dtype, + A.device, + streamk=False, + ) + + # Auto-calculate optimal m_tiles_per_batch if not provided + if m_tiles_per_batch is None: + m_tiles_per_batch = _auto_m_tiles_per_batch(selector, M_local, N) + + # Calculate number of tiles based on selector + num_tiles_m = (M_local + selector.block_m - 1) // selector.block_m + num_tiles_n = (N + selector.block_n - 1) // selector.block_n + + num_tiles = num_tiles_m * num_tiles_n + + # Calculate number of batches + num_batches = (num_tiles_m + m_tiles_per_batch - 1) // m_tiles_per_batch + + ws = FusedWorkspace( + operation="matmul_all_gather_copy_engine", + shape=(M, N, K), + dtype=A.dtype, + world_size=world_size, + prepared=True, + ) + + # Allocate one readiness counter per M-batch. + ws.locks = shmem.zeros((num_batches,), dtype=torch.int32) + ws.completion_signals = shmem.zeros((world_size,), dtype=torch.int32) + + # Store metadata for later use + ws.selector = selector + ws.num_tiles_m = num_tiles_m + ws.num_tiles_n = num_tiles_n + ws.num_batches = num_batches + ws.m_tiles_per_batch = m_tiles_per_batch + + return ws + + +def matmul_all_gather_copy_engine( + shmem, + output_tensor: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + bias: Optional[torch.Tensor] = None, + async_op: bool = False, + workspace: Optional[FusedWorkspace] = None, + flag_iteration: int = 0, + verbose: bool = False, +) -> FusedWorkspace: + """ + Fused matrix multiplication and all-gather using SDMA (copy engine) for scatter. + + Computes: output = all_gather(A @ B + bias) along M dimension + + Each rank has A of shape (M_local, K) where M_local = M / world_size. + The operation computes C_local = A @ B on each rank and uses SDMA hardware + to scatter the tiles to all ranks (all-gather pattern). + + Args: + shmem: Iris shmem context + output_tensor: Output tensor C of shape (M, N) where M = M_local * world_size + A: Input matrix A of shape (M_local, K) + B: Input matrix B of shape (K, N) + bias: Optional bias vector (M_local,) + async_op: If False, performs barrier at end + workspace: Optional pre-allocated workspace + flag_iteration: Launch generation for cumulative batch counters. + Batch readiness counters are not reset each iteration; + the poster waits for the generation-adjusted target. + verbose: If True, print poster/main/quiet timing breakdown in sync mode + + Returns: + FusedWorkspace object + """ + # Fail fast: without the SDMA engine the GEMM launched below would wait forever on flags. + shmem._require_copy_engine() + + M_local, K = A.shape + K2, N = B.shape + world_size = shmem.get_num_ranks() + rank = shmem.get_rank() + + assert K == K2, f"Inner dimensions must match: A has K={K}, B has K={K2}" + + M = M_local * world_size + assert output_tensor.shape == (M, N), f"Output must be ({M}, {N}), got {output_tensor.shape}" + + timing_events = None + cpu_timing = None + if verbose: + current_stream = torch.cuda.current_stream(device=A.device) + timing_events = { + "stream": current_stream, + "poster_start": torch.cuda.Event(enable_timing=True), + "poster_end": torch.cuda.Event(enable_timing=True), + "main_start": torch.cuda.Event(enable_timing=True), + "main_end": torch.cuda.Event(enable_timing=True), + "quiet_start": torch.cuda.Event(enable_timing=True), + "quiet_end": torch.cuda.Event(enable_timing=True), + } + cpu_timing = { + "total_start": time.perf_counter(), + "poster_launch_ms": 0.0, + "main_launch_ms": 0.0, + "quiet_cpu_ms": 0.0, + "sync_wait_ms": 0.0, + } + + # Allocate workspace if not provided + if workspace is None: + workspace = matmul_all_gather_copy_engine_preamble(shmem, A, B) + + stride_cm, stride_cn = output_tensor.stride() + + selector = workspace.selector + + launch_wave_plan = _ensure_tritonblas_launch_wave_workspace( + shmem, + workspace, + selector, + A.device, + M_local, + N, + world_size, + output_tensor.element_size(), + ) + + # Get metadata from workspace after any schedule planning. + num_batches = workspace.num_batches + + if timing_events is not None: + timing_events["poster_start"].record(timing_events["stream"]) + poster_launch_start = time.perf_counter() if cpu_timing is not None else None + poster_grid = (world_size,) + c_local_base = output_tensor[rank * M_local :, :] + _launch_wave_wait_poster_kernel[poster_grid]( + c_local_base, + workspace.locks, + workspace.completion_signals, + flag_iteration, + workspace.wave_tile_counts, + workspace.wave_transfer_offsets, + workspace.wave_transfer_counts, + workspace.transfer_row_offsets, + workspace.transfer_col_offsets, + workspace.transfer_width_bytes, + workspace.transfer_heights, + shmem.get_heap_bases(), + shmem.get_copy_engine_ctx(), + rank, + world_size, + workspace.num_waves, + workspace.max_rects_per_wave, + stride_cm, + stride_cm, + stride_cn, + ) + if cpu_timing is not None: + cpu_timing["poster_launch_ms"] = (time.perf_counter() - poster_launch_start) * 1000.0 + if timing_events is not None: + timing_events["poster_end"].record(timing_events["stream"]) + + # Launch GEMM after poster submission so SDMA can wait autonomously. + if timing_events is not None: + timing_events["main_start"].record(timing_events["stream"]) + main_launch_start = time.perf_counter() if cpu_timing is not None else None + if bias is not None: + import warnings + + warnings.warn( + "Bias is not yet supported in the tritonBLAS SignalView path for " + "matmul_all_gather_copy_engine. Ignoring bias for this launch." + ) + + epilogue_config = EpilogueConfig( + workspace.locks, + map_type="launch_wave", + block_group_m=launch_wave_plan.wave_size, + ) + c_local_view = output_tensor[rank * M_local : (rank + 1) * M_local, :] + persistent_matmul_lt( + A, + B, + c_local_view, + selector, + bias=None, + work_stealing=False, + epilogue_config=epilogue_config, + ) + if cpu_timing is not None: + cpu_timing["main_launch_ms"] = (time.perf_counter() - main_launch_start) * 1000.0 + if timing_events is not None: + timing_events["main_end"].record(timing_events["stream"]) + + if not async_op: + wait_cpu_start = time.perf_counter() if cpu_timing is not None else None + if timing_events is not None: + timing_events["quiet_start"].record(timing_events["stream"]) + _wait_completion_signals_kernel[(world_size,)]( + workspace.completion_signals, + flag_iteration + 1, + rank, + world_size, + ) + if timing_events is not None: + timing_events["quiet_end"].record(timing_events["stream"]) + sync_wait_start = time.perf_counter() if cpu_timing is not None else None + torch.cuda.synchronize() + if cpu_timing is not None: + cpu_timing["sync_wait_ms"] = (time.perf_counter() - sync_wait_start) * 1000.0 + cpu_timing["quiet_cpu_ms"] = (time.perf_counter() - wait_cpu_start) * 1000.0 + + if verbose and rank == 0: + poster_ms = timing_events["poster_start"].elapsed_time(timing_events["poster_end"]) + main_ms = timing_events["main_start"].elapsed_time(timing_events["main_end"]) + quiet_ms = timing_events["quiet_start"].elapsed_time(timing_events["quiet_end"]) + gpu_total_ms = poster_ms + main_ms + quiet_ms + cpu_launch_total_ms = cpu_timing["poster_launch_ms"] + cpu_timing["main_launch_ms"] + cpu_total_ms = (time.perf_counter() - cpu_timing["total_start"]) * 1000.0 + tile_transfer_count = (world_size - 1) * getattr(workspace, "num_transfers", num_batches) + shmem.info( + f"[Rank {rank}] Device copy-engine GPU timing. " + f"Poster: {poster_ms:.2f}ms, Main: {main_ms:.2f}ms, Wait: {quiet_ms:.2f}ms, " + f"GPU total: {gpu_total_ms:.2f}ms" + ) + shmem.info( + f"[Rank {rank}] Device copy-engine CPU timing. " + f"Poster launch: {cpu_timing['poster_launch_ms']:.2f}ms, " + f"Main launch: {cpu_timing['main_launch_ms']:.2f}ms, " + f"Launch total: {cpu_launch_total_ms:.2f}ms, " + f"Wait: {cpu_timing['quiet_cpu_ms']:.2f}ms, " + f"Barrier wait: {cpu_timing['sync_wait_ms']:.2f}ms, " + f"CPU total: {cpu_total_ms:.2f}ms, " + f"transfers={tile_transfer_count}" + ) + + return workspace diff --git a/iris/ops/matmul_all_gather_host_copy_engine.py b/iris/ops/matmul_all_gather_host_copy_engine.py new file mode 100644 index 000000000..a8ae23967 --- /dev/null +++ b/iris/ops/matmul_all_gather_host_copy_engine.py @@ -0,0 +1,356 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved. + +""" +Fused GEMM + All-Gather operation using host-initiated SDMA with POLL packets. + +Each rank has a row-sharded input A_local (M_local x K) and computes C_local = A_local @ B. +The host pre-queues SDMA POLL+COPY packets for scatter, then the device kernel just stores +tiles to local HBM and sets flags to trigger the pre-queued transfers. + +This is more efficient than device-initiated SDMA because: +- SDMA queue setup happens once on host (not per-tile) +- Device kernel is lightweight (store + set flag) +- SDMA hardware automatically performs scatter when flags are set + +This implementation supports two backends: +- Custom Triton kernel (legacy, controlled by use_tritonblas=False) +- tritonBLAS with SignalView (default, use_tritonblas=True) +""" + +from typing import Optional +import torch +import triton +import triton.language as tl +import time + +from .workspace import FusedWorkspace + +# Import tritonBLAS +from tritonblas.matmul import persistent_matmul_lt +from tritonblas.config import EpilogueConfig +from tritonblas.matmul import _make_matmul_selector +from .tritonblas_launch_wave_schedule import build_launch_wave_plan + + +@triton.jit() +def wait_cnt(): + tl.inline_asm_elementwise("s_waitcnt vmcnt(0)", "=r", [], dtype=tl.int32, is_pure=False, pack=1) + + +@triton.jit() +def _wait_completion_signals_kernel( + completion_signals, + expected_value, + cur_rank: tl.constexpr, + world_size: tl.constexpr, +): + src_rank = tl.program_id(0) + if src_rank >= world_size or src_rank == cur_rank: + return + # while tl.atomic_add(completion_signals + src_rank, 0, sem="acquire", scope="sys") < expected_value: + while tl.load(completion_signals + src_rank, cache_modifier=".cv", volatile=True) < expected_value: + pass + + +def _auto_m_tiles_per_batch(selector, M_local: int, N: int) -> int: + """Auto-calculate optimal m_tiles_per_batch based on selector and shape.""" + num_tiles_m = (M_local + selector.block_m - 1) // selector.block_m + num_tiles_n = (N + selector.block_n - 1) // selector.block_n + active_cus = getattr(selector, "_ACTIVE_CU", None) + if active_cus is None or active_cus <= 0: + active_cus = torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count + tiles_per_group = max(1, selector.group_m * num_tiles_n) + groups_per_wave = max(1, int(active_cus) // tiles_per_group) + return max(1, min(num_tiles_m, groups_per_wave * selector.group_m)) + + +def matmul_all_gather_host_copy_engine_preamble( + shmem, + A: torch.Tensor, + B: torch.Tensor, + m_tiles_per_batch: Optional[int] = None, + trace: bool = False, + selector=None, +) -> FusedWorkspace: + """Allocate workspace for matmul_all_gather_host_copy_engine. + + ``trace`` is not supported yet and must be False. + """ + shmem._require_copy_engine() + if trace: + raise NotImplementedError("trace=True is not supported for matmul_all_gather_host_copy_engine yet") + + M_local, K = A.shape + K2, N = B.shape + world_size = shmem.get_num_ranks() + + assert K == K2, f"Inner dimensions must match: A has K={K}, B has K={K2}" + + M = M_local * world_size + + ws = FusedWorkspace( + operation="matmul_all_gather_host_copy_engine", + shape=(M, N, K), + dtype=A.dtype, + world_size=world_size, + prepared=True, + ) + + if selector is None: + selector = _make_matmul_selector( + M_local, + N, + K, + A.dtype, + B.dtype, + A.dtype, + A.device, + streamk=False, + ) + + # Auto-calculate optimal m_tiles_per_batch if not provided + if m_tiles_per_batch is None: + m_tiles_per_batch = _auto_m_tiles_per_batch(selector, M_local, N) + num_tiles_m = triton.cdiv(M_local, selector.block_m) + num_tiles_n = triton.cdiv(N, selector.block_n) + launch_wave_plan = build_launch_wave_plan( + num_tiles_m=num_tiles_m, + num_tiles_n=num_tiles_n, + group_size_m=selector.group_m, + launch_grid=num_tiles_m * num_tiles_n, + wave_size=selector._ACTIVE_CU, + num_xcds=selector.num_sms, + ) + ws.selector = selector + ws.launch_wave_plan = launch_wave_plan + ws.locks = shmem.zeros((launch_wave_plan.num_waves,), dtype=torch.int32) + ws.num_tiles_m = num_tiles_m + ws.num_tiles_n = num_tiles_n + ws.num_batches = launch_wave_plan.num_waves + ws.num_waves = launch_wave_plan.num_waves + ws.m_tiles_per_batch = m_tiles_per_batch + + ws.completion_signals = shmem.zeros((world_size,), dtype=torch.int32) + + return ws + + +def matmul_all_gather_host_copy_engine( + shmem, + output_tensor: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + bias: Optional[torch.Tensor] = None, + async_op: bool = False, + workspace: Optional[FusedWorkspace] = None, + flag_iteration: int = 0, + trace: bool = False, + verbose: bool = False, +) -> FusedWorkspace: + """ + Fused matrix multiplication and all-gather using host-initiated SDMA with POLL packets. + + Computes: output = all_gather(A @ B + bias) along M dimension + + Each rank has A of shape (M_local, K) where M_local = M / world_size. + The host pre-queues SDMA POLL+COPY packets for all tiles and ranks. + The device kernel computes tiles, stores to local HBM, then sets flags. + SDMA hardware automatically performs scatter when flags are set. + + Args: + shmem: Iris shmem context + output_tensor: Output tensor C of shape (M, N) where M = M_local * world_size + A: Input matrix A of shape (M_local, K) + B: Input matrix B of shape (K, N) + bias: Optional bias vector (M_local,) + async_op: If False, performs barrier at end + workspace: Optional pre-allocated workspace + trace: Not supported yet; must be False. + + Returns: + FusedWorkspace object + """ + # Fail fast: without the SDMA engine the GEMM launched below would wait forever on flags. + shmem._require_copy_engine() + if trace: + raise NotImplementedError("trace=True is not supported for matmul_all_gather_host_copy_engine yet") + + M_local, K = A.shape + K2, N = B.shape + world_size = shmem.get_num_ranks() + rank = shmem.get_rank() + + assert K == K2, f"Inner dimensions must match: A has K={K}, B has K={K2}" + + M = M_local * world_size + assert output_tensor.shape == (M, N), f"Output must be ({M}, {N}), got {output_tensor.shape}" + + # Allocate workspace if not provided + if workspace is None: + workspace = matmul_all_gather_host_copy_engine_preamble(shmem, A, B, trace=trace) + + stride_cm, stride_cn = output_tensor.stride() + device = A.device + + selector = workspace.selector + launch_wave_plan = getattr(workspace, "launch_wave_plan", None) + num_tiles_m = workspace.num_tiles_m + num_tiles_n = workspace.num_tiles_n + num_batches = workspace.num_batches + if launch_wave_plan is None: + raise ValueError("workspace.launch_wave_plan must be initialized in preamble") + + # ═══════════════════════════════════════════════════════════════════════ + # Device Phase: Compute GEMM + store + set flags + # ═══════════════════════════════════════════════════════════════════════ + # tritonBLAS path with SignalView + selector = workspace.selector + + # Create a view of output_tensor at this rank's position + C_local_view = output_tensor[rank * M_local : (rank + 1) * M_local, :] + + epilogue_config = EpilogueConfig( + workspace.locks, + map_type="launch_wave", + block_group_m=launch_wave_plan.wave_size, + ) + + # Use work-stealing if enabled + tritonblas_config = None + + # Warn about bias + if bias is not None: + import warnings + + warnings.warn("Bias is not yet supported in tritonBLAS integration. Consider adding bias manually after GEMM.") + + # Launch tritonBLAS GEMM with SignalView + persistent_matmul_lt( + A, + B, + C_local_view, + selector, + config=tritonblas_config, + bias=None, + work_stealing=False, + epilogue_config=epilogue_config, + ) + + # ═══════════════════════════════════════════════════════════════════════ + # Host Phase: Enqueue SDMA POLL+COPY packets for all tiles + # (While kernel is running in parallel on device) + # ═══════════════════════════════════════════════════════════════════════ + element_size = output_tensor.element_size() + + if verbose and rank == 0: + shmem.info( + f"[Rank {rank}] Starting SDMA loop (launch-wave transfers)... " + f"num_m_tiles={num_tiles_m}, num_tiles_n={num_tiles_n}, " + f"wave_size={launch_wave_plan.wave_size}" + ) + shmem.info( + f"[Rank {rank}] Will transfer in {launch_wave_plan.num_waves} waves " + f"across {len(launch_wave_plan.transfers)} rects" + ) + + sdma_start_time = time.perf_counter() + tile_transfer_count = 0 + + # Get block size from workspace selector + block_size_m = workspace.selector.block_m + block_size_n = workspace.selector.block_n + + signal_ptr_local = workspace.completion_signals.data_ptr() + rank * workspace.completion_signals.element_size() + transfers_by_wave = [[] for _ in range(launch_wave_plan.num_waves)] + for transfer in launch_wave_plan.transfers: + transfers_by_wave[transfer.wave_id].append(transfer) + + for wave_id, wave_transfers in enumerate(transfers_by_wave): + if not wave_transfers: + continue + + expected_flag_value = (flag_iteration + 1) * launch_wave_plan.wave_tile_counts[wave_id] + wait_flag_ptr = workspace.locks.data_ptr() + wave_id * workspace.locks.element_size() + is_last_wave = wave_id == (launch_wave_plan.num_waves - 1) + + src_ptrs = [] + dst_ptrs_local = [] + tile_heights = [] + tile_widths = [] + src_strides = [] + dst_strides = [] + + for transfer in wave_transfers: + m_start = transfer.m_tile_start * block_size_m + n_start = transfer.n_tile_start * block_size_n + batch_height = min(transfer.m_tile_count * block_size_m, M_local - m_start) + batch_width = min(transfer.n_tile_count * block_size_n, N - n_start) + + src_offset = (m_start + rank * M_local) * stride_cm + n_start * stride_cn + src_ptrs.append(output_tensor.data_ptr() + src_offset * element_size) + dst_offset_local = (m_start + rank * M_local) * stride_cm + n_start * stride_cn + + dst_ptrs_local.append(output_tensor.data_ptr() + dst_offset_local * element_size) + tile_heights.append(batch_height) + tile_widths.append(batch_width) + src_strides.append(stride_cm) + dst_strides.append(stride_cm) + + for remote_rank in range(world_size): + if remote_rank == rank: + continue + + dst_ptrs_remote = [ + shmem.heap.translate(dst_ptr_local, rank, remote_rank) for dst_ptr_local in dst_ptrs_local + ] + signal_ptr_remote = None + if is_last_wave: + signal_ptr_remote = shmem.heap.translate(signal_ptr_local, rank, remote_rank) + + shmem.put_tiles( + src_ptrs, + to_rank=remote_rank, + dst_addrs=dst_ptrs_remote, + data_size=element_size, + tile_heights=tile_heights, + tile_widths=tile_widths, + src_strides=src_strides, + dst_strides=dst_strides, + wait_flag=wait_flag_ptr, + wait_value=expected_flag_value, + signal_flag=signal_ptr_remote, + signal_value=1, + async_op=True, + channel=0, + ) + tile_transfer_count += len(wave_transfers) + + if verbose and wave_id < 2 and remote_rank == (rank + 1) % world_size: + shmem.info( + f"[Rank {rank}→{remote_rank}] Queued wave={wave_id} " + f"transfers={len(wave_transfers)} tiles={launch_wave_plan.wave_tile_counts[wave_id]}" + ) + + sdma_end_post_time = time.perf_counter() + + if not async_op: + _wait_completion_signals_kernel[(world_size,)]( + workspace.completion_signals, + flag_iteration + 1, + rank, + world_size, + ) + sdma_end_time = time.perf_counter() + torch.cuda.synchronize() + if verbose and rank == 0: + post_ms = (sdma_end_post_time - sdma_start_time) * 1000.0 + quiet_ms = (sdma_end_time - sdma_end_post_time) * 1000.0 + total_ms = (sdma_end_time - sdma_start_time) * 1000.0 + shmem.info( + f"[Rank {rank}] SDMA complete. " + f"Post: {post_ms:.2f}ms, Wait: {quiet_ms:.2f}ms, Total: {total_ms:.2f}ms, " + f"transfers={tile_transfer_count}" + ) + + return workspace diff --git a/iris/ops/tritonblas_launch_wave_schedule.py b/iris/ops/tritonblas_launch_wave_schedule.py new file mode 100644 index 000000000..e52a68438 --- /dev/null +++ b/iris/ops/tritonblas_launch_wave_schedule.py @@ -0,0 +1,301 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. + +""" +Helpers for describing the tritonBLAS launch-wave tile schedule. + +The current tritonBLAS GEMM path launches the full logical grid, applies the +chunked XCD remap to each launched program id, and each launched program +produces one tile. This helper mirrors that launch order and coalesces the +tiles issued by each hardware wave of ``wave_size`` launched programs into one +or more rectangular transfers. +""" + +from __future__ import annotations + +from dataclasses import dataclass + + +def ceil_div(a: int, b: int) -> int: + return (a + b - 1) // b + + +def chiplet_transform_chunked(pid: int, num_workgroups: int, num_xcds: int, chunk_size: int) -> int: + if num_xcds <= 1 or chunk_size <= 0: + return pid + if pid > (num_workgroups // (num_xcds * chunk_size)) * (num_xcds * chunk_size): + return pid + + local_pid = pid // num_xcds + chunk_idx = local_pid // chunk_size + pos_in_chunk = local_pid % chunk_size + xcd = pid % num_xcds + return chunk_idx * num_xcds * chunk_size + xcd * chunk_size + pos_in_chunk + + +def default_chunk_size(total_tiles: int, group_size_m: int, num_xcds: int) -> int: + chunk_size = group_size_m * group_size_m + if num_xcds > 0: + chunk_size = min(chunk_size, max(1, total_tiles // num_xcds)) + return max(1, chunk_size) + + +def grouped_tile_coords(tile_id: int, num_tiles_m: int, num_tiles_n: int, group_size_m: int) -> tuple[int, int, int]: + num_pid_in_group = group_size_m * num_tiles_n + group_id = tile_id // num_pid_in_group + first_pid_m = group_id * group_size_m + actual_group_size_m = min(num_tiles_m - first_pid_m, group_size_m) + pid_m = first_pid_m + ((tile_id % num_pid_in_group) % actual_group_size_m) + pid_n = (tile_id % num_pid_in_group) // actual_group_size_m + return pid_m, pid_n, group_id + + +@dataclass(frozen=True) +class LaunchWaveTransfer: + wave_id: int + group_id: int + m_tile_start: int + n_tile_start: int + m_tile_count: int + n_tile_count: int + tile_count: int + + +@dataclass(frozen=True) +class LaunchWavePlan: + num_tiles_m: int + num_tiles_n: int + total_tiles: int + launch_grid: int + wave_size: int + num_xcds: int + chunk_size: int + num_waves: int + wave_tile_counts: tuple[int, ...] + transfers: tuple[LaunchWaveTransfer, ...] + + +def _coalesce_group_columns( + wave_id: int, + group_id: int, + first_pid_m: int, + columns: dict[int, set[int]], +) -> list[LaunchWaveTransfer]: + transfers: list[LaunchWaveTransfer] = [] + merged: list[tuple[int, int, int, int]] = [] + + for pid_n in sorted(columns): + local_ms = sorted(columns[pid_n]) + if not local_ms: + continue + + seg_start = local_ms[0] + seg_prev = local_ms[0] + for local_m in local_ms[1:]: + if local_m == seg_prev + 1: + seg_prev = local_m + continue + merged.append((pid_n, seg_start, seg_prev - seg_start + 1, 1)) + seg_start = local_m + seg_prev = local_m + merged.append((pid_n, seg_start, seg_prev - seg_start + 1, 1)) + + for pid_n, local_m_start, m_tile_count, n_tile_count in merged: + if transfers: + prev = transfers[-1] + if ( + prev.group_id == group_id + and prev.n_tile_start + prev.n_tile_count == pid_n + and prev.m_tile_start == first_pid_m + local_m_start + and prev.m_tile_count == m_tile_count + ): + transfers[-1] = LaunchWaveTransfer( + wave_id=prev.wave_id, + group_id=prev.group_id, + m_tile_start=prev.m_tile_start, + n_tile_start=prev.n_tile_start, + m_tile_count=prev.m_tile_count, + n_tile_count=prev.n_tile_count + n_tile_count, + tile_count=prev.tile_count + m_tile_count * n_tile_count, + ) + continue + + transfers.append( + LaunchWaveTransfer( + wave_id=wave_id, + group_id=group_id, + m_tile_start=first_pid_m + local_m_start, + n_tile_start=pid_n, + m_tile_count=m_tile_count, + n_tile_count=n_tile_count, + tile_count=m_tile_count * n_tile_count, + ) + ) + + return transfers + + +def _coalesce_group_rows( + wave_id: int, + group_id: int, + first_pid_m: int, + columns: dict[int, set[int]], +) -> list[LaunchWaveTransfer]: + rows: dict[int, set[int]] = {} + for pid_n, local_ms in columns.items(): + for local_m in local_ms: + rows.setdefault(local_m, set()).add(pid_n) + + transfers: list[LaunchWaveTransfer] = [] + merged: list[tuple[int, int, int, int]] = [] + + for local_m in sorted(rows): + ns = sorted(rows[local_m]) + if not ns: + continue + + seg_start = ns[0] + seg_prev = ns[0] + for pid_n in ns[1:]: + if pid_n == seg_prev + 1: + seg_prev = pid_n + continue + merged.append((local_m, seg_start, 1, seg_prev - seg_start + 1)) + seg_start = pid_n + seg_prev = pid_n + merged.append((local_m, seg_start, 1, seg_prev - seg_start + 1)) + + for local_m_start, n_tile_start, m_tile_count, n_tile_count in merged: + if transfers: + prev = transfers[-1] + if ( + prev.group_id == group_id + and prev.m_tile_start + prev.m_tile_count == first_pid_m + local_m_start + and prev.n_tile_start == n_tile_start + and prev.n_tile_count == n_tile_count + ): + transfers[-1] = LaunchWaveTransfer( + wave_id=prev.wave_id, + group_id=prev.group_id, + m_tile_start=prev.m_tile_start, + n_tile_start=prev.n_tile_start, + m_tile_count=prev.m_tile_count + m_tile_count, + n_tile_count=prev.n_tile_count, + tile_count=prev.tile_count + m_tile_count * n_tile_count, + ) + continue + + transfers.append( + LaunchWaveTransfer( + wave_id=wave_id, + group_id=group_id, + m_tile_start=first_pid_m + local_m_start, + n_tile_start=n_tile_start, + m_tile_count=m_tile_count, + n_tile_count=n_tile_count, + tile_count=m_tile_count * n_tile_count, + ) + ) + + return transfers + + +def build_launch_wave_plan( + num_tiles_m: int, + num_tiles_n: int, + group_size_m: int, + launch_grid: int, + wave_size: int, + num_xcds: int, + chunk_size: int | None = None, + merge_order: str = "column", +) -> LaunchWavePlan: + total_tiles = num_tiles_m * num_tiles_n + if launch_grid <= 0: + raise ValueError("launch_grid must be positive") + if wave_size <= 0: + raise ValueError("wave_size must be positive") + if total_tiles <= 0: + raise ValueError("tile grid must be non-empty") + if merge_order not in {"column", "row"}: + raise ValueError("merge_order must be 'column' or 'row'") + + if chunk_size is None: + chunk_size = default_chunk_size(total_tiles, group_size_m, num_xcds) + + launch_grid = max(launch_grid, total_tiles) + num_waves = ceil_div(launch_grid, wave_size) + wave_tile_counts: list[int] = [] + transfers: list[LaunchWaveTransfer] = [] + + launched_tile_ids = [ + chiplet_transform_chunked(pid, launch_grid, num_xcds, chunk_size) for pid in range(launch_grid) + ] + + for wave_id in range(num_waves): + groups: dict[int, tuple[int, dict[int, set[int]]]] = {} + tiles_in_wave = 0 + + wave_program_start = wave_id * wave_size + wave_program_end = min(wave_program_start + wave_size, launch_grid) + for launch_pid in range(wave_program_start, wave_program_end): + tile_id = launched_tile_ids[launch_pid] + if tile_id >= total_tiles: + continue + + pid_m, pid_n, group_id = grouped_tile_coords(tile_id, num_tiles_m, num_tiles_n, group_size_m) + first_pid_m = group_id * group_size_m + columns = groups.setdefault(group_id, (first_pid_m, {}))[1] + local_m = pid_m - first_pid_m + columns.setdefault(pid_n, set()).add(local_m) + tiles_in_wave += 1 + + wave_tile_counts.append(tiles_in_wave) + if tiles_in_wave == 0: + continue + + for group_id in sorted(groups): + first_pid_m, columns = groups[group_id] + if merge_order == "row": + transfers.extend(_coalesce_group_rows(wave_id, group_id, first_pid_m, columns)) + else: + transfers.extend(_coalesce_group_columns(wave_id, group_id, first_pid_m, columns)) + + return LaunchWavePlan( + num_tiles_m=num_tiles_m, + num_tiles_n=num_tiles_n, + total_tiles=total_tiles, + launch_grid=launch_grid, + wave_size=wave_size, + num_xcds=num_xcds, + chunk_size=chunk_size, + num_waves=num_waves, + wave_tile_counts=tuple(wave_tile_counts), + transfers=tuple(transfers), + ) + + +def build_launch_wave_plan_for_shape( + m: int, + n: int, + block_m: int, + block_n: int, + group_size_m: int, + launch_grid: int, + wave_size: int, + num_xcds: int, + chunk_size: int | None = None, + merge_order: str = "column", +) -> LaunchWavePlan: + num_tiles_m = ceil_div(m, block_m) + num_tiles_n = ceil_div(n, block_n) + return build_launch_wave_plan( + num_tiles_m=num_tiles_m, + num_tiles_n=num_tiles_n, + group_size_m=group_size_m, + launch_grid=launch_grid, + wave_size=wave_size, + num_xcds=num_xcds, + chunk_size=chunk_size, + merge_order=merge_order, + ) diff --git a/iris/ops/workspace.py b/iris/ops/workspace.py index a9c7cb616..ffa0bb505 100644 --- a/iris/ops/workspace.py +++ b/iris/ops/workspace.py @@ -42,6 +42,18 @@ class FusedWorkspace: aux_buffer: Optional[torch.Tensor] = None # Generic buffer for intermediate results locks: Optional[torch.Tensor] = None # Synchronization primitives + # Copy-engine workspace + completion_signals: Optional[torch.Tensor] = None # (world_size,) sender-completion slots + + # Launch-wave copy-engine metadata + wave_tile_counts: Optional[torch.Tensor] = None + wave_transfer_offsets: Optional[torch.Tensor] = None + wave_transfer_counts: Optional[torch.Tensor] = None + transfer_row_offsets: Optional[torch.Tensor] = None + transfer_col_offsets: Optional[torch.Tensor] = None + transfer_width_bytes: Optional[torch.Tensor] = None + transfer_heights: Optional[torch.Tensor] = None + prepared: bool = False def matches( @@ -82,4 +94,12 @@ def clear(self): """Free all allocated buffers.""" self.aux_buffer = None self.locks = None + self.completion_signals = None + self.wave_tile_counts = None + self.wave_transfer_offsets = None + self.wave_transfer_counts = None + self.transfer_row_offsets = None + self.transfer_col_offsets = None + self.transfer_width_bytes = None + self.transfer_heights = None self.prepared = False diff --git a/pyproject.toml b/pyproject.toml index 9e9e0b598..80f53cf1d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,7 +21,7 @@ dependencies = [ "numpy", "requests", "ruff<0.18", # 0.16 formats code blocks in Markdown; see lint.yml - "tritonblas @ git+https://github.com/ROCm/tritonBLAS.git@muhaawad/iris", + "tritonblas @ git+https://github.com/ROCm/tritonBLAS.git@dev/dasidler/merge-main-into-muhaawad-iris", ] [project.urls] @@ -39,7 +39,7 @@ amd = [ ] sdma = [ # SDMA copy engine (use_copy_engine=True, host-initiated put/quiet). Install via `pip install iris[sdma]`. - "rocm-xio @ git+https://github.com/ROCm/rocm-xio.git@cbe97e6392066bef7901121965ffadad19404da4", + "rocm-xio @ git+https://github.com/ROCm/rocm-xio.git@97fa1663a57a5c401f66a9e3ef677fa226ec1163", ] dev = [ "pytest", diff --git a/tests/ops/test_all_gather_matmul.py b/tests/ops/test_all_gather_matmul.py index ce2b6e3ce..d8cf2ffdd 100644 --- a/tests/ops/test_all_gather_matmul.py +++ b/tests/ops/test_all_gather_matmul.py @@ -9,11 +9,15 @@ Covers both the baseline pull kernel and the HBM-buffered kernel. """ +import gc import pytest import torch import torch.distributed as dist +from tritonblas.config import matmul_preamble +from tritonblas.matmul import _make_matmul_selector, matmul_lt import iris +import os from iris.ops.all_gather_matmul_hbm_buffer import ( _auto_config, _CHAMPION_CONFIGS, @@ -23,8 +27,91 @@ from iris.ops.config import FusedConfig -def _make_reference(rank, world_size, M, K_local, N, dtype): - """Build a torch reference output for all_gather + matmul.""" +@pytest.fixture(autouse=True) +def cleanup_gpu_memory(): + """Fixture to clean up GPU memory before and after each test.""" + # Cleanup before test starts + gc.collect() + torch.cuda.empty_cache() + torch.cuda.synchronize() + yield # Run the test + # Cleanup after test completes (pass or fail) + gc.collect() + torch.cuda.empty_cache() + torch.cuda.synchronize() + + +def _param_shapes(): + if "IRIS_TEST_M" in os.environ: + return [ + ( + int(os.environ["IRIS_TEST_M"]), + int(os.environ["IRIS_TEST_K_LOCAL"]), + int(os.environ["IRIS_TEST_N"]), + ) + ] + return [ + (128, 32, 64), + (256, 64, 128), + ] + + +def _heap_size() -> int: + return int(os.environ.get("IRIS_TEST_HEAP_SIZE", 1 << 33)) + + +def _full_validation_threshold_bytes() -> int: + return int(os.environ.get("IRIS_TEST_FULL_VALIDATION_THRESHOLD_BYTES", 2 << 30)) + + +def _validation_mode() -> str: + return os.environ.get("IRIS_TEST_VALIDATION_MODE", "auto").lower() + + +def _validation_tile_rows(M: int) -> int: + return max(1, min(M, int(os.environ.get("IRIS_TEST_VALIDATION_TILE_ROWS", 128)))) + + +def _validation_tile_cols(N: int) -> int: + return max(1, min(N, int(os.environ.get("IRIS_TEST_VALIDATION_TILE_COLS", 256)))) + + +def _validation_num_row_tiles() -> int: + return max(1, int(os.environ.get("IRIS_TEST_VALIDATION_NUM_ROW_TILES", 4))) + + +def _validation_num_col_tiles() -> int: + return max(1, int(os.environ.get("IRIS_TEST_VALIDATION_NUM_COL_TILES", 4))) + + +def _sample_starts(size: int, tile_size: int, num_tiles: int) -> list[int]: + if size <= tile_size: + return [0] + + max_start = size - tile_size + if num_tiles == 1: + return [0] + + starts = {0, max_start} + for tile_idx in range(1, num_tiles - 1): + starts.add((max_start * tile_idx) // (num_tiles - 1)) + return sorted(starts) + + +def _should_sample_validation(M: int, N: int, dtype: torch.dtype) -> bool: + mode = _validation_mode() + if mode == "full": + return False + if mode == "sampled": + return True + if mode != "auto": + raise ValueError(f"Unknown IRIS_TEST_VALIDATION_MODE={mode!r}; expected auto, full, or sampled") + + element_size = torch.tensor([], dtype=dtype).element_size() + return M * N * element_size > _full_validation_threshold_bytes() + + +def _make_inputs(rank, world_size, M, K_local, N, dtype): device = f"cuda:{rank}" K = K_local * world_size @@ -34,12 +121,95 @@ def _make_reference(rank, world_size, M, K_local, N, dtype): torch.manual_seed(123) B = torch.randn(K, N, dtype=dtype, device=device) - A_gathered_list = [torch.zeros(M, K_local, dtype=dtype, device=device) for _ in range(world_size)] + return A_sharded, B + + +def _make_full_reference(A_sharded, B, world_size): + """Build a full torch reference output for small all_gather + matmul cases.""" + A_gathered_list = [torch.empty_like(A_sharded) for _ in range(world_size)] dist.all_gather(A_gathered_list, A_sharded) A_gathered_ref = torch.cat(A_gathered_list, dim=1) ref_output = torch.matmul(A_gathered_ref, B) torch.cuda.synchronize() - return A_sharded, B, ref_output + return ref_output + + +def _make_reference(rank, world_size, M, K_local, N, dtype): + """Build inputs and the dense all-gather + matmul reference output.""" + A_sharded, B = _make_inputs(rank, world_size, M, K_local, N, dtype) + return A_sharded, B, _make_full_reference(A_sharded, B, world_size) + + +def _assert_close_tile(output_tile, ref_tile, atol, rtol, rank, row_start, col_start, context): + close = torch.isclose(output_tile, ref_tile, atol=atol, rtol=rtol) + if torch.all(close): + return + + mismatch_idx = torch.nonzero(~close, as_tuple=False)[0] + local_row = int(mismatch_idx[0].item()) + local_col = int(mismatch_idx[1].item()) + global_row = row_start + local_row + global_col = col_start + local_col + abs_diff = torch.abs(output_tile - ref_tile) + max_diff = torch.nan_to_num(abs_diff, nan=float("inf")).max().item() + output_val = output_tile[local_row, local_col].item() + ref_val = ref_tile[local_row, local_col].item() + pytest.fail( + f"Rank {rank}: sampled validation mismatch in {context} at row={global_row}, col={global_col}: " + f"output={output_val}, ref={ref_val}, max_diff={max_diff}, expected within atol={atol}, rtol={rtol}" + ) + + +def _assert_full_output_matches(output, ref_output, atol, rtol, rank, context): + if torch.allclose(output, ref_output, atol=atol, rtol=rtol): + return + + max_diff = (output - ref_output).abs().max().item() + pytest.fail(f"Rank {rank}: Max diff {max_diff}, expected < {atol} ({context})") + + +def _assert_sampled_output_matches(output, A_sharded, B, rank, world_size, atol, rtol, context, bias=None): + M, N = output.shape + rows_per_tile = _validation_tile_rows(M) + cols_per_tile = _validation_tile_cols(N) + row_starts = _sample_starts(M, rows_per_tile, _validation_num_row_tiles()) + col_starts = _sample_starts(N, cols_per_tile, _validation_num_col_tiles()) + + for row_start in row_starts: + row_end = min(row_start + rows_per_tile, M) + local_a = A_sharded[row_start:row_end].contiguous() + gathered_a_parts = [torch.empty_like(local_a) for _ in range(world_size)] + dist.all_gather(gathered_a_parts, local_a) + gathered_a_rows = torch.cat(gathered_a_parts, dim=1) + + for col_start in col_starts: + col_end = min(col_start + cols_per_tile, N) + ref_tile = torch.matmul(gathered_a_rows, B[:, col_start:col_end]) + if bias is not None: + ref_tile = ref_tile + bias[row_start:row_end, None] + output_tile = output[row_start:row_end, col_start:col_end] + _assert_close_tile(output_tile, ref_tile, atol, rtol, rank, row_start, col_start, context) + + del gathered_a_rows, gathered_a_parts, local_a + + torch.cuda.synchronize() + + +def _assert_output_matches_reference(output, A_sharded, B, rank, world_size, atol, rtol, context, bias=None): + if _should_sample_validation(output.shape[0], output.shape[1], output.dtype): + _assert_sampled_output_matches(output, A_sharded, B, rank, world_size, atol, rtol, context, bias=bias) + return + + ref_output = _make_full_reference(A_sharded, B, world_size) + if bias is not None: + ref_output = ref_output + bias[:, None] + _assert_full_output_matches(output, ref_output, atol, rtol, rank, context) + + +def _hbm_buffer_test_config(M: int, K_local: int, N: int) -> FusedConfig | None: + if M <= 256 or K_local <= 64 or N <= 128: + return FusedConfig(block_size_m=64, block_size_n=64, block_size_k=32) + return None @pytest.mark.parametrize( @@ -51,17 +221,14 @@ def _make_reference(rank, world_size, M, K_local, N, dtype): ) @pytest.mark.parametrize( "M,K_local,N", - [ - (128, 32, 64), - (256, 64, 128), - ], + _param_shapes(), ) def test_all_gather_matmul_baseline(dtype, atol, rtol, M, K_local, N): """Test baseline all_gather_matmul against torch all_gather + matmul.""" if not dist.is_initialized(): pytest.skip("torch.distributed not initialized") - heap_size = 2**33 + heap_size = _heap_size() ctx = iris.iris(heap_size) rank = ctx.get_rank() world_size = ctx.get_num_ranks() @@ -72,8 +239,7 @@ def test_all_gather_matmul_baseline(dtype, atol, rtol, M, K_local, N): if M < min_block_size or K_local < min_block_size or N < min_block_size: pytest.skip(f"Problem too small for min block size {min_block_size}") - A_sharded, B, ref_output = _make_reference(rank, world_size, M, K_local, N, dtype) - device = f"cuda:{rank}" + A_sharded, B = _make_inputs(rank, world_size, M, K_local, N, dtype) A_sharded_shmem = ctx.zeros((M, K_local), dtype=dtype) A_sharded_shmem.copy_(A_sharded) @@ -98,27 +264,94 @@ def test_all_gather_matmul_baseline(dtype, atol, rtol, M, K_local, N): torch.cuda.synchronize() ctx.barrier() - max_diff = (output - ref_output).abs().max().item() - assert torch.allclose(output, ref_output, atol=atol, rtol=rtol), ( - f"Rank {rank}: Max diff {max_diff}, expected < {atol}" + _assert_output_matches_reference( + output, + A_sharded, + B, + rank, + world_size, + atol, + rtol, + "all_gather_matmul_baseline", ) + # Clean up to prevent OOM in subsequent tests + del A_sharded, B, A_sharded_shmem, B_shmem, output + torch.cuda.empty_cache() + torch.cuda.synchronize() + @pytest.mark.parametrize( "dtype, atol, rtol", [ (torch.float16, 1e-2, 1e-2), - (torch.bfloat16, 1e-2, 1e-2), ], ) @pytest.mark.parametrize( "M,K_local,N", + _param_shapes(), +) +def test_tritonblas_rccl_all_gather_matmul(dtype, atol, rtol, M, K_local, N): + """Test RCCL all_gather + tritonBLAS matmul against torch reference.""" + if not dist.is_initialized(): + pytest.skip("torch.distributed not initialized") + + heap_size = _heap_size() + ctx = iris.iris(heap_size) + rank = ctx.get_rank() + world_size = ctx.get_num_ranks() + device = torch.device(f"cuda:{rank}") + + K = K_local * world_size + A_sharded, B = _make_inputs(rank, world_size, M, K_local, N, dtype) + + A_gathered_parts = [torch.empty((M, K_local), dtype=dtype, device=device) for _ in range(world_size)] + output = ctx.zeros((M, N), dtype=dtype) + selector = _make_matmul_selector( + M, + N, + K, + A_sharded.dtype, + B.dtype, + output.dtype, + device, + streamk=False, + ) + config = matmul_preamble(selector) + + dist.all_gather(A_gathered_parts, A_sharded) + A_gathered = torch.cat(A_gathered_parts, dim=1) + matmul_lt(A_gathered, B, output, selector, config) + + torch.cuda.synchronize() + + _assert_output_matches_reference( + output, + A_sharded, + B, + rank, + world_size, + atol, + rtol, + f"tritonblas+rccl, M={M}, K_local={K_local}, N={N}", + ) + + # Clean up to prevent OOM in subsequent tests + del A_sharded, B, output, A_gathered_parts, A_gathered + torch.cuda.empty_cache() + torch.cuda.synchronize() + + +@pytest.mark.parametrize( + "dtype, atol, rtol", [ - (128, 32, 64), - (256, 64, 128), - (512, 64, 128), + (torch.float16, 1e-2, 1e-2), ], ) +@pytest.mark.parametrize( + "M,K_local,N", + _param_shapes(), +) @pytest.mark.parametrize( "staged_a_layout", [ @@ -131,14 +364,14 @@ def test_all_gather_matmul_hbm_buffer(dtype, atol, rtol, M, K_local, N, staged_a if not dist.is_initialized(): pytest.skip("torch.distributed not initialized") - heap_size = 2**33 + heap_size = _heap_size() ctx = iris.iris(heap_size) rank = ctx.get_rank() world_size = ctx.get_num_ranks() K = K_local * world_size - A_sharded, B, ref_output = _make_reference(rank, world_size, M, K_local, N, dtype) + A_sharded, B = _make_inputs(rank, world_size, M, K_local, N, dtype) A_sharded_shmem = ctx.zeros((M, K_local), dtype=dtype) A_sharded_shmem.copy_(A_sharded) @@ -148,16 +381,10 @@ def test_all_gather_matmul_hbm_buffer(dtype, atol, rtol, M, K_local, N, staged_a ctx.barrier() - config = FusedConfig(block_size_m=64, block_size_n=64, block_size_k=32) - - # k_per_flag must divide num_k_blocks = K // block_size_k; use 1 for small shapes - num_k_blocks = K // config.block_size_k - k_per_flag = 1 - while k_per_flag * 2 <= 8 and num_k_blocks % (k_per_flag * 2) == 0: - k_per_flag *= 2 + config = _hbm_buffer_test_config(M, K_local, N) workspace = all_gather_matmul_hbm_buffer_preamble( - ctx, A_sharded_shmem, B_shmem, config=config, staged_a_layout=staged_a_layout, k_per_flag=k_per_flag + ctx, A_sharded_shmem, B_shmem, config=config, staged_a_layout=staged_a_layout ) all_gather_matmul_hbm_buffer( @@ -167,7 +394,6 @@ def test_all_gather_matmul_hbm_buffer(dtype, atol, rtol, M, K_local, N, staged_a B_shmem, config=config, workspace=workspace, - k_per_flag=k_per_flag, staged_a_layout=staged_a_layout, trace=False, ) @@ -175,44 +401,50 @@ def test_all_gather_matmul_hbm_buffer(dtype, atol, rtol, M, K_local, N, staged_a torch.cuda.synchronize() ctx.barrier() - max_diff = (output - ref_output).abs().max().item() - assert torch.allclose(output, ref_output, atol=atol, rtol=rtol), ( - f"Rank {rank}: Max diff {max_diff}, expected < {atol} " - f"(staged_a_layout={staged_a_layout}, M={M}, K_local={K_local}, N={N})" + _assert_output_matches_reference( + output, + A_sharded, + B, + rank, + world_size, + atol, + rtol, + f"staged_a_layout={staged_a_layout}, M={M}, K_local={K_local}, N={N}", ) + # Clean up to prevent OOM in subsequent tests + del A_sharded, B, A_sharded_shmem, B_shmem, output, workspace + torch.cuda.empty_cache() + torch.cuda.synchronize() + @pytest.mark.parametrize( "dtype, atol, rtol", [ (torch.float16, 1e-2, 1e-2), - (torch.bfloat16, 1e-2, 1e-2), ], ) @pytest.mark.parametrize( "M,K_local,N", - [ - (128, 32, 64), - ], + _param_shapes(), ) def test_all_gather_matmul_hbm_buffer_with_bias(dtype, atol, rtol, M, K_local, N): """Test all_gather_matmul_hbm_buffer with a bias vector.""" if not dist.is_initialized(): pytest.skip("torch.distributed not initialized") - heap_size = 2**33 + heap_size = _heap_size() ctx = iris.iris(heap_size) rank = ctx.get_rank() world_size = ctx.get_num_ranks() K = K_local * world_size - A_sharded, B, ref_output_no_bias = _make_reference(rank, world_size, M, K_local, N, dtype) + A_sharded, B = _make_inputs(rank, world_size, M, K_local, N, dtype) device = f"cuda:{rank}" torch.manual_seed(77) bias = torch.randn(M, dtype=dtype, device=device) - ref_output = ref_output_no_bias + bias[:, None] A_sharded_shmem = ctx.zeros((M, K_local), dtype=dtype) A_sharded_shmem.copy_(A_sharded) @@ -224,13 +456,7 @@ def test_all_gather_matmul_hbm_buffer_with_bias(dtype, atol, rtol, M, K_local, N ctx.barrier() - config = FusedConfig(block_size_m=64, block_size_n=64, block_size_k=32) - - # k_per_flag must divide num_k_blocks = K // block_size_k; use 1 for small shapes - num_k_blocks = K // config.block_size_k - k_per_flag = 1 - while k_per_flag * 2 <= 8 and num_k_blocks % (k_per_flag * 2) == 0: - k_per_flag *= 2 + config = _hbm_buffer_test_config(M, K_local, N) all_gather_matmul_hbm_buffer( ctx, @@ -239,25 +465,36 @@ def test_all_gather_matmul_hbm_buffer_with_bias(dtype, atol, rtol, M, K_local, N B_shmem, bias=bias_shmem, config=config, - k_per_flag=k_per_flag, trace=False, ) torch.cuda.synchronize() ctx.barrier() - max_diff = (output - ref_output).abs().max().item() - assert torch.allclose(output, ref_output, atol=atol, rtol=rtol), ( - f"Rank {rank}: Max diff {max_diff}, expected < {atol} (with bias)" + _assert_output_matches_reference( + output, + A_sharded, + B, + rank, + world_size, + atol, + rtol, + "with bias", + bias=bias, ) + # Clean up to prevent OOM in subsequent tests + del A_sharded, B, bias, A_sharded_shmem, B_shmem, bias_shmem, output + torch.cuda.empty_cache() + torch.cuda.synchronize() + def test_all_gather_matmul_hbm_buffer_auto_workspace(): """Test all_gather_matmul_hbm_buffer with workspace=None (auto preamble).""" if not dist.is_initialized(): pytest.skip("torch.distributed not initialized") - heap_size = 2**33 + heap_size = _heap_size() ctx = iris.iris(heap_size) rank = ctx.get_rank() world_size = ctx.get_num_ranks() @@ -310,7 +547,7 @@ def test_all_gather_matmul_hbm_buffer_workspace_reuse(): if not dist.is_initialized(): pytest.skip("torch.distributed not initialized") - heap_size = 2**33 + heap_size = _heap_size() ctx = iris.iris(heap_size) rank = ctx.get_rank() world_size = ctx.get_num_ranks() @@ -368,7 +605,7 @@ def test_all_gather_matmul_hbm_buffer_trace(): if not dist.is_initialized(): pytest.skip("torch.distributed not initialized") - heap_size = 2**33 + heap_size = _heap_size() ctx = iris.iris(heap_size) rank = ctx.get_rank() world_size = ctx.get_num_ranks() diff --git a/tests/ops/test_all_gather_matmul_copy_engine.py b/tests/ops/test_all_gather_matmul_copy_engine.py new file mode 100644 index 000000000..599708965 --- /dev/null +++ b/tests/ops/test_all_gather_matmul_copy_engine.py @@ -0,0 +1,293 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. + +""" +Tests for all_gather_matmul_copy_engine. + +Each rank owns A_sharded (M, K_local), gathers the K dimension across ranks, +and computes C = A_gathered @ B. This file exercises both the host-initiated +and device-initiated copy-engine paths against a torch reference. +""" + +import gc +import pytest +import torch +import torch.distributed as dist + +import iris +import os +from iris.ops.all_gather_matmul_copy_engine import ( + all_gather_matmul_copy_engine, + all_gather_matmul_copy_engine_preamble, +) +from tritonblas.matmul import _make_matmul_selector + + +@pytest.fixture(autouse=True) +def cleanup_gpu_memory(): + """Fixture to clean up GPU memory before and after each test.""" + # Cleanup before test starts + gc.collect() + torch.cuda.empty_cache() + torch.cuda.synchronize() + yield # Run the test + # Cleanup after test completes (pass or fail) + gc.collect() + torch.cuda.empty_cache() + torch.cuda.synchronize() + + +def _param_shapes(): + if "IRIS_TEST_M" in os.environ: + return [ + ( + int(os.environ["IRIS_TEST_M"]), + int(os.environ["IRIS_TEST_K_LOCAL"]), + int(os.environ["IRIS_TEST_N"]), + ) + ] + # (1024, 512, 1024) stays divisible by any selector block size at 1-8 ranks (K_local % block_k, M % block_m) + return [(256, 128, 256), (1024, 512, 1024)] + + +def _device_initiated_modes(): + mode = os.environ.get("IRIS_TEST_COPY_ENGINE_MODE") + if mode == "host": + return [False] + if mode == "device": + return [True] + return [False, True] + + +def _heap_size() -> int: + return int(os.environ.get("IRIS_TEST_HEAP_SIZE", 1 << 33)) + + +def _full_validation_threshold_bytes() -> int: + return int(os.environ.get("IRIS_TEST_FULL_VALIDATION_THRESHOLD_BYTES", 2 << 30)) + + +def _validation_mode() -> str: + return os.environ.get("IRIS_TEST_VALIDATION_MODE", "auto").lower() + + +def _validation_tile_rows(M: int) -> int: + return max(1, min(M, int(os.environ.get("IRIS_TEST_VALIDATION_TILE_ROWS", 128)))) + + +def _validation_tile_cols(N: int) -> int: + return max(1, min(N, int(os.environ.get("IRIS_TEST_VALIDATION_TILE_COLS", 256)))) + + +def _validation_num_row_tiles() -> int: + return max(1, int(os.environ.get("IRIS_TEST_VALIDATION_NUM_ROW_TILES", 4))) + + +def _validation_num_col_tiles() -> int: + return max(1, int(os.environ.get("IRIS_TEST_VALIDATION_NUM_COL_TILES", 4))) + + +def _sample_starts(size: int, tile_size: int, num_tiles: int) -> list[int]: + if size <= tile_size: + return [0] + + max_start = size - tile_size + if num_tiles == 1: + return [0] + + starts = {0, max_start} + for tile_idx in range(1, num_tiles - 1): + starts.add((max_start * tile_idx) // (num_tiles - 1)) + return sorted(starts) + + +def _should_sample_validation(M: int, N: int, dtype: torch.dtype) -> bool: + mode = _validation_mode() + if mode == "full": + return False + if mode == "sampled": + return True + if mode != "auto": + raise ValueError(f"Unknown IRIS_TEST_VALIDATION_MODE={mode!r}; expected auto, full, or sampled") + + element_size = torch.tensor([], dtype=dtype).element_size() + return M * N * element_size > _full_validation_threshold_bytes() + + +def _make_inputs(rank, world_size, M, K_local, N, dtype): + device = f"cuda:{rank}" + K = K_local * world_size + + torch.manual_seed(42 + rank) + A_sharded = torch.randn(M, K_local, dtype=dtype, device=device) + + torch.manual_seed(123) + B = torch.randn(K, N, dtype=dtype, device=device) + + return A_sharded, B + + +def _make_full_reference(A_sharded, B, world_size): + """Build a full torch reference output for small all_gather + matmul cases.""" + A_gathered_list = [torch.empty_like(A_sharded) for _ in range(world_size)] + dist.all_gather(A_gathered_list, A_sharded) + A_gathered_ref = torch.cat(A_gathered_list, dim=1) + ref_output = torch.matmul(A_gathered_ref, B) + torch.cuda.synchronize() + return ref_output + + +def _assert_close_tile(output_tile, ref_tile, atol, rtol, rank, row_start, col_start, context): + close = torch.isclose(output_tile, ref_tile, atol=atol, rtol=rtol) + if torch.all(close): + return + + mismatch_idx = torch.nonzero(~close, as_tuple=False)[0] + local_row = int(mismatch_idx[0].item()) + local_col = int(mismatch_idx[1].item()) + global_row = row_start + local_row + global_col = col_start + local_col + abs_diff = torch.abs(output_tile - ref_tile) + max_diff = torch.nan_to_num(abs_diff, nan=float("inf")).max().item() + output_val = output_tile[local_row, local_col].item() + ref_val = ref_tile[local_row, local_col].item() + pytest.fail( + f"Rank {rank}: sampled validation mismatch in {context} at row={global_row}, col={global_col}: " + f"output={output_val}, ref={ref_val}, max_diff={max_diff}, expected within atol={atol}, rtol={rtol}" + ) + + +def _assert_full_output_matches(output, ref_output, atol, rtol, rank, context): + if torch.allclose(output, ref_output, atol=atol, rtol=rtol): + return + + max_diff = (output - ref_output).abs().max().item() + pytest.fail(f"Rank {rank}: Max diff {max_diff}, expected < {atol} ({context})") + + +def _assert_sampled_output_matches(output, A_sharded, B, rank, world_size, atol, rtol, context): + M, N = output.shape + rows_per_tile = _validation_tile_rows(M) + cols_per_tile = _validation_tile_cols(N) + row_starts = _sample_starts(M, rows_per_tile, _validation_num_row_tiles()) + col_starts = _sample_starts(N, cols_per_tile, _validation_num_col_tiles()) + + for row_start in row_starts: + row_end = min(row_start + rows_per_tile, M) + local_a = A_sharded[row_start:row_end].contiguous() + gathered_a_parts = [torch.empty_like(local_a) for _ in range(world_size)] + dist.all_gather(gathered_a_parts, local_a) + gathered_a_rows = torch.cat(gathered_a_parts, dim=1) + + for col_start in col_starts: + col_end = min(col_start + cols_per_tile, N) + ref_tile = torch.matmul(gathered_a_rows, B[:, col_start:col_end]) + output_tile = output[row_start:row_end, col_start:col_end] + _assert_close_tile(output_tile, ref_tile, atol, rtol, rank, row_start, col_start, context) + + del gathered_a_rows, gathered_a_parts, local_a + + torch.cuda.synchronize() + + +def _assert_output_matches_reference(output, A_sharded, B, rank, world_size, atol, rtol, context): + if _should_sample_validation(output.shape[0], output.shape[1], output.dtype): + _assert_sampled_output_matches(output, A_sharded, B, rank, world_size, atol, rtol, context) + return + + ref_output = _make_full_reference(A_sharded, B, world_size) + _assert_full_output_matches(output, ref_output, atol, rtol, rank, context) + + +def _make_selector(M, N, K, dtype, device): + return _make_matmul_selector( + M, + N, + K, + dtype, + dtype, + dtype, + device, + streamk=False, + ) + + +@pytest.mark.parametrize("dtype, atol, rtol", [(torch.float16, 5e-2, 5e-2)]) +@pytest.mark.parametrize("device_initiated", _device_initiated_modes()) +@pytest.mark.parametrize("M,K_local,N", _param_shapes()) +def test_all_gather_matmul_copy_engine(dtype, atol, rtol, device_initiated, M, K_local, N): + """Test all_gather_matmul_copy_engine against torch all_gather + matmul.""" + if not dist.is_initialized(): + pytest.skip("torch.distributed not initialized") + + heap_size = _heap_size() + ctx = iris.iris(heap_size, engines=("sdma",)) + rank = ctx.get_rank() + world_size = ctx.get_num_ranks() + K = K_local * world_size + + A_sharded, B = _make_inputs(rank, world_size, M, K_local, N, dtype) + selector = _make_selector(M, N, K, dtype, B.device) + + if M % selector.block_m != 0: + pytest.skip(f"M={M} must be divisible by block_m={selector.block_m}") + if K % selector.block_k != 0: + pytest.skip(f"K={K} must be divisible by block_k={selector.block_k}") + if K_local % selector.block_k != 0: + pytest.skip(f"K_local={K_local} must be divisible by block_k={selector.block_k}") + + A_sharded_shmem = ctx.zeros((M, K_local), dtype=dtype) + A_sharded_shmem.copy_(A_sharded) + B_shmem = ctx.zeros((K, N), dtype=dtype) + B_shmem.copy_(B) + output = ctx.zeros((M, N), dtype=dtype) + + workspace = all_gather_matmul_copy_engine_preamble( + ctx, + A_sharded_shmem, + B_shmem, + selector=selector, + ) + + ctx.barrier() + + all_gather_matmul_copy_engine( + ctx, + output, + A_sharded_shmem, + B_shmem, + workspace=workspace, + device_initiated=device_initiated, + trace=False, + ) + + torch.cuda.synchronize() + ctx.barrier() + + _assert_output_matches_reference( + output, + A_sharded, + B, + rank, + world_size, + atol, + rtol, + f"device_initiated={device_initiated}, M={M}, K_local={K_local}, N={N}", + ) + + # Clean up to prevent OOM in subsequent tests + del A_sharded, B, A_sharded_shmem, B_shmem, output, workspace + torch.cuda.empty_cache() + torch.cuda.synchronize() + + +if __name__ == "__main__": + import sys + + if not dist.is_initialized(): + print("Run with: torchrun --nproc_per_node=2 tests/ops/test_all_gather_matmul_copy_engine.py") + sys.exit(1) + + rank = dist.get_rank() + torch.cuda.set_device(rank) + print(f"[Rank {rank}] Tests in this file require pytest + torchrun. See tests/run_tests_distributed.py") diff --git a/tests/ops/test_matmul_all_gather.py b/tests/ops/test_matmul_all_gather.py index ad2f37c4d..1ed3033e8 100644 --- a/tests/ops/test_matmul_all_gather.py +++ b/tests/ops/test_matmul_all_gather.py @@ -8,10 +8,156 @@ Install with: pip install git+https://github.com/ROCm/tritonBLAS.git """ +import gc import pytest import torch import torch.distributed as dist +import tritonblas import iris +import os + + +@pytest.fixture(autouse=True) +def cleanup_gpu_memory(): + """Release Python and CUDA allocations between parametrized tests.""" + gc.collect() + torch.cuda.empty_cache() + torch.cuda.synchronize() + yield + gc.collect() + torch.cuda.empty_cache() + torch.cuda.synchronize() + + +def _param_shapes(): + if "IRIS_TEST_M" in os.environ: + return [ + ( + int(os.environ["IRIS_TEST_M"]), + int(os.environ["IRIS_TEST_N"]), + int(os.environ["IRIS_TEST_K"]), + ) + ] + return [ + (64, 64, 32), + (512, 256, 512), + (1024, 2048, 1024), + ] + + +def _heap_size() -> int: + return int(os.environ.get("IRIS_TEST_HEAP_SIZE", 1 << 33)) + + +def _full_validation_threshold_bytes() -> int: + return int(os.environ.get("IRIS_TEST_FULL_VALIDATION_THRESHOLD_BYTES", 2 << 30)) + + +def _should_stream_validation(rows_per_rank: int, n: int, dtype: torch.dtype) -> bool: + element_size = torch.tensor([], dtype=dtype).element_size() + local_ref_bytes = rows_per_rank * n * element_size + return local_ref_bytes > _full_validation_threshold_bytes() + + +def _validation_rows_per_chunk(rows_per_rank: int, n: int, dtype: torch.dtype) -> int: + if "IRIS_TEST_VALIDATION_ROWS_PER_CHUNK" in os.environ: + return max(1, min(rows_per_rank, int(os.environ["IRIS_TEST_VALIDATION_ROWS_PER_CHUNK"]))) + element_size = torch.tensor([], dtype=dtype).element_size() + target_bytes = 2 << 30 + rows_per_chunk = max(1, target_bytes // max(1, n * element_size)) + return max(1, min(rows_per_rank, rows_per_chunk)) + + +def _validation_cols_per_chunk(rows_per_chunk: int, n: int, dtype: torch.dtype) -> int: + if "IRIS_TEST_VALIDATION_COLS_PER_CHUNK" in os.environ: + return max(1, min(n, int(os.environ["IRIS_TEST_VALIDATION_COLS_PER_CHUNK"]))) + element_size = torch.tensor([], dtype=dtype).element_size() + target_bytes = 128 << 20 + cols_per_chunk = max(1, target_bytes // max(1, rows_per_chunk * element_size * 4)) + return max(1, min(n, cols_per_chunk)) + + +def _assert_close_chunked(output_chunk, ref_chunk, atol, rtol, src_rank, row_start): + cols_per_chunk = _validation_cols_per_chunk(output_chunk.shape[0], output_chunk.shape[1], output_chunk.dtype) + + for col_start in range(0, output_chunk.shape[1], cols_per_chunk): + col_end = min(col_start + cols_per_chunk, output_chunk.shape[1]) + output_slice = output_chunk[:, col_start:col_end] + ref_slice = ref_chunk[:, col_start:col_end] + abs_diff = torch.abs(output_slice - ref_slice) + tolerance = atol + rtol * torch.abs(ref_slice) + mismatch = abs_diff > tolerance + + if torch.any(mismatch): + mismatch_idx = torch.nonzero(mismatch, as_tuple=False)[0] + local_row = int(mismatch_idx[0].item()) + local_col = int(mismatch_idx[1].item()) + global_row = row_start + local_row + global_col = col_start + local_col + max_diff = torch.max(abs_diff).item() + output_val = output_slice[local_row, local_col].item() + ref_val = ref_slice[local_row, local_col].item() + pytest.fail( + f"Mismatch in gathered rows from src_rank={src_rank} at row={global_row}, col={global_col}: " + f"output={output_val}, ref={ref_val}, max_diff={max_diff}, expected within atol={atol}, rtol={rtol}\n" + f"Rank validation failed for shmem.ops.matmul_all_gather" + ) + + +def _assert_gathered_rows_match_dense(output, local_ref, rank, world_size, atol, rtol): + recv_chunk = None + rows_per_rank = local_ref.shape[0] + + for src_rank in range(world_size): + if rank == src_rank: + ref_chunk = local_ref + else: + if recv_chunk is None: + recv_chunk = torch.empty_like(local_ref) + ref_chunk = recv_chunk + + dist.broadcast(ref_chunk, src=src_rank) + row_start = src_rank * rows_per_rank + row_end = row_start + rows_per_rank + output_chunk = output[row_start:row_end] + + if not torch.allclose(output_chunk, ref_chunk, atol=atol, rtol=rtol): + max_diff = torch.max(torch.abs(output_chunk - ref_chunk)).item() + pytest.fail( + f"Max difference in gathered rows from src_rank={src_rank}: {max_diff}, expected < {atol}\n" + f"Rank {rank}: shmem.ops.matmul_all_gather output doesn't match reference" + ) + + +def _assert_gathered_rows_match_streamed(output, A_local, B, rank, world_size, atol, rtol): + rows_per_rank = A_local.shape[0] + rows_per_chunk = _validation_rows_per_chunk(rows_per_rank, output.shape[1], output.dtype) + + for src_rank in range(world_size): + for local_row_start in range(0, rows_per_rank, rows_per_chunk): + local_row_end = min(local_row_start + rows_per_chunk, rows_per_rank) + chunk_rows = local_row_end - local_row_start + + if rank == src_rank: + ref_chunk = torch.matmul(A_local[local_row_start:local_row_end], B) + else: + ref_chunk = torch.empty((chunk_rows, output.shape[1]), dtype=output.dtype, device=output.device) + + dist.broadcast(ref_chunk, src=src_rank) + global_row_start = src_rank * rows_per_rank + local_row_start + global_row_end = global_row_start + chunk_rows + output_chunk = output[global_row_start:global_row_end] + _assert_close_chunked(output_chunk, ref_chunk, atol, rtol, src_rank, global_row_start) + + +def _assert_gathered_rows_match(output, A_local, B, rank, world_size, atol, rtol): + if _should_stream_validation(A_local.shape[0], output.shape[1], output.dtype): + _assert_gathered_rows_match_streamed(output, A_local, B, rank, world_size, atol, rtol) + return + + local_ref = torch.matmul(A_local, B) + torch.cuda.synchronize() + _assert_gathered_rows_match_dense(output, local_ref, rank, world_size, atol, rtol) @pytest.mark.parametrize( @@ -24,18 +170,14 @@ ) @pytest.mark.parametrize( "M, N, K", - [ - (64, 64, 32), - (512, 256, 512), - (1024, 2048, 1024), - ], + _param_shapes(), ) def test_matmul_all_gather(dtype, atol, rtol, M, N, K): """Test matmul_all_gather using shmem.ops API with proper config.""" if not dist.is_initialized(): pytest.skip("torch.distributed not initialized") - heap_size = 2**33 # 8GB + heap_size = _heap_size() shmem = iris.iris(heap_size) rank = shmem.get_rank() world_size = shmem.get_num_ranks() @@ -61,15 +203,6 @@ def test_matmul_all_gather(dtype, atol, rtol, M, N, K): B = shmem.randn((K, N), dtype=dtype) output = shmem.zeros((M, N), dtype=dtype) - # Reference: compute local GEMM, then all-gather along M dimension - A_ref = A_local.clone() - B_ref = B.clone() - C_local_ref = torch.matmul(A_ref, B_ref) - C_gathered_list = [torch.zeros(M_local, N, dtype=dtype, device=f"cuda:{rank}") for _ in range(world_size)] - dist.all_gather(C_gathered_list, C_local_ref) - pytorch_output = torch.cat(C_gathered_list, dim=0) # Concatenate along M dimension - torch.cuda.synchronize() - shmem.barrier() # Use appropriate block sizes based on problem size @@ -102,18 +235,77 @@ def test_matmul_all_gather(dtype, atol, rtol, M, N, K): torch.cuda.synchronize() shmem.barrier() - max_diff = torch.abs(output - pytorch_output).max().item() - - assert torch.allclose(output, pytorch_output, atol=atol, rtol=rtol), ( - f"Max difference: {max_diff}, expected < {atol}\n" - f"Rank {rank}: shmem.ops.matmul_all_gather output doesn't match reference" - ) + _assert_gathered_rows_match(output, A_local, B, rank, world_size, atol, rtol) if rank == 0: print(f"✓ matmul_all_gather test passed: {dtype}, M={M}, N={N}, K={K}") shmem.barrier() + del output + del B + del A_local del shmem import gc gc.collect() + torch.cuda.empty_cache() + + +@pytest.mark.parametrize( + "dtype, atol, rtol", + [ + (torch.float16, 0.5, 0.01), + (torch.bfloat16, 0.5, 0.01), + ], +) +@pytest.mark.parametrize( + "M, N, K", + _param_shapes(), +) +def test_tritonblas_rccl_matmul_all_gather(dtype, atol, rtol, M, N, K): + """Test tritonBLAS matmul + RCCL all_gather against a dense local reference.""" + if not dist.is_initialized(): + pytest.skip("torch.distributed not initialized") + + heap_size = _heap_size() + shmem = iris.iris(heap_size) + rank = shmem.get_rank() + world_size = shmem.get_num_ranks() + + if M % world_size != 0: + pytest.skip(f"M={M} not divisible by world_size={world_size}") + + M_local = M // world_size + min_block_size = 32 + if M_local < min_block_size: + pytest.skip(f"M_local={M_local} too small for world_size={world_size} (need >= {min_block_size})") + if K < min_block_size: + pytest.skip(f"K={K} too small (need >= {min_block_size})") + if N < min_block_size: + pytest.skip(f"N={N} too small (need >= {min_block_size})") + + torch.manual_seed(123 + rank) + A_local = shmem.randn((M_local, K), dtype=dtype) + torch.manual_seed(456) + B = shmem.randn((K, N), dtype=dtype) + C_local = shmem.zeros((M_local, N), dtype=dtype) + output = shmem.zeros((M, N), dtype=dtype) + selector = tritonblas.OrigamiMatmulSelector( + M_local, + N, + K, + A_local.dtype, + B.dtype, + C_local.dtype, + A_local.device, + ) + config = tritonblas.matmul_preamble(selector) + + shmem.barrier() + tritonblas.matmul_lt(A_local, B, C_local, selector, config) + dist.all_gather_into_tensor(output, C_local) + + torch.cuda.synchronize() + shmem.barrier() + + _assert_gathered_rows_match(output, A_local, B, rank, world_size, atol, rtol) diff --git a/tests/ops/test_matmul_all_gather_copy_engine.py b/tests/ops/test_matmul_all_gather_copy_engine.py new file mode 100644 index 000000000..c01c0cafb --- /dev/null +++ b/tests/ops/test_matmul_all_gather_copy_engine.py @@ -0,0 +1,272 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. + +""" +Tests for matmul_all_gather_copy_engine. + +Each rank computes C_local = A_local @ B and the copy engine scatters the +result tiles so every rank observes the gathered output C. +""" + +import pytest +import torch +import torch.distributed as dist + +import iris +import os +from iris.ops.matmul_all_gather_host_copy_engine import ( + matmul_all_gather_host_copy_engine, + matmul_all_gather_host_copy_engine_preamble, +) +from iris.ops.matmul_all_gather_copy_engine import ( + matmul_all_gather_copy_engine, + matmul_all_gather_copy_engine_preamble, +) +from tritonblas.matmul import _make_matmul_selector + + +def _param_shapes(): + if "IRIS_TEST_M" in os.environ: + return [ + ( + int(os.environ["IRIS_TEST_M"]), + int(os.environ["IRIS_TEST_N"]), + int(os.environ["IRIS_TEST_K"]), + ) + ] + return [(1024, 256, 256)] + + +def _copy_engine_modes(): + mode = os.environ.get("IRIS_TEST_COPY_ENGINE_MODE") + if mode in {"host", "device"}: + return [mode] + return ["host", "device"] + + +def _heap_size() -> int: + return int(os.environ.get("IRIS_TEST_HEAP_SIZE", 1 << 33)) + + +def _full_validation_threshold_bytes() -> int: + return int(os.environ.get("IRIS_TEST_FULL_VALIDATION_THRESHOLD_BYTES", 2 << 30)) + + +def _should_stream_validation(rows_per_rank: int, n: int, dtype: torch.dtype) -> bool: + element_size = torch.tensor([], dtype=dtype).element_size() + local_ref_bytes = rows_per_rank * n * element_size + return local_ref_bytes > _full_validation_threshold_bytes() + + +def _validation_rows_per_chunk(rows_per_rank: int, n: int, dtype: torch.dtype) -> int: + if "IRIS_TEST_VALIDATION_ROWS_PER_CHUNK" in os.environ: + return max(1, min(rows_per_rank, int(os.environ["IRIS_TEST_VALIDATION_ROWS_PER_CHUNK"]))) + element_size = torch.tensor([], dtype=dtype).element_size() + target_bytes = 2 << 30 + rows_per_chunk = max(1, target_bytes // max(1, n * element_size)) + return max(1, min(rows_per_rank, rows_per_chunk)) + + +def _validation_cols_per_chunk(rows_per_chunk: int, n: int, dtype: torch.dtype) -> int: + if "IRIS_TEST_VALIDATION_COLS_PER_CHUNK" in os.environ: + return max(1, min(n, int(os.environ["IRIS_TEST_VALIDATION_COLS_PER_CHUNK"]))) + element_size = torch.tensor([], dtype=dtype).element_size() + target_bytes = 128 << 20 + cols_per_chunk = max(1, target_bytes // max(1, rows_per_chunk * element_size * 4)) + return max(1, min(n, cols_per_chunk)) + + +def _assert_close_chunked(output_chunk, ref_chunk, atol, rtol, copy_engine_mode, src_rank, row_start): + cols_per_chunk = _validation_cols_per_chunk(output_chunk.shape[0], output_chunk.shape[1], output_chunk.dtype) + + for col_start in range(0, output_chunk.shape[1], cols_per_chunk): + col_end = min(col_start + cols_per_chunk, output_chunk.shape[1]) + output_slice = output_chunk[:, col_start:col_end] + ref_slice = ref_chunk[:, col_start:col_end] + abs_diff = torch.abs(output_slice - ref_slice) + tolerance = atol + rtol * torch.abs(ref_slice) + mismatch = abs_diff > tolerance + + if torch.any(mismatch): + mismatch_idx = torch.nonzero(mismatch, as_tuple=False)[0] + local_row = int(mismatch_idx[0].item()) + local_col = int(mismatch_idx[1].item()) + global_row = row_start + local_row + global_col = col_start + local_col + max_diff = torch.max(abs_diff).item() + output_val = output_slice[local_row, local_col].item() + ref_val = ref_slice[local_row, local_col].item() + pytest.fail( + f"Mismatch in gathered rows from src_rank={src_rank} at row={global_row}, col={global_col}: " + f"output={output_val}, ref={ref_val}, max_diff={max_diff}, expected within atol={atol}, rtol={rtol}\n" + f"Rank validation failed for matmul_all_gather_copy_engine (mode={copy_engine_mode})" + ) + + +def _assert_gathered_rows_match_dense(output, local_ref, rank, world_size, atol, rtol, copy_engine_mode): + recv_chunk = None + rows_per_rank = local_ref.shape[0] + + for src_rank in range(world_size): + if rank == src_rank: + ref_chunk = local_ref + else: + if recv_chunk is None: + recv_chunk = torch.empty_like(local_ref) + ref_chunk = recv_chunk + + dist.broadcast(ref_chunk, src=src_rank) + row_start = src_rank * rows_per_rank + row_end = row_start + rows_per_rank + output_chunk = output[row_start:row_end] + + if not torch.allclose(output_chunk, ref_chunk, atol=atol, rtol=rtol): + max_diff = torch.max(torch.abs(output_chunk - ref_chunk)).item() + pytest.fail( + f"Max difference in gathered rows from src_rank={src_rank}: {max_diff}, expected < {atol}\n" + f"Rank validation failed for matmul_all_gather_copy_engine (mode={copy_engine_mode})" + ) + + +def _assert_gathered_rows_match_streamed(output, A_local, B, rank, world_size, atol, rtol, copy_engine_mode): + rows_per_rank = A_local.shape[0] + rows_per_chunk = _validation_rows_per_chunk(rows_per_rank, output.shape[1], output.dtype) + + for src_rank in range(world_size): + for local_row_start in range(0, rows_per_rank, rows_per_chunk): + local_row_end = min(local_row_start + rows_per_chunk, rows_per_rank) + chunk_rows = local_row_end - local_row_start + + if rank == src_rank: + ref_chunk = torch.matmul(A_local[local_row_start:local_row_end], B) + else: + ref_chunk = torch.empty((chunk_rows, output.shape[1]), dtype=output.dtype, device=output.device) + + dist.broadcast(ref_chunk, src=src_rank) + global_row_start = src_rank * rows_per_rank + local_row_start + global_row_end = global_row_start + chunk_rows + output_chunk = output[global_row_start:global_row_end] + _assert_close_chunked( + output_chunk, + ref_chunk, + atol, + rtol, + copy_engine_mode, + src_rank, + global_row_start, + ) + + +def _assert_gathered_rows_match(output, A_local, B, rank, world_size, atol, rtol, copy_engine_mode): + if _should_stream_validation(A_local.shape[0], output.shape[1], output.dtype): + _assert_gathered_rows_match_streamed(output, A_local, B, rank, world_size, atol, rtol, copy_engine_mode) + return + + local_ref = torch.matmul(A_local, B) + torch.cuda.synchronize() + _assert_gathered_rows_match_dense(output, local_ref, rank, world_size, atol, rtol, copy_engine_mode) + + +def _make_selector(M_local, N, K, dtype, device): + return _make_matmul_selector( + M_local, + N, + K, + dtype, + dtype, + dtype, + device, + streamk=False, + ) + + +@pytest.mark.parametrize("dtype, atol, rtol", [(torch.float16, 5e-2, 5e-2)]) +@pytest.mark.parametrize("copy_engine_mode", _copy_engine_modes()) +@pytest.mark.parametrize("M,N,K", _param_shapes()) +def test_matmul_all_gather_copy_engine(dtype, atol, rtol, copy_engine_mode, M, N, K): + if not dist.is_initialized(): + pytest.skip("torch.distributed not initialized") + + heap_size = _heap_size() + shmem = iris.iris(heap_size, engines=("sdma",)) + rank = shmem.get_rank() + world_size = shmem.get_num_ranks() + device = torch.device(f"cuda:{rank}") + + if M % world_size != 0: + pytest.skip(f"M={M} not divisible by world_size={world_size}") + + M_local = M // world_size + selector = _make_selector(M_local, N, K, dtype, device) + + if M_local % selector.block_m != 0: + pytest.skip(f"M_local={M_local} must be divisible by block_size_m={selector.block_m}") + if K % selector.block_k != 0: + pytest.skip(f"K={K} must be divisible by block_size_k={selector.block_k}") + + A_local = shmem.randn((M_local, K), dtype=dtype) + B = shmem.randn((K, N), dtype=dtype) + output = shmem.zeros((M, N), dtype=dtype) + + if copy_engine_mode == "device": + workspace = matmul_all_gather_copy_engine_preamble( + shmem, + A_local, + B, + selector=selector, + ) + else: + workspace = matmul_all_gather_host_copy_engine_preamble( + shmem, + A_local, + B, + trace=False, + selector=selector, + ) + + shmem.barrier() + + if copy_engine_mode == "device": + matmul_all_gather_copy_engine( + shmem, + output, + A_local, + B, + workspace=workspace, + ) + else: + matmul_all_gather_host_copy_engine( + shmem, + output, + A_local, + B, + workspace=workspace, + trace=False, + ) + + torch.cuda.synchronize() + shmem.barrier() + + _assert_gathered_rows_match(output, A_local, B, rank, world_size, atol, rtol, copy_engine_mode) + + shmem.barrier() + del output + del B + del A_local + del shmem + import gc + + gc.collect() + torch.cuda.empty_cache() + + +if __name__ == "__main__": + import sys + + if not dist.is_initialized(): + print("Run with: torchrun --nproc_per_node=2 tests/ops/test_matmul_all_gather_copy_engine.py") + sys.exit(1) + + rank = dist.get_rank() + torch.cuda.set_device(rank) + print(f"[Rank {rank}] Tests in this file require pytest + torchrun. See tests/run_tests_distributed.py")