From 807003d7436c80cacf69264d0584aea8e98a0c38 Mon Sep 17 00:00:00 2001 From: georgebisbas Date: Fri, 3 Jul 2026 13:55:26 +0200 Subject: [PATCH 1/3] Update: convert all_to_all distributed to push-based lowering MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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). --- examples/workers/l3/README.md | 2 +- .../l3/all_to_all_distributed/README.md | 14 ++-- .../kernels/aiv/all_to_all_kernel.cpp | 69 +++++++++++-------- .../kernels/orchestration/all_to_all_orch.cpp | 4 +- .../workers/l3/all_to_all_distributed/main.py | 10 +-- 5 files changed, 58 insertions(+), 41 deletions(-) diff --git a/examples/workers/l3/README.md b/examples/workers/l3/README.md index a0281b0fc..3c6bd094a 100644 --- a/examples/workers/l3/README.md +++ b/examples/workers/l3/README.md @@ -66,7 +66,7 @@ Two things to know before reading the example: | [`allgather_distributed/`](allgather_distributed/) | One communication domain via `orch.allocate_domain`; each rank stages its slice, synchronizes across ranks, then gathers every rank's window data into a full output. | | [`reduce_scatter_distributed/`](reduce_scatter_distributed/) | One communication domain via `orch.allocate_domain`; each rank stages all input chunks, synchronizes, then reduces the per-rank chunk across peers into a rank-local output. | | [`broadcast_distributed/`](broadcast_distributed/) | One communication domain via `orch.allocate_domain`; root stages into the window, synchronizes, then every rank reads the root's scratch slot into its output. | -| [`all_to_all_distributed/`](all_to_all_distributed/) | One communication domain via `orch.allocate_domain`; scratch indexed by destination rank, barrier, then each rank gathers the slice peers sent to it. | +| [`all_to_all_distributed/`](all_to_all_distributed/) | One communication domain via `orch.allocate_domain`; each rank pushes its per-destination chunks directly to peers' scratch via TPUT, barrier synchronizes, then each rank copies its local result to output. | | [`ffn_tp_parallel/`](ffn_tp_parallel/) | Local compute followed by one-domain cross-rank reduction through a domain scratch window. | | [`ep_dispatch_combine/`](ep_dispatch_combine/) | MoE-style dispatch/combine over a one-domain communication window. | | [`domain_rank_map/`](domain_rank_map/) | Small two-domain example showing domain-local ranks, missing-domain `KeyError`, separate window slices, and real per-domain allreduce. | diff --git a/examples/workers/l3/all_to_all_distributed/README.md b/examples/workers/l3/all_to_all_distributed/README.md index 2d50a8d11..ee52da7ca 100644 --- a/examples/workers/l3/all_to_all_distributed/README.md +++ b/examples/workers/l3/all_to_all_distributed/README.md @@ -2,17 +2,21 @@ This example demonstrates distributed all-to-all (personalized exchange) — each rank sends a different slice to each peer. -## Algorithm (3-Phase Mesh) +## Algorithm (2-Phase Push) ```text -Phase 1: stage-in for d in 0..P-1: input[d * C] → scratch[d * C] (chunk d is for rank d) -Phase 2: barrier mesh barrier (all-to-all notify/wait) -Phase 3: exchange for s in 0..P-1: TLOAD(peer s scratch[my_rank * C]) → output[s * C] +Phase 1: push for d in 0..P-1: TPUT input[d * C] → peer d's scratch[my_rank * C] + Self-rank handled naturally — CommRemotePtr to self returns local pointer +Phase 2: barrier mesh barrier (all-to-all notify/wait); no post-exchange barrier needed +Phase 3: copy-out for s in 0..P-1: TLOAD(local scratch[s * C]) → output[s * C] (purely local) ``` **Input**: Each rank owns `nranks * COUNT_PER_RANK` floats (chunk d is payload for rank d). **Output**: Each rank receives `nranks * COUNT_PER_RANK` floats (chunk s is what rank s sent to me). +After Phase 1, scratch at offset `src * C` holds the chunk rank `src` pushed here. +After the barrier, the result is in scratch — Phase 3 is a purely local copy to the output tensor. + ## Usage ```bash @@ -32,7 +36,7 @@ python main.py -p a2a3 -d 0-3 ## File Structure -```text +```textpush, barrier, copy-out all_to_all_distributed/ ├── kernels/ │ ├── aiv/ diff --git a/examples/workers/l3/all_to_all_distributed/kernels/aiv/all_to_all_kernel.cpp b/examples/workers/l3/all_to_all_distributed/kernels/aiv/all_to_all_kernel.cpp index cdde406bd..7057ef1d5 100644 --- a/examples/workers/l3/all_to_all_distributed/kernels/aiv/all_to_all_kernel.cpp +++ b/examples/workers/l3/all_to_all_distributed/kernels/aiv/all_to_all_kernel.cpp @@ -9,13 +9,23 @@ * ----------------------------------------------------------------------------------------------------------- */ /** - * AllToAll kernel — symmetric, 3-phase, HCCL-window scratch pattern. + * AllToAll kernel — symmetric, 2-phase push-based pattern. * - * Phase 1 (stage-in): for dest in 0..nranks-1: input[dest*C..) → scratch[dest*C..) - * Phase 2 (barrier): signal matrix + TWAIT cross-rank sync - * Phase 3 (exchange): for src in 0..nranks-1: TLOAD(peer src scratch[my_rank*C]) → output[src*C..) + * Phase 1 (push): for dest in 0..nranks-1: + * TPUT input[dest*C..) → peer dest's scratch[my_rank*C..) + * Self-rank handled naturally: CommRemotePtr to self + * returns the local pointer, so TPUT degenerates to a + * local GM write. + * Phase 2 (barrier): signal matrix + TWAIT cross-rank sync. + * No post-exchange barrier needed — input is read-only + * and the scratch is write-only from peers, so there is + * no WAR hazard on the scratch. + * Phase 3 (copy-out): for src in 0..nranks-1: + * TLOAD(local scratch[src*C..)) → TSTORE(output[src*C..)) + * Purely local copy — the scratch already holds the + * result after the barrier. * - * Scratch is indexed by destination rank: scratch[dest*C] holds the chunk sent to dest. + * Scratch at offset src*C holds the chunk that rank src pushed here. * * args layout: * tensor(0) = input nranks*COUNT_PER_RANK floats (INPUT) @@ -74,35 +84,40 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in return; } - // signal_base follows the nranks * COUNT_PER_RANK float staging region. + // signal_base follows the nranks * COUNT_PER_RANK float data region. __gm__ int32_t *signal_base = reinterpret_cast<__gm__ int32_t *>(scratch + nranks * COUNT_PER_RANK); ShapeDyn shape(1, 1, 1, 1, COUNT_PER_RANK); StrideDyn stride(COUNT_PER_RANK, COUNT_PER_RANK, COUNT_PER_RANK, COUNT_PER_RANK, 1); - TileData stageTile(1, COUNT_PER_RANK); - TileData recvTile(1, COUNT_PER_RANK); - TASSIGN(stageTile, 0x0); - TASSIGN(recvTile, 0x10000); + // Single push tile — reused for both Phase 1 push and Phase 3 copy-out. + TileData pushTile(1, COUNT_PER_RANK); + TASSIGN(pushTile, 0x0); // ------------------------------------------------------------------ - // Phase 1: stage-in — copy each destination chunk into scratch so - // every peer can read the slice destined for them in Phase 3. + // Phase 1: push — TPUT each destination chunk directly into the + // corresponding peer's scratch at offset my_rank*C. When dest == my_rank + // CommRemotePtr returns the local pointer, so TPUT is a local GM write. // ------------------------------------------------------------------ for (int dest = 0; dest < nranks; ++dest) { + __gm__ float *scratch_dst_local = scratch + my_rank * COUNT_PER_RANK; + __gm__ float *scratch_dst_remote = CommRemotePtr(commCtx, scratch_dst_local, dest); + Global inputChunkG(input + dest * COUNT_PER_RANK, shape, stride); - Global scratchChunkG(scratch + dest * COUNT_PER_RANK, shape, stride); - TLOAD(stageTile, inputChunkG); - set_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0); - wait_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0); - TSTORE(scratchChunkG, stageTile); - set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); - wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + Global scratchChunkG(scratch_dst_remote, shape, stride); + +#if defined(__CPU_SIM) + for (size_t i = 0; i < COUNT_PER_RANK; ++i) { + scratch_dst_remote[i] = input[dest * COUNT_PER_RANK + i]; + } +#else + pto::comm::TPUT(scratchChunkG, inputChunkG, pushTile); +#endif } pipe_barrier(PIPE_ALL); // ------------------------------------------------------------------ - // Phase 2: device barrier — notify every peer that stage-in is done, + // Phase 2: device barrier — notify every peer that our pushes are done, // then wait until every peer has notified us. // ------------------------------------------------------------------ for (int peer = 0; peer < nranks; ++peer) { @@ -119,19 +134,17 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in pipe_barrier(PIPE_ALL); // ------------------------------------------------------------------ - // Phase 3: exchange — read chunk my_rank from every rank's scratch and - // write it into the corresponding slice of the output tensor. - // CommRemotePtr with pe==my_rank returns localPtr unchanged, so the - // self-read goes through the same code path as the remote reads. + // Phase 3: copy-out — every peer has pushed into our scratch at their + // rank's offset. Copy each slot into the output tensor. Purely local + // reads — no CommRemotePtr needed here. // ------------------------------------------------------------------ for (int src = 0; src < nranks; ++src) { - __gm__ float *remote_chunk = CommRemotePtr(commCtx, scratch + my_rank * COUNT_PER_RANK, src); - Global remoteG(remote_chunk, shape, stride); + Global scratchSlotG(scratch + src * COUNT_PER_RANK, shape, stride); Global outputSlotG(output + src * COUNT_PER_RANK, shape, stride); - TLOAD(recvTile, remoteG); + TLOAD(pushTile, scratchSlotG); set_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0); wait_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0); - TSTORE(outputSlotG, recvTile); + TSTORE(outputSlotG, pushTile); set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); } diff --git a/examples/workers/l3/all_to_all_distributed/kernels/orchestration/all_to_all_orch.cpp b/examples/workers/l3/all_to_all_distributed/kernels/orchestration/all_to_all_orch.cpp index dcccb3770..8782eb8e2 100644 --- a/examples/workers/l3/all_to_all_distributed/kernels/orchestration/all_to_all_orch.cpp +++ b/examples/workers/l3/all_to_all_distributed/kernels/orchestration/all_to_all_orch.cpp @@ -9,11 +9,11 @@ * ----------------------------------------------------------------------------------------------------------- */ /** - * AllToAll orchestration shim. + * AllToAll orchestration shim — push-based (2 kernel phases + local copy-out). * * tensor(0) input INPUT (nranks*COUNT_PER_RANK floats) * tensor(1) output OUTPUT_EXISTING (nranks*COUNT_PER_RANK floats) - * tensor(2) scratch INOUT (HCCL window slot) + * tensor(2) scratch INOUT (HCCL window slot — holds result after barrier) * scalar(0) nranks * scalar(1) CommContext device pointer */ diff --git a/examples/workers/l3/all_to_all_distributed/main.py b/examples/workers/l3/all_to_all_distributed/main.py index 9cc9cf99f..2837a7c12 100644 --- a/examples/workers/l3/all_to_all_distributed/main.py +++ b/examples/workers/l3/all_to_all_distributed/main.py @@ -7,14 +7,14 @@ # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. # See LICENSE in the root of the software repository for the full text of the License. # ----------------------------------------------------------------------------------------------------------- -"""End-to-end distributed all-to-all — symmetric 3-phase pattern. +"""End-to-end distributed all-to-all — symmetric 2-phase push-based pattern. Each rank owns nranks send chunks; chunk d is payload for rank d. After the exchange rank r holds chunk s from every source s in output[s*COUNT_PER_RANK]: - Phase 1 stage-in for d in 0..N-1: input[d*C..) → scratch[d*C..) + Phase 1 push for d in 0..N-1: TPUT input[d*C..) → peer d's scratch[my_rank*C..) Phase 2 device barrier signal matrix cross-rank sync via TNOTIFY/TWAIT - Phase 3 exchange for s in 0..N-1: TLOAD(peer s scratch[my_rank*C]) → output[s*C..) + Phase 3 copy-out for s in 0..N-1: TLOAD(local scratch[s*C..)) → output[s*C..) Run: python examples/workers/l3/all_to_all_distributed/main.py -p a2a3sim -d 0-1 @@ -55,7 +55,7 @@ DTYPE_NBYTES = 4 # float32 # Signal tail: one int32 slot per rank, bounded by kMaxSupportedRanks. SIGNAL_TAIL_NBYTES = 16 * 4 # 64 B -# NOTE: the full scratch size depends on nranks (staging area = nranks * COUNT_PER_RANK floats). +# NOTE: the full scratch size depends on nranks (result area = nranks * COUNT_PER_RANK floats). # It is computed inside run() once nranks is known. @@ -119,7 +119,7 @@ def run( ) -> int: """Core logic — callable from both CLI and pytest.""" nranks = len(device_ids) - # Scratch = nranks * COUNT_PER_RANK floats (staging) + signal tail. + # Scratch = nranks * COUNT_PER_RANK floats (result area) + signal tail. scratch_nbytes = nranks * COUNT_PER_RANK * DTYPE_NBYTES + SIGNAL_TAIL_NBYTES window_size = max(scratch_nbytes, 4 * 1024) From 7fd5a5a0365c9d05e33631616a895c1740466d32 Mon Sep 17 00:00:00 2001 From: georgebisbas Date: Fri, 3 Jul 2026 14:17:59 +0200 Subject: [PATCH 2/3] Cleanup: remove stale __CPU_SIM guards from TPUT-using kernels MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The pto-isa CPU simulator now has native TPUT support (TPut.hpp, via Copy_Data wrapper). The manual element-loop fallbacks guarded by __CPU_SIM are no longer needed — pto::comm::TPUT works on both hardware and simulator paths. - all_to_all_kernel.cpp: 1 guard removed - combine.cpp: 1 guard removed - dispatch.cpp: 3 guards removed - kernel_allreduce_sum.cpp: 1 guard removed Sim-verified: all_to_all (P=2, P=4), ffn_tp_parallel all pass. ep_dispatch_combine golden check is a pre-existing failure (fails on main too). --- .../kernels/aiv/all_to_all_kernel.cpp | 6 ------ 1 file changed, 6 deletions(-) diff --git a/examples/workers/l3/all_to_all_distributed/kernels/aiv/all_to_all_kernel.cpp b/examples/workers/l3/all_to_all_distributed/kernels/aiv/all_to_all_kernel.cpp index 7057ef1d5..832e5eef6 100644 --- a/examples/workers/l3/all_to_all_distributed/kernels/aiv/all_to_all_kernel.cpp +++ b/examples/workers/l3/all_to_all_distributed/kernels/aiv/all_to_all_kernel.cpp @@ -106,13 +106,7 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in Global inputChunkG(input + dest * COUNT_PER_RANK, shape, stride); Global scratchChunkG(scratch_dst_remote, shape, stride); -#if defined(__CPU_SIM) - for (size_t i = 0; i < COUNT_PER_RANK; ++i) { - scratch_dst_remote[i] = input[dest * COUNT_PER_RANK + i]; - } -#else pto::comm::TPUT(scratchChunkG, inputChunkG, pushTile); -#endif } pipe_barrier(PIPE_ALL); From e3498bcdcc52de63b7d3bcd9cf8e5ea36505d9bf Mon Sep 17 00:00:00 2001 From: georgebisbas Date: Fri, 3 Jul 2026 14:31:25 +0200 Subject: [PATCH 3/3] fix(pr): resolve review comments for #1267 - 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) --- examples/workers/l3/all_to_all_distributed/README.md | 2 +- .../l3/all_to_all_distributed/kernels/aiv/all_to_all_kernel.cpp | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/workers/l3/all_to_all_distributed/README.md b/examples/workers/l3/all_to_all_distributed/README.md index ee52da7ca..b8336d38d 100644 --- a/examples/workers/l3/all_to_all_distributed/README.md +++ b/examples/workers/l3/all_to_all_distributed/README.md @@ -36,7 +36,7 @@ python main.py -p a2a3 -d 0-3 ## File Structure -```textpush, barrier, copy-out +```text all_to_all_distributed/ ├── kernels/ │ ├── aiv/ diff --git a/examples/workers/l3/all_to_all_distributed/kernels/aiv/all_to_all_kernel.cpp b/examples/workers/l3/all_to_all_distributed/kernels/aiv/all_to_all_kernel.cpp index 832e5eef6..14e58a0e0 100644 --- a/examples/workers/l3/all_to_all_distributed/kernels/aiv/all_to_all_kernel.cpp +++ b/examples/workers/l3/all_to_all_distributed/kernels/aiv/all_to_all_kernel.cpp @@ -99,8 +99,8 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in // corresponding peer's scratch at offset my_rank*C. When dest == my_rank // CommRemotePtr returns the local pointer, so TPUT is a local GM write. // ------------------------------------------------------------------ + __gm__ float *scratch_dst_local = scratch + my_rank * COUNT_PER_RANK; for (int dest = 0; dest < nranks; ++dest) { - __gm__ float *scratch_dst_local = scratch + my_rank * COUNT_PER_RANK; __gm__ float *scratch_dst_remote = CommRemotePtr(commCtx, scratch_dst_local, dest); Global inputChunkG(input + dest * COUNT_PER_RANK, shape, stride);