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. .. code-block:: python # 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 .. code-block:: python # 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 :doc:`shmem_host` and :doc:`shmem_device`). Kernel Design Patterns ---------------------- AllGather + GEMM Overlapping ~~~~~~~~~~~~~~~~~~~~~~~~~~~~ .. code-block:: text 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 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ .. code-block:: text 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**: .. code-block:: python # 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: .. code-block:: python # 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 :doc:`shmem_host`) - Access remote memory using ``remote_ptr()`` in kernels (see :doc:`shmem_device`) - 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 :doc:`../developer-guide/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 :doc:`../developer-guide/kernel_debugging`) Example: Ring AllReduce with Overlap ------------------------------------- Here's a conceptual example of how tile-centric semantics enable efficient ring AllReduce: .. code-block:: python 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 ---------- - `TileLink: Generating Efficient Compute-Communication Overlapping Kernels using Tile-Centric Primitives (MLSys 2025) `_ - `Triton-distributed: Programming Overlapping Kernels on Distributed AI Systems with the Triton Compiler `_ - :doc:`triton_dist_language` — API reference for distributed primitives - :doc:`shmem_device` — ACLSHMEM device-side operations - :doc:`shmem_host` — ACLSHMEM host-side initialization and memory management