Technical Summary: Triton for MTIA

Problem Statement

The rapid diversification of machine learning workloads (from LLMs to recommendation systems) has driven the proliferation of custom AI accelerators, such as Meta's MTIA-2i. While these custom ASICs offer performance and cost benefits tailored to specific workloads, they introduce significant programming challenges. Unlike GPUs, which utilize a SIMT (Single Instruction, Multiple Threads) model with hardware-managed caches and thread schedulers, custom accelerators like MTIA-2i often employ distinct programming models. MTIA-2i, for instance, relies on asynchronous command issuance to fixed-function units (FFUs), software-managed on-chip memory via Circular Buffers (CBs), and explicit software pipelining to hide latency.

This architectural divergence creates a "programming model gap." Developers must hand-optimize kernels in low-level languages (e.g., C++) for each hardware generation, a process that is slow, error-prone, and hinders the rapid iteration required for modern model development. While Domain-Specific Languages (DSLs) like Triton have successfully abstracted GPU complexities, their viability on non-GPU architectures with fundamentally different execution and memory models had not been established at production scale.

Methodology

The authors present the first production-scale application of Triton on the MTIA-2i accelerator. The approach involves developing a new compiler backend, enhancing the TorchInductor integration, and introducing minimal language extensions to bridge the gap between high-level Triton syntax and MTIA's low-level constraints.

1. Compiler Architecture

The Triton-MTIA compiler is built on MLIR and follows a four-stage flow:

• Parsing: Leverages the upstream Triton parser to generate Triton IR (ttir), ensuring compatibility with the evolving Triton language.
• Middle-End (Mapping): Maps operations to MTIA hardware engines. It distinguishes between operations suitable for Fixed-Function Units (FFUs) like the Dot Product Engine (DPE) and SIMD Engine (SE), versus those requiring the RISC-V vector cores. Crucially, it analyzes memory access patterns to determine if they can be lowered to Direct Memory Access (DMA) transfers (which require structured access) or must fall back to scalar loads/stores.
• Backend (Lowering): Handles architecture-specific details, including tensor bufferization. Since Triton does not natively model on-chip memory hierarchies, the compiler allocates and manages Circular Buffers (CBs) in Local Memory (LS), enforcing read/write pointer dependencies to simulate FIFO behavior.
• Code Generation: Supports both LLVM IR and human-readable C++ generation, allowing developers to inspect and manually tune generated code before final linking.

2. Optimization Strategies

To achieve performance competitive with expert-tuned C++, the compiler employs three key strategies:

• Software Pipelining: The compiler utilizes the asynchronous nature of MTIA's RISC-V cores and FFUs. It performs "greedy CB sizing" to allocate the minimum necessary on-chip memory while ensuring the pipeline remains full. It also implements "effective shape propagation" to skip unnecessary computations and buffer initializations for masked data.
• Vector Codegen: For "long-tail" operations not supported by FFUs, the compiler leverages the RISC-V Vector Extensions (RVV). It optimizes vector-scalar operations by avoiding unnecessary broadcasts and fuses back-to-back vector loops to eliminate redundant memory transfers between LS and registers.
• Load Balancing: The compiler supports flexible PID-to-PE (Processing Element) mapping strategies (linear, round-robin, zigzag). For workloads with uneven distributions (e.g., causal attention masks), a zigzag distribution is used to balance the load across the 8x8 PE grid, yielding significant speedups over linear mapping.

3. Language Extensions

To address cases where automatic inference is insufficient, the authors introduced minimal extensions:

• Architecture Exposure: New primitives like tl.mtia.get_pe_id and tl.mtia.re_reduce allow kernel authors to explicitly query hardware topology and utilize the Reduction Engine network.
• Memory Reuse: The tl.extract_slice extension enables explicit fine-grained memory reuse, allowing authors to load data once into LS and reuse slices without re-fetching from main memory.
• Custom Types: Integration with C++ structures via ctypes.Structure allows Triton kernels to interact with existing MTIA software stack metadata (e.g., TensorViews).

4. Inductor Integration

The authors tailored the TorchInductor code generation for MTIA. This includes generating analyzable memory access patterns (tiling scalar indices) to ensure DMA lowering and expanding the use of Triton Templates for complex kernels like LayerNorm, enabling automatic fusion of activation functions.

Key Contributions

1. Production-Scale Deployment: The paper demonstrates that high-level DSLs can be successfully applied to custom accelerators with non-GPU programming models. Triton kernels were deployed in production across approximately 60 model types, covering 50% of layers and 47% of non-GEMM execution time.
2. Compiler and Language Innovations: The work introduces a new compiler backend that lowers Triton to MTIA, along with specific optimizations (software pipelining, CB management) and minimal language extensions that expose hardware features without sacrificing portability.
3. Performance Validation: The study provides empirical evidence that Triton-generated kernels achieve performance competitive with expert-tuned C++ implementations, often reaching over 80% of the architectural roofline for GEMM operations.

Results

• Performance: Triton GEMM kernels consistently achieve >80% of the architecture roofline, matching expert-tuned C++ kernels. Fused FlashAttention kernels show significant speedups, particularly for generative models with long sequence lengths.
• Development Velocity: Within a single quarter, the team tripled the number of model types utilizing Triton and doubled the runtime spent in Triton kernels on MTIA-2i.
• Coverage: The Inductor-Triton flow achieved a 94.5% success rate for inference and 94.7% for training across a set of internal models.
• Portability: 11 out of 16 ported GPU kernels were directly portable with minimal changes, demonstrating the viability of cross-architecture code reuse, provided undefined behaviors are avoided.

Significance

The paper argues that the diversity of ML hardware and models poses a significant barrier to developer productivity. By successfully bridging the programming model gap between ML frameworks (PyTorch), high-level DSLs (Triton), and custom accelerators (MTIA), the authors demonstrate that it is possible to maintain high development velocity without sacrificing performance. The work suggests that flexible DSLs, augmented with targeted language extensions and compiler optimizations, can serve as a unifying layer for heterogeneous AI hardware, enabling rapid innovation and efficient deployment at scale. The authors position this as a step toward a stronger community around standardized DSLs to lower the adoption barrier for future custom accelerators.