Triton is a compiler and language for writing performance-portable AI kernels with a productive blocked programming model. This FBTriton repository is Meta's development branch, where compiler-managed features such as automatic warp specialization are developed alongside lower-level TLX kernels that provide an expert-written performance reference.
The central goal is to keep ordinary Triton kernels productive while allowing the compiler to discover schedules that use modern GPU pipelines efficiently.
Automatic warp specialization
Warp specialization assigns different groups of warps to different roles inside one kernel. A data partition can issue asynchronous loads, compute partitions can drive tensor cores, and epilogue or correction partitions can run concurrently when dependencies allow.
Specialization helps by:
Triton's AutoWS pipeline makes these decisions at JIT compilation time rather than requiring every partition and barrier to be handwritten.
The compiler pipeline today
The current design separates decision-making passes from the transformations that materialize those decisions:
| Stage | Responsibility |
|---|---|
| Data partitioning | Expose independent GEMMs or operation chains that can execute concurrently |
| Loop scheduler | Build a software-pipeline schedule and attach scheduling decisions to the IR |
| Partition scheduler | Assign operations to data, compute, correction, reduction, or epilogue warp partitions |
| Buffer creation | Identify cross-partition values and choose SMEM or TMEM communication buffers |
| Memory planner | Choose buffer depth and reuse based on liveness, dependencies, and hardware limits |
| Code partitioner | Create channels, barriers, buffer rotation, and physically outlined warp-group regions |
Partition scheduling
The partition scheduler propagates task decisions as operation attributes. Current strategies recognize roles such as TMA data movement, tensor-core computation, softmax correction, reductions, and epilogue stores. Once operations are assigned, cross-partition SSA dependencies become explicit communication channels.
Software pipelining
The loop scheduler reorders independent chains from adjacent iterations to maximize producer-consumer distance. In attention, for example, a dot product from one iteration can overlap softmax and value accumulation from another. Data partitioning gives the scheduler more independent work to interleave.
Channels and memory planning
A channel represents both the storage and synchronization needed to move a value between partitions. The memory planner decides:
Code partitioning
The code partitioner turns abstract channels into executable regions. It derives rotating buffer indices and barrier phases from loop iteration counts, reuses hardware completion barriers where possible, and performs local instruction ordering to reduce live ranges and register pressure.
TLX and AutoWS
TLX and AutoWS solve related problems at different levels:
| TLX | Triton AutoWS |
|---|---|
| Kernel author explicitly defines warp-group tasks and synchronization | Compiler derives partitions, communication, and synchronization from Triton IR |
| Best for expert control, new hardware mechanisms, and hand-tuned reference kernels | Best for applying warp specialization broadly without rewriting kernels in a lower-level DSL |
| Exposes the intended pipeline directly in Python | Keeps the source close to ordinary Triton and specializes during compilation |
The TLX documentation covers the explicit programming model used to validate and refine many of these compiler strategies.
Roadmap
Profile-guided partition scheduling
Near-term work uses measured operation and region costs to rank partition choices. Communication overhead between partitions must be weighed against the concurrency gained by separating them.
Model-based global optimization
Scheduling, partitioning, synchronization, layouts, and memory planning form a joint combinatorial search problem. The roadmap calls for hardware cost models—specified statically or learned from profiling—to guide a global planner that can choose software-pipeline schedules, channel depths, buffer reuse, partitions, and memory placement together.
Kernel fusion and megakernels
IR-level fusion could combine handwritten and generated Triton kernels at tile granularity. The long-term goal is to optimize dependent and independent kernels together, reducing launch overhead and global-memory traffic while retaining provable correctness for algorithmic transformations.
Broader kernel and hardware coverage
Current AutoWS focuses on Hopper and Blackwell. Continued work aims to generalize beyond attention-style independent dot chains, reduce target-specific heuristics, and make measured hardware models the main ingredient needed to support new architectures.
Read the design article
This page is a condensed guide to the published design. See Warp Specialization in Triton: Design and Roadmap for the full discussion.