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 producer_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 consumer_kernel(...):
    # Wait for tile to be ready
    token = tdl.wait(signal_ptr + tile_id, 1)
    # 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 + 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 + 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 + 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)
data = tl.load(data_ptr)  # BUG: might execute before wait!

# With consume_token: explicit dependency
token = tdl.wait(signal_ptr, 1)
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

  • UB (Unified Buffer): A local scratchpad/register-like memory used by the Vector Core (AIV) for element-wise, reduction, and vector operations. Size: 192 KB on Atlas 800T/I A2 series

  • L0 Buffers: Specialized, ultra-fast input and accumulation buffers dedicated to the Cube Core (AIC) for matrix multiplication:

    • L0A: Input buffer A for matrix multiplication

    • L0B: Input buffer B for matrix multiplication

    • L0C: Accumulation buffer for matrix multiplication results

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

Architecture Overview:

Each AI Core consists of one Cube Core (AIC) paired with two Vector Cores (AIV):

  • Cube Core (AIC): Executes tl.dot operations using L0A, L0B, L0C buffers

  • 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 A2 series, further reduced with doublebuffer enabled)

  • Tensor tail axis must be aligned: 32 bytes for Vector operations, 512 bytes for Cube+Vector operations

  • 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.

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 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)
        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)
        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