Triton: The Compiler That Pretends to Be a Library

Triton is a compiler with a Python frontend. The @triton.jit decorator does not decorate a function. It parses the function’s AST, runs it through an MLIR pipeline, and emits a GPU binary. The Python function never runs as Python.

This matters because it sets what Triton can and cannot do. It tiles a GEMM for you, stages data through shared memory, maps blocks to warps, and picks tensor core instructions. You never write a line of PTX or manage a thread index. But you also cannot reach past the abstraction when it gets in the way.

This post walks through how Triton compiles code, what the compiler decides on your behalf, where the abstraction falls short, and where Triton fits in the ML systems stack.

The Programming Model: Blocks, Not Threads

In CUDA you write code for one thread and the hardware runs it across thousands. Triton flips this. You write code for a block of data and the compiler splits it across threads.

import triton
import triton.language as tl

@triton.jit
def softmax_kernel(
    input_ptr, output_ptr,
    n_cols,
    BLOCK_SIZE: tl.constexpr,
):
    # 1. Each "program" handles one row. No thread indexing.
    row_idx = tl.program_id(0)
    col_offsets = tl.arange(0, BLOCK_SIZE)
    mask = col_offsets < n_cols

    # 2. Load an entire row as a block. The compiler picks
    #    coalescing, vector width, and predication.
    row = tl.load(input_ptr + row_idx * n_cols + col_offsets, mask=mask, other=-float('inf'))

    # 3. Block-level reductions. The compiler turns these into
    #    warp shuffles and shared memory reductions.
    row_max = tl.max(row, axis=0)
    numerator = tl.exp(row - row_max)
    denominator = tl.sum(numerator, axis=0)
    result = numerator / denominator

    tl.store(output_ptr + row_idx * n_cols + col_offsets, result, mask=mask)

The key operations:

• tl.load / tl.store: Tile-level memory access. The compiler turns these into coalesced global memory instructions with predication from the mask argument.
• tl.dot: Block-level matrix multiply. The compiler emits mma.sync (Ampere), wgmma (Hopper), or tcgen05.mma (Blackwell) depending on the target.
• tl.max, tl.sum: Block-level reductions. The compiler lowers these to warp shuffles (SHFL.BFLY) or shared memory reductions depending on block size.
• tl.constexpr: Compile-time constants. BLOCK_SIZE gets baked into the binary. Different values produce different kernels — hence the autotuning step.

You never write threadIdx.x. There is no __syncthreads(). There is no shared memory declaration. The compiler handles all of it. That is both the selling point and the ceiling.

The Compilation Pipeline

Triton lowers code through four IRs:¹

graph LR
    PY["Python AST"] --> TTIR["Triton IR<br/>(TTIR)"]
    TTIR --> TTGIR["Triton GPU IR<br/>(TTGIR)"]
    TTGIR --> LLVM["LLVM IR"]
    LLVM --> PTX["PTX"]
    PTX --> CUBIN["cubin"]

    style PY fill:#fce4ec,color:#1a1a1a
    style TTIR fill:#e8eaf6,color:#1a1a1a
    style TTGIR fill:#e0f2f1,color:#1a1a1a
    style LLVM fill:#fff3e0,color:#1a1a1a
    style PTX fill:#f3e5f5,color:#1a1a1a
    style CUBIN fill:#e8f5e9,color:#1a1a1a

Stage 1: Python AST → Triton IR (TTIR)

On first call with concrete arguments, Triton parses the Python AST and traces it into TTIR — a hardware-independent MLIR dialect (tt namespace). Operations stay abstract here: tl.load becomes tt.load, tl.dot becomes tt.dot. Standard compiler passes run — constant folding, CSE, dead code removal.

TTIR knows nothing about threads, warps, or shared memory. It works on tensors whose shapes come from the constexpr parameters.

Stage 2: Triton IR → Triton GPU IR (TTGIR)

This is where the compiler makes its hard calls. TTGIR adds GPU-specific structure:

• Thread-to-data mapping. The compiler decides how to spread a block across threads. A BLOCK_SIZE=128 vector might become 4 elements per thread across 32 threads (one warp), or 2 per thread across 64 (two warps).
• Shared memory allocation. When a tl.dot operand gets reused (say, in a GEMM inner loop), TTGIR inserts shared memory allocation and async_copy ops to stage data from global memory.
• Layout propagation. TTGIR tracks tensor layouts — blocked, shared, slice, dot-operand — and inserts conversions when an op needs a layout its input does not provide. A tl.dot needs operands in a specific “dot-operand” layout that matches the hardware MMA instruction.
• Software pipelining. For loops with known trip counts, TTGIR overlaps memory loads from iteration N+1 with compute from iteration N.

This stage is where most of Triton’s value lives. It is also where most of its performance bugs come from. A few real examples from the Triton issue tracker:

• The TritonGPUCoalesce pass sometimes picks a default blocked layout that does not match the downstream tl.dot layout. The compiler then inserts a convert_layout inside the inner loop — each one going through shared memory — and the kernel slows down 2–5x. Manually forcing a better layout recovers the performance. (#6206)²
• A change to the convert_layout swizzling algorithm caused an 18% TFLOPs regression in flex_attention on B200 hardware. The fix: fold the layout conversion for TMEM stores so the swizzle does not introduce extra shared memory traffic. (#8328)
• Layout propagation through ReshapeOp and DotOp has caused regressions in memory-bound kernels across multiple releases (v3.4–v3.6). The community is addressing this with a rewrite of the layout system using “linear layouts” — modeling tensor layouts as linear algebra over GF(2) — to replace the case-by-case heuristics that keep breaking. (#9640)

The pattern is consistent: TTGIR layout decisions are the single biggest lever on Triton kernel performance, and the compiler does not always get them right. Diagnosing it means reading TTGIR dumps.

Stage 3: TTGIR → LLVM IR

The GPU-specific MLIR gets lowered to plain LLVM IR. By now the tile ops have been broken into per-thread scalar ops. LLVM handles register allocation, instruction selection, and scheduling. The target is nvptx64.

Stage 4: LLVM IR → PTX → cubin

LLVM’s NVPTX backend emits PTX. Triton then calls ptxas to produce the cubin. The binary gets cached on disk, keyed by a hash of the source, the constexpr values, and the target architecture. Later calls with the same inputs skip the whole pipeline and load the cached binary.

What the IR actually looks like

The pipeline description above is abstract. Here is what it looks like concretely. Take a simpler kernel — a vector add — and trace it through TTIR and TTGIR.

The Python source:

@triton.jit
def add_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offs < n
    x = tl.load(x_ptr + offs, mask=mask)
    y = tl.load(y_ptr + offs, mask=mask)
    tl.store(out_ptr + offs, x + y, mask=mask)

TTIR (simplified, from TRITON_KERNEL_DUMP=1):

// 1. Everything is tensor-typed. No threads, no warps.
// 2. tt.load and tt.store operate on tensor<1024xf32> — the full block.
// 3. The mask is a tensor<1024xi1>, not a scalar branch.
module {
  tt.func @add_kernel(%x: !tt.ptr<f32>, %y: !tt.ptr<f32>,
                       %out: !tt.ptr<f32>, %n: i32) {
    %pid  = tt.get_program_id {axis = 0} : i32
    %offs = tt.make_range {start = 0, end = 1024} : tensor<1024xi32>
    %base = tt.splat %pid : tensor<1024xi32>
    %idx  = arith.addi %base, %offs : tensor<1024xi32>
    %mask = arith.cmpi slt, %idx, %n_splat : tensor<1024xi1>
    %xv   = tt.load %x_ptrs, %mask : tensor<1024xf32>   // ← block load
    %yv   = tt.load %y_ptrs, %mask : tensor<1024xf32>
    %sum  = arith.addf %xv, %yv : tensor<1024xf32>
    tt.store %out_ptrs, %sum, %mask : tensor<1024xf32>   // ← block store
    tt.return
  }
}

Notice: no shared memory, no thread indices, no layout annotations. TTIR is pure math on tensors.

TTGIR (simplified):

// 1. Tensors now carry layout attributes: #blocked<sizePerThread=[4],
//    threadsPerWarp=[32], warpsPerCTA=[8]>.
//    This means: 4 elements per thread, 32 threads per warp, 8 warps.
//    Total = 4 × 32 × 8 = 1024 elements. That is the full BLOCK.
// 2. tt.load becomes ttg.load — same op, but now the compiler knows
//    which thread loads which 4 elements.
// 3. No layout conversions needed here (element-wise add).
//    For tl.dot, you would see ttg.convert_layout ops inserted
//    to reshape data into the dot-operand layout.
module {
  ttg.func @add_kernel(%x: !tt.ptr<f32>, ...) {
    // ... same index math, but tensors carry #blocked layout ...
    %xv = ttg.load %x_ptrs, %mask
           {layout = #blocked<sPerT=[4], tPerW=[32], wPerCTA=[8]>}
           : tensor<1024xf32, #blocked>
    %yv = ttg.load %y_ptrs, %mask ... : tensor<1024xf32, #blocked>
    %sum = arith.addf %xv, %yv : tensor<1024xf32, #blocked>
    ttg.store %out_ptrs, %sum, %mask : tensor<1024xf32, #blocked>
    ttg.return
  }
}

The key difference: every tensor now has a layout. For this element-wise add, the layouts match throughout, so no conversions are needed. For a GEMM, TTGIR would insert ttg.convert_layout ops between the load layout (#blocked) and the dot-operand layout (#dot_op) — and those conversions go through shared memory. That is where bank conflicts and performance bugs hide.

What the Compiler Decides for You

The point of Triton’s abstraction is that the compiler handles the decisions that eat most of a CUDA author’s time:

Decision	CUDA Programmer	Triton Compiler
Thread block dimensions	Manual (dim3)	Derived from BLOCK_SIZE constexprs
Shared memory size and layout	Manual (__shared__, bank conflict avoidance)	Automatic (TTGIR layout propagation)
Memory coalescing	Manual (stride analysis)	Automatic (vectorized load/store lowering)
Warp synchronization	Manual (__syncwarp, __syncthreads)	Automatic (barrier insertion in TTGIR)
Tensor core instruction selection	Manual (wmma / mma.sync / PTX)	Automatic (tl.dot → MMA/WGMMA)
Software pipelining	Manual (multi-stage buffering)	Automatic (TTGIR pipelining pass)
Boundary predication	Manual (if/else on thread index)	Automatic (mask on loads/stores)

For many workloads — fused elementwise ops, reductions, small-to-medium GEMMs, attention variants — these choices are good enough. “Good enough” here means within 10–20% of hand-tuned CUDA, at a fraction of the development time.

Where Triton Wins

Operator fusion

This is Triton’s best case. Take a sequence like GEMM → GeLU → Dropout → LayerNorm. In CUDA, each step is usually a separate kernel launch. Each launch reads from and writes to HBM. Bandwidth is the bottleneck, not compute.

A Triton kernel fuses the whole chain. Intermediate values stay in registers or shared memory and never touch HBM. For bandwidth-bound workloads — which covers most LLM inference at small batch sizes — fusion delivers 2–4x speedups over cuBLAS plus separate elementwise kernels. The Triton tutorials repository includes benchmarks for fused softmax and matrix multiplication that show these gains on A100 hardware.³

This is why PyTorch’s TorchInductor uses Triton as its default codegen backend.⁴ When torch.compile traces a model graph, it finds groups of ops it can fuse and emits Triton kernels for them.

Fast iteration

A Triton kernel runs 30–50 lines of Python. The same kernel in CUDA runs 200–500 lines of C++. When you are trying out a new attention variant or a quantized GEMM for a new model, that gap is the difference between trying an idea and skipping it.

Multi-backend portability

Because Triton’s pipeline is MLIR-based, it supports non-NVIDIA targets. The AMD ROCm backend lowers TTGIR to HIP and targets AMD’s MFMA instructions on MI300 hardware.⁵ The abstraction boundary at TTIR means the same kernel source can run on both vendors. The hardware-specific work happens in TTGIR lowering, not in user code.

A caveat on maturity: the AMD backend works and is in production at several companies, but it is not at parity with NVIDIA. Some TTGIR passes (notably software pipelining and certain layout optimizations) are tuned for NVIDIA’s memory hierarchy and warp size. On MI300, expect 10–30% lower performance than the same kernel on a comparable NVIDIA GPU for compute-bound workloads, though AMD is closing the gap with each ROCm release.

Where Triton Loses

No warp-level control

Triton hides warps. You cannot call __shfl_sync, __ballot_sync, or __match_any_sync. You cannot assign work to specific warps.

This hurts when your algorithm needs direct thread-to-thread communication: warp-level sorting, custom reduction trees, cooperative group patterns. If performance depends on controlling warp topology, Triton cannot express it.

Irregular memory access

Triton’s shared memory and layout passes are built for regular, tiled access. Scatter/gather workloads — graph neural networks, sparse attention, hash table lookups — have access patterns the compiler cannot predict at compile time. The result is either wasted shared memory (conservative allocation) or worse performance than plain CUDA.

Debugging

When a Triton kernel gives wrong results or runs slow, debugging is hard. Three IR transforms sit between your Python and the final PTX, so the source mapping is loose at best.

In CUDA, cuda-gdb, compute-sanitizer, and ncu (Nsight Compute) map straight to the source. Triton has no equivalent. Here is the workflow I use when something goes wrong:

Step 1: Dump the IR. Set TRITON_KERNEL_DUMP=1 to write the TTIR and TTGIR for every compiled kernel to /tmp/triton_dump/. For more detail, MLIR_ENABLE_DUMP=1 prints every pass to stderr.

Step 2: Look for convert_layout. In the TTGIR dump, search for ttg.convert_layout ops. Each one is a data reshape that usually goes through shared memory. If you see a convert_layout between two ops that should share a layout (e.g., two elementwise ops), something is forcing a layout change. That is your performance bug.

// This is the red flag. A convert_layout between a load and an add
// means the load produced #blocked but the add wants #blocked<different>.
// The fix is usually adjusting BLOCK dimensions so the layouts agree.
%xv = ttg.load ... : tensor<1024xf32, #blocked<sPerT=[4], ...>>
%xv_cvt = ttg.convert_layout %xv : ... -> tensor<1024xf32, #blocked<sPerT=[8], ...>>
%sum = arith.addf %xv_cvt, %yv : tensor<1024xf32, #blocked<sPerT=[8], ...>>

Step 3: Check register pressure. Run the compiled kernel through ncu (Nsight Compute). Look at the “Occupancy” section. If register count per thread is above 128, Triton allocated too many registers — often because a tile is too large or the compiler failed to pipeline a loop. Shrinking BLOCK_SIZE or splitting the kernel usually helps.

Step 4: Compare PTX. For correctness bugs, dump the PTX (TRITON_KERNEL_DUMP=1 writes it alongside the IR) and look for predication errors. A common mistake: the mask tensor has the wrong shape, so the compiler predicates the wrong lanes. The PTX will show @!p0 guards on loads/stores — check that they match your intent.

This workflow is not elegant. But it is the one that works today.

The Hopper and Blackwell problem

Each GPU generation ships features that need new abstractions, not just new instruction selection:

• Hopper added TMA for async multi-dimensional copies and WGMMA that needs 128 threads to issue cooperatively. Triton could not express either at launch. Support came later through experimental APIs and heuristics that detect GEMM patterns and insert TMA loads and WGMMA instructions. But the abstraction leaks: block size choices that work fine on Ampere can stop the compiler from picking the WGMMA path on Hopper.

• Blackwell adds TMEM, tcgen05 single-thread issue, and FP4 microscaling. The ISA changed in ways that matter. Triton’s backend needs a deep rework to target these features, because the thread-to-data mapping in TTGIR was built around the warp-level model that Blackwell has left behind.

This is the core tension in Triton’s design. Each generation breaks the abstraction the compiler tries to hold together. The team adds heuristics and special cases to patch it, but the abstraction gains weight without getting cleaner.

Common pitfalls

A few things that trip up most Triton authors:

1. Block sizes that are not powers of two. Triton’s layout math assumes power-of-two dimensions. Non-power-of-two BLOCK sizes compile but often produce poor vectorization. Use triton.next_power_of_2() and mask the excess.

2. Too-large tiles on Hopper. A BLOCK_M=256, BLOCK_N=256 tile that runs well on A100 might miss the WGMMA path on H100 because the compiler cannot map it to a supported WGMMA shape. Stick to tiles that divide cleanly into 64×64 or 128×128 warp-group chunks.

3. Dynamic control flow inside the kernel. if/else branches that depend on runtime values (not constexpr) cause warp divergence. Triton cannot optimize this away. Use tl.where for branchless selection instead.

4. Forgetting that autotuning is not optional. Triton’s performance is sensitive to BLOCK_SIZE, num_warps, and num_stages. A kernel that is 3x slower than CUDA might just need a @triton.autotune decorator with 5–10 configurations. The compiler does not pick good defaults for all shapes.

5. Assuming loads are free. tl.load with a large block and a sparse mask still issues memory transactions for the entire block. The mask only controls which lanes write to registers. If your access is truly sparse, Triton is the wrong tool.

Triton’s Place in the Stack

graph TD
    subgraph User["User-Facing"]
        PT["PyTorch<br/>torch.compile"]
        JAX["JAX/XLA"]
        TF["TensorFlow"]
    end

    subgraph Codegen["Code Generation"]
        IND["TorchInductor"]
        XLA["XLA Compiler"]
        TRI["Triton"]
    end

    subgraph Runtime["Runtime Libraries"]
        CUBLAS["cuBLAS"]
        CUDNN["cuDNN"]
        CUTLASS["CUTLASS"]
        FLASH["FlashAttention"]
    end

    subgraph HW["Hardware"]
        GPU["NVIDIA GPU<br/>(PTX/SASS)"]
        AMD["AMD GPU<br/>(GCN/CDNA)"]
    end

    PT --> IND
    IND --> TRI
    IND --> CUBLAS
    IND --> CUDNN
    JAX --> XLA
    TRI --> GPU
    TRI --> AMD
    CUBLAS --> GPU
    CUTLASS --> GPU
    FLASH --> GPU

    style User fill:#e8eaf6,color:#1a1a1a
    style Codegen fill:#e0f2f1,color:#1a1a1a
    style Runtime fill:#fff3e0,color:#1a1a1a
    style HW fill:#fce4ec,color:#1a1a1a

Triton fills a specific role: code generation for custom fused kernels. It does not replace cuBLAS for dense GEMMs (cuBLAS still wins for large square matrices). It does not replace CUTLASS for library-grade template code. What it replaces is writing 500-line CUDA kernels every time you need a fused op that cuBLAS does not ship.

In production:

• TorchInductor uses Triton for fused subgraphs and falls back to cuBLAS/cuDNN for standard ops.⁴
• vLLM uses Triton for PagedAttention, quantized GEMM kernels (AWQ, GPTQ), and custom activations.⁶
• FlashAttention ships both CUDA and Triton versions. The CUDA one is faster. The Triton one lets you experiment with attention variants (sliding window, block-sparse, grouped-query) quickly.

Picking the right tool

Workload	Best tool	Why
Standard dense GEMM	cuBLAS	Tuned by NVIDIA, autotuned per GPU
Custom fused operator	Triton	10x faster to write, within 10–20% of CUDA
Library-grade GEMM template	CUTLASS	Full control over tiling, pipelining, epilogue
Custom attention variant	Triton	Iterate in hours, not weeks
Sparse / irregular kernel	CUDA	Triton’s model does not fit
Multi-backend portability	Triton	Same source targets NVIDIA and AMD

The Abstraction Tax

You trade control for speed of development. Triton makes that trade clear: give up thread-level control, get block-level programming with automatic shared memory and tensor core use. For workloads that fit, this is a good deal.

But the tax grows each hardware generation. Each new GPU ships features — TMA, WGMMA, TMEM, tcgen05 — built for a lower level of control than Triton offers. The compiler team adds support, but the lag between hardware launch and Triton readiness is typically 6–12 months. In that window, only CUDA and CUTLASS can use the new features.

There is a deeper limit too. Triton generates code for a single GPU. It has no notion of multi-GPU communication (NCCL), host-device data movement, or memory spaces. A tl.load reads from device memory. There is no tl.load_from_host or tl.transfer. Where the data came from, which device runs the kernel, whether the pointer is even valid on this device — that is the caller’s problem.

For a single-GPU kernel, this is fine. For the multi-device world that real ML training lives in, Triton only covers the innermost loop. Placement, data movement, memory space management — all of that stays untyped, unchecked, and handled by Python framework code that no compiler can see or verify.⁷

References

---

Disclaimer: This article was generated using the Gemini 3.1 Pro and Claude Opus 4.8 models.

1. Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations. Tillet, P., Kung, H. T., Cox, D. MAPL 2019 / PLDI 2019. The original paper on block-level GPU programming and the Triton compilation pipeline. (Link) ↩︎

2. Triton Layout Conversion Issues. The TritonGPUCoalesce pass and convert_layout mechanism are a recurring source of performance regressions. Issues #6206, #8328, and #9640 document the pattern across multiple releases. ↩︎

3. Triton Tutorials: Fused Softmax and Matrix Multiplication. Triton project. Includes benchmarks comparing Triton kernels against cuBLAS and PyTorch native ops on A100. (Link) ↩︎

4. TorchInductor: A PyTorch-Native Compiler. PyTorch team, 2022. TorchInductor uses Triton as its default GPU codegen backend for fused operator graphs. (Link) ↩︎ ↩︎²

5. AMD ROCm Support for Triton. Triton community, 2024. The AMD backend lowers Triton GPU IR to HIP and targets MFMA instructions on MI300 hardware. (Link) ↩︎

6. vLLM: Easy, Fast, and Cheap LLM Serving. Kwon, W. et al. SOSP 2023. vLLM uses custom Triton kernels for PagedAttention and quantized serving. (Link) ↩︎

7. MLIR: A Compiler Infrastructure for the End of Moore’s Law. Lattner, C. et al. 2020. The multi-level IR framework that Triton’s pipeline is built on. (Link) ↩︎