How We Built TLX: A Hardware-Native GPU Compiler for Production ML at Scale

Triton Low-Level Extensions (TLX) , is an extension to Triton that helps engineers optimize GPU performance while preserving Triton’s productive programming model. TLX introduces explicit warp-group and cluster-level orchestration, letting developers adopt new hardware capabilities without sacrificing velocity. TLX has demonstrated performance improvements in RecSys training, RecSys inference, and large-scale LLM pretraining. Modern GPU architectures increasingly rely on execution-oriented features such as warp specialization, asynchronous memory movement, tensor memory, and cluster-level execution to optimize their performance. But these same features also expose a growing gap between low-level programming models for control and high-level programming models for productivity. To bridge this gap, we built Triton Low-Level Extensions (TLX), an extension that gives engineers explicit control over GPU execution while preserving Triton’s productive programming model. Today, TLX powers production workloads across Meta, from AI training to large-scale LLM pretraining, enabling engineers to adopt new hardware capabilities quickly without sacrificing developer velocity. Across NVIDIA and AMD GPUs, TLX is pushing the frontier of GPU performance — combining state-of-the-art kernels with explicit execution orchestration and model-aware fusion to deliver end-to-end improvements beyond what standalone libraries can achieve. From AI Models to GPU Kernels To understand where TLX fits, it helps to trace the path from an AI model to the GPU instructions that ultimately execute it. At Meta, machine learning (ML) workloads are built on a common software stack centered around PyTorch. Model code is transformed through graph capture, optimization, compilation, and lowering before eventually becoming GPU kernels running across large accelerator fleets. Over the past several years, Triton has evolved from an experimental kernel language into a foundational component of this stack. A significant share of production GPU kernels — both compiler-generated and hand-optimized — are now implemented in Triton. Improvements at this layer can therefore propagate broadly across training and inference workloads, making kernel performance an increasingly important lever for improving end-to-end system efficiency. This shift has also changed who writes kernels. Kernel optimization is no longer the exclusive domain of GPU specialists. ML engineers routinely profile workloads, identify bottlenecks, and optimize the hottest kernels through autotuning, fusion, memory-layout improvements, or algorithmic restructuring. Triton’s productivity allows many of these iterations to happen in hours rather than days, enabling engineers to stay close to the model-development workflow. As a result, GPU kernel development has increasingly become part of the standard optimization toolkit for ML engineers. Triton already provides a productive foundation for many memory-bound workloads through fusion and compiler-managed memory-access optimizations. Building on these strengths, TLX addresses a complementary challenge: coordinating execution groups, asynchronous data movement, and shared resources on modern GPUs. The Problem: When Hardware Evolves Faster Than Programming Models While Triton dramatically improved developer productivity, modern GPU architectures have evolved in ways that place increasing emphasis on execution orchestration rather than computation alone. NVIDIA’s Hopper introduced warp specialization and the Tensor Memory Accelerator (TMA). NVIDIA’s Blackwell GPU expanded this trend further with Tensor Memory (TMEM), cluster-level execution, paired-CTA execution, and increasingly asynchronous execution mechanisms. These capabilities are not simply new instructions that can be hidden behind compiler optimizations, they fundamentally change how high-performance kernels are organized. For example, a modern FlashAttention kernel may dedicate one warp group to issuing TMA transfers while another executes tensor-core computation. A Blackwell GEMM kernel may coordinate multiple warp groups, TMEM allocations, and cluster resources simultaneously. In both cases, performance depends not only on the computation being performed, but also on how multiple execution agents cooperate. As we adopted Hopper and Blackwell across Meta’s ML infrastructure, we repeatedly encountered the same challenge: CUDA could express these execution structures, but doing so often required large amounts of architecture-specific orchestration code. Triton provided a dramatically more productive programming model, but many of the execution relationships required by modern hardware remained hidden behind compiler abstractions. The result was a growing compiler catch-up problem. Each new hardware generation introduced new execution capabilities, while programmers either waited for compiler support to arrive or dropped into increasingly specialized implementations. We realized that the challenge was not simply missing compiler features, it was the absence of an appropriate programming abstraction. The Missing Middle: Why SIMT and SIMB Weren’t Enough The capabilities that matter most on recent GPUs operate neither at thread granularity nor at block granularity. Warp specialization partitions work across cooperating warp groups. TMA introduces dedicated agents for memory movement. TMEM enables new communication patterns between execution groups. Cluster execution lets multiple CTAs cooperate within a shared execution context. Organizing a kernel around these features means reasoning about how several execution agents cooperate, communicate, and synchronize — a level of structure that sits between CUDA and Triton. This exposed what we came to view as a missing abstraction. CUDA’s SIMT model operates at too fine a granularity: It provides complete control, but requires developers to express every execution relationship manually. Triton’s blocked programming model operates at too coarse a granularity: It improves productivity, but turns many execution structures into implicit compiler decisions. Modern hardware increasingly requires something in between. We refer to this model as Multiple Instructions, Multiple Warps (MIMW). MIMW elevates the warp group to the fundamental unit of orchestration. Rather than reasoning about individual threads or entire blocks, programmers express how multiple execution agents cooperate, communicate, synchronize, and share resources. This is the natural level at which many Hopper and Blackwell optimizations are expressed, and it is the abstraction layer TLX implements. Figure 2 illustrates the evolution of GPU programming models. CUDA’s SIMT model exposed execution directly at thread granularity, giving developers complete control over scheduling, synchronization, and memory movement. Triton moved toward a SIMB model, raising the abstraction to blocked tensor computations and improving productivity. MIMW occupies the space between these two extremes. Rather than treating threads or tensor blocks as the primary execution abstraction, it elevates warp groups to first-class execution agents. The emergence of MIMW is not a software-design preference but a consequence of how hardware has evolved. Volta and Ampere exposed tensor cores through thread- and warp-level programming. Hopper made asynchronous data movement (enabled byTMA) and warp-group execution (enabled by WGMMA) central to high-performance kernels. Blackwell extends this trend further with TMEM, register reallocation, cluster execution, and paired-CTA MMA. Each generation pushes more cooperation across warp groups — and sometimes across CTAs — into territory that the programming model, not just the compiler, must address. Making MIMW practical requires expressing cooperating warp or wave groups as first-class execution tasks, with explicit control over shared-memory buffer access, synchronization, and handoff between them. TLX brings these capabiliti