Skip to content

Triggered plans: one counter per chunk, per-peer completion counters - #616

Open
ariscript wants to merge 3 commits into
mainfrom
arprakas/plan-gate-counters
Open

ariscript wants to merge 3 commits into
mainfrom
arprakas/plan-gate-counters

Conversation

@ariscript

@ariscript ariscript commented Oct 29, 2026 •

Copy link
Copy Markdown
Member

First of three stacked follow-ups to #573, in the order agreed on that PR: gate layout (this PR),
then rounds (#617), then destinations and arrival groups (#618). They address @dsidler's review points
1 and 4.

What changes

  • Gates are one int32 symmetric tensor: a counter and an arrival counter per chunk, then a completion
    counter per peer.
  • The copy engine polls a chunk's counter against writers × round, so mark_ready is a single
    release atomic. There is no last-writer election and no separate ready word.
  • plan.counters is contiguous int32, so a kernel that signals per tile can increment it directly
    instead of calling TriggeredView.mark_ready. That includes a tritonBLAS EpilogueConfig(counter_buffer=...)
    (Add SignalView and WaitView tritonBLAS#81).
  • Every round, each peer queue's chains end with one copy-engine add on the sending rank's completion
    counter for that peer:
    • plan.wait() now reads only local memory; the progress kernel's remote reads are gone.
    • Kernels can wait inline with the new TriggeredView.sent() / wait_sent(), for example before
      overwriting their source.
  • Limit: writers × round must stay at or below 2³¹ − 1 (over 7M rounds at 304 writers).

Testing (MI300X, torch allocator)

  • tests/unittests/test_plan.py at 2 and 4 ranks. New tests:
    • a kernel that increments plan.counters with a plain gpu-scope atomic_add after
      debug_barrier, as the tritonBLAS epilogue does: 8 writers on one 4 MiB chunk, 10 rounds,
      no stale data;
    • wait_sent before reusing the source;
    • the round limit.
  • View tests (Triton and Gluon) on one GPU.
  • Docs build.

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

4 open issues
What changed in this PR

Updates triggered-plan gate layout and device/host synchronization to use contiguous int32 counters (per-chunk marks + arrivals, and per-peer completion), enabling single-atomic mark_ready and local-only wait() progress checks.

Changes:

  • Replace per-chunk “ready word + counter” gate rows with int32 counters/arrivals and per-peer completion counters (ABI v2).
  • Add device APIs TriggeredView.sent() / wait_sent() and host-side completion signaling per peer queue.
  • Update unit tests and docs to reflect the new gate tensor layout and semantics.
File Description
iris/​experimental/​plan/​triton.py Switches device ABI/layout to int32 gates and adds sent()/wait_sent() logic.
iris/​experimental/​plan/​gluon.py Mirrors Triton TriggeredView changes for Gluon device code.
iris/​experimental/​plan/​triggered.py Allocates/slices new gate tensor, adds round limit check, and switches wait logic to local-only gates.
iris/​experimental/​plan/​kernels.py Replaces progress kernel with local-only wait kernel; updates release semantics.
iris/​experimental/​plan/​armer.py Arms chains polling counter thresholds and signals int32 arrivals + per-peer completion.
tests/​unittests/​test_plan.py Adds coverage for external counter increments, wait_sent, and round limit; updates arrivals assertions.
tests/​unittests/​test_plan_view.py Updates view tests for new counter/arrival layout and adds sent()/wait_sent() coverage.
tests/​unittests/​test_plan_view_gluon.py Updates Gluon view tests similarly and adds sent()/wait_sent() coverage.
docs/​conceptual/​triggered-plans.md Updates conceptual model and memory-ordering text for counter-based gating and completion counters.
docs/​reference/​host/​plan.md Exposes new device methods sent/wait_sent in API docs.
iris/​experimental/​plan/​__init__.py Updates module docstring to mention new device APIs.

Review effort: Lite

Comment thread iris/experimental/plan/gluon.py Resolved
Comment thread iris/experimental/plan/triton.py Resolved
Comment thread tests/unittests/test_plan.py Resolved
Comment thread iris/experimental/plan/triggered.py Resolved
ariscript and others added 3 commits October 30, 2026 16:42
…iggered plans

Gates become one int32 symmetric tensor: a counter and an arrival counter per chunk, then a
completion counter per peer.

- Chains poll the chunk's counter against writers x round, so mark_ready is a single release
  atomic: no last-writer election and no ready word. plan.counters is contiguous int32, so a kernel
  that signals per tile (a tritonBLAS epilogue) can increment it directly.
- Every round, a rank's chains to each peer end with an ATOMIC ADD32 on the sending rank's
  completion counter for that peer. wait() then reads only local words; the progress kernel's
  remote reads and armed table are gone. TriggeredView gains sent() / wait_sent().
- Arrival signals are 32-bit (ADD32, 0x0f). release_kernel raises the counter to writers x round
  with atomic_max, so close() opens partially marked gates.
- Round limit: writers x round <= 2^31 - 1.

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

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
The device view keeps the peers a rank sends to as bits of one int64 word; at 64 ranks building it
failed with an int64 overflow from torch.

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

Assigning a Python int to an element of the CUDA handle tensor copies from the host and blocks
until the stream drains, so start() (and close()'s release) waited for all queued GPU work.
fill_ on a one-element view is queued like any kernel.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
@ariscript
ariscript force-pushed the arprakas/plan-gate-counters branch from 53d8704 to 8accc21 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