Skip to content

SDMA AG-GEMM and GEMM-AG benchmarks - #496

Open
dsidler wants to merge 209 commits into
mainfrom
dev/dasidler/sdma-benchmark
Open

dsidler wants to merge 209 commits into
mainfrom
dev/dasidler/sdma-benchmark

Conversation

@dsidler

@dsidler dsidler commented Apr 7, 2026 •

Copy link
Copy Markdown
Collaborator

Summary

Adds copy-engine (SDMA) variants of the fused all-gather + GEMM ops. The SDMA engines move the remote tiles while the compute units run a persistent tritonBLAS GEMM, so no workgroups are spent on fetching.

What's in it

Ops (iris.ops)

  • all_gather_matmul_copy_engine: host-initiated or device-initiated (device_initiated=True) transfers. Remote K-shards are posted in M-batches, and the GEMM waits per batch on PrologueConfig flags.
  • matmul_all_gather_copy_engine (device-initiated) and matmul_all_gather_host_copy_engine: the GEMM signals tile completion through EpilogueConfig counters, and SDMA gathers launch-wave by launch-wave (tritonblas_launch_wave_schedule.py).
  • All three require iris.iris(..., engines=("sdma",)) and fail fast with a clear error otherwise.

Host and device API

  • Iris.put_tile / put_tiles take strided 2D copies, optional signal flags and wait-then-put.
  • New device ops in iris.mem.triton: put_signal, put_signal_rect, wait_then_put_rect(s), wait_then_put_signal_rect(s).

Benchmarks and tests

  • Copy-engine benchmarks for both ops. Plus baselines: RCCL + tritonBLAS for all-gather + matmul, and torch.mm / tritonBLAS + RCCL for matmul + all-gather.
  • iris.bench has a new --engines sdma flag. The copy-engine benchmarks skip with a message without it.
  • Tests: new copy-engine tests plus the baseline test updates. Large-shape validation is sampled and controlled through IRIS_TEST_* env vars. The default heap is 8 GB.

Fixes to existing ops

  • int64 pointer offsets for C and staged_a in all_gather_matmul and all_gather_matmul_hbm_buffer.
  • all_gather_matmul_hbm_buffer halves the default k_per_flag until it divides the K-blocks, so small and single-rank problems work.

🤖 Generated with Claude Code

dsidler and others added 12 commits October 5, 2026 23:51
…rologueConfig

- tritonblas: pin to the tritonBLAS branch that merges main (SignalView/WaitView,
  EpilogueConfig/PrologueConfig) into muhaawad/iris, installs origami from PyPI
  (rocm-origami), and includes the gfx950 LDS estimate fix.
- rocm-xio: pin to current main (97fa166), 24 commits past the previous pin.
- iris.ops copy-engine variants: construct EpilogueConfig / PrologueConfig
  directly and pass epilogue_config / prologue_config to persistent_matmul_lt,
  replacing the removed create_counter_config / create_wait_config helpers.

Co-Authored-By: Claude <noreply@anthropic.com>
@dsidler dsidler changed the title Copy Engine AG-GEMM and GEMM-AG benchmarks SDMA AG-GEMM and GEMM-AG benchmarks Oct 23, 2026
Comment thread pyproject.toml
"requests",
"ruff<0.17", # 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",

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

revert once the two tritonBLAS PR merged:
ROCm/tritonBLAS#100
ROCm/tritonBLAS#101

@dsidler
dsidler marked this pull request as ready for review October 23, 2026 22:00

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

in-progress We are working on it iris Iris project issue

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants