SDMA AG-GEMM and GEMM-AG benchmarks

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])

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()

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()

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()