Triton programming model
Map an output tile to Triton program instances, pointer blocks, masks, and configuration parameters
Before you begin
Follow a linked prerequisite when you cannot explain it from memory. Background skills without links are stated as abilities rather than hidden terminology.
- Blocked tensor operations
- Masks and strides
- Python decorators
Launch the tutorial vector-add kernel on 1000 FP32 elements with BLOCK=256. Triton creates four program instances because ceiling(1000/256) equals four. Printing or testing the computed offsets shows that the fourth instance describes positions 768 through 1023, so a Boolean mask must prevent loads and stores at positions 1000 through 1023. A comparison with torch.add verifies the output, while a profiler shows one generated GPU kernel.
A Triton program instance is one invocation of the decorated function for one program identifier. The function operates on a block of offsets rather than on one explicitly named CUDA thread. Triton's compiler maps those blocked operations onto GPU lanes and instructions. Generated Triton IR, PTX, and Nsight Compute data reveal that mapping when performance differs from the intended block design.
Triton changes the unit of programming from one scalar CUDA thread to a program instance operating on vectors or tiles of indices. The source expresses which block of a tensor is loaded, transformed, reduced, and stored. The compiler maps that block to warps, vector instructions, shared memory, and native code. Block-level source makes common AI kernels concise without pretending hardware constraints have disappeared.
A Triton program describes tiled dataflow. A program ID chooses an output region, index vectors describe elements, masks handle boundaries, and loads and stores define memory access. Block sizes and warp counts determine resource use. Use traffic, roofline, and occupancy analysis to select and evaluate these parameters.
A learner coming from CUDA should not equate one Triton program instance with one CUDA thread. One instance describes a block of work. The compiler distributes that block across a configured number of warps. The difference explains both the concise source and the need to reason about the complete block's live values.
Program instances and block values
In vector addition, program id p owns offsets p·BLOCK through p·BLOCK+BLOCK−1. tl.arange creates that vector of offsets. Loads return a Triton block value with one element per offset. Arithmetic applies across the block, and the store writes the block results. The source describes data parallelism without writing a scalar thread index for each lane.
The block value is not a conventional array allocated in global memory. During lowering, portions can reside in registers, use vector instructions, participate in reductions, or move through shared memory. A very large block can therefore create high register pressure even though the source contains only one variable name.
Masks and boundary values
The final program instance often covers offsets beyond the logical tensor. A load mask prevents invalid memory access and supplies an alternative value for inactive positions. The alternative value must match the operation. Zero is appropriate for addition, while negative infinity is appropriate for a maximum used by softmax. Stores use a mask so invalid output positions are not written.
A mask preserves memory safety, but a mask does not eliminate all cost for inactive positions. Block shapes are compile-time fixed, and some operations still execute across the padded block. Triton's fused-softmax tutorial pads row width to a power of two and masks the extra elements. The method works well for a bounded width family but can waste resources when one configuration covers widely different widths.
Pointer expressions and memory order
Memory behavior follows the addresses created by the program. Consecutive offsets along the fastest-changing axis should usually map to consecutive memory locations so hardware can combine lane requests. A tensor can be logically contiguous along one axis and strided along another; the pointer expression must incorporate the actual strides.
Compiler hints about alignment and divisibility can permit wider or more efficient operations, but an incorrect hint creates an invalid program assumption. First express correct pointers and masks. Then add only properties guaranteed by the operator's eligibility checks. Confirm actual transaction sizes and requested bytes with profiler evidence.
Reductions, dot operations, and resources
A block reduction such as tl.max or tl.sum requires communication across the elements represented by one program instance. The compiler selects a warp and shared-memory reduction schedule. A tl.dot operation exposes a tiled matrix product and can lower to hardware matrix instructions when shapes, layouts, datatypes, and target support permit.
Block size, number of warps, and number of pipeline stages affect instruction parallelism, register count, and shared memory. Their effects are not monotonic. More warps can help a large reduction but add overhead to a small row. More stages can cover memory latency but increase capacity use. Inspect generated IR and native code when the measured result conflicts with the expected schedule.
Autotuning and specialization
Autotuning defines a bounded set of configurations and measures them for keys such as row width or matrix dimensions. The key should change when the best schedule can change. The key should not include every incidental runtime value, because each unique key can trigger compilation and benchmarking.
Production systems need a plan for first use, cache persistence, and rare shapes. Tune representative families offline or during controlled warm-up when request latency cannot absorb the search. Retain a deterministic default and cap the candidate set. Benchmark selection noise can choose unstable winners when candidates differ by less than timing variance.
Worked example: Designing a row-wise softmax program
Each program instance handles one row whose width varies up to a known range.
- Choose a power-of-two block size covering the row and create an index vector. Mask indices beyond the actual width and load them as negative infinity for the maximum.
- For width 1000, the block contains 1024 positions. The last 24 loads return negative infinity, so they cannot increase the maximum or contribute a positive exponential.
- Compute row maximum, subtract, exponentiate, reduce the sum, and normalize while values remain in the program's fast storage.
- Select warp count and block size based on width. A huge block for every row can waste lanes and registers on small cases, so dispatch or autotuning may define width families.
- Test fully masked behavior, extreme logits, odd widths, and gradients if needed. Inspect register use and achieved bandwidth across the shape distribution.
- Compare with a width near the next power-of-two boundary, such as 1025. Padding to 2048 can double the represented block and may justify a multi-pass or separate configuration.
Conclusion: The Triton source mirrors the mathematical row algorithm while its performance still depends on physical block size, memory order, reduction lowering, and resources.
Common errors
- Assuming Triton removes the need for GPU architecture knowledge. Triton relocates decisions, but tile shape, memory access, warps, and registers remain decisive.
- Autotuning an enormous search space without workload keys. Compilation cost and noisy winners can outweigh small gains.
- Using one padded power-of-two block for every width. Convenience can create severe waste and register pressure on irregular shapes.
- Using a masked load with the wrong inactive value. A zero tail in a maximum reduction changes an all-negative row.
Reference summary
@triton.jit
def add_kernel(x, y, out, n, BLOCK: tl.constexpr):
offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
mask = offsets < n
x_block = tl.load(x + offsets, mask=mask, other=0.0)
y_block = tl.load(y + offsets, mask=mask, other=0.0)
tl.store(out + offsets, x_block + y_block, mask=mask)Triton exposes program ids, block pointers or explicit offsets, masked loads/stores, reductions, dot operations, and compile-time meta-parameters. Autotuning searches discrete configurations, but the search space must remain meaningful and compilation cost must be accounted for. Inspect generated IR/PTX and profiler behavior when performance diverges from the block model.
- Use masks for arbitrary tails and safe out-of-bounds behavior.
- Tune block size, warps, stages, and layout-relevant parameters by shape family.
- Keep numerical reductions and approximate math explicit.
- Cache compiled variants and avoid unbounded specialization.
Knowledge check
Why can autotuning over dozens of configurations be counterproductive in an online system with many rare shapes?
Show the answer guide
- Autotuning pays an up-front cost to compile, launch, and time candidate configurations. A rare shape may execute too few times to recover that cost through later speedups.
- Dozens of candidates can delay the request that first encounters the shape, violating an online latency objective even if the selected kernel is faster afterward.
- Many rare shapes can expand compiled-code and tuning-result caches, consume host and device memory, and evict entries for common shapes. Concurrent tuning can also compete with serving work.
- An online system should restrict the search space, tune offline, bucket shapes, or use a robust default unless expected executions times per-execution savings exceed compilation, benchmarking, cache, and tail-latency costs.
Calculate whether online autotuning repays its cost
Create a Triton kernel with at least 24 legal configurations and feed it a declared trace of 10,000 calls: 80% from four common shapes and 20% distributed across 200 rare shapes. Compare exhaustive first-use autotuning, a four-candidate search, shape bucketing, and one fixed configuration. Measure cold and warm calls separately.
Evidence to keep: Submit the program and workload trace, a table of compile or tuning time, p50 and p99 call latency, steady-state kernel time, cache entries, and total trace time for each policy, plus a plot of cumulative time. Give a causal conclusion that identifies when tuning cost is or is not amortized.
A fast kernel becomes useful to model code only after it satisfies the framework's operator contract, including shapes, devices, compilation, and gradients.
Sources and further reading
The lesson explanations synthesize primary papers and official references. Follow the originals for proofs, exact APIs, architecture details, and version-sensitive behavior.