Skip to content

Triggered plans: other destinations, some peers, arrival groups - #618

Open
ariscript wants to merge 3 commits into
arprakas/plan-roundsfrom
arprakas/plan-destinations
Open

ariscript wants to merge 3 commits into
arprakas/plan-roundsfrom
arprakas/plan-destinations

Conversation

@ariscript

@ariscript ariscript commented Oct 29, 2026 •

Copy link
Copy Markdown
Member

Stacked on #617 (review only the last commit). With these, plans can express reduce-scatter inboxes, all-gather into a staging buffer (as
AG-GEMM needs), and two-shot all-reduce.

What changes

  • TriggeredPlan(destination=...) lands peers' chunks in another symmetric tensor.
    Region(dst_offset_bytes=, dst_stride_bytes=) places each region there at its own offset and row
    pitch. The schedule builders take dst= / dst_base=.
  • Chunk(peers=...) sends a chunk to only some ranks. Other ranks neither receive it nor wait for it.
  • Chunk(group=...): a group's chunks signal one arrival counter at each receiver.
    TriggeredView.wait_group() / group_arrived() wait for all of them at once.
  • The overlap check runs per receiver, on destination bytes. When the destination is the buffer
    itself, what lands on a rank must also miss the bytes that rank sends; that is the previous rule,
    generalized.
  • Each chunk record names the gate word that shows it on this rank and how much that word grows per
    round. arrived() and plan.wait() therefore treat own, grouped, ungrouped and not-addressed
    chunks alike.
  • Not supported: sending a chunk to its own rank.

Testing (MI300X, torch allocator)

  • test_plan.py at 2 and 4 ranks. New tests:
    • tiles all-gathered into a staging buffer with a wider row pitch;
    • a reduce-scatter-shaped exchange into per-source inbox slots, with one wait_group in the
      consumer kernel;
    • a chunk sent to one peer only.
  • Schedule unit tests: destination offsets in the builders, validation, the per-receiver overlap check.
  • View tests: recipients and groups.
  • Docs build.

@github-actions github-actions Bot added in-progress We are working on it iris Iris project issue labels Oct 29, 2026
@ariscript
ariscript added this pull request to stack #619 October 29, 2026 19:15
@ariscript
ariscript force-pushed the arprakas/plan-destinations branch from 2ce10da to c1bd950 Compare October 29, 2026 22:26
Copilot AI lite review requested due to automatic review settings October 29, 2026 22:26

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Warning

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.

Copilot review overview

5 open issues
What changed in this PR

This PR extends triggered plans to support non-default destination placement, selective peer delivery, and grouped arrivals to enable patterns like reduce-scatter inboxes, all-gather into staging buffers, and two-shot all-reduce.

Changes:

  • Add destination placement (TriggeredPlan(destination=...), Region(dst_offset_bytes=..., dst_stride_bytes=...)) and propagate it through schedule builders and SDMA arming.
  • Add selective delivery (Chunk(peers=...)) and grouped arrivals (Chunk(group=...), TriggeredView.group_arrived() / wait_group()).
  • Update disjointness validation to check destination bytes per receiver (and protect own sent bytes when destination is the buffer).
File Description
tests/​unittests/​test_plan_view_gluon.py Updates view gate partitioning and adds tests for recipients/groups in Gluon.
tests/​unittests/​test_plan_view.py Updates view gate partitioning and adds tests for recipients/groups in Triton view.
tests/​unittests/​test_plan_schedule.py Adds unit tests for destination defaults, builders landing in dst, dst pitch, peer filtering, and per-receiver disjointness checks.
tests/​unittests/​test_plan.py Adds integration tests for tiling into staging w/ pitch, reduce-scatter inbox + wait_group, and chunk-to-some-peers.
iris/​experimental/​plan/​triton.py Bumps ABI, expands view record layout, and implements group arrival/wait logic.
iris/​experimental/​plan/​triggered.py Adds destination support, per-peer queue sizing, group gate allocation, updated view encoding, and per-receiver overlap validation.
iris/​experimental/​plan/​schedule.py Extends Region with destination placement, Chunk with peers/groups, updates builders, validate(), and check_disjoint() logic.
iris/​experimental/​plan/​kernels.py Updates wait_kernel to use per-chunk arrival word + expected-per-round derived from records.
iris/​experimental/​plan/​gluon.py Mirrors Triton TriggeredView changes (arrival word/expected + group wait APIs).
iris/​experimental/​plan/​armer.py Adds dst_pitch to Copy, updates per-peer queue sizing for selective peers, and uses destination pitch for tiled copies.
docs/​reference/​host/​plan.md Exposes group_arrived/wait_group in API reference.
docs/​conceptual/​triggered-plans.md Documents destinations/peers/groups and updates constraints/limits accordingly.

Review effort: Lite

Comment thread iris/experimental/plan/schedule.py Resolved
Comment thread iris/experimental/plan/gluon.py Resolved
Comment thread iris/experimental/plan/triggered.py Resolved
Comment thread iris/experimental/plan/triton.py Resolved
Comment thread iris/experimental/plan/armer.py Resolved
@ariscript
ariscript force-pushed the arprakas/plan-destinations branch from ffb6dfc to e7f3d33 Compare October 30, 2026 21:22
ariscript and others added 3 commits October 30, 2026 17:21
…in arrival groups

- TriggeredPlan(destination=...) lands peers' chunks in another symmetric tensor, and
  Region(dst_offset_bytes=, dst_stride_bytes=) places each region at its own offset and row pitch
  there (staging buffers, reduce-scatter inboxes). The schedule builders take dst= and dst_base=.
- Chunk(peers=...) sends a chunk to only some ranks; ranks it does not go to neither receive nor
  wait for it, and completion counts only the peers a rank sends to.
- Chunk(group=...): chunks of a group signal one arrival counter at each receiver, and
  TriggeredView.wait_group() / group_arrived() wait for all of them at once.
- The overlap check is per receiver, on destination bytes; with the buffer as destination, what
  lands on a rank must also miss the bytes it sends (the previous rule, generalized).
- Each chunk record now names the gate word that shows it on this rank and how much it grows per
  round, so arrived() and the wait kernel treat own, grouped, ungrouped and not-addressed chunks
  alike.

Tests pass at 2 and 4 ranks on MI300X (torch allocator).

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
…work

- The schedule builders now require dst_base whenever dst is given: without it, a slice of the
  destination silently landed at offset 0.
- arrived()'s docstring says groups apply to other ranks' chunks; a rank's own chunk always reads its
  counter.
- The round-limit error names group counts too, not only writers.
- Each plan computes its per-peer queue bytes once instead of every round.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
prepare_all_gather (#609) checks a caller's schedule with the plan's _rows(schedule); keep that helper
and give the per-region version its own name. An all-gather sends every chunk in place to every rank,
so it now rejects chunks with peers or destination offsets.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
@ariscript
ariscript force-pushed the arprakas/plan-destinations branch from e7f3d33 to 883de57 Compare October 30, 2026 22:43

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.

2 participants