7.1 Introduction to the Triton programming framework

This article uses the official Fused Softmax operator to introduce Triton’s programming model and basic syntax. Unlike CUDA, where you usually write compute logic from a thread’s point of view, Triton lets you describe load, compute, and store operations in tiles (data blocks), and leaves most thread- and warp-level details to the compiler.

Softmax is a common operator in deep learning and large models, with implementations in frameworks like PyTorch and TensorFlow. Here we implement row-wise Softmax in Triton and compare it with naive_softmax built from several PyTorch ops and with native torch.sofrmax, to see how operator fusion cuts intermediate tensor memory traffic and kernel launch overhead.

1-Triton programming model

Before covering Triton’s programming model, let’s recall how we wrote CUDA operators before.

Take Softmax: we usually write the compute logic from a single thread’s view, as in the warp-sum reduce below:

// Sum reduce within a warp
__device__ __forceinline__ float warpReduceSum(float val){
    for(int i = 1 ; i < 32 ; i <<= 1)
        val += __shfl_xor_sync(0xffffffff, val ,i) ;
    return val ;
}

__device__ __forceinline__ float blockReduceSum(float val){
    __shared__ float smem[32] ;

    int lane_ID = threadIdx.x % 32 ;
    int warp_ID = threadIdx.x / 32 ;

    val = warpReduceSum(val) ;
    if(lane_ID == 0) smem[warp_ID] = val ;
    __syncthreads() ;

    val = (lane_ID < blockDim.x / 32) ? smem[lane_ID] : 0.0f ;
    if(warp_ID == 0) val = warpReduceSum(val) ;

    __shared__ float block_res ;
    if(threadIdx.x == 0) block_res = val ;
    __syncthreads() ;
    return block_res ;
}

Each thread first finishes the warp reduce via shuffle, then each warp writes its partial result to shared memory, and the first warp finishes the block-level reduce. You have to handle inter-thread communication, shared memory, and sync yourself. That CUDA style gives fine control over memory and data, but makes implementation and debugging harder. Triton takes a lot of that messy work off your plate and lets the compiler handle most of the details.

Triton was first proposed by Philippe Tillet, H. T. Kung, and David Cox in 2019, aiming at a GPU programming style for deep learning operators that sits above CUDA while still keeping strong performance. In 2021, OpenAI released Python-based Triton 1.0 and kept pushing it forward.

Unlike CUDA’s SIMT model centered on threads and warps, Triton uses a programming model built around tiles (data blocks). You mainly describe the data block a Triton program instance should handle, plus the load, compute, and store ops on that block, without explicitly managing every CUDA thread. The compiler then maps those block-level ops onto concrete threads and warps, and optimizes for coalesced memory access, shared-memory allocation, needed thread sync, and instruction layout based on the access pattern. Which SM runs a given block, and when a warp’s instructions are issued, is still decided by the GPU hardware scheduler. The compiler also cannot promise that every access pattern will coalesce well.

Triton’s goal is not to hide GPU parallelism entirely. It sits between algorithm expression and the hardware: you still choose how to tile the operator, the program grid, and the overall algorithm structure, but you can leave a lot of thread-level implementation detail to the compiler. That lowers the cost of writing and debugging custom GPU operators while keeping performance close to hand-written CUDA when possible.

2-Torch Softmax operator implementation

Consider an input matrix $X \in \mathbb{R}^{M \times N}$ , where $M$ and $N$ are the number of rows and columns. Softmax is applied to each row:

y_{i,j} = \operatorname{softmax}(X_{i,:})_j = \frac{e^{x_{i,j}}}{\sum_{k=0}^{N-1} e^{x_{i,k}}}

Here $x_{i,j}$ is the element at row $i$ , column $j$ , and $k$ is an independent index used for the denominator sum. In practice you first subtract each row’s max $m_i$ , which gives a numerically stabler form that is mathematically equivalent:

The naive_softmax below uses that stabilization:

import torch

def naive_softmax(x) :

    x_max = x.max(dim=1)[0]

    z = x - x_max[:, None]

    numerator = torch.exp(z)

    denominator = numerator.sum(dim=1)

    ret = numerator / denominator[:, None]

    return ret

M, N = 10, 10
x = torch.rand(size=(M, N,), dtype=torch.float32, device='cuda')
y_torch = naive_softmax(x)

3-Basic Triton syntax

Most ops in a Triton kernel handle a whole data block, not a single scalar. When reading Triton code, it helps to follow this order:

Use tl.program_id to decide which data region the current program owns;
Use tl.arange to build element offsets inside a data block;
Build a set of memory addresses with pointer arithmetic;
Use tl.load to load data into a Triton tensor;
Run elementwise ops, reduces, or matrix multiplies on the tensor;
Use tl.store to write results back to device memory.

One important constraint: a Triton kernel cannot use native Python data structures (list, dict, and so on) or PyTorch ops (torch.max, torch.relu, and so on). The sections below cover basic syntax commonly used inside Triton kernels:

tl.program_id

tl.program_id returns the current program instance’s index on a given launch-grid axis, similar to CUDA’s blockIdx:

# Triton's launch grid has at most three dimensions, so axis is 0, 1, or 2
# Similar to CUDA's blockIdx.x / blockIdx.y / blockIdx.z 
pid = tl.program_id(axis=0)

tl.num_programs

tl.num_programs(axis) returns how many programs are on that launch-grid axis, similar to CUDA’s gridDim:


# Likewise, axis is 0, 1, or 2
# Matching CUDA's gridDim.x / gridDim.y / gridDim.z
num_programs = tl.num_programs(axis=0)

tl.arange

tl.arange(start, end) builds contiguous logical indices:


offsets = tl.arange(0, BLOCK_SIZE)

BLOCK_SIZE is usually a power of two passed in via tl.constexpr.

tl.range

tl.range works like Python’s native range iterator and is often used in loops:


@triton.jit
def kernel(M, num_stages):
    for idx in tl.range(0, M, num_stages=num_stages) :
        ...

num_stages is the software-pipeline depth; see the Triton Softmax implementation later for a concrete case.

tl.load

tl.load reads data from device memory given a pointer or a set of pointers:


values = tl.load(ptrs)

If ptrs is a pointer tensor of shape [BLOCK_SIZE], the returned values have the same shape.
For positions that may go out of bounds, use a mask:


values = tl.load(
    ptrs,
    mask=offsets < n_elements,
    other=0.0,
)

Where mask=False, the code does not touch that memory location and fills the return with other. If you use mask, you should usually set other explicitly. Otherwise the masked positions have undefined return values and can pollute later compute.

tl.store

tl.store writes a scalar or Triton tensor to the given memory locations, and can also use a mask to avoid out-of-bounds writes:


tl.store(output_ptrs, output, mask=mask)

tl.where

tl.where(condition, x, y) picks results elementwise from a condition:


output = tl.where(condition, x, y)

Semantically it is like: output[i] = x[i] if condition[i] else y[i]. A common example is ReLU:


output = tl.where(x > 0.0, x, 0.0)

Important: both the x and y branches of tl.where are evaluated. It only chooses the final result by condition; it does not skip compute or memory access in either branch. For example:


# Unsafe branch: may hit divide-by-zero
result = tl.where(denominator != 0, numerator / denominator, 0.0)

# Safe branch
safe_denominator = tl.where(denominator != 0, denominator, 1.0)
result = numerator / safe_denominator
result = tl.where(denominator != 0, result, 0.0)

So tl.where is good for elementwise selection, but do not treat it as control flow that can skip a branch.

tl.dot

Triton provides many elementwise math functions in the style of NumPy and PyTorch, for example:

y = tl.exp(x)
y = tl.log(x)
y = tl.sqrt(x)
y = tl.abs(x)
y = tl.maximum(x, 0.0)
y = tl.dot(x, y)

tl.dot computes matrix multiplication of two 2D or 3D data blocks. In a GEMM kernel, tl.dot multiplies and accumulates along the $K$ dimension tile by tile:


# Create an all-zero tensor with the given shape and dtype
accumulator = tl.zeros(
    (BLOCK_SIZE_M, BLOCK_SIZE_N),
    dtype=tl.float32,
)

for k in range(0, K, BLOCK_SIZE_K):
    a = tl.load(a_ptrs, mask=a_mask, other=0.0)
    b = tl.load(b_ptrs, mask=b_mask, other=0.0)

    accumulator = tl.dot(a, b, accumulator)

    a_ptrs += BLOCK_SIZE_K * stride_ak
    b_ptrs += BLOCK_SIZE_K * stride_bk

The accumulator argument adds this matrix-multiply result into an existing tensor.

4-Triton Softmax operator implementation

A Triton kernel launches many program instances on a grid. You can think of one program instance as similar to a CUDA thread block: several threads/warps cooperate to process one or more data blocks.

As with a CUDA kernel, you declare a Triton kernel function; Triton defines kernels with @triton.jit. The next three code blocks concatenate into the full _softmax function; the two blocks after the signature are its body:

import triton
import triton.language as tl
from triton.runtime import driver

# Declare this as a Triton kernel
@triton.jit
def _softmax(
    input_ptr,                    # base pointer of the input matrix
    output_ptr,                   # base pointer of the output matrix
    input_row_stride,             # stride between the starts of two adjacent input rows
    output_row_stride,            # stride between the starts of two adjacent output rows
    n_rows,                       # number of matrix rows
    n_cols,                       # number of matrix columns
    BLOCK_SIZE: tl.constexpr,     # logical elements per loop iteration; compile-time constant
    num_stages: tl.constexpr      # software-pipeline depth passed to tl.range
) :

Here BLOCK_SIZE is how many logical elements each loop iteration handles in this example. tl.constexpr means the value is known at compile time, so the compiler can fix tensor shapes and emit specialized code. Think of it like a C++ compile-time constant.

In this article’s example, one program loops over several rows of Softmax, so we need the current program id and the number of launched programs:

    # A Triton program is similar to a CUDA thread block
    # tl.program_id(0) is like blockIdx.x; axis=0/1/2 maps to the grid's x/y/z axes
    row_start = tl.program_id(axis=0)

    # Like CUDA's gridDim.x: number of programs on grid axis 0
    row_step = tl.num_programs(axis=0)

Suppose $P$ programs are launched. The program numbered $p$ processes rows $p, p+P, p+2P, \ldots$ until the row index reaches $M$ , computing one row per loop iteration.

For each row, first build the per-column pointers, then tl.load the data, finish the reduce and elementwise ops, and finally tl.store the result:


    # Stride is row_step; row_idx is the row index this program is handling
    for row_idx in tl.range(row_start, n_rows, row_step, num_stages=num_stages) :
        # Base pointer of the row currently being processed
        row_start_ptr = input_ptr + row_idx * input_row_stride

        # col_offset: [0, 1, 2, ..., BLOCK_SIZE-1]
        # BLOCK_SIZE here must be a compile-time-known power of two
        col_offset = tl.arange(0, BLOCK_SIZE)

        # Build pointers for BLOCK_SIZE logical positions; out-of-bounds ones are masked below
        input_ptrs = row_start_ptr + col_offset

        # Matrix column count is not always a power of two; use a mask to bound the reads
        mask = col_offset < n_cols

        # mask=False positions skip memory reads and fill -inf; when the row max is finite, their exp is 0
        row = tl.load(input_ptrs, mask=mask, other=-float('inf'))

        # Softmax compute
        row_minus_max = row - tl.max(row, axis=0)
        numerator = tl.exp(row_minus_max)
        denominator = tl.sum(numerator, axis=0)
        softmax_output = numerator / denominator

        # Base pointer of the output row
        output_row_ptr = output_ptr + row_idx * output_row_stride
        # Per-element pointer offsets
        output_ptrs = output_row_ptr + col_offset
        tl.store(output_ptrs, softmax_output, mask=mask)

Note that row here is a 1D Triton tensor of length BLOCK_SIZE, so tl.max(row, axis=0) and tl.sum(..., axis=0) both reduce along the logical column positions of the current row and yield a scalar.

tl.range(..., num_stages=num_stages) hints a software-pipeline depth to the compiler so loads and compute from different loop iterations can overlap and improve resource use. Higher depth means more pressure on shared memory and registers.

As with CUDA, a Triton kernel also needs a launch grid. The approach below follows the official tutorial: precompile first, query resource use, then estimate how many programs to launch. Programs keep processing later rows in a loop; that style is often called a persistent kernel.

DEVICE = torch.cuda.current_device()
properties = driver.active.utils.get_device_properties(DEVICE)
NUM_SM = properties['multiprocessor_count']   # number of SMs
NUM_REGS = properties['max_num_regs']         # register capacity budget, counted in 32-bit registers
SIZE_SMEM = properties['max_shared_mem']      # shared-memory capacity budget, in bytes
WARP_SIZE = properties['warpSize']            # 32 on NVIDIA GPUs
kernels = {}

def softmax(x):
    assert x.is_cuda

    n_rows, n_cols = x.shape
    y = torch.empty_like(x)

    # Round up to the smallest power of two that is >= n_cols
    BLOCK_SIZE = triton.next_power_of_2(n_cols)
    num_warps = 8       # one program uses 8 warps, i.e. 256 CUDA threads
    num_stages = 4 if SIZE_SMEM > 200000 else 2       # pipeline depth
    kernel, num_programs = kernels.get(BLOCK_SIZE, (None, 0))

    if kernel is None:
        # warmup only compiles; it does not really run. Returns a compiled kernel object
        kernel = _softmax.warmup(
            y, x, 
            x.stride(0), y.stride(0), 
            n_rows, n_cols, BLOCK_SIZE=BLOCK_SIZE, 
            num_stages=num_stages, num_warps=num_warps, 
            grid=(1,))

        # Initialize CUDA kernel handles so we can query resource use of the compiled artifact.
        kernel._init_handles()
        n_regs = kernel.n_regs
        size_smem = kernel.metadata.shared
        occupancy = NUM_REGS // (n_regs * WARP_SIZE * num_warps)
        occupancy = min(occupancy, SIZE_SMEM // size_smem)
        num_programs = NUM_SM * occupancy
        kernels[BLOCK_SIZE] = (kernel, num_programs)

    num_programs = min(num_programs, n_rows)

    # Launch the Softmax kernel
    # x and y are implicitly converted to the matrix base addresses
    kernel[(num_programs, 1, 1)](
        x,
        y,
        x.stride(0),   # address difference between adjacent elements in the same column
        y.stride(0),
        n_rows,
        n_cols,
    )

    return y

5-Performance comparison

Triton ships handy benchmark helpers. By setting different configs and benchmark implementations, you can quickly compare operator performance across methods.

Following the previous article, Softmax with large enough data is often memory-bound. So we use effective bandwidth as a shared metric for both implementations:

B_{\mathrm{eff}} = \frac{2MN \cdot s}{t_{\mathrm{ms}}} \times 10^{-6} \quad (\mathrm{GB/s})

where $s$ is bytes per element and $t_{\mathrm{ms}}$ is latency in milliseconds.


# Triton's benchmark decorator
@triton.testing.perf_report(
    triton.testing.Benchmark(
        x_names=['N'],                               # parameter N under different configs
        x_vals=[128 * i for i in range(2, 17)],      # N=256, 384, ..., 2048
        x_log=True,
        line_arg='provider',                         # pick the implementation for each curve via provider
        line_vals=['triton', 'torch'],               # two implementations under test
        line_names=['Triton', 'Torch'],              # names of the comparison curves
        styles=[('blue', '-'), ('green', '-')],      # colors of the comparison curves
        ylabel='Effective bandwidth (GB/s)',
        plot_name='fused-softmax-performance',
        args={'M': 64},                              # fix M at 64
    )
)
def benchmark(M, N, provider):
    # Build input and run the measured op on the same current CUDA stream
    x = torch.rand(M, N, device=DEVICE, dtype=torch.float32)

    # Return median, 20th percentile, and 80th percentile latency in order
    quantiles = [0.5, 0.2, 0.8]

    # Compare against native torch
    if provider == 'torch':
        ms, min_ms, max_ms = triton.testing.do_bench(
            lambda: torch.softmax(x, axis=-1), 
            quantiles=quantiles)

    # Benchmark the Triton Softmax operator
    if provider == 'triton':
        ms, min_ms, max_ms = triton.testing.do_bench(
            lambda: softmax(x), 
            quantiles=quantiles)

    # Compute GB/s from the same effective read/write byte count
    gbps = lambda ms: 2 * x.numel() * x.element_size() / ms * 1e-6
    return gbps(ms), gbps(max_ms), gbps(min_ms)

if __name__ == '__main__':
    torch.manual_seed(0)
    # Cover small sizes, non-power-of-two column counts, and different shapes under the same BLOCK_SIZE
    for M, N in [(1, 1), (128, 256), (129, 255), (1823, 781)]:
        x = torch.randn(M, N, device=DEVICE, dtype=torch.float32)
        y_torch = naive_softmax(x)
        y_triton = softmax(x)
        torch.testing.assert_close(y_triton, y_torch, rtol=1e-5, atol=1e-7)
    print('Correctness checks passed.')

    # Compare effective bandwidth of Triton vs Torch
    benchmark.run(print_data=True, show_plots=True)

do_bench warms up and times repeatedly. With quantiles=[0.5, 0.2, 0.8], the returns are median latency, the 20th percentile, and the 80th percentile.

Results from a run:

the max error is 1.862645149230957e-09
fused-softmax-performance:
         N      Triton      Torch
0    256.0   18.285714   6.736842
1    384.0   27.428571   9.142857
2    512.0   42.666665  12.190476
3    640.0   53.333331  15.238095
4    768.0   54.857142  17.454545
5    896.0   63.999999  18.666666
6   1024.0   73.142856  20.480001
7   1152.0   82.285713  21.735849
8   1280.0  106.666663  24.615385
9   1408.0  117.333329  26.074074
10  1536.0  127.999995  26.482758
11  1664.0  104.000000  28.689654
12  1792.0  149.333328  28.000000
13  1920.0  159.999994  30.967742
14  2048.0  146.285712  30.117647

From those results:

As data size grows, Triton’s Softmax pulls further ahead of native Torch, which suggests the fused Triton kernel in this test set really does cut overhead.
Triton’s memory efficiency is higher, and the operator code is simpler and easier to maintain.

📝 Summary

From CUDA’s thread-level implementation to Triton’s block-level expression, the core of this article comes down to three threads:

Describe parallel compute in data blocks: move from explicitly writing thread/warp reduce, communication, and sync to describing the data and ops a program should handle. You still choose tiling, grid, and addressing; the compiler maps block-level ops onto threads and warps.
Use basic syntax to get a correct operator: assign matrix rows with tl.program_id, tl.num_programs, and a loop; build column offsets with tl.arange; handle padding with mask and -inf; then finish a stable Softmax with tl.max, tl.exp, and tl.sum.
Cut overhead with fusion, and judge the gain by measurement: fuse several ops into one kernel to cut intermediate global-memory traffic and launch count; pick launch size from resource use; after correctness checks, compare performance against a clear baseline with an effective-bandwidth metric.

Triton simplifies thread-level details; operator performance still depends on the algorithm, data layout, resource pressure, and input size. After this example, you can use the official docs to keep learning vector add, LayerNorm, and matrix multiply, and gradually practice more complex tiling and tuning.

🎯 Self-check list

Explain what Triton’s tile, program, and grid each mean, and under what conditions a program is analogous to a CUDA thread block;
Describe the rough flow of the Triton Softmax operator, from resolving program_id, computing data pointers, loading and computing, through writing back to memory, and understand the code for each step;
Distinguish what BLOCK_SIZE means in Triton’s model vs CUDA’s in this article (in Triton: how many data elements one program handles; in CUDA: how many threads one block contains)
Understand how tl.where works and how its branch selection behaves;
From row_start, row_step, and tl.range, list the row indices each program actually processes;
Explain why input_row_stride is passed in, and how address calculation works in a Triton kernel;
Understand how num_programs is computed dynamically during Triton kernel launch;
Explain why column count is padded up to a power of two, and why both load and store need a mask;
Explain why different Softmax implementations diverge in efficiency as data size grows, and where the optimization bottleneck sits;
Try implementing other operators in Triton (LayerNorm, GEMM, and so on) and compare them with native Torch;

Search