How We Built TLX: A Hardware-Native GPU Compiler for Production ML at Scale
- Source
- engineering.fb.com
- Date

Context
Meta's TLX, short for Triton Low-Level Extensions, gives engineers direct control over how recent GPUs schedule their work without leaving Triton, the Python-based language in which a significant share of Meta's production GPU kernels are now written (Meta's engineering post). A kernel is the small program that performs one piece of a model's math on the GPU. Triton lets an engineer describe that math as operations on blocks of data and leaves the compiler to decide how the hardware carries it out. CUDA, the lower-level alternative, exposes every thread but requires large amounts of hardware-specific code. Meta says newer NVIDIA chips, Hopper and Blackwell, get much of their speed from groups of threads called warp groups doing different jobs at the same time, for example one group copying the next tile of data into fast on-chip memory while another runs the matrix multiplication. Triton hid those arrangements inside compiler decisions, so engineers either waited for compiler support or fell back to specialized code.
TLX keeps the computation in ordinary Triton and adds a second layer that says how it is organized. An engineer assigns roles to warp groups, such as a producer that loads data, a consumer that runs tensor-core math and an epilogue that writes results, and coordinates them with lightweight barriers. TLX also exposes the Tensor Memory Accelerator, a hardware unit for asynchronous memory copies, along with Blackwell's tensor memory and clusters of thread blocks that can share on-chip memory. Unlike CUDA, it does not require the engineer to specify how data is laid out in memory; the compiler infers a layout unless someone overrides it. Meta says this makes it much easier to fuse several operations into one kernel, which saves the memory traffic between them. TLX currently targets NVIDIA Hopper and Blackwell and AMD's MI300 and MI350x.
The evidence is Meta's own production use over several months. Meta reports a TLX kernel with up to a 2x forward speedup in training its Generative Ads Model, a 4x speedup for one recommendation- kernel on the H100, and up to a two-thirds cut in compute-intensive latency for recommendation models built together with their kernels. A TLX kernel for layers also runs continuously in Meta's large language model . Meta says TLX typically matches hand-tuned CUDA while letting engineers iterate faster, citing weeks rather than months for one project. Several of these gains came from redesigning models and kernels together, so they are not attributable to TLX alone. The post does not link a code repository or say how outside developers can use TLX; Meta's next steps include making it a backend for torch.compile so common PyTorch code gets TLX kernels without manual annotation.
- Meta argues CUDA's thread-level SIMT model is too fine-grained and Triton's blocked model too coarse for Hopper and Blackwell features like TMA and TMEM, so TLX implements a middle model it calls MIMW, where warp groups are the unit programmers orchestrate.
- In TLX, Triton still expresses the computation while annotations assign producer, consumer and epilogue warp groups and connect them with mbarriers, so a matmul overlaps TMA loads, tensor-core math and stores without leaving Python or Triton autotuning.
- Because TLX infers layouts unless a programmer overrides them, Meta says operations can be fused without manually reconciling tiling and swizzling between stages, the layout bookkeeping that tends to dominate fusion work in CUDA-style workflows.
- Meta reports TLX in production: 20-25% end-to-end MFU training its GEM ads model on thousands of Blackwell GPUs, a 4x speedup for a RecSys compression kernel on H100, and up to two-thirds lower compute-intensive latency for co-designed inference models.
- Meta says its Triton kernel count grew about 3.8x versus 2.0x for CUDA over 10 months, that TLX typically matches hand-tuned CUDA performance, and that its in-kernel broadcast work took weeks with TLX rather than months.
- attention — The mechanism that lets a model weigh which earlier words matter for the word it's currently processing — the core operation of a transformer.
- inference — Running a trained model to get answers — the phase where AI is actually used, as opposed to trained.
- mixture-of-experts — A model built from many specialist sub-networks where only a few activate per token, giving big-model capability at small-model running cost.
- pretraining — The first, biggest phase of building a model: training it on enormous amounts of text so it learns language, facts, and reasoning in general.
- articleLiger Kernel Open Source Ecosystem For Efficient Llm TrainingLinkedIn Engineering
articleRealisticTritonBench: A Benchmark for Triton-Kernel Generation in Real-World AI FrameworksJinjun Huang, Zhongzhen Wen, Tongtong Xu, Meng Yan, Xin Xia, Zhongxin Liu
articleFrom CUDA to MLX: How K-Search Brings Decades of Kernel Expertise to Apple Siliconbair.berkeley.edu
Checking sign-in…
Loading comments…