Technical Summary: Triton for MTIA

Problem Statement

The rapid proliferation of custom AI accelerators, such as Meta's MTIA-2i, has created a significant programming model gap. Unlike GPUs, which utilize a SIMT (Single Instruction, Multiple Threads) model with hardware-managed caches and thread schedulers, custom ASICs like MTIA-2i often employ distinct architectures featuring fixed-function units (FFUs), software-managed on-chip memory (Circular Buffers), and asynchronous command issuance.

This architectural divergence forces developers to hand-optimize kernels in low-level languages (e.g., C++) for each hardware generation, a process that is labor-intensive, error-prone, and hinders rapid iteration on new model architectures. While Domain-Specific Languages (DSLs) like Triton have successfully abstracted GPU complexity, their viability on non-GPU accelerators with fundamentally different programming models had not been established at production scale. The core challenge is to determine if a high-level DSL can support the full range of ML accelerators in modern datacenters without sacrificing developer productivity or performance.

Methodology

The authors developed a production-scale application of Triton on Meta's second-generation custom ASIC, MTIA-2i. The methodology involved three primary technical pillars:

1. A New Triton Compiler Backend for MTIA-2i

Since the existing open-source Triton compiler targets GPU-style programming models, a new backend was built from the ground up using MLIR. The compilation flow includes:

• Parsing: Leveraging the upstream Triton parser to generate Triton IR (ttir), enhanced to support MTIA-specific language extensions.
• Middle-End Mapping: Mapping operations to MTIA hardware engines. The compiler 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. It specifically analyzes memory access patterns to determine if they can be lowered to high-bandwidth DMA transfers (requiring structured access) or must fall back to scalar loads/stores.
• Backend Optimization: Implementing software pipelining to hide memory and computation latencies. This involves:
  • Effective Shape Propagation: Analyzing mask-based boundary conditions to skip unnecessary buffer initialization and computation.
  • Greedy Circular Buffer (CB) Sizing: Allocating on-chip Local Memory (LS) to maximize pipeline depth while fitting within hardware constraints.
  • Instruction Scheduling: Distributing commands between the two RISC-V cores per PE to avoid bottlenecks, utilizing a static analysis algorithm to synchronize shared CBs.
• Code Generation: Supporting both LLVM IR and human-readable C++ generation to facilitate debugging and manual optimization.

2. Integration with TorchInductor

The authors tailored the TorchInductor code generation flow for MTIA to improve operator coverage and performance:

• Strided Memory Access: Modified Inductor to generate pointer arithmetic patterns that are analyzable by the Triton-MTIA compiler, ensuring successful lowering to DMAs.
• Triton Templates: Expanded the use of Triton Templates beyond GEMM to include complex kernels like LayerNorm, allowing for automatic fusion of operations and elimination of intermediate tensor storage.
• Resource Modeling: Enhanced Inductor to model MTIA's specific memory hierarchy (15 CBs, 384 KB LS per PE) to prevent register spilling and optimize fusion decisions.

3. Language Extensions

To bridge the gap between Triton's high-level abstraction and MTIA's specific requirements, minimal language extensions were introduced:

• Architecture Exposure: Added tl.mtia.get_pe_id and tl.mtia.get_pe_grid_size to allow explicit control over Processing Element (PE) mapping, crucial for topology-aware communication. Added tl.mtia.re_reduce to leverage hardware reduction engines.
• Fine-Grained Memory Reuse: Introduced tl.extract_slice to explicitly express data reuse patterns within Circular Buffers, allowing the compiler to avoid redundant loads.
• Custom Types: Enabled the use of ctypes.Structure to define custom data types (e.g., TensorViews), facilitating integration with the existing C++-based MTIA software stack.
• Libdevice: Adapted the libdevice mechanism to allow seamless execution of GPU Triton kernels while applying MTIA-specific optimizations (e.g., lowering to FFUs vs. RISC-V).

Key Results

The authors successfully deployed both manually written and Inductor-generated Triton kernels in production across approximately 60 different model types.

• Performance: Triton-MTIA kernels achieved performance competitive with expert-tuned C++ implementations.
  • GEMM: Triton kernels consistently achieved over 80% of the architecture roofline, matching expert-tuned C++ performance.
  • FlashAttention: Fused kernels showed significant speedups, particularly for generative models with long sequence lengths.
  • Long-Tail Kernels: 17 production long-tail kernels performed on par with baseline C++ solutions, with the added benefit of rapid development for new scenarios.
  • Inductor Fusion: Automatically fused kernels (e.g., LayerNorm-Sigmoid) were competitive with handwritten baselines. In a recommendation model training scenario, Inductor-generated kernels were 31% faster than the handwritten kernels they replaced, contributing to a 4% improvement in Queries Per Second (QPS).
• Coverage: The Triton-MTIA flow accounted for 50% of layers and 47% of non-GEMM execution time across the deployed models.
• Development Velocity: Within a single quarter, the number of model types employing Triton tripled, and the runtime spent in Triton kernels doubled.

Significance and Claims

The paper claims to present the first production-scale application of Triton on a custom ML accelerator with a programming model distinct from GPUs.

The primary significance lies in demonstrating that high-level DSLs can effectively bridge the gap between ML frameworks, kernels, and custom hardware. The authors argue that:

1. Productivity vs. Performance: It is possible to achieve high performance on highly customized accelerators without reverting to low-level, hand-optimized code, thereby accelerating the development velocity of new models.
2. Portability: A significant portion of existing GPU Triton kernels can be ported to MTIA with minimal changes, provided developers avoid undefined behaviors and adhere to specific coding styles (e.g., ensuring structured memory access for DMA).
3. Scalability: The integration of Inductor with Triton on custom hardware enables automatic operator coverage and optimization, which is essential for supporting the growing diversity of ML models.

The work suggests that composable DSLs, which offer a baseline of portability with optional, architecture-specific extensions, are a viable path forward for the heterogeneous landscape of AI accelerators.