Reverse All2All (EP MoE Distributed Kernel)

The Reverse All2All kernel implements the backward direction of the All2All communication pattern in EP (Expert Parallel) MoE models. In standard MoE All2All, tokens are dispatched from their home rank to the rank holding their assigned expert (forward All2All). Reverse All2All is the second half: after expert computation, results must be sent back from the expert’s home rank to the token’s original rank.

This document covers two implementation versions:

• MTE Version (general purpose, see “MTE Implementation” section): The kernel uses a Producer-Consumer pipeline pattern on Vector Cores, with even Vector Cores acting as Producers (remote cross-card stores via tl.store + dl.symm_at symmetric address mapping), and odd Vector Cores acting as Consumers (local reads and output writes).

• UDMA Optimized Version (Ascend 950, see “UDMA Optimized Implementation” section): Leverages the new UDMA link on 950 for higher cross-card bandwidth, splitting “cross-card data exchange” and “local data reordering” into two separate phases for respective processing.

Reverse All2All Operator Semantics

In an EP MoE model with world_size ranks, each rank holds a subset of experts. The communication pattern is:

Forward All2All (dispatch):
  Input:  (S, H, D) per rank           # Tokens grouped by target expert rank
  Output: (S/world_size, H * world_size, D) per rank  # Tokens from all ranks

Expert Computation:
  Each rank processes tokens assigned to its experts

Reverse All2All (combine):
  Input:  (S * world_size, H, D) per rank   # Local expert outputs
  Output: (S, H * world_size, D) per rank    # Outputs returned to original token ranks

The Reverse All2All kernel specifically handles the final step. The input tensor shape is (S * world_size, H, D) — containing expert outputs that need to be sent back to all other ranks. After reverse All2All, each rank receives outputs from all ranks, reconstructing a complete output with shape (S, H * world_size, D).

The key difference from forward All2All is the dimension swap: the world_size factor moves from the H dimension to the S dimension (or vice versa), essentially transposing the cross-rank communication pattern.

EP MoE Scenario

In a typical EP MoE layer (such as DeepSeek-V3 with 256 experts):

1. Forward All2All (Dispatch): Tokens are sent to the rank holding their assigned expert. Each rank’s input (S, H, D) becomes (S/world_size, H * world_size, D).

2. Expert GEMM: Each rank computes GEMM for local experts using GroupedGEMM or Fused MoE.

3. Reverse All2All (Combine): Expert outputs are sent back to original token ranks. Each rank’s input (S * world_size, H, D) becomes (S, H * world_size, D).

The Reverse All2All kernel is the third step. A Producer-Consumer pipeline based on per-task signal synchronization achieves overlap between remote cross-card stores (Producer) and local reads/output writes (Consumer), obtaining higher bandwidth utilization than barrier-based approaches.

---

MTE Implementation

Producer-Consumer Pipeline

The Reverse All2All kernel uses a Vector Core-level Producer-Consumer pipeline based on program_id, launching the kernel at Vector Core granularity (not AICore granularity):

Even Vector Cores (role = pid % 2 == 0):  Producer
  - Read local input tensor data
  - Write data to remote symmetric memory via dl.symm_at
  - Wait for signal with dl.wait (waitValue=0, acquire) before remote writes
  - Use dl.consume_token to enforce dependency
  - Call libshmem_device.fence() after all remote writes complete
  - Send dl.notify to each target rank to signal completion

Odd Vector Cores (role = pid % 2 == 1):  Consumer
  - Wait for incoming data signal with dl.wait (acquire)
  - Use dl.consume_token to enforce dependency
  - Read data from local symmetric memory
  - Write gathered data to output tensor
  - Call libshmem_device.fence() after all reads complete
  - Send dl.notify (rank, signal=0) to self to reset signal for next iteration

This differs from the previous sub_vec_id() sub-block model. In the Vector Core PID model:

• logical_core_id = pid // 2 — Each Producer-Consumer pair shares a logical_core_id

• n_role_cores = ncore // 2 — Number of cores allocated to each role

• Even pids handle all cross-card remote store operations (Producer)

• Odd pids handle all local symmetric memory reads and output writes (Consumer)

The kernel launches with the total number of Vector Cores rather than the number of AICores:

from triton.backends.ascend.driver import NPUUtils

vec_num = NPUUtils().get_aivector_core_num()  # Vector Core count (not AICore count)
kernel_hccl_reverse_a2a_pipelined[vec_num, 1, 1](...)

This allows each Vector Core to independently act as a Producer or Consumer, maximizing parallelism for communication operations.

Signal-Based Synchronization: dl.wait / dl.notify / fence

Unlike barrier-based synchronization (barrier_all() / barrier_all_vec()), the Reverse All2All kernel uses fine-grained per-task signal synchronization:

Producer side:

1. dl.wait(remote_signal_ptr, 1, "gpu", "acquire", waitValue=0) — Wait for the target rank’s signal slot to be cleared (value == 0), indicating the target rank’s consumer has finished reading the previous buffer

2. dl.consume_token(remote_ptr, token) — Enforce dependency, ensuring remote writes only occur after wait completes

3. tl.store(remote_ptr + offset, data, mask=mask) — Write data to remote symmetric memory

4. libshmem_device.fence() — Ensure all remote writes are visible before sending notification

5. dl.notify(signal_mem_ptr, target_rank) — Notify target rank that data is available

Consumer side:

1. dl.wait(signal_mem_ptr, 1, "gpu", "acquire") — Wait for signal from source rank (data available)

2. dl.consume_token(local_peer_ptr, token) — Enforce dependency, ensuring local reads only occur after wait completes

3. tl.load(local_peer_ptr + offset, mask=mask) — Read data from local symmetric memory

4. tl.store(output_ptr + offset, data, mask=mask) — Write to output tensor

5. libshmem_device.fence() — Ensure all reads complete before resetting signal

6. dl.notify(signal_mem_ptr, rank, 0) — Reset signal (value=0) to release next iteration’s producer

This per-task signal mechanism achieves finer-grained overlap between Producer and Consumer than coarse-grained barrier approaches. Each Producer task can proceed as soon as its dedicated signal slot is cleared, without waiting for all tasks to complete a global barrier.

Signal Memory Layout

Signal memory is organized as a per-task slot array:

signal_mem shape: [H * world_size * buffer_num * 8]  (int64, 8 bytes per slot)

Layout: [buffer_id][num_blocks_d][H * rank_size][rank_factor]
        - buffer_id: Identifies double buffer slot
        - num_blocks_d * H * rank_size: Total tasks per buffer iteration
        - rank_factor: Identifies producer/consumer pair (producer uses rank, consumer uses r)

Producer uses: remote_signal_ptr + buffer_id * num_blocks_d * H * rank_size * 8
               + num_blocks_d * H * rank * 8 + task_idx * 8

Consumer uses: signal_mem_ptr + buffer_id * num_blocks_d * H * rank_size * 8
               + num_blocks_d * H * r * 8 + task_idx * 8

Triton-Distributed Reverse All2All Operator

import triton
import triton.language as tl
import triton_dist.language as dl
from triton_dist.language.extra import libshmem_device
from triton.backends.ascend.driver import NPUUtils

@triton.jit
def kernel_hccl_reverse_a2a_pipelined(
    a_ptr,                           # Input tensor: shape (S_total, H, D)
    c_ptr,                           # Output tensor: shape (S, H * world_size, D)
    peer_mem_ptr,                    # Symmetric memory pointer
    signal_mem_ptr,                  # Signal memory pointer (int64)
    rank,                            # Current rank ID (host side)
    rank_size,                       # Total number of ranks (host side)
    buffer_num,                      # Number of double buffers
    S, H, D,                         # Tensor dimensions (S = S_total / rank_size)
    stride_as, stride_ah, stride_ad, # Input strides
    stride_cs, stride_ch, stride_cd, # Output strides
    COMM_BLOCK_S: tl.constexpr,      # Tile size in S dimension
    COMM_BLOCK_D: tl.constexpr,      # Tile size in D dimension
):
    ncore = tl.num_programs(axis=0)
    pid = tl.program_id(axis=0)

    # Producer (even PID) / Consumer (odd PID) role assignment
    role = pid % 2
    logical_core_id = pid // 2
    n_role_cores = ncore // 2

    num_blocks_s = tl.cdiv(S, COMM_BLOCK_S)
    num_blocks_d = tl.cdiv(D, COMM_BLOCK_D)

    # Symmetric memory layout: [buffer_num, S_block, H * rank_size, D]
    buffer_chunk_size = (H * rank_size) * D
    stride_ps = H * rank_size * D
    stride_ph = D
    stride_pd = 1

    # Outer loop: Iterate over S dimension blocks
    for global_id_s in range(0, num_blocks_s):
        buffer_id = global_id_s % buffer_num

        # ---- Producer Stage (Even Vector Cores): Cross-card remote store ----
        if role == 0:
            total_prod_tasks = num_blocks_d * H * rank_size
            for task_idx in range(logical_core_id, total_prod_tasks, n_role_cores):
                tmp = task_idx
                rank_loop_id = tmp % rank_size
                tmp //= rank_size
                h_id = tmp % H
                block_id_d = tmp // H

                target_rank = (rank + rank_loop_id) % rank_size

                # Boundary mask for non-aligned S and D dimensions
                offs_s = global_id_s * COMM_BLOCK_S + tl.arange(0, COMM_BLOCK_S)
                offs_d = block_id_d * COMM_BLOCK_D + tl.arange(0, COMM_BLOCK_D)
                mask = (offs_s < S)[:, None] & (offs_d < D)[None, :]

                # Read from local A tensor (row offset = target_rank * S)
                a_s = target_rank * S + offs_s
                a_offs = a_s[:, None] * stride_as + h_id * stride_ah + offs_d[None, :] * stride_ad
                a_data = tl.load(a_ptr + a_offs, mask=mask, other=0.0)

                # Wait for signal: target rank's consumer has completed previous read
                remote_ptr = dl.symm_at(peer_mem_ptr, target_rank)
                remote_signal_ptr = dl.symm_at(signal_mem_ptr, target_rank)
                token = dl.wait(
                    remote_signal_ptr
                    + buffer_id * num_blocks_d * H * rank_size * 8
                    + num_blocks_d * H * rank * 8
                    + tmp * 8,
                    1, "gpu", "acquire", waitValue=0,
                )
                remote_ptr_dummy = dl.consume_token(remote_ptr, token)

                # Write to remote symmetric memory
                peer_h_write = rank * H + h_id
                peer_offs_write = (
                    buffer_id * (COMM_BLOCK_S * buffer_chunk_size)
                    + tl.arange(0, COMM_BLOCK_S)[:, None] * stride_ps
                    + peer_h_write * stride_ph
                    + offs_d[None, :] * stride_pd
                )
                tl.store(remote_ptr_dummy + peer_offs_write, a_data, mask=mask)

            # Ensure all remote writes are visible before notify
            libshmem_device.fence()
            for task_idx in range(logical_core_id, total_prod_tasks, n_role_cores):
                tmp = task_idx
                rank_loop_id = tmp % rank_size
                tmp //= rank_size
                target_rank = (rank + rank_loop_id) % rank_size
                dl.notify(
                    signal_mem_ptr
                    + buffer_id * num_blocks_d * H * rank_size * 8
                    + num_blocks_d * H * rank * 8
                    + tmp * 8,
                    target_rank,
                )

        # ---- Consumer Stage (Odd Vector Cores): Read from local Shmem and store to C ----
        if role == 1:
            local_peer_ptr = dl.symm_at(peer_mem_ptr, rank)
            total_cons_tasks = num_blocks_d * H * rank_size
            for task_idx in range(logical_core_id, total_cons_tasks, n_role_cores):
                tmp = task_idx
                r = tmp % rank_size
                tmp //= rank_size
                h_id = tmp % H
                block_id_d = tmp // H

                # Boundary mask
                offs_s = global_id_s * COMM_BLOCK_S + tl.arange(0, COMM_BLOCK_S)
                offs_d = block_id_d * COMM_BLOCK_D + tl.arange(0, COMM_BLOCK_D)
                mask = (offs_s < S)[:, None] & (offs_d < D)[None, :]

                # Wait for signal: source rank's producer has completed write
                peer_h_read = r * H + h_id
                peer_offs_read = (
                    buffer_id * (COMM_BLOCK_S * buffer_chunk_size)
                    + tl.arange(0, COMM_BLOCK_S)[:, None] * stride_ps
                    + peer_h_read * stride_ph
                    + offs_d[None, :] * stride_pd
                )
                token = dl.wait(
                    signal_mem_ptr
                    + buffer_id * num_blocks_d * H * rank_size * 8
                    + num_blocks_d * H * r * 8
                    + tmp * 8,
                    1, "gpu", "acquire",
                )
                local_peer_ptr = dl.consume_token(local_peer_ptr, token)
                peer_data = tl.load(local_peer_ptr + peer_offs_read, mask=mask, other=0.0)

                # Write to local C tensor (transpose: r * H + h_id becomes column index)
                c_h_idx = r * H + h_id
                c_offs = offs_s[:, None] * stride_cs + c_h_idx * stride_ch + offs_d[None, :] * stride_cd
                tl.store(c_ptr + c_offs, peer_data, mask=mask)

            # Ensure all reads complete before resetting signal
            libshmem_device.fence()
            for task_idx in range(logical_core_id, total_cons_tasks, n_role_cores):
                tmp = task_idx
                r = tmp % rank_size
                tmp //= rank_size
                h_id = tmp % H
                block_id_d = tmp // H
                dl.notify(
                    signal_mem_ptr
                    + buffer_id * num_blocks_d * H * rank_size * 8
                    + num_blocks_d * H * r * 8
                    + tmp * 8,
                    rank, 0,
                )

# Launch with Vector Core count
def hccl_reverse_a2a_kernel_launcher(A, C, peer_mem, signal_mem, rank, rank_size, buffer_num, COMM_BLOCK_S, COMM_BLOCK_D):
    S_total, H, D = A.shape
    S = S_total // rank_size
    vec_num = NPUUtils().get_aivector_core_num()
    kernel_hccl_reverse_a2a_pipelined[vec_num, 1, 1](
        A, C, peer_mem, signal_mem, rank, rank_size, buffer_num, S, H, D,
        A.stride(0), A.stride(1), A.stride(2),
        C.stride(0), C.stride(1), C.stride(2),
        COMM_BLOCK_S=COMM_BLOCK_S, COMM_BLOCK_D=COMM_BLOCK_D,
    )

Symmetric Memory Layout

For Reverse All2All, the symmetric memory layout is organized at block level (rather than complete tensor level):

peer_mem layout per buffer slot: [COMM_BLOCK_S, H * rank_size, D]
  - Each S dimension block has COMM_BLOCK_S rows
  - H * rank_size columns: Each rank's H experts each occupy a column slot
  - D dimension depth

Complete peer_mem size: COMM_BLOCK_S * (H * world_size) * D * buffer_num
  (Allocated as a flat 1D tensor of this total size)

stride_ps = H * rank_size * D   (stride along S-block rows within buffer slot)
stride_ph = D                    (stride along H columns)
stride_pd = 1                    (stride along D depth)

When the Producer writes to a remote rank’s symmetric memory, the H dimension offset is rank * H + h_id (writing to the slot reserved for this rank’s data). When the Consumer reads from local symmetric memory, the offset is r * H + h_id (reading the slot where rank r wrote data).

This block-level layout supports pipelined double buffering: Each buffer slot holds only COMM_BLOCK_S rows at a time, rather than the entire S * world_size rows, significantly reducing symmetric memory usage.

Performance Considerations

• Per-task signals vs global barrier: The dl.wait/dl.notify/fence mechanism achieves finer-grained synchronization than barrier_all_vec(). Each task can proceed as soon as its own signal is satisfied, without waiting for all tasks to complete, reducing idle time and improving throughput.

• Dynamic masking: The kernel uses boundary masks (offs_s < S)[:, None] & (offs_d < D)[None, :] to correctly handle non-aligned S and D dimensions, avoiding out-of-bounds reads/writes.

• Double buffering: The Producer-Consumer pattern with buffer_num=2 allows the Producer to write data for iteration i+1 while the Consumer reads data from iteration i, overlapping the two stages.

---

UDMA Optimized Implementation (Ascend 950)

Ascend 950 introduces UDMA links, which provide higher cross-card bandwidth compared to the original MTE links (tl.load/tl.store + dl.symm_at symmetric address mapping), but UDMA can only be accessed through RMA primitives like libshmem_device.putmem/getmem. This section covers the UDMA-optimized version of Reverse All2All (EP MoE combine phase) on 950.

Core idea: Split “cross-card data exchange” and “local data reordering” into two phases with different characteristics, handling each with the most suitable link — cross-card bulk transfers are handled by UDMA (one-shot, large blocks, one core per pe), while local data movement (including the special case of “sending to self” which UDMA cannot express) is handled by MTE’s fine-grained double-buffered pipeline.

Two-Phase Data Copy

Phase 1 — UDMA cross-card bulk exchange (one-shot, core ↔ target rank one-to-one mapping):
  Each rank putmems the entire segment of (S, H, D) data in local A that belongs to target rank r to r in one shot.
  One core handles only one target rank, and moves all data for that rank in a single call (no tiling by s/h/d).

Phase 2 — Local MTE self-loop movement (fine-grained, double-buffered, handles only the "send to self" special case):
  UDMA putmem cannot "send data to self" (dst must be remote shmem address),
  so the r == rank (self) portion of data uses local MTE tiled copy to write to the local slot of symmetric memory,
  using S-block double buffering, overlapping with the subsequent Consumer read phase.

Single barrier_all() synchronization:
  Wait for both Phase 1 cross-card UDMA writes and Phase 2 local MTE writes to complete.

Consumer phase (fine-grained, MTE read + output write):
  For r != rank: Read directly from symmetric memory where Phase 1 bulk data has arrived (no need to wait for double buffering, as bulk data has arrived all at once).
  For r == rank: Read from the double-buffered slot written by Phase 2 (corresponding to Phase 2's double buffering rhythm).

Phase 1 uses UDMA for higher cross-card bandwidth; Phase 2 and Consumer phase still use MTE, because they require fine-grained tile-level parallelism and double-buffered pipeline overlap, which are precisely what UDMA’s constraints do not allow and MTE supports (on-card data copy) scenarios.

Vector Core Role Division

The kernel continues to use the Vector Core-level sub-block programming model (sub_vec_id()/sub_vec_num()), and further divides even sub-blocks (global_vec_id % 2 == 0) into two role groups by physical core id:

role A (even sub-blocks, implemented by range(pid, rank_size, ncore) striping): UDMA sender
  When ncore >= rank_size (common case, e.g., on 950 ncore is much larger than rank_size),
  this loop is equivalent to cores with pid < rank_size each getting a unique target_rank == pid;
  When ncore < rank_size, the same core may handle multiple target_ranks, but each target_rank
  is still handled by only one core (Constraint 2's one-core-one-pe semantics is guaranteed by stride ncore).
  Each target_rank sends the entire segment of data in one putmem.

role B (even sub-blocks, when pid >= rank_size the UDMA loop is empty, then participate): Local MTE self-loop Producer
  Only handles r == rank local data, writing to the local slot of symmetric memory in fine-grained (s_block, h, d) tile double buffering.

role C (odd sub-blocks): Consumer
  Reads symmetric memory in fine-grained (s_block, h, d, r) tiles (r==rank uses double-buffered slot, r!=rank uses Phase 1 bulk-arrived slot),
  writing to output tensor.

This role division allows UDMA sending (role A) and local MTE filling (role B) to execute in parallel within the same even sub-block group without conflict (the (rank, s_block) combinations each handles do not overlap), then unified synchronization with Consumer (role C) using a single barrier_all().

Triton-Distributed Operator Implementation

import triton
import triton.language as tl
import triton_dist.language as dl
from triton_dist.language.extra import libshmem_device
from triton.language.extra.cann.extension import sub_vec_id, sub_vec_num
from triton.backends.ascend.driver import NPUUtils

@triton.jit
def kernel_reverse_a2a_udma(
    a_ptr, c_ptr, peer_mem_ptr,
    rank, rank_size, buffer_num,
    S, H, D,
    stride_as, stride_ah, stride_ad,
    stride_cs, stride_ch, stride_cd,
    COMM_BLOCK_S: tl.constexpr,
    COMM_BLOCK_D: tl.constexpr,
    elem_bytes: tl.constexpr,
):
    vec_num_per_aicore = 2
    ncore = tl.num_programs(axis=0) * sub_vec_num() // vec_num_per_aicore
    pid = tl.program_id(axis=0) * sub_vec_num() // vec_num_per_aicore
    global_vec_id = pid * sub_vec_num() + sub_vec_id()

    num_blocks_s = tl.cdiv(S, COMM_BLOCK_S)
    num_blocks_d = tl.cdiv(D, COMM_BLOCK_D)
    stride_ps = H * D  # per-rank slot layout: [S, H, D]

    # ---- Phase 1 (role A): one core per target rank, one-shot UDMA putmem ----
    # Constraint 2: stride must be ncore, ensuring same target_rank is handled by only one pid
    # (range(pid, rank_size) without stride would cause multiple pids to repeatedly handle same target_rank, see udma-programming.md Constraint 2)
    if global_vec_id % 2 == 0:
        for target_rank in range(pid, rank_size, ncore):
            if target_rank != rank:
                # Constraint 3: one call moves entire segment (S, H, D) for this target, no tiling
                libshmem_device.putmem(
                    peer_mem_ptr + rank * S * H * D,      # dst: target_rank's symmetric memory, slot "from rank"
                    a_ptr + target_rank * S * H * D,       # src: entire segment in local A belonging to target_rank
                    S * H * D * elem_bytes,
                    target_rank,                            # pe: target rank for putmem
                )

    # ---- Phase 2 (role B): local self-copy via tiled MTE double buffer ----
    for global_id_s in range(0, num_blocks_s):
        buffer_id = global_id_s % buffer_num
        if global_vec_id % 2 == 0 and pid >= rank_size:
            total_tasks = num_blocks_d * H
            logical_pid = pid - rank_size
            logical_ncores = ncore - rank_size
            for task_idx in range(logical_pid, total_tasks, logical_ncores):
                h_id = task_idx % H
                block_id_d = task_idx // H
                offs_s = global_id_s * COMM_BLOCK_S + tl.arange(0, COMM_BLOCK_S)
                offs_d = block_id_d * COMM_BLOCK_D + tl.arange(0, COMM_BLOCK_D)
                mask = (offs_s < S)[:, None] & (offs_d < D)[None, :]

                a_s = rank * S + offs_s  # "send to self" segment in local data
                a_offs = a_s[:, None] * stride_as + h_id * stride_ah + offs_d[None, :] * stride_ad
                a_data = tl.load(a_ptr + a_offs, mask=mask, other=0.0)

                peer_offs = (
                    rank * S * H * D
                    + buffer_id * (COMM_BLOCK_S * stride_ps)
                    + tl.arange(0, COMM_BLOCK_S)[:, None] * stride_ps
                    + h_id * D
                    + offs_d[None, :]
                )
                tl.store(peer_mem_ptr + peer_offs, a_data, mask=mask)

        # ---- Sync: wait for both UDMA bulk write + local MTE write to complete for current S-block ----
        libshmem_device.barrier_all()

        # ---- Consumer (role C): read symmetric memory and write output ----
        if global_vec_id % 2 == 1:
            total_tasks = num_blocks_d * H * rank_size
            for task_idx in range(pid, total_tasks, ncore):
                tmp = task_idx
                r = tmp % rank_size
                tmp //= rank_size
                h_id = tmp % H
                block_id_d = tmp // H
                offs_s = global_id_s * COMM_BLOCK_S + tl.arange(0, COMM_BLOCK_S)
                offs_d = block_id_d * COMM_BLOCK_D + tl.arange(0, COMM_BLOCK_D)
                mask = (offs_s < S)[:, None] & (offs_d < D)[None, :]

                if r == rank:
                    # Local self-loop data: read Phase 2 written double-buffered slot
                    peer_offs = (
                        rank * S * H * D
                        + buffer_id * (COMM_BLOCK_S * stride_ps)
                        + tl.arange(0, COMM_BLOCK_S)[:, None] * stride_ps
                        + h_id * D
                        + offs_d[None, :]
                    )
                else:
                    # Remote data: Phase 1 has bulk-arrived all at once, read directly by s_block index
                    peer_offs = (
                        r * S * H * D
                        + global_id_s * COMM_BLOCK_S * stride_ps
                        + tl.arange(0, COMM_BLOCK_S)[:, None] * stride_ps
                        + h_id * D
                        + offs_d[None, :]
                    )
                peer_data = tl.load(peer_mem_ptr + peer_offs, mask=mask, other=0.0)

                c_h_idx = r * H + h_id
                c_offs = offs_s[:, None] * stride_cs + c_h_idx * stride_ch + offs_d[None, :] * stride_cd
                tl.store(c_ptr + c_offs, peer_data, mask=mask)


def launch_reverse_a2a_udma(A, C, peer_mem, rank, rank_size, buffer_num, COMM_BLOCK_S, COMM_BLOCK_D):
    S_total, H, D = A.shape
    S = S_total // rank_size
    aicore_num = NPUUtils().get_aicore_num()
    kernel_reverse_a2a_udma[aicore_num, 1, 1](
        A, C, peer_mem, rank, rank_size, buffer_num, S, H, D,
        A.stride(0), A.stride(1), A.stride(2),
        C.stride(0), C.stride(1), C.stride(2),
        COMM_BLOCK_S=COMM_BLOCK_S, COMM_BLOCK_D=COMM_BLOCK_D,
        elem_bytes=A.element_size(),
    )

Symmetric Memory Layout

Unlike the pure MTE version (symmetric memory organized as [buffer_num, COMM_BLOCK_S, H*rank_size, D], tile-by-tile double buffering), the UDMA version’s symmetric memory is organized by rank slots, with double buffering only within the “local self-loop” slot:

peer_mem total size: rank_size * S * H * D

per-rank slot: [S, H, D], slot starting offset = r * S * H * D
  - r != rank slots: Written by Phase 1 UDMA putmem in one-shot bulk, Consumer reads directly by s_block index, no buffering needed
  - r == rank slot: Written by Phase 2 local MTE, only the prefix interval
                    [0, buffer_num * COMM_BLOCK_S * H * D) within the slot does double-buffered cyclic reuse

This layout separates “bulk one-shot arrival, no buffering needed” (remote UDMA data) from “block-by-block pipelined, needs double buffering” (local self-loop data), avoiding allocating buffer_num times the space for remote data that doesn’t need buffering, saving symmetric memory compared to the pure MTE version.

Performance Considerations

• UDMA bandwidth advantage concentrated in Phase 1: Cross-card data volume is (rank_size - 1) × S × H × D × elem_bytes, completed with rank_size - 1 large-block putmem calls, fully utilizing UDMA’s bandwidth advantage over MTE; larger tiles and fewer calls get closer to UDMA peak bandwidth.

• Local self-loop still needs MTE fine-grained double buffering: The r == rank portion of data volume is S × H × D, accounting for 1/rank_size of total data volume, handled with MTE tiled double buffering, overlapping with Consumer reads, its relative overhead decreases as rank_size increases.

• Core count allocation tradeoff: Cores with pid < rank_size are dedicated to UDMA sending, cores with pid >= rank_size are dedicated to local MTE filling — when rank_size is large, fewer cores participate in local MTE filling (ncore - rank_size), potentially making the local self-loop phase a bottleneck.

• Single barrier_all() covers both links: Phase 1 (UDMA) and Phase 2 (MTE) writes share the same barrier_all(), meaning Consumer must wait for the slower of the two links to complete. If UDMA and MTE latencies differ significantly, setting independent signals for the two links and having Consumer wait for the corresponding link separately based on data source (r == rank vs r != rank), rather than being dragged down by the slowest one, could improve performance.

• Memory alignment: Same as pure MTE version, (H, D) layout must satisfy 32B alignment; UDMA’s bulk transfer has no additional alignment requirements for total bytes (S*H*D*elem_bytes), but tile-level tl.load/tl.store in Phase 2/Consumer are still subject to MTE alignment constraints.

Common Issues

1. Directly applying MTE version’s tile loop to putmem

If you copy the pure MTE version’s tile-by-tile (s_block, h, d) splitting and call putmem tile-by-tile, it will cause precision issues, because currently UDMA only supports single QP communication, and when multiple cores concurrently read/write the same rank, the single QP’s signal slots are insufficient to guarantee all multi-core communication completes. Solution: Phase 1 must be redesigned as “one core per pe, one call for entire segment data”, it cannot be a line-by-line replacement of the MTE tile loop.

2. Trying to use putmem to handle “send to self” data

putmem’s dst must be a remote shmem address, it cannot express “send to self”. Solution: This portion of data always goes through local MTE (tl.load/tl.store), as shown in Phase 2.

3. Still applying double-buffer indexing to bulk-arrived remote data

After remote data arrives in one shot via UDMA as a complete [S, H, D], Consumer should read directly using global_id_s * COMM_BLOCK_S indexing into the interior of that bulk, rather than using buffer_id for double-buffer offset like local self-loop data — the symmetric memory addressing methods for the two are different, mixing them will read wrong offsets. Solution: Consumer selects offset formula separately for r == rank vs r != rank (as shown in code example).

Wait/Notify Optimization

The above UDMA + MTE dual-link implementation uses a single barrier_all() to synchronize Phase 1 (UDMA cross-card writes), Phase 2 (local MTE self-loop writes), and Consumer reads, which means Consumer must wait for all Producer tasks (including all ranks’ UDMA sends + local MTE self-loop) to complete before starting reads, creating a global synchronization bottleneck.

Optimization Goal

Replace barrier_all() with fine-grained wait/notify signal synchronization, allowing each Consumer task to start reading as soon as the Producer signal it depends on is ready, without being blocked by the slowest Producer.

Core Changes

1. Dual Signal Memory Region Design

signal_mem: flat int64 array, divided into two independent regions

Region A — Cross-card UDMA signals [num_blocks_s, rank_size]:
  Slot semantics: (global_id_s, sender_rank) indicates "sender_rank has completed UDMA send of S-block global_id_s"
  Total slots: num_blocks_s * rank_size
  Slot address: signal_mem_ptr + (global_id_s * rank_size + sender_rank) * SIG_STRIDE
  Notifier: UDMA Sender (even cores, logical_core_id < rank_size)
  Waiter: Consumer (odd cores), waits for r != rank remote data
  Buffer reuse: None (each (block, rank) pair uses unique address, no reset needed)
  Credit return: Not needed

Region B — Local MTE self-loop double buffer signals [buffer_num, tasks_per_block]:
  Slot semantics: (buffer_id, task_idx) indicates "local MTE has completed task_idx write of buffer_id"
  Total slots: buffer_num * num_blocks_d * H
  Slot address: signal_mem_ptr + region_b_base + (buffer_id * tasks_per_block + task_idx) * SIG_STRIDE
  Notifier: MTE Producer (even cores, logical_core_id >= rank_size)
  Waiter: Consumer (odd cores), waits for r == rank local self-loop data
  Buffer reuse: Yes (num_blocks_s / buffer_num rounds)
  Credit return: Required (Consumer notifies(..., 0) to reset signal after read complete)

Where SIG_STRIDE = 8 (one cache line), avoiding false sharing.
region_b_base = num_blocks_s * rank_size * SIG_STRIDE.

Key design principles:

• Independent signals for two links: UDMA and MTE each have independent signal regions, no interference, Consumer waits for corresponding signal separately based on data source (r == rank vs r != rank)

• UDMA signals need no reset: Each (block, sender) pair uses a unique address, write once use once, unlimited waiters

• MTE signals need bidirectional handshake: Producer waits for signal clear (wait(waitValue=0)) → write → notify(1); Consumer waits for signal set (wait()) → read → notify(0) to return credit

2. UDMA Sender Stage

No need to wait for signal, notify immediately after send

# Phase 1 — UDMA cross-card bulk send (even cores, logical_core_id < rank_size)
if global_vec_id % 2 == 0 and logical_core_id < rank_size:
    for udma_target_rank in range(logical_core_id, rank_size, n_role_cores):
        if udma_target_rank != rank:
            # Step 1: UDMA putmem (synchronous semantics, returns when data has reached remote)
            libshmem_device.putmem(
                peer_mem_ptr + rank * S * H * D + global_id_s * COMM_BLOCK_S * stride_ps,
                a_ptr + udma_target_rank * S * H * D + global_id_s * COMM_BLOCK_S * stride_as,
                cur_block_s * H * D * elem_bytes,
                udma_target_rank,
            )
    # Step 2: Send completion signal to each target rank (UDMA synchronous, no fence needed)
    for udma_target_rank in range(logical_core_id, rank_size, n_role_cores):
        if udma_target_rank != rank:
            dl.notify(
                signal_mem_ptr + (global_id_s * rank_size + rank) * SIG_STRIDE,
                udma_target_rank.to(tl.int32),
            )

Key points:

• UDMA putmem is synchronous (call returns when data has reached remote), no fence() needed before notify

• Each sender_rank sends signal to target rank’s signal slot (global_id_s, rank), indicating “this rank’s this S-block has finished sending”

3. MTE Producer Stage

Bidirectional handshake: wait for clear → write → send signal

# Phase 2 — Local MTE self-loop write (even cores, logical_core_id >= rank_size)
if global_vec_id % 2 == 0 and logical_core_id >= rank_size:
    logical_mte_pid = logical_core_id - rank_size
    logical_mte_ncores = n_role_cores - rank_size
    for task_idx in range(logical_mte_pid, tasks_per_block, logical_mte_ncores):
        h_id = task_idx % H
        block_id_d = task_idx // H
        # ... calculate offsets and mask ...

        # Step 1: Wait for Consumer to finish reading previous round (signal value is 0)
        self_sig_ptr = (
            signal_mem_ptr + region_b_base
            + (buffer_id * tasks_per_block + task_idx) * SIG_STRIDE
        )
        token = dl.wait(self_sig_ptr, 1, "gpu", "acquire", waitValue=0)
        store_ptr = dl.consume_token(peer_mem_ptr, token)

        # Step 2: Load local data and write to symmetric memory
        a_data = tl.load(a_ptr + a_offs, mask=mask, other=0.0)
        peer_offs_write = ( ... )
        tl.store(store_ptr + peer_offs_write, a_data, mask=mask)

    # Step 3: fence ensures all MTE writes are visible
    libshmem_device.fence()

    # Step 4: Send completion signal for each task
    for task_idx in range(logical_mte_pid, tasks_per_block, logical_mte_ncores):
        self_sig_ptr = (
            signal_mem_ptr + region_b_base
            + (buffer_id * tasks_per_block + task_idx) * SIG_STRIDE
        )
        dl.notify(self_sig_ptr, rank)

Key points:

• MTE needs fence() to ensure writes are visible, and fence() must be before notify()

• Bidirectional handshake: Producer first wait(waitValue=0) to ensure Consumer has finished reading previous round, then write, then notify(1)

4. Consumer Stage

Wait for signal separately by data source + credit return

# Consumer (odd cores)
if global_vec_id % 2 == 1:
    total_cons_tasks = tasks_per_block * rank_size
    for task_idx in range(logical_core_id, total_cons_tasks, n_role_cores):
        tmp = task_idx
        r = tmp % rank_size
        tmp //= rank_size
        h_id = tmp % H
        block_id_d = tmp // H
        # ... calculate offsets and mask ...

        if r == rank:
            # Local self-loop data: wait for MTE Producer signal
            self_sig_ptr = (
                signal_mem_ptr + region_b_base
                + (buffer_id * tasks_per_block + tmp) * SIG_STRIDE
            )
            token = dl.wait(self_sig_ptr, 1, "gpu", "acquire")
            load_ptr = dl.consume_token(peer_mem_ptr, token)
            peer_offs_read = ( ... )  # Double buffer address
            peer_data = tl.load(load_ptr + peer_offs_read, mask=mask, other=0.0)
        else:
            # Remote data: wait for UDMA Sender signal
            remote_sig_ptr = (
                signal_mem_ptr + (global_id_s * rank_size + r) * SIG_STRIDE
            )
            token = dl.wait(remote_sig_ptr, 1, "gpu", "acquire")
            load_ptr = dl.consume_token(peer_mem_ptr, token)
            peer_offs_read = ( ... )  # Bulk address
            peer_data = tl.load(load_ptr + peer_offs_read, mask=mask, other=0.0)

        # Write to output tensor
        c_offs = ( ... )
        tl.store(c_ptr + c_offs, peer_data, mask=mask)

    # Step: fence ensures all reads complete
    libshmem_device.fence()

    # Step: Credit return (only for local self-loop data)
    for task_idx in range(logical_core_id, total_cons_tasks, n_role_cores):
        tmp = task_idx
        r = tmp % rank_size
        tmp //= rank_size
        if r == rank:
            self_sig_ptr = (
                signal_mem_ptr + region_b_base
                + (buffer_id * tasks_per_block + tmp) * SIG_STRIDE
            )
            dl.notify(self_sig_ptr, rank, 0)  # Reset to 0, release next round Producer

Performance Gains

Synchronization mode	Consumer wait time	Global sync points
barrier_all() (baseline)	max(all UDMA sends, all MTE self-loops)	One per S-block
wait/notify (optimized)	t_UDMA[r] (remote data) or t_MTE[task] (local data)	No global sync

Host Side Modifications

1. Signal Memory Allocation

num_blocks_s = -(-S // COMM_BLOCK_S)
num_blocks_d = -(-D // COMM_BLOCK_D)
signal_mem_size = (
    num_blocks_s * rank_size + buffer_num * num_blocks_d * H
) * 8  # int64, 8 bytes per slot
signal_mem = ash.aclshmem_create_tensor(
    [signal_mem_size], dtype=torch.int64, device_id=rank
)

2. Reset signal before kernel launch

for _ in range(iters):
    signal_mem.fill_(0)  # Reset all signals to 0
    dist.barrier()        # Ensure all ranks have cleared
    launcher(...)
    torch.npu.synchronize()
    dist.barrier()        # Ensure all ranks have completed

Key points:

• Must fill_(0) before each launch, otherwise residual signals will cause next wait to pass through immediately (race condition)

• One dist.barrier() before and after clearing: former prevents this rank’s clear from erasing peer’s signals, latter prevents next round’s clear from erasing current round’s tail signals

/home/vincent/dev/huawei/Triton-Distributed-Ascend/docs/ascend/en/tutorial/reverse_all2all.md