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..b8336d38d 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 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..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 @@ -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,34 @@ 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. // ------------------------------------------------------------------ + __gm__ float *scratch_dst_local = scratch + my_rank * COUNT_PER_RANK; for (int dest = 0; dest < nranks; ++dest) { + __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); + + pto::comm::TPUT(scratchChunkG, inputChunkG, pushTile); } 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 +128,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)