update(distributed): convert all_to_all distributed to push-based lowering#1267
update(distributed): convert all_to_all distributed to push-based lowering#1267georgebisbas wants to merge 3 commits into
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 Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughThe all-to-all distributed example is refactored from a symmetric 3-phase stage-in/exchange pattern to a symmetric 2-phase push-based pattern. The kernel now performs direct TPUT pushes into peer scratch, a barrier via signal notify/wait, and local copy-out. Documentation and comments across README files, orchestration shim, and main.py are updated accordingly. Changes2-Phase Push Refactor
Estimated code review effort: 2 (Simple) | ~12 minutes 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 |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the distributed all-to-all example to use a more efficient 2-phase push-based algorithm (push, barrier, and local copy-out) instead of the previous 3-phase pattern. Feedback on these changes includes fixing a markdown formatting typo in the README and optimizing the kernel loop by hoisting a loop-invariant pointer calculation outside of the loop.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
examples/workers/l3/all_to_all_distributed/README.md (1)
39-43: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winMalformed code fence and stale kernel description.
Line 39 concatenates the fence language tag with content:
textpush, barrier, copy-out instead of ````text ```` on its own line — this breaks the code fence/language identifier. It looks likepush, barrier, copy-outwas intended to replace the stale `(stage-in, barrier, exchange)` description of the kernel on Line 43, which still describes the old 3-phase design.📝 Proposed fix
-```textpush, barrier, copy-out +```text all_to_all_distributed/ ├── kernels/ │ ├── aiv/ -│ │ └── all_to_all_kernel.cpp # AIV kernel (stage-in, barrier, exchange) +│ │ └── all_to_all_kernel.cpp # AIV kernel (push, barrier, copy-out)🤖 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 `@examples/workers/l3/all_to_all_distributed/README.md` around lines 39 - 43, Fix the malformed fenced block in the all_to_all_distributed README and update the stale kernel phase description. In the README’s tree snippet, separate the code fence language tag from the content so the fence starts correctly as a standalone text block, and update the all_to_all_kernel.cpp comment to use the current “push, barrier, copy-out” wording instead of the old “stage-in, barrier, exchange” description. Use the README tree section and the all_to_all_kernel.cpp entry as the anchors when editing.
🤖 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.
Outside diff comments:
In `@examples/workers/l3/all_to_all_distributed/README.md`:
- Around line 39-43: Fix the malformed fenced block in the
all_to_all_distributed README and update the stale kernel phase description. In
the README’s tree snippet, separate the code fence language tag from the content
so the fence starts correctly as a standalone text block, and update the
all_to_all_kernel.cpp comment to use the current “push, barrier, copy-out”
wording instead of the old “stage-in, barrier, exchange” description. Use the
README tree section and the all_to_all_kernel.cpp entry as the anchors when
editing.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: e8bf3e22-0c52-406a-b6ea-86964d314f96
📒 Files selected for processing (5)
examples/workers/l3/README.mdexamples/workers/l3/all_to_all_distributed/README.mdexamples/workers/l3/all_to_all_distributed/kernels/aiv/all_to_all_kernel.cppexamples/workers/l3/all_to_all_distributed/kernels/orchestration/all_to_all_orch.cppexamples/workers/l3/all_to_all_distributed/main.py
georgebisbas
changed the title
- Fix markdown code block language specifier typo in README.md (textpush,barrier,copy-out -> text) - Hoist loop-invariant scratch_dst_local pointer outside the Phase 1 loop (gemini-code-assist review)
- Fix markdown code block language specifier typo in README.md (textpush,barrier,copy-out -> text) - Hoist loop-invariant scratch_dst_local pointer outside the Phase 1 loop (gemini-code-assist review)
georgebisbas
force-pushed
the
perf/all-to-all-push-based
branch
2 times, most recently
from
e3498bc to
4d82a92
Compare
|
@ChaoWao this is rebased and ready to review again, fails on onboard-a2a3 for unrelated reasons right now. |
Replace the 3-phase pull-based all_to_all kernel (stage-in → barrier → remote TLOAD) with a 2-phase push-based kernel (TPUT to peers → barrier → local copy-out). - Phase 1: TPUT each input chunk directly to the corresponding peer's scratch at offset my_rank*C via pto::comm::TPUT. Self-rank handled naturally — CommRemotePtr to self returns the local pointer. - Phase 2: unchanged notify/wait barrier. - Phase 3: purely local TLOAD from scratch → TSTORE to output (no CommRemotePtr needed here — scratch already holds the result). - Single pushTile reused for both push and copy-out phases (was two tiles: stageTile + recvTile). No post-exchange barrier needed (no WAR hazard — input read-only, scratch write-only from peers). - CPU_SIM fallback: manual element loop instead of TPUT. Scratch naming preserved for consistency with all other L3 examples (allreduce, allgather, broadcast, reduce_scatter). Sim-verified: P=2 and P=4 golden checks pass (0.000e+00 max diff).
Align naming to push → barrier → local copy-out, update the collectives index, and trim phase-banner comments after the ST relocation rebase.
ChaoWao
force-pushed
the
perf/all-to-all-push-based
branch
from
d7da9e0 to
346d2c9
Compare
1 participant
Summary
Replace the pull-based all_to_all kernel (stage-in → barrier → remote TLOAD)
with a push-based kernel (TPUT → barrier → local copy-out), mirroring the
approach in pypto#1937.
Collective demos were relocated to
tests/st/worker/collectives/in #1294;this PR applies the push rewrite in that ST home (not under
examples/workers/l3/all_to_all_distributed/).Benefits
TPUTinstead ofstaging locally for remote reads.
the result — only local
TLOAD/TSTORE.pushTile(wasstageTile+recvTile).distinct scratch slots (no WAR on scratch).
What changed
tests/st/worker/collectives/all_to_all/kernels/aiv/all_to_all_kernel.cpppto::comm::TPUT…/orchestration/all_to_all_orch.cpp…/all_to_all/README.mdtests/st/worker/collectives/README.mdHow it works
pto::comm::TPUTfollows the pattern inexamples/workers/l3/ep_dispatch_combine/kernels/aiv/combine.cpp.Verification
Dockerfile.simpler.sim.ubuntu22.04):TestAllToAllP2PASSTestAllToAllP4PASSReferences
ep_dispatch_combine/kernels/aiv/combine.cpp: existing TPUT usagetests/st/worker/collectives/