[ariscript]

A plan used every iteration of a fused op should not cost a host barrier and two allgathers per
round.

What changes

• prepare_next() prepares the next round on this rank while the current one runs. It needs a plan
  created with prepare_ahead=True, which holds queue space for two rounds instead of one. Its
  copies wait behind gates that no producer can open before every rank has passed next_round().

• next_round() is a round fence on the current stream, so the host never waits. The fence is a
  kernel with counters of the plan's own: it waits at most the plan's timeout, and a peer that never
  arrives is recorded as an error rather than hanging the stream. next_round() prepares the next
  round first if prepare_next() wasn't called.

• wait(blocking=False) queues the wait on the stream. Errors are recorded on the GPU (the first
  one wins) and raised once, by the next blocking wait() or by close():

  • a timeout, including too few marks (with the same TimeoutError as a blocking wait());
  • extra marks;
  • a fence that a peer never reached;
  • a peer that could not prepare the round.

  Once an error is recorded, later stream-ordered waits return at once.

• If preparing fails on one rank, that rank still takes part in the fence and marks an abort on its
  peers. Their wait() for that round then raises, even if the failed rank's close() has flushed
  old data to them; then the failed rank's next_round() raises.

• The queue ledger tracks each prepared round separately, and start() and close() check order per
  round. A blocking wait() releases its round, so host put() works again before next_round().

• Chains are built once and re-armed every round with a new threshold, instead of translating and
  wrapping every region again. Most of the host time saved comes from this.

• prepare_all_gather takes prepare_ahead, and examples/34_sdma_all_gather gains a
  stream-ordered loop (prepare_next(), wait(blocking=False), a consumer kernel, next_round()),
  checked once after the loop. It prints the host time per round.

Host time per round (MI300X, tile_chunks of 64×64 tiles, 304 tiles per chunk, 4096 tiles per

rank; the producer only opens the gates):

	#616	this PR, blocking wait()	prepare_next()	fully stream-ordered
2 ranks	15.3 ms	2.7 ms	2.3 ms	2.1 ms
4 ranks	44.4 ms	3.6 ms	2.5 ms	2.1 ms

Wall time per round over 100 rounds, re-measured after the review changes (the #616 column is from
before them).

Until #624 merges, CI may show false timeouts from the round fence or a stream-ordered wait
(read_realtime() can exit a clock-bounded loop early on CDNA3/CDNA4). CI tests the merge with
main, so rerunning it once #624 has merged picks up the fix.

Testing (MI300X, torch allocator)

• test_plan.py at 2 and 4 ranks. New tests:
  • prepare_next() before wait();
  • a loop with no host synchronization;
  • errors from a stream-ordered wait, raised once;
  • a round fence that a peer never reaches;
  • closing with the next round prepared;
  • two plans sharing queues per round;
  • prepare_ahead being opt-in;
  • host put() between wait() and next_round();
  • a prepare failure on one rank, after part of the round was armed or with only the completion
    signal failing.
• Ledger unit tests; copy-engine, tracing and SDMA all-gather tests; example 34 at 2 ranks.
• Docs build.

[Copilot]

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

2 open issues

• tl.constexpr is an annotation used in Triton function signatures, not a value constructor;… · New
• kernels.WRONG_MARKS is referenced with .value, but the PR introduces it in kernels.py as a… · New

What changed in this PR

This PR reduces per-round host overhead for TriggeredPlan by preparing the next round early and making round transitions/waits stream-ordered, while improving queue accounting to support multiple prepared rounds safely.

Changes:

• Add prepare_next() and stream-ordered wait(blocking=False)/next_round() behavior with deferred error reporting.
• Introduce per-plan queue “holds” and per-round queue ledger tracking to avoid repeated capacity checks and barriers.
• Rework copy-chain arming to build chains once and re-arm them each round with updated thresholds.

File	Description
tests/​unittests/​test_queue_ledger.py	Adds unit coverage for holds, per-round ordering, and release semantics in the queue ledger.
tests/​unittests/​test_plan.py	Adds new multi-round/stream-ordered tests for prepare_next(), deferred errors, and close behavior.
iris/​host/​queue_ledger.py	Implements held-capacity accounting and per-round ordering checks for prepared copies.
iris/​experimental/​plan/​triggered.py	Adds overlapping-round API and stream-ordered barriers/waits; defers errors to later blocking operations/close.
iris/​experimental/​plan/​kernels.py	Adds GPU-side error recording for stream-ordered waits and replaces budgeted polling with time-based polling.
iris/​experimental/​plan/​armer.py	Introduces Chain objects to reuse per-chain fixed arguments and re-arm each round with a new threshold.
docs/​reference/​host/​plan.md	Exposes prepare_next in docs for TriggeredPlan.
docs/​conceptual/​triggered-plans.md	Documents overlapping rounds and updated semantics for device-barrier round transitions and deferred errors.

_{Review effort: Lite}

[nirvedhmeshram]

I will review more thoroughly in a bit, just leaving one comment that I noticed while reviewing #616

[nirvedhmeshram]

Thanks, this is great progress towards using triggered plans performantly.

[nirvedhmeshram]

next_round() now fences with ctx.device_barrier(), and that kernel has no real timeout: the spin loop in _device_barrier_kernel only has a tl.device_assert on MAX_SPINS, which is compiled out unless debug is on, so it spins forever. Before this PR the fence was a host torch.distributed barrier, which at least has a process-group timeout.

Concrete case: rank 0 raises anywhere between wait() and next_round() (a TimeoutError from a blocking wait that only rank 0 hit near the deadline, or any exception in the user's consumer code). The with block sends rank 0 into close(). Rank 1 has already called next_round(), so its stream now holds a device barrier that rank 0 will never enter. Rank 1 then calls start() and wait(), and wait() immediately does self._error.tolist(), which syncs the stream and blocks forever behind the spinning barrier. Rank 0 sits in close() at ctx.barrier() waiting for rank 1. The job hangs with a kernel spinning on rank 1's GPU. It also leaves the pairwise device-barrier counters off by one, so every later ctx.device_barrier() in the process is out of step.

Could the plan use its own barrier kernel bounded by the plan timeout (read_realtime, same as wait_kernel) that records TIMED_OUT into self._error and exits, instead of ctx.device_barrier()? Then the next blocking wait() or close() would raise instead of hanging, and the shared barrier counters would not be touched.

[ariscript]

Agreed, fixed. next_round() no longer calls ctx.device_barrier(). The plan has its own fence
kernel with one counter per peer, kept in the plan's gates. Each rank adds 1 to its counter on every
peer, then waits up to the plan's timeout (read_realtime, as in wait_kernel) for every peer's
counter to reach the round. A peer that never arrives is recorded in the error word as
TimeoutError("round r: rank i did not reach next_round()"), which the next blocking wait() or
close() raises. Because the counters belong to the plan, a fence that times out leaves no shared
barrier out of step. A rank whose plan already holds an error skips the fence, so its peers record a
timeout instead of waiting forever. test_round_fence_times_out covers your case: one rank leaves
before next_round(), the other rank's next wait() raises, and ctx.device_barrier() still
works afterwards.

While testing this I found that read_realtime() could make these loops exit at once. On gfx9,
Triton's memrealtime() waits on vmcnt after s_memrealtime, which completes on lgkmcnt, so in the
fence kernel the next instruction read the clock register before the value had landed, and the fence
recorded a timeout in under a second. Triton fixed this in triton-lang/triton#12179, but no release
or Iris pin has it yet, so #624 works around it in Iris. Until it merges, this PR's CI can show
those false timeouts.

[nirvedhmeshram]

When this plan is queued behind another plan that has not started, close() returns here with the plan CLOSING and never looks at self._error. The plan is later finished by the other plan's close() through _drain_closing and _mark_closed, but that path only raises the closing plan's own deferred error, so a timeout or wrong mark count recorded by this plan's stream-ordered wait() is silently lost. Could _drain_closing collect the deferred errors of the plans it finishes and have close() raise them too (or at least warn), or could the early return check self._error once the queue is flushed?

[ariscript]

Fixed. When close() finishes plans that were waiting as CLOSING, it now reads each one's error
word and warns with the error. It can't raise them, since they belong to other plans, and the call
that closes one plan shouldn't fail for another's error.

[nirvedhmeshram]

In a loop that never synchronizes, once a chunk goes missing (a peer died, a dropped mark) every later round's wait_kernel also spins for the whole plan timeout, 30 s by default, and the host can't notice because it never syncs. 100 queued rounds means close to an hour of GPU time before close() reports the first error. Since the error word is sticky anyway, could the kernel load error[0] at entry and exit immediately (with status 0) when it is already non-zero? Then the stream reaches close() in roughly one timeout.

[ariscript]

Done. In stream-ordered mode the wait kernel checks the error word first, and if it is set, writes
status 0 and returns. Only the stream-ordered path does this: the flush in close() uses the same
kernel to wait for this rank's queues, and it has to wait whatever was recorded.

[nirvedhmeshram]

The stream-ordered loop (prepare_next, wait(blocking=False), next_round on the device) is the main use case for this PR, but outside this docs snippet it only appears in tests, and nothing in the stack up to #618 adds an example. Could examples/34_sdma_all_gather get a variant that uses it, with a consumer kernel reading out instead of a host check and the results checked once after the loop? That would also give the host-time numbers in the description something reproducible to point at. Fine as a follow-up once the stack lands if that is easier.

[ariscript]

Added it in this PR rather than as a follow-up. examples/34_sdma_all_gather has a third part that
runs the same plan with prepare_ahead=True: prepare_next() after the producer,
wait(blocking=False), a consumer kernel that adds each round's output into a float64 sum, then
next_round(). The sum is checked once after the loop, and rank 0 prints the host time per round
(excluding round 1, which compiles the kernels). On MI300X at 2 ranks with the defaults it printed
0.2 to 0.4 ms per round. The table in the description, from a longer benchmark, is updated too.