Add iris.concurrent (concurrent, independent, fused comm+gemm) - #544
Open
ryanswann-amd wants to merge 9 commits into
Open
Add iris.concurrent (concurrent, independent, fused comm+gemm)#544ryanswann-amd wants to merge 9 commits into
ryanswann-amd wants to merge 9 commits into
Conversation
Introduces iris.concurrent: run an independent GEMM concurrently with a collective (all_gather / all_reduce / reduce_scatter / all_to_all / broadcast) via persistent work-stealing kernels (fused single-kernel or two-stream modes), with an optional autotuner. - concurrent/_kernels.py, gemm.py: work-stealing GEMM+collective kernels and host launchers; the GEMM/comm CU split (gemm_wgs) and GEMM tile are the knobs. - concurrent/predictor.py: analytical origami-based CU-split cost model (origami is an optional dependency; import-safe without it). - concurrent/autotune.py: pass tune=True to any op to predict -> benchmark the top-k configs -> cache the winner to a JSON tuning DB (~/.cache/iris/...). Always benchmarks the static default too, so tuning is monotonic (never worse than the untuned baseline). - concurrent/tuning_data/: committed measurement corpus (MI300X/gfx942 and MI350X/gfx950, world 2/4/8; 4 collectives x 15 shapes) plus per-arch derived candidate sets and the collect/derive tooling, so the tuner has good defaults with or without the cost model. - tests: predictor + autotuner (CPU-only, pass) and fused correctness (2-GPU). Co-authored-by: Cursor <cursoragent@cursor.com>
The per-shape grid measurements (data/<arch>_world<W>.json) are research provenance, not a runtime dependency, and grow per arch x world. Keep them in the origami_comms workbench instead of this repo. iris ships only the derived candidate_set.json (loaded at runtime) plus the collect/derive tooling; derive_candidates.py gains a --data flag to point at the external corpus. The regenerated candidate_set.json is byte-identical, so this is data-location only. Co-authored-by: Cursor <cursoragent@cursor.com>
Replace the hand-rolled MFMA K-loop in _gemm_tile with tritonblas.GemmContext.reduce_axis (the same building block iris.ops uses), keeping the GROUP_M work-stealing tile swizzle and the local .wt store. All four kernels that call _gemm_tile (the three fused dual-queue kernels + ws_gemm) inherit it; comm tiles are unchanged. cache_modifier_a/b="" is required: GemmContext defaults to ".cg" (bypass-L1 streaming) which discards GEMM operand reuse and is ~1.3-1.6x slower here; with default L1/L2 caching the swap is perf-neutral-to-faster (measured 1.5-2.7% faster in the fused work-stealing kernel, bit-correct). Validated 22/22 tests/ccl/test_concurrent_fused (all 5 collectives x fused/concurrent). Note: this couples iris.concurrent to tritonblas + a Triton>=3.7 aggregate API (same dependency as iris.ops). Co-authored-by: Cursor <cursoragent@cursor.com>
Replace the manual iris.put peer loop in _all_gather_tile with the shared iris device-context tile collective ctx.all_gather (the same primitive iris.ops uses); its row-offset layout matches the old hand-rolled push exactly. The two all-gather kernels (fused_ws_gemm_all_gather, ws_all_gather) and the gemm.py all_gather launcher now pass the device-context tensor (shmem.get_device_context()) instead of heap_bases; the other four collectives are unchanged (their device-ctx methods use lock-based ready-flags / a different partition axis, so they stay hand-rolled on heap_bases). Validated 22/22 tests/ccl/test_concurrent_fused (all 5 collectives x fused/ concurrent). Same put/store traffic as before, so perf is expected-neutral. Co-authored-by: Cursor <cursoragent@cursor.com>
…s/tile_layout Extract the copy-pasted GROUP_M work-stealing swizzle + rm/rn/mask preamble from the four hand-rolled comm tiles (_all_reduce_tile, _reduce_scatter_tile, _all_to_all_tile, _broadcast_tile) into one _grouped_coords helper that reuses the shared iris.mem.triton.types.tile_layout (with its vectorization hints). Behavior-preserving: the collective algorithms (iris.put/load/store on heap_bases, block/chunk partitioning) are unchanged; only the duplicated coordinate boilerplate is shared. tile_layout drops the "% M" address wrap in favor of a real bounds mask, equivalent for the block-divisible shapes these kernels target (same equivalence already validated by the GEMM tile swap). Validated 22/22 tests/ccl/test_concurrent_fused (all 5 collectives x fused/concurrent). Co-authored-by: Cursor <cursoragent@cursor.com>
…e primitives Move the per-tile @triton.jit device functions (_grouped_coords, _gemm_tile, _all_gather_tile, _all_reduce_tile, _reduce_scatter_tile, _all_to_all_tile, _broadcast_tile, _comm_tile_ext) into a new iris/concurrent/_tiles.py; _kernels.py now holds only the 7 launched work-stealing kernels and imports the tiles it calls. Pure code organization -- no behavior change (emitted kernels are identical). Validated: imports resolve across modules and 22/22 tests/ccl/test_concurrent_fused (all 5 collectives x fused/concurrent). Co-authored-by: Cursor <cursoragent@cursor.com>
ryanswann-amd
requested review from
BKP,
artulab,
mawad-amd and
neoblizz
as code owners
July 23, 2026 20:54
Contributor
There was a problem hiding this comment.
Pull request overview
Note
Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.
Adds iris.concurrent support for overlapping an independent GEMM with GPU-resident collective communication while explicitly splitting GPU workgroups between the two workloads, including optional top‑k autotuning and an analytical CU-allocation predictor.
Changes:
- Introduces
iris.concurrent.gemmAPIs forall_gather,all_reduce,reduce_scatter,all_to_all, andbroadcastwithmode="fused"andmode="concurrent". - Adds predictor + autotuner modules (with a committed per-arch candidate set) and supporting tuning-data tooling.
- Adds unit tests for predictor/autotune logic and distributed correctness tests for fused/concurrent execution.
Reviewed changes
Copilot reviewed 15 out of 16 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/unittests/test_concurrent_predictor.py | Unit tests for predictor cost/scheduler logic, with origami-dependent tests skipped when unavailable |
| tests/unittests/test_concurrent_autotune.py | Unit tests for autotuner cache path, key stability, JSON DB roundtrip, and cache-hit behavior |
| tests/ccl/test_concurrent_fused.py | Distributed correctness tests validating GEMM+collective outputs in fused and concurrent modes |
| iris/concurrent/tuning_data/derive_candidates.py | Tool to derive a greedy top‑k candidate ordering from a measurement corpus |
| iris/concurrent/tuning_data/collect_corpus.py | Tool to collect on-device tuning corpus (grid sweep) for candidate derivation |
| iris/concurrent/tuning_data/candidate_set.json | Committed per-arch derived candidate ordering and coverage curves for the tuner |
| iris/concurrent/tuning_data/init.py | Loader for the committed candidate set JSON |
| iris/concurrent/tuning_data/README.md | Documentation for tuning corpus schema, regeneration, and coverage status |
| iris/concurrent/predictor.py | Analytical predictor and work-stealing makespan simulator for split/tile ranking |
| iris/concurrent/gemm.py | Public concurrent GEMM+collective APIs with fused and concurrent kernel launch paths + tuning hook |
| iris/concurrent/autotune.py | Top‑k autotuning logic with JSON cache and predictor/heuristic candidate ranking |
| iris/concurrent/_tiles.py | Triton per-tile device primitives for GEMM and collectives (shared by fused and concurrent kernels) |
| iris/concurrent/_kernels.py | Triton persistent work-stealing kernels (fused dual-queue and two-kernel variants) |
| iris/concurrent/init.py | Package init exposing autotune and gemm namespaces |
| iris/init.py | Exposes iris.concurrent at the top-level package namespace |
| .gitignore | Ensures the committed candidate_set.json is not accidentally ignored |
Comment on lines
+154
to
+171
| def _comm_sat_channels(collective, world, base=COMM_SAT_CHANNELS): | ||
| """CUs comm needs to SATURATE xGMI, scaled by per-tile transfer count (the | ||
| physics, not a per-collective fit): a tile that pushes/pulls more remote | ||
| blocks needs more CUs to hit link bandwidth. Relative to all_gather's N-1 | ||
| transfers/tile: all_reduce = 2(N-1) (read-all + write-all) -> 2x; all_to_all | ||
| = 1 (one peer) -> ~1/(N-1)x; reduce_scatter/broadcast = N-1 -> 1x. This | ||
| penalizes starving a heavy-comm collective (AR) of CUs -- the fix for the | ||
| measured gw=272/comm_wgs=32 cliff the model otherwise walks into -- WITHOUT | ||
| over-flooring the light A2A (which a blanket cap did).""" | ||
| # Only all_reduce needs the floor: its comm is FEW tiles (full/W) each HEAVY | ||
| # (2(N-1) read+write transfers), so starving it to ~32 CUs is catastrophic | ||
| # (measured 8-13x cliff at gw=272). AG/RS/BC/A2A have many light comm tiles, | ||
| # so a floor there only mis-shifts their (already good) picks -- an A/B on the | ||
| # fine grid showed a blanket floor drops AG exact 78->65%, BC 75->63%. So | ||
| # floor AR at 2x base; leave the rest unfloored. | ||
| if collective == "all_reduce": | ||
| return max(1, min(int(round(base * 2.0)), 304)) | ||
| return 1 |
Comment on lines
+375
to
+388
| def predict_split( | ||
| shape, | ||
| num_wgs=304, | ||
| candidates=None, | ||
| mt=(256, 256, 64), | ||
| cb=(256, 64), | ||
| c_atomic_gemm=C_ATOMIC_GEMM_MS, | ||
| c_atomic_comm=C_ATOMIC_COMM_MS, | ||
| oneshot=True, | ||
| comm_sat=COMM_SAT_CHANNELS, | ||
| mt_candidates=None, | ||
| mt_prefer=(256, 256, 64), | ||
| mt_margin=0.20, | ||
| ): |
Comment on lines
+45
to
+66
| def _load_comm_model(): | ||
| global _comm | ||
| if _comm is not None: | ||
| return _comm | ||
| cand = [] | ||
| if os.environ.get("ORIGAMI_COMM_MODEL"): | ||
| cand.append(os.environ["ORIGAMI_COMM_MODEL"]) | ||
| cand += [p for p in sys.path] | ||
| for base in cand: | ||
| try: | ||
| if base and base not in sys.path: | ||
| sys.path.insert(0, base) | ||
| from model.collective import predict_row # noqa: F401 | ||
| from model.hardware import MI300X, MI300X_COMM # noqa: F401 | ||
| import model.collective as mc | ||
| import model.hardware as mh | ||
|
|
||
| _comm = (mc, mh) | ||
| return _comm | ||
| except Exception: | ||
| continue | ||
| raise ImportError("comm model not found; set ORIGAMI_COMM_MODEL to the origami_comms repo root") |
Comment on lines
+90
to
+99
| def _save_db(db): | ||
| path = cache_path() | ||
| try: | ||
| os.makedirs(os.path.dirname(path), exist_ok=True) | ||
| tmp = f"{path}.{os.getpid()}.tmp" | ||
| with open(tmp, "w") as f: | ||
| json.dump(db, f, indent=2, sort_keys=True) | ||
| os.replace(tmp, path) | ||
| except OSError: | ||
| pass |
Comment on lines
+57
to
+63
| def _arch(device): | ||
| """Short device-arch string for the tuning-db cache key.""" | ||
| try: | ||
| props = torch.cuda.get_device_properties(device) | ||
| return getattr(props, "gcnArchName", None) or props.name | ||
| except Exception: | ||
| return "unknown" |
Make predictor tuning parameters and hardware limits effective, harden cache persistence across processes, and stabilize architecture cache keys. Co-authored-by: Cursor <cursoragent@cursor.com>
Collaborator
|
Do the fused variants added here replace the ones inside https://github.com/ROCm/iris/tree/main/iris/ops ? |
Collaborator
Author
No, the two ops here are specifically independent, so there's no dependencies between them. iris/ops is ops that have dependencies between each other. If I drew the DAG for ops it's for concurrent it's Long running comm overlapped with long running GEMM is (I believe) something that happens in training |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
Add support for overlapping an independent GEMM with collective communication while explicitly dividing GPU resources between the workloads.
Technical Details
Introduces
iris.concurrent.gemmwithall_gather,all_reduce,reduce_scatter,all_to_all, andbroadcast.fused: one persistent kernel with two work-stealing queues.concurrent: separate persistent kernels on different streams.See the full comparison.
Test Plan
Run predictor/autotuner unit tests and distributed correctness tests for all collectives in both execution modes.
Test Result
All unit tests and 22 distributed tests passed. GEMM and collective outputs matched their PyTorch references.
Submission Checklist