feat(collectives): ring allreduce TPUT push model on HOST builtin + InCore composite (requires PTOAS >= 0.55) - #2280
Open
georgebisbas wants to merge 5 commits into
Open
Conversation
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughChangesRing allreduce TPUT conversion
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Rank
participant NeighborBarrier
participant RightNeighbor
Rank->>Rank: Read own receive slot
Rank->>NeighborBarrier: Wait for ready generation
Rank->>RightNeighbor: TPUT partial chunk
Rank->>NeighborBarrier: Wait for push-done generation
Rank->>Rank: Reduce and store partial
Rank->>RightNeighbor: TPUT finalized chunk
Rank->>NeighborBarrier: Wait for allgather completion
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (2)
src/ir/transforms/lower_composite_ops_pass.cpp (1)
1533-1537: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winCorrect the allgather ready-barrier rationale.
The comment states the ready barrier is safe because "the cumulative generation count includes the previous step's push-done notifies". Each round uses its own signal row (
ag_round), so rowk's counters start at zero and carry no credit from rowk-1. The actual guarantee comes from the previous round's push-done barrier, which every rank must pass before it enters roundk. Restate the rationale so a future change does not remove the previous round's push-done barrier on a false premise. The same wording appears at Lines 1571-1573.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/ir/transforms/lower_composite_ops_pass.cpp` around lines 1533 - 1537, Update the ready-barrier rationale comment around lines 1533-1537 to remove the incorrect explanation about cumulative generation counts including previous push-done notifies. Replace it with the correct rationale: the guarantee comes from the previous round's push-done barrier that every rank must pass before entering round k. Apply the same correction to the identical wording at lines 1571-1573 to ensure consistency and prevent future maintainers from removing the previous round's push-done barrier based on the false premise.python/pypto/runtime/builtins/collectives/allreduce_ring/templates/kernel.cpp.in (1)
121-144: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low valueAdd a comment documenting the two-rank
NeighborBarrierbehavior.For
nranks == 2, both neighbours map to rank 1, so each row cell receives twoAtomicAddwrites and is observed twice via the sameGE 1waits. A short comment keeps the two-rank safety invariant explicit; the existing row allocation path already uses a fresh zeroed signal buffer per allreduce call.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/pypto/runtime/builtins/collectives/allreduce_ring/templates/kernel.cpp.in` around lines 121 - 144, Add a concise comment above NeighborBarrier documenting the nranks == 2 case: both neighbours resolve to rank 1, each signal cell receives two AtomicAdd notifications, and the corresponding GE 1 waits observe the shared cell twice. Keep the existing implementation unchanged and note that the signal buffer is freshly zeroed per allreduce call.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In
`@python/pypto/runtime/builtins/collectives/allreduce_ring/templates/kernel.cpp.in`:
- Around line 146-149: Update the stale StepBarrier comment near the barrier
configuration to reflect that kUseNeighborBarrier is already enabled and
NeighborBarrier has been verified on 910B2 for P=2 and P=4. Remove the outdated
instructions to land with RoundBarrier first and defer the swap; keep the
comment focused on the current implementation and its verification status.
In `@src/ir/transforms/lower_composite_ops_pass.cpp`:
- Around line 1439-1454: The tile.set_validshape operation returns a new
narrowed TileType and does not mutate put_stage in-place. Update the
pld.tile.put call to pass the narrowed alias rs_stage_valid instead of
put_stage, so the push uses the correct narrowed staging tile extent. Apply the
same fix to the allgather push path by using ag_stage_valid instead of put_stage
in its corresponding pld.tile.put call.
In `@tests/ut/ir/transforms/test_lower_composite_ops.py`:
- Around line 2102-2109: Add a non-empty assertion for the stage_creates
collection in the FP16 test, matching the sibling test’s assertion and message.
Keep the existing chunk_shapes loop and validation unchanged so the test
requires FP16 staging-tile creation before checking alignment and size
constraints.
---
Nitpick comments:
In
`@python/pypto/runtime/builtins/collectives/allreduce_ring/templates/kernel.cpp.in`:
- Around line 121-144: Add a concise comment above NeighborBarrier documenting
the nranks == 2 case: both neighbours resolve to rank 1, each signal cell
receives two AtomicAdd notifications, and the corresponding GE 1 waits observe
the shared cell twice. Keep the existing implementation unchanged and note that
the signal buffer is freshly zeroed per allreduce call.
In `@src/ir/transforms/lower_composite_ops_pass.cpp`:
- Around line 1533-1537: Update the ready-barrier rationale comment around lines
1533-1537 to remove the incorrect explanation about cumulative generation counts
including previous push-done notifies. Replace it with the correct rationale:
the guarantee comes from the previous round's push-done barrier that every rank
must pass before entering round k. Apply the same correction to the identical
wording at lines 1571-1573 to ensure consistency and prevent future maintainers
from removing the previous round's push-done barrier based on the false premise.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 73efeda0-1ee1-418f-ae6b-d8ff15d8b379
📒 Files selected for processing (7)
docs/en/dev/passes/12-lower_composite_ops.mddocs/en/dev/passes/40-lower_host_tensor_collectives.mddocs/zh/dev/passes/12-lower_composite_ops.mddocs/zh/dev/passes/40-lower_host_tensor_collectives.mdpython/pypto/runtime/builtins/collectives/allreduce_ring/templates/kernel.cpp.insrc/ir/transforms/lower_composite_ops_pass.cpptests/ut/ir/transforms/test_lower_composite_ops.py
HOST builtin (builtin.tensor.allreduce_ring): - Convert reduce-scatter + allgather from TLOAD/TSTORE pull to pto::comm::TPUT push: TPUT<AtomicAdd> remote-accumulate for RS, non-atomic TPUT for AG. - Ordering: pipe_barrier(PIPE_ALL) around every TPUT + dsb(DSB_DDR) before TNOTIFY (mirrors the allgather/all_to_all host builtins; NOT a GM fence). - Enable O(1) NeighborBarrier (notify/wait left+right only) — NPU-safe because the TPUT write pipeline orders data ahead of the signal, which the pull model lacked. NPU-verified at P=2/4 (both RoundBarrier baseline and NeighborBarrier). InCore composite (LowerTensorRingAllReduceRule): - Replace pld.tile.remote_load pulls with pld.tile.put pushes (non-atomic TPUT + local reduce, preserving Sum/Max/Min/Prod). - Race-free protocol: own-value read before the ready barrier, push after it, push-done barrier, then local read+reduce+store; credits stay 2 per subchunk. - Requires PTOAS >= v0.55 (dynamic partition-view shapes for tput, issue hw-native-sys#1069). - Ragged/arbitrary lengths and FP16 preserved via balanced segments + valid shapes; stage tile narrowed with tile.set_validshape per transfer. Docs: ring builtin lowering notes (en/zh) updated to describe the push model. NOTE: ring composite UT tests (test_lower_composite_ops.py ring section) still pin the old pull IR and will be updated in a follow-up commit.
The ring allreduce composite now lowers to pld.tile.put pushes (TPUT) with a tile.create staging tile instead of pld.tile.remote_load pulls. Update the ring UT assertions to pin the push structure: - required ops: pld.tile.remote_load -> pld.tile.put - staging tile: tile.create shape [1, chunk_cols] is the UB-bounded aligned chunk (stage is created once per ring call) - puts carry 7 args (dst, peer, src, stage, dst_offsets, src_offsets, shape) - control flow: if_count 13 -> 14 (the extra If guards each push phase) - FP16: aligned 16-element stage instead of allow_physical_tail_padding - round-trip: pld.tile.put print survives parse All 130 composite UTs pass.
Update the composite lowering pass docs (en/zh): the ring allreduce rule now moves data with pld.tile.put (TPUT push) instead of pld.tile.remote_load pulls — own-value read before the ready barrier, push to the right neighbour, push-done barrier, then local read+reduce+store. The VEC staging tile is narrowed with tile.set_validshape per transfer; PTOAS >= v0.55 dynamic partition-view shapes (hw-native-sys/PTOAS#1069) keep ragged/FP16 tails exact. Non-atomic push + local reduce preserves Sum/Max/Min/Prod.
Pre-commit formatting for the TPUT push-model ring changes (matches the upstream pre-commit CI hooks).
CodeRabbit review: tile.set_validshape returns the narrowed alias (memory inherited from the input) — bind it and pass it as the pld.tile.put stage in both the reduce-scatter and allgather pushes so the TPUT single-shot path reads exactly the transfer's valid width by construction (matches the mesh store pattern). Also address review nits: - kernel.cpp.in: update the stale StepBarrier comment (NeighborBarrier is the active, silicon-verified default). - test: assert stage_creates is non-empty in the FP16 aligned-schedule test.
georgebisbas
force-pushed
the
ring-allreduce-tput-push
branch
from
August 5, 2026 06:50
5f1e8c2 to
83c4b26
Compare
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
1 participant
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.
Summary
Replaces the pull-model engine of the ring allreduce on both rails with a TPUT push model (remote write), enabling O(1)
NeighborBarrieron the HOST builtin and eliminating the pull-model NPU memory-ordering gap.builtin.tensor.allreduce_ring): reduce-scatter + allgather converted fromTLOAD/TSTOREpull topto::comm::TPUTpush —TPUT<AtomicAdd>remote-accumulate for RS, non-atomicTPUTfor AG. Ordering ispipe_barrier(PIPE_ALL)around every transfer +dsb(DSB_DDR)beforeTNOTIFY(mirrors the in-tree allgather/all_to_all host builtins; not a GM fence). The O(P²)RoundBarrieris replaced by the O(1)NeighborBarrier(notify/wait the two ring neighbours only), which is NPU-safe because the TPUT write pipeline orders the data ahead of the signal — the pull model could not provide that.LowerTensorRingAllReduceRule): replacespld.tile.remote_loadpulls withpld.tile.putpushes (non-atomic TPUT + local reduce, preserving Sum/Max/Min/Prod). Race-free per-subchunk protocol: own-value read → ready barrier → push to right neighbour → push-done barrier → local read+reduce+store; barrier credits stay 2 per subchunk (signal shape unchanged). Ragged/arbitrary lengths and FP16 are preserved via balanced segments + valid shapes, with the shared VEC staging tile narrowed per transfer viatile.set_validshape.This PR depends on PTOAS v0.55 (release: hw-native-sys/PTOAS#1069, fixed in PR #1079).
The InCore composite's
pld.tile.puttransfers carry the exact raggedvalid_colsas the partition-view extent. PTOAS ≤ v0.54 rejects dynamic partition-view shapes forpto.comm.tput('pto.comm.tput' op expects dst to have a positive static shape), so the pure push model cannot compile below v0.55. The HOST builtin does not depend on this (its kernel is hand-written), but the composite rail does.CI / developers must use PTOAS v0.55+ (install via
ptoas-0.55wheel orptoas-bin-aarch64.tar.gz, setPTOAS_ROOTaccordingly). The UT tests now pin the push structure (pld.tile.put+tile.createstaging tile instead ofpld.tile.remote_load).Verification (NPU silicon, 910B2, PTOAS v0.55)
All on real NPUs (8x 910B2), P=2 and P=4:
tests/st/distributed/test_l3_host_tensor_allreduce_ring.py— HOST ring, P=2/4 ✅ (withNeighborBarrierenabled)tests/st/distributed/collectives/test_l3_tensor_allreduce_ring_intrinsic.py— InCore ring, P=2/4, sizes {1, 17, 4097, 65537} (ragged + >UB), Sum/Max/Min/Prod, FP16 ✅tests/st/distributed/collectives/test_l3_allreduce_ring.py+test_l3_ring_sizing_prewarm.py— no regression ✅test_lower_host_tensor_collectives.py,test_host_orch_distributed.py,test_lower_composite_ops.py(+ numerical) all green ✅Total: 25/25 ST + 190 UT passed.
Trade-off note (ReduceOp)
The HOST builtin is
ReduceOp::kSumonly by construction, so itsTPUT<AtomicAdd>RS is fine. The composite keeps non-atomic push + local reduce to preserve Sum/Max/Min/Prod; only a remote-atomicTPUT<AtomicAdd>variant would be Sum-only (AtomicTypehas noAtomicMax/Min).