Triton-distributed Semantics

This document describes the design philosophy and semantic model of Triton-distributed-ascend.

Design Philosophy

Triton-distributed is built on a Tile-Centric design philosophy (as described in the MLSys 2025 paper), which means:

  1. Tile as the Unit of Work: Both computation and communication are organized around tiles (blocks of data). Each tile is a self-contained unit that can be computed, transferred, and synchronized independently.

  2. Decoupled Computation and Communication: Communication (data transfer) and computation (GEMM, etc.) are explicitly separated and can be performed by different AI cores or different parts of the same kernel.

  3. Fine-grained Overlapping: By organizing work around tiles, we can achieve fine-grained overlapping where computation on tiles that are already available can proceed while other tiles are still being transferred.

Core Semantic Concepts

Producer-Consumer Model

Triton-distributed uses a producer-consumer model for overlapping computation with communication:

  • Producer: Responsible for data transfer (e.g., AllGather, All-to-All). The producer copies data to a shared buffer and signals when each tile is ready.

  • Consumer: Responsible for computation (e.g., GEMM). The consumer waits for tiles to be ready, then computes on them immediately.

# Producer: Transfer data and signal
@triton.jit
def overlap_kernel(...):
    # Transfer tile to shared buffer
    tl.store(remote_buffer_ptr + tile_offset, data)
    # Signal that tile is ready
    tdl.notify(signal_ptr, peer_rank, signal=tile_id, sig_op="set")

# Consumer: Wait for data and compute
@triton.jit
def overlap_kernel(...):
    # Wait for tile to be ready
    token = tdl.wait(signal_ptr + tile_id, 1, "sys", "acquire")
    # Consume token to establish data dependency
    data_ptr = tdl.consume_token(data_ptr, token)
    # Now safe to compute on the tile
    result = compute(tl.load(data_ptr))

Signal-Based Synchronization

The synchronization model is based on signals rather than barriers:

  • wait(ptr, n): Wait until n signals at ptr reach expected values

  • notify(ptr, rank, signal, sig_op): Send a signal to a specific rank

  • consume_token(value, token): Establish data dependency between wait and memory access

Key Insight: Unlike global barriers that synchronize all ranks, signal-based synchronization allows fine-grained tile-level coordination. This enables:

  • Different tiles to proceed independently

  • Maximum overlapping between communication and computation

  • Avoiding the “straggler effect” where slow ranks block everyone

Symmetric Memory Model

Triton-distributed-ascend uses symmetric memory (ACLSHMEM) for cross-rank communication:

  • All PEs (Processing Elements) allocate memory at the same virtual address

  • symm_at(ptr, rank): Maps a local pointer to the corresponding address on another PE

  • Enables direct load/store across PEs without explicit send/receive

# Access remote PE's memory directly
remote_ptr = tdl.symm_at(local_ptr, peer_pe)
tl.store(remote_ptr + offset, data)  # Write to peer's memory

Note: On Ascend NPUs, symmetric memory is provided by ACLSHMEM (see SHMEM Host API and SHMEM Device API).

Kernel Design Patterns

AllGather and GEMM Overlapping

Time →
┌─────────────────────────────────────────────────────────────┐
│ Producer (AllGather)                                        │
│ [Tile 0] → [Tile 1] → [Tile 2] → [Tile 3] → ...             │
│    ↓          ↓          ↓          ↓                       │
│  signal    signal     signal     signal                     │
│    ↓          ↓          ↓          ↓                       │
│ Consumer (GEMM)                                             │
│         [Tile 0] → [Tile 1] → [Tile 2] → [Tile 3] → ...     │
└─────────────────────────────────────────────────────────────┘

  1. Producer transfers tiles from other PEs via AllGather

  2. After each tile transfer completes, producer signals the consumer

  3. Consumer waits for each tile, then immediately starts GEMM computation

  4. Computation and communication overlap, hiding communication latency

GEMM and ReduceScatter Overlapping

Time →
┌─────────────────────────────────────────────────────────────┐
│ Producer (GEMM)                                             │
│ [Tile 0] → [Tile 1] → [Tile 2] → [Tile 3] → ...             │
│    ↓          ↓          ↓          ↓                       │
│  signal    signal     signal     signal                     │
│    ↓          ↓          ↓          ↓                       │
│ Consumer (ReduceScatter)                                    │
│         [Tile 0] → [Tile 1] → [Tile 2] → [Tile 3] → ...     │
└─────────────────────────────────────────────────────────────┘

  1. Producer (GEMM) computes output tiles

  2. After each tile is computed, producer signals the consumer

  3. Consumer (ReduceScatter) waits for tiles and performs reduction and scatter

  4. Communication happens as computation produces results, maximizing overlap

Threadblock Swizzling

To maximize overlap and minimize synchronization, Triton-distributed uses threadblock swizzling:

# Swizzle tile assignment so each PE starts with its local data
pid_m = (pid_m + rank * tiles_per_pe) % total_tiles

This ensures:

  • Each PE starts computing on locally available data (no wait needed)

  • By the time a PE needs remote data, it’s likely already transferred

  • Reduces synchronization overhead and improves cache efficiency

Token-Based Data Dependency

The consume_token primitive establishes explicit data dependencies:

# Without consume_token: compiler might reorder load before wait
token = tdl.wait(signal_ptr, 1, "sys", "acquire")
data = tl.load(data_ptr)  # BUG: might execute before wait!

# With consume_token: explicit dependency
token = tdl.wait(signal_ptr, 1, "sys", "acquire")
data_ptr = tdl.consume_token(data_ptr, token)  # Establishes dependency
data = tl.load(data_ptr)  # Guaranteed to execute after wait

This is essential for correctness because:

  1. Compilers aggressively reorder instructions for performance

  2. The wait and load have no syntactic dependency

  3. consume_token creates an explicit data dependency that prevents reordering

Ascend-Specific Considerations

Memory Hierarchy

Ascend NPUs have a distinct memory hierarchy optimized for AI workloads:

  • Global Memory (GM/HBM): Main device memory with high bandwidth

  • L1 (Unified Buffer/UB): Local scratchpad memory for element-wise, reduction, and vector operations. Size: 192 KB on Atlas 800T A2 training servers and Atlas 800I A2 inference servers.

  • Symmetric Memory: Allocated via ACLSHMEM from Global Memory, accessible across PEs for distributed communication

Architecture Overview:

Each AI Core consists of 1 Cube core paired with 2 Vector cores:

  • Cube core (AIC): Executes tl.dot matrix multiplication operations

  • Vector cores (AIV): Execute element-wise operations, reductions, and gather/scatter using the UB

When using symmetric memory for distributed operations:

  • Allocate buffers with aclshmem_malloc() on the host (see SHMEM Host API)

  • Access remote memory using remote_ptr() in kernels (see SHMEM Device API)

  • Use symm_at() from triton_dist.language for pointer translation

Memory Constraints:

  • UB size is limited (192 KB on Atlas 800T A2 training servers and Atlas 800I A2 inference servers, further reduced with doublebuffer enabled)

  • Use tiling to fit data within UB constraints

Synchronization Implementation

On Ascend NPUs, synchronization primitives map to ACLSHMEM operations:

  • wait() → Uses ACLSHMEM signal waiting mechanisms

  • notify() → Uses ACLSHMEM putmem_signal or signal_op

  • Barriers → ACLSHMEM barrier_all() or barrier()

For debugging synchronization issues, see Kernel Debugging.

Runtime Implementation of Communication Primitives

Mechanism Summary

Operation

Underlying mechanism

signal SET (P2P / MTE)

Ordinary scalar write to remote GM followed by a cacheline writeback instruction, not hardware atomic and not a doorbell

signal ADD (P2P)

Hardware atomic add combined with UB→GM DataCopy, performing read-modify-write at the write end

signal (cross-node RoCE)

RDMA write followed by quiet, or RDMA atomic

wait / wait_until

Local spin busy-wait, invalidating the cacheline before each read

barrier

Software algorithm (centralized pull), composed of signal_set and spin

UDMA

True doorbell: fill WQE/SGE into the SQ ring → write the doorbell register

signal / notify

The SET path consists of two steps:

  1. Ordinary scalar store *addr = val;

  2. dcci_cacheline(addr) — flush the data cache back to GM. The comment explicitly states that this is to “flush data cache to GM after signal to ensure it is visible to other ranks”.

The ADD path goes through UB: place the value in UB → set/wait for the MTE3 flag → enable hardware atomic-add mode → copy_ub2gm → disable atomic mode.

The mechanism of aclshmem_<T>_p (point write) is isomorphic to signal SET: aclshmem_ptr obtains the address, followed by a scalar store and cacheline writeback.

Wait Polling Mechanism

The six comparison predicates (_eq / _ne / _gt / _ge / _lt / _le) have the same form:

do { dcci_cacheline((__gm__ uint8_t *)sig_addr); } while (*sig_addr != cmp_val);

There is no hardware blocking or interrupt mechanism; correctness depends entirely on invalidating the L2 cacheline on every loop iteration. dcci_cacheline internally uses DataCacheCleanAndInvalid, with an empty __asm__ __volatile__("") before and after to prevent compiler optimization.

This also explains a characteristic of fine-grained synchronization: dl.wait has non-consuming semantics—it only polls and does not modify the signal value, so multiple waiters on the same slot are safe. The constraint falls on the reset side (which party is responsible for writing the signal back to 0).

barrier_all Implementation

aclshmem_barrier_all() ultimately uses a centralized pull barrier:

  1. Each PE writes only its own flag (signal_set);

  2. It then directly spin-reads the peer’s flag via aclshmem_ptr(sync_pool, remote_pe)—the peer does not need to actively push;

  3. Work is distributed across the available Vector Cores, with O(N/K) complexity.

A key detail preventing deadlock: The barrier-specific wait predicate accepts not only == cmp_val but also == cmp_val + 1. The comment explains that this is “in case when peer pe enters next barrier”, preventing a fast PE that has already entered the next barrier from making a slow PE wait forever.

UDMA

UDMA is not a new set of dl.* operations; it is a transport engine selected during host initialization. The same libshmem_device.putmem() call takes completely different underlying paths under MTE and UDMA configurations.

The UDMA device-side send flow is a true doorbell mechanism:

QP fetch → slot calculation → fill SQE → fill SGE → dcci_cachelines → ring the doorbell register

Best Practices

  1. Minimize Synchronization Granularity: Use per-tile signals instead of global barriers

  2. Overlap Aggressively: Start computation as soon as any tile is ready, don’t wait for all tiles

  3. Use Threadblock Swizzling: Arrange work so local data is processed first

  4. Batch Signals When Possible: Multiple tiles can share a signal if they’re always accessed together

  5. Consider Communication Topology: On Ascend clusters, leverage the network topology (NVLink, RoCE) for efficient data placement

  6. Profile and Tune: Use profiling tools to identify overlap efficiency and synchronization bottlenecks

  7. Use ACLSHMEM Efficiently:

    • Prefer non-blocking operations (*_nbi) when possible

    • Use quiet() or fence() judiciously to enforce ordering

    • Consider using team-based operations for subset communication

  8. Test with Barrier Sync First: When debugging, use --enable-hivm-inject-barrier-all-sync=true to rule out fine-grained synchronization bugs (see Kernel Debugging)

Example: Ring AllReduce with Overlap

Here’s a conceptual example of how tile-centric semantics enable efficient ring AllReduce:

import triton
import triton.language as tl
import triton_dist.language as tdl

@triton.jit
def ring_allreduce_kernel(
    local_data_ptr,
    remote_data_ptr,
    signal_ptr,
    rank: tl.constexpr,
    world_size: tl.constexpr,
    num_tiles: tl.constexpr,
):
    pid = tl.program_id(0)

    # Phase 1: Reduce-Scatter (send to next, receive from prev)
    for step in range(world_size - 1):
        send_rank = (rank + 1) % world_size
        recv_rank = (rank - 1 + world_size) % world_size

        # Determine which tile to send/receive in this step
        tile_id = (rank - step + world_size) % world_size

        if pid == tile_id:
            # Send tile to next rank
            remote_ptr = tdl.symm_at(remote_data_ptr, send_rank)
            data = tl.load(local_data_ptr + tile_id * TILE_SIZE)
            tl.store(remote_ptr + tile_id * TILE_SIZE, data)
            tdl.notify(signal_ptr, send_rank, signal=step, sig_op="set")

        # Wait for incoming tile from previous rank
        token = tdl.wait(signal_ptr + step, 1, "sys", "acquire")
        recv_ptr = tdl.consume_token(remote_data_ptr, token)
        incoming_data = tl.load(recv_ptr + tile_id * TILE_SIZE)

        # Reduce with local data
        local_data = tl.load(local_data_ptr + tile_id * TILE_SIZE)
        reduced = local_data + incoming_data
        tl.store(local_data_ptr + tile_id * TILE_SIZE, reduced)

    # Phase 2: AllGather (propagate reduced tiles)
    for step in range(world_size - 1):
        send_rank = (rank + 1) % world_size
        recv_rank = (rank - 1 + world_size) % world_size

        # Determine which tile to send in this step
        # Each rank sends the tile it just finished reducing
        send_tile_id = (rank - step + world_size) % world_size

        # Determine which tile to receive in this step
        recv_tile_id = (rank - step - 1 + world_size) % world_size

        if pid == send_tile_id:
            # Send reduced tile to next rank
            remote_ptr = tdl.symm_at(remote_data_ptr, send_rank)
            data = tl.load(local_data_ptr + send_tile_id * TILE_SIZE)
            tl.store(remote_ptr + send_tile_id * TILE_SIZE, data)
            tdl.notify(signal_ptr, send_rank, signal=step + world_size - 1, sig_op="set")

        # Wait for incoming tile from previous rank
        token = tdl.wait(signal_ptr + step + world_size - 1, 1, "sys", "acquire")
        recv_ptr = tdl.consume_token(remote_data_ptr, token)
        incoming_data = tl.load(recv_ptr + recv_tile_id * TILE_SIZE)

        # Store the received reduced tile
        tl.store(local_data_ptr + recv_tile_id * TILE_SIZE, incoming_data)

Key Points:

  • Each tile is independently synchronized with wait()/notify()

  • Tiles can be reduced as soon as they arrive (no barrier needed)

  • Uses consume_token() to ensure correct load ordering

  • Overlaps communication (sending next tile) with computation (reducing current tile)

References