Triton — GPU Programming for Neural Networks

1. Why Triton?

Deep Neural Networks (DNNs) demand immense computational power, making GPUs the accelerators of choice. However, writing efficient GPU code directly in CUDA or OpenCL is notoriously difficult, requiring deep expertise to manage parallelism and memory hierarchies effectively. Existing solutions like polyhedral compilers or scheduling languages often struggle to match the performance of hand-tuned vendor libraries (like cuBLAS/cuDNN) or lack flexibility for novel operations, especially sparse ones.

Triton offers a different path. It’s a Python-based language and compiler designed to bridge this gap.

1.1 The Problem

GPU libraries like cuBLAS/cuDNN are fast but rigid; CUDA/OpenCL are flexible but fragile and verbose. Most DL researchers need that middle ground: custom ops that are as easy as Python yet as fast as handcrafted CUDA.

1.2 The Triton Idea

Triton revives an old HPC trope — Single Program, Multiple Data (SPMD) — but flips the casting: instead of a scalar program running on many threads, you write a blocked program that operates on an entire tile of data. The compiler then sprinkles the tile across threads, warps, tensor‑cores, and cache lines.

CUDA Model: Scalar program, blocked threads.

// Scalar Program, Blocked Threads
// Traditional CUDA (Conceptual): Many threads, each computing one output element.
for(m = 0..M) for(n = 0..N) { C[m,n] = dot(A[m,:], B[:,n]); }

Triton Model: Blocked program, scalar threads

# Blocked Program, Scalar Threads 
# Triton (Conceptual): Fewer program instances, each computing a tile of the output. 
for(m = 0..M step MB) for(n = 0..N step NB) {
    # Operates on tiles: A, B
    C_tile = block_dot(A_tile, B_tile);
    C = C_tile;
}

This block‑structured IR gives Triton two killer properties:

• Performance: Near‑peak TFLOPS on A100/H100/MI300 without manual shared mem juggling.
• Flexibility: Handles irregular or sparse iteration spaces — something polyhedral and schedule‑based DSLs struggle with.

1.3 Why This Matters

This “Blocked Program” approach allows the Triton compiler to automatically:

• Optimize Memory Access: Coalesce loads/stores and manage data movement through fast shared memory.
• Leverage Hardware: Generate instructions for specialized units like Tensor Cores (NVIDIA) or MFMA (AMD).
• Simplify Programming: You write Python code operating on tiles; the compiler handles the low-level complexity.
• Increase Flexibility: Easily implement custom operations, including those with irregular memory access patterns (like sparse kernels), which are challenging for other systems.

The result? Code that feels like NumPy but performs like optimized libraries.

1.4 A 60‑Second Detour: What exactly is SPMD?

SPMD means every GPU thread executes the same instruction stream but over a different slice of the input. Picture a choir reading identical sheet music, each singer projecting to a distinct corner of the hall.

Traditional CUDA assigns one thread per singer — fine for scalar notes, painful for chords. Triton hands each singer a block of notes (say a 16×16 matrix tile). This ‘blocked SPMD’ style gives the compiler room to:

• fuse loads/stores for coalesced memory,
• stage data in shared memory without manual index algebra,
• emit tensor‑core MMA instructions automatically.

The result is code that sounds like NumPy but performs like hand‑tuned cuBLAS.

2. Installation

Triton 3 requires Python ≥3.8 and CUDA ≥11.8 (or ROCm ≥5.6).

pip install --upgrade triton  # identical on CUDA & ROCm

No nvcc, no environment variables—it bundles its own PTX/LLVM toolchain.

3. Your First Kernel: Vector Addition

Let’s implement element-wise vector addition (output = x + y). This demonstrates basic Triton concepts.

import torch
import triton
import triton.language as tl

# Determine the active PyTorch device (CPU or GPU)
DEVICE = triton.runtime.driver.active.get_active_torch_device()

@triton.jit
def add_kernel(
    x_ptr,  # Pointer to the first input vector
    y_ptr,  # Pointer to the second input vector
    output_ptr,  # Pointer to the output vector
    n_elements,  # Total number of elements in the vectors
    BLOCK_SIZE: tl.constexpr,  # Size of the block each program handles
                               # tl.constexpr indicates it's a compile-time constant
):
    # 1. Identify the program instance and its data range
    pid = tl.program_id(axis=0)  # ← Get the 1D program ID
    block_start = pid * BLOCK_SIZE
    # Create offsets for the elements this program will process
    offsets = block_start + tl.arange(0, BLOCK_SIZE) # ← + start

    # 2. Create a mask to handle potential out-of-bounds elements
    # (if n_elements is not perfectly divisible by BLOCK_SIZE)
    mask = offsets < n_elements # ← Boolean tensor for safe access

    # 3. Load data safely using the mask
    # tl.load reads a block of data from memory (DRAM) into registers/SRAM
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)

    # 4. Perform the computation
    output = x + y

    # 5. Store the result back to memory safely using the mask
    # tl.store writes a block of data back to DRAM
    tl.store(output_ptr + offsets, output, mask=mask)

# Helper function to launch the kernel
def add(x: torch.Tensor, y: torch.Tensor):
    output = torch.empty_like(x)  # Pre-allocate output tensor
    assert x.is_contiguous() and y.is_contiguous(), "Input tensors must be contiguous"
    n_elements = output.numel()

    # Define the launch grid: 1D grid with one program per block
    # triton.cdiv(a, b) computes ceil(a / b)
    grid = lambda meta: (triton.cdiv(n_elements, meta),) # ← Grid size calculation

    # Launch the kernel
    add_kernel[grid](
        x, y, output,  # Pass tensors (implicitly converted to pointers)
        n_elements,    # Pass size argument
        BLOCK_SIZE=1024 # Pass compile-time constant as keyword argument
    )
    return output

# --- Verification ---
torch.manual_seed(0)
size = 98432
x = torch.rand(size, device=DEVICE)
y = torch.rand(size, device=DEVICE)

output_triton = add(x, y)
output_torch = x + y

print(f"Input device: {x.device}")
print("Triton Output (first 5):", output_triton[:5])
print("PyTorch Output (first 5):", output_torch[:5])
print(f'Maximum difference: {torch.max(torch.abs(output_torch - output_triton))}')

assert torch.allclose(output_triton, output_torch, atol=1e-5, rtol=0)
print("✅ Triton and PyTorch match!")

Key Concepts Introduced:

• @triton.jit: Decorator to define a Triton kernel.
• tl.program_id(axis=...): Gets the ID of the current program instance.
• tl.arange(start, end): Creates a 1D tensor sequence (like torch.arange).
• tl.constexpr: Marks compile-time constants (needed for block shapes).
• tl.load(ptr, mask=..., other=...): Loads data from memory, handling boundaries.
• tl.store(ptr, value, mask=...): Stores data to memory, handling boundaries.
• grid = lambda meta: (...): Defines the launch grid size dynamically based on meta-parameters.
• kernel[grid](...): Launches the kernel with the specified grid.

4. Debugging Triton Kernels

Triton ships four built‑ins:

• Compile‑time : static_print, static_assert
• Run‑time : device_print, device_assert (execute only with TRITON_DEBUG=1)

Need a full Pythonic stack trace? Enable the interpreter:

TRITON_INTERPRET=1 python my_script.py

This runs every kernel on CPU with NumPy fallbacks; you can sprinkle pdb.set_trace() inside the JITed function to single‑step through the tensorized IR. For memory bugs on GPU, prepend compute‑sanitizer (CUDA) or use ROCm’s AddressSanitizer.

5. Kernel Fusion: Fused Softmax

A naive softmax involves multiple steps (max, subtract, exp, sum, divide), each potentially requiring a slow DRAM roundtrip. Triton excels at fusing these into one kernel, drastically reducing memory bandwidth usage.

import torch
import triton
import triton.language as tl
from triton.runtime import driver # To get device properties

DEVICE = triton.runtime.driver.active.get_active_torch_device()

@triton.jit
def softmax_kernel(
    output_ptr, input_ptr,
    input_row_stride, output_row_stride, # Strides to move between rows
    n_rows, n_cols, # Matrix dimensions
    BLOCK_SIZE: tl.constexpr, # Power-of-2 block size >= n_cols
    # num_stages: tl.constexpr # Optional hint for software pipelining
):
    # Each program handles a subset of rows
    row_idx = tl.program_id(0) # ← Each program starts processing row `row_idx`

    # --- Load a row ---
    row_start_ptr = input_ptr + row_idx * input_row_stride
    col_offsets = tl.arange(0, BLOCK_SIZE)
    input_ptrs = row_start_ptr + col_offsets
    # Mask to handle columns beyond n_cols (since BLOCK_SIZE >= n_cols)
    mask = col_offsets < n_cols
    # Load the row, using -inf for padding to ensure max works correctly
    row = tl.load(input_ptrs, mask=mask, other=-float('inf')) # ← Safe load

    # --- Compute softmax ---
    # 1. Subtract max for numerical stability (max reduction)
    row_minus_max = row - tl.max(row, axis=0) # ← Max reduction
    # 2. Exponentiate
    numerator = tl.exp(row_minus_max)
    # 3. Sum numerators for denominator (sum reduction)
    denominator = tl.sum(numerator, axis=0) # ← Sum reduction
    # 4. Divide
    softmax_output = numerator / denominator

    # --- Store the result row ---
    output_row_start_ptr = output_ptr + row_idx * output_row_stride
    output_ptrs = output_row_start_ptr + col_offsets
    tl.store(output_ptrs, softmax_output, mask=mask) # ← Safe store

# Helper function to launch the softmax kernel
def softmax(x: torch.Tensor):
    n_rows, n_cols = x.shape
    # BLOCK_SIZE must be >= n_cols and a power of 2
    BLOCK_SIZE = triton.next_power_of_2(n_cols)

    y = torch.empty_like(x) # Allocate output

    # Launch grid: Simple 1D grid with one program per row
    grid = (n_rows,) # ← Launch one program per row

    softmax_kernel[grid](
        y, x,
        x.stride(0), y.stride(0),
        n_rows, n_cols,
        BLOCK_SIZE=BLOCK_SIZE,
        # num_stages=... # Optional tuning parameter
    )
    return y

# --- Verification ---
torch.manual_seed(0)
x = torch.randn(1823, 781, device=DEVICE) # Irregular shape
y_triton = softmax(x)
y_torch = torch.softmax(x, axis=1)

print("Softmax Triton Output (sample):", y_triton[0, :5])
print("Softmax PyTorch Output (sample):", y_torch[0, :5])
assert torch.allclose(y_triton, y_torch, atol=1e-2, rtol=0), "Softmax outputs don't match!"
print("✅ Softmax Triton and PyTorch match!")

# (Benchmark code from tutorial shows >4x speedup vs PyTorch JIT on V100)

Key Concepts Introduced:

• Kernel Fusion: Combining multiple logical steps into one kernel launch.
• Reductions: tl.max, tl.sum operate across an axis of a block.
• Handling Arbitrary Shapes: Using triton.next_power_of_2 and masking.

6. Compute Powerhouse: Matmul (FP16/FP8)

Matrix multiplication (GEMM) is compute-bound. Triton can implement high-performance GEMM kernels rivaling vendor libraries.

6.1 Core Concepts

• Blocked Algorithm: Each program computes one BLOCK_M x BLOCK_N tile of the output C.
• Tiled Accumulation: It iterates through BLOCK_K slices of the K dimension, loading tiles of A (BLOCK_M x BLOCK_K) and B (BLOCK_K x BLOCK_N), performing a tl.dot operation, and accumulating the result in FP32 for precision.
• Pointer Arithmetic: Calculates pointers for 2D blocks using strides and tl.arange with broadcasting ([:, None], [None, :]).
• L2 Cache Optimization (Grouped Order): Launches programs in a specific order (grouping by GROUP_SIZE_M rows) to improve data reuse in the L2 cache, boosting performance.
• Masking: Handles boundaries where M, N, or K are not multiples of block sizes.

6.2 Matmul Kernel Code

import torch
import triton
import triton.language as tl

DEVICE = triton.runtime.driver.active.get_active_torch_device()

# --- Autotune Configurations (Placeholder) ---
# Define different block sizes, stages, warps to test.
# The optimal configs depend heavily on the GPU architecture (CUDA vs ROCm)
# and data types (FP16 vs FP8).
# See the official Triton matmul tutorial for detailed, platform-specific configs.
AUTOTUNE_CONFIGS =

# Define leaky_relu if used in fusion
@triton.jit
def leaky_relu(x):
    return tl.where(x >= 0, x, 0.01 * x)

@triton.autotune(
    configs=AUTOTUNE_CONFIGS,
    key=['M', 'N', 'K'], # Retune when these dimensions change
)
@triton.jit
def matmul_kernel(
    # Pointers to matrices
    a_ptr, b_ptr, c_ptr,
    # Matrix dimensions
    M, N, K,
    # Strides for matrices
    stride_am, stride_ak,
    stride_bk, stride_bn,
    stride_cm, stride_cn,
    # Meta-parameters (tuned)
    BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr,
    GROUP_SIZE_M: tl.constexpr,
    # Optional: Activation function fused (passed as constexpr string)
    ACTIVATION: tl.constexpr = "" # Default: no activation
):
    """Computes C = Activation(A x B)."""
    # === Program ID Calculation (Grouped Order for L2 Cache) ===
    pid = tl.program_id(axis=0)
    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
    num_pid_in_group = GROUP_SIZE_M * num_pid_n
    group_id = pid // num_pid_in_group
    first_pid_m = group_id * GROUP_SIZE_M
    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
    pid_m = first_pid_m + (pid % num_pid_in_group) % group_size_m
    pid_n = (pid % num_pid_in_group) // group_size_m

    # === Pointer Setup for A and B Tiles ===
    offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
    offs_k = tl.arange(0, BLOCK_SIZE_K)
    a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
    b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)

    # === Accumulation Loop ===
    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) # ← FP32 accumulator
    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
        a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0)
        b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
        accumulator = tl.dot(a, b, accumulator) # ← Block matrix multiplication
        a_ptrs += BLOCK_SIZE_K * stride_ak
        b_ptrs += BLOCK_SIZE_K * stride_bk

    # === Optional Activation Fusion ===
    if ACTIVATION == "leaky_relu":
        accumulator = leaky_relu(accumulator) # ← Apply activation before storing

    c = accumulator.to(tl.float16) # ← Cast to output type

    # === Write Back Output Tile ===
    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) # ← 2D mask
    tl.store(c_ptrs, c, mask=c_mask)

# --- Python Wrapper ---
def matmul(a, b, activation=""):
    assert a.shape[1] == b.shape, "Incompatible dimensions"
    assert a.is_contiguous(), "Matrix A must be contiguous"
    M, K = a.shape; _, N = b.shape
    c = torch.empty((M, N), device=a.device, dtype=torch.float16) # Assuming FP16 output
    grid = lambda META: (triton.cdiv(M, META) * triton.cdiv(N, META),)
    matmul_kernel[grid](
        a, b, c, M, N, K,
        a.stride(0), a.stride(1), b.stride(0), b.stride(1), c.stride(0), c.stride(1),
        ACTIVATION=activation
    )
    return c

# --- Verification (FP16) ---
torch.manual_seed(0)
a = torch.randn((512, 512), device=DEVICE, dtype=torch.float16)
b = torch.randn((512, 512), device=DEVICE, dtype=torch.float16)
triton_output = matmul(a, b)
torch_output = torch.matmul(a, b)
print("Matmul Triton Output (FP16 sample):", triton_output[0, :5])
print("Matmul PyTorch Output (FP16 sample):", torch_output[0, :5])
# Adjust tolerance for specific hardware if needed (e.g., AMD CDNA2)
is_hip_cdna2 = triton.runtime.driver.active.get_current_target().backend == 'hip' and \
               triton.runtime.driver.active.get_current_target().arch == 'gfx90a'
rtol = 1e-2 if is_hip_cdna2 else 0
assert torch.allclose(triton_output, torch_output, atol=1e-2, rtol=rtol), "Matmul FP16 outputs don't match!"
print("✅ Matmul Triton and PyTorch match (FP16)!")

# --- FP8 Note ---
# FP8 execution requires compatible hardware (NVIDIA Hopper/Blackwell+) and
# potentially kernel modifications for FP8 data loading and tl.dot usage.
# The autotune configs for FP8 also differ significantly.
# See official Triton tutorials for FP8 matmul specifics.

Key Concepts Introduced:

• tl.dot(a, b, acc): Blocked matrix multiplication primitive.
• FP32 Accumulation: Improves numerical stability for FP16/FP8 inputs.
• Grouped Launch Order: Optimizes L2 cache usage.
• Activation Fusion: Applying activations directly in the kernel saves memory bandwidth.

7. Autotuning 101

Wrap any kernel with @triton.autotune and provide a list of triton.Configs. At first invocation Triton will profile every config on actual tensor shapes and memoize the winner. You get cuDNN‑style autotuning without writing a CMake plug‑in.

configs = [triton.Config({'BLOCK_SIZE_M':128,'BLOCK_SIZE_N':256,'BLOCK_SIZE_K':64}, num_warps=8, num_stages=3), ...]
@triton.autotune(configs, key=['M','N','K'])
@triton.jit
...

8. Memory Magic: Seeded Dropout

Standard dropout requires storing a large boolean mask, consuming significant memory during training. Triton’s built-in pseudo-random number generator (tl.rand) allows recomputing the mask on-the-fly using just a seed and offsets, saving memory.

import torch
import triton
import triton.language as tl
import tabulate # For nice printing

DEVICE = triton.runtime.driver.active.get_active_torch_device()

@triton.jit
def _seeded_dropout(
    x_ptr,
    output_ptr,
    n_elements,
    p, # Dropout probability (probability to zero out)
    seed, # Single seed for the entire operation
    BLOCK_SIZE: tl.constexpr,
):
    # Calculate offsets for this program instance
    pid = tl.program_id(axis=0)
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE) # ← Element indices

    # Boundary mask
    mask = offsets < n_elements

    # Load input data
    x = tl.load(x_ptr + offsets, mask=mask)

    # --- Generate dropout mask on the fly ---
    # tl.rand needs a seed and offsets to generate unique random numbers
    # for each element across all programs.
    random = tl.rand(seed, offsets) # ← Generate random numbers),)
    _seeded_dropout[grid](
        x, output, n_elements, p, seed,
        BLOCK_SIZE=1024 # Example block size
    )
    return output

# --- Verification ---
x = torch.randn(size=(10,), device=DEVICE)
p = 0.5
seed1 = 123
seed2 = 512

output1_seed1 = seeded_dropout(x, p=p, seed=seed1)
output2_seed1 = seeded_dropout(x, p=p, seed=seed1) # Should be identical
output3_seed2 = seeded_dropout(x, p=p, seed=seed2) # Should be different

print("Seeded Dropout Example:")
# Using tabulate for cleaner output
table = ["Input"] + [f"{val:.2f}" for val in x.tolist()],
    [f"Output (seed={seed1}, run 1)"] + [f"{val:.2f}" for val in output1_seed1.tolist()],
    [f"Output (seed={seed1}, run 2)"] + [f"{val:.2f}" for val in output2_seed1.tolist()],
    [f"Output (seed={seed2})"] + [f"{val:.2f}" for val in output3_seed2.tolist()]
print(tabulate.tabulate(table))


# Verify reproducibility with the same seed
assert torch.allclose(output1_seed1, output2_seed1)
# Verify difference with different seeds (highly likely)
assert not torch.allclose(output1_seed1, output3_seed2)
print("\n✅ Seeded dropout reproducibility verified.")

Key Concepts Introduced:

• tl.rand(seed, offsets): Parallel pseudo-random number generation.
• Stateless Mask Generation: Avoids storing the mask, saving memory.
• tl.where(condition, val_if_true, val_if_false): Conditional selection.

9. Beyond the Basics: Layer Norm & FlashAttention

Triton isn’t limited to simple operations. The official tutorials include implementations for more complex, performance-critical kernels:

• Layer Normalization: Provides fused forward and backward pass implementations, including parallel reductions for weight/bias gradients, often outperforming libraries like Apex FusedLayerNorm.
• FlashAttention v2: A Triton implementation reproducing the highly optimized FlashAttention v2 algorithm is available, crucial for efficient LLM training and inference. It leverages tiling, shared memory, and careful scheduling to minimize DRAM access for the attention mechanism.

These examples demonstrate Triton’s capability to handle complex algorithms and achieve state-of-the-art performance. Check the official documentation for the full code and explanations.

10. Triton in the Wild: LLMs, PyTorch & ROCm

Triton is increasingly integral to the AI ecosystem:

• LLM Training & Inference: Used in libraries like TensorRT-LLM, Unsloth, and Liger-Kernel to implement fused kernels, optimized attention (FlashAttention), quantization (FP8/INT8), and efficient fine-tuning techniques.
• PyTorch Integration: Serves as a key backend for torch.compile() (via TorchInductor), automatically generating Triton kernels from PyTorch code for significant speedups.
• Inference Serving: Triton Inference Server (TIS) often uses backends (like TensorRT-LLM) built with Triton/CUDA kernels to serve models efficiently.
• Cross-Platform: Supports both NVIDIA CUDA and AMD ROCm, enabling high-performance kernels on different GPU architectures, although platform-specific tuning might be needed.

Compatibility Note: The tight integration with PyTorch (torch.compile) means specific PyTorch versions often require specific Triton versions. Check PyTorch release notes or use isolated environments (venv, conda) to manage dependencies carefully. Mismatches can cause errors.

11. Best Practices

• Start Simple: Use the interpreter (TRITON_INTERPRET=1) for initial logic debugging.
• Use tl.constexpr: For block sizes and other parameters affecting kernel structure or compile-time logic.
• Profile: Use compute-sanitizer, Nsight/rocprof, and triton.testing.perf_report to identify bottlenecks and verify correctness.
• Autotune: Leverage @triton.autotune for performance portability, especially when tensor shapes vary. Consult official tutorials for good starting configs.
• Check Compatibility: Ensure your Triton version matches your PyTorch version if using torch.compile.
• Read the Docs: The official Triton documentation is the best resource for detailed explanations, advanced features, and troubleshooting.

Triton empowers developers to write high-performance GPU code with Pythonic ease. By understanding its core concepts and leveraging its tools, you can unlock significant performance gains for your deep learning workloads.

Useful Links:

1. https://openai.com/index/triton/
2. https://triton-lang.org/main/index.html

triton/python/tutorials at main · triton-lang/triton