ML-Triton: Multi-Level GPU Compiler

• ML-Triton is a multi-level compiler and language extension for GPU programming that decouples workgroup, warp, and SIMD levels to optimize dense ML kernels.
• It implements a hierarchical lowering pipeline mirroring modern GPU architecture, enabling precise tiling and hardware-aware optimizations that achieve within 5% of hand-optimized code.
• Key features include compiler hints and a warp-level API for explicit control over tiling and memory operations, making it versatile across architectures and DSLs.

ML-Triton is a multi-level compiler and language extension for Triton, a Python-based domain-specific language (DSL) for GPU programming. It implements a hierarchical compilation pipeline that closely mirrors the physical and logical structure of modern GPUs, enabling fine-grained performance optimizations and enhanced programmability for dense ML operations such as GEMM and multi-head attention. By decoupling workgroup, warp, and SIMD/intrinsic levels in both the compiler and programming model, ML-Triton achieves performance within 5% of hand-optimized expert code for critical ML kernels across architectures, while supporting advanced language features and explicit hardware-aware optimizations (Wang et al., 19 Mar 2025).

1. Motivation and Limitations of Flat-Lowering Compilers

Traditional Triton and similar DSLs lower directly from the workgroup (threadblock) abstraction to per-thread LLVM intermediate representation (IR) in a single monolithic step. This approach fails to expose the three-level hierarchy that characterizes modern GPUs:

• Workgroup (CTA/threadblock)
• Warp/subgroup (SIMD units operating in lock-step)
• Per-lane SIMD execution (hardware intrinsics such as DPAS/WMMA/MFMA)

The lack of explicit warp-level partitioning obscures opportunities for leveraging specialized hardware instructions (e.g., blocked load, blocked matrix-multiply accumulate), entangles compiler logic for inter- and intra-warp control, and forces kernel developers to write around abstractions to achieve optimal performance. ML-Triton corrects these deficits with a true multi-level lowering pipeline (Wang et al., 19 Mar 2025).

2. Multi-Level Lowering Pipeline: Stages and Representation

ML-Triton’s compiler pipeline consists of four stages, each corresponding to a hardware abstraction, and propagates a layout encoding throughout to ensure consistent tiling:

1. Triton IR (workgroup-level): High-level kernel written in Python, expressed in Triton-specific IR.
2. convert-triton-to-tritongpu-warp: Introduces a BlockedEncoding that captures user-specified workgroup and warp partitioning, recording tile shapes and their mapping to CTAs and warps.
3. distribute-to-warps: Slices each encoded tensor into warp-local fragments, computes per-warp tile start indices, and converts global operations (tt.dot, tt.load, tt.store) into warp-scoped instances.
4. match-target-size: Further decomposes warp-local tiles into sub-blocks aligned to hardware vector/matrix-instruction constraints (e.g., 8x16 for DPAS on Intel PVC), inserting extraction/insertion ops.
5. convert-tritongpu-to-llvm: Each local operation is lowered to a hardware intrinsic (SIMT or SIMD) suitable for the target device.

This pipeline enables a precise mapping of user and hardware constraints (workgroup sizes, warp counts, blocked layouts) to efficient emulator- or hardware-level code, providing a principled foundation for both correctness and performance (Wang et al., 19 Mar 2025).

3. Language Extensions: Compiler Hints and Warp-Level API

ML-Triton exposes both high-level and low-level extensions for user control:

• Compiler Hints: New tiling hints on operations such as tl.dot, e.g., tiling="horizontal", allow explicit control over warp-level tiling strategy (square, horizontal, vertical). This supports, for example, row-wise partitioning for FlashAttention-2 kernels.
• Warp-Level API: New constructs for:
  • Explicit per-warp control (@warp_level)
  • Accessing warp identifiers (tl.warp_id())
  • Allocating shared-local memory (tl.alloc)
  • Warp-scope reduction operations with control over cross-warp behavior (tl.reduce(..., cross_warp=Bool, dst_warps=Mask)).

These extensions allow concise expression of advanced usage such as warp-synchronous tiled loads, cross-warp reductions, and hardware-optimized attention kernels. Empirically, implementing paged attention with explicit warp logic required only ~10 additional lines compared to the pure workgroup style (Wang et al., 19 Mar 2025).

4. Hardware Mapping and Instruction Utilization

The explicit multi-level (workgroup, warp, SIMD/intrinsic) approach enables direct emission of specialized hardware instructions:

• Blocked Load: For example, 2DBlockRead.v64i16 to read a 32x32 fp16 tile with SIMD16.
• Blocked Matrix Multiply (DPAS): Direct dispatch of dpas instructions to hardware matrix units, provided the local tiles are of the aligned shape.

Thread-lane mapping formulas and tiling choices ensure that accumulator and fragment register allocation is tight and each lane performs a regular sub-block of the tile. This matches the optimal data movement and compute mapping required to saturate hardware (Wang et al., 19 Mar 2025).

5. Performance Evaluation and Empirical Benchmarks

On Intel Ponte Vecchio PVC (max 1550), ML-Triton achieves the following, measured as geometric mean throughput versus expert-tuned XeTLA kernels:

Benchmark	XeTLA (Gops/s)	ML-Triton (Gops/s)	Ratio (%)
GEMM (compute)	12,000	11,520	96
GEMM (memory)	8,300	7,802	94
FlashAttention-2 (64)	4,500	4,275	95
FlashAttention-2 (128)	4,200	4,050	96
Paged Attention	3,800	3,630	96

For all kernel categories (GEMM, memory-bound, attention), ML-Triton produces code that is within 5% of the best-in-class expert-tuned C++ and assembly (Wang et al., 19 Mar 2025).

6. Generalization Across Architectures and DSLs

The three-stage lowering and layout-encoding propagation are applicable to:

• NVIDIA (e.g., via WMMA), AMD (e.g., MFMA), and other hardware with explicit warp- or group-level intrinsics,
• DSLs built atop MLIR (e.g., TVM, Hidet) with an intermediate warp dialect,
• Any domain where hierarchical hardware is present (distinct load/compute units at multiple levels).

Combined with compiler hints and explicit APIs, this strikes a balance between productivity and near-peak performance, facilitating maintainable high-performance GPU kernel development for both research and production settings (Wang et al., 19 Mar 2025).

7. Significance and Impact

ML-Triton enables ML kernel authors to:

• Compose kernels in Python with high-level semantics,
• Deploy hardware-specific optimizations with minimal boilerplate,
• Realize performance that consistently approaches or matches expert code,
• Future-proof kernels against hardware evolution by decoupling logical tiling and codegen concerns,
• Share kernels across architectures with minimal adaptation.

Its approach is foundational for high-productivity, near-optimal code synthesis in research and production ML frameworks targeting diverse modern GPUs (Wang et al., 19 Mar 2025).