Triton and when to write a kernel

Most inference engineers spend far more time choosing and configuring library kernels than writing new ones. Writing one matmul and one fused elementwise in Triton is still the fastest way to make the roofline intuitive. This page is the decision rule, plus a sketch of how Triton thinks.

The big idea

Triton is a Python DSL for tiled GPU kernels. You think in program IDs and block pointers. The compiler handles much of the indexing, coalescing and shared-memory choreography. You still own the tile shape and the algorithm.

A food processor versus a chef's knife: Triton is the processor. Raw CUDA is the knife. cuBLAS is the restaurant's existing line. Do not open a new kitchen to boil water.

Walking through the choice

Sketch of a fused softmax (illustrative shape, not production code):

@triton.jit
def softmax_kernel(out_ptr, in_ptr, stride, n_cols, BLOCK: tl.constexpr):
    row = tl.program_id(0)
    cols = tl.arange(0, BLOCK)
    mask = cols < n_cols
    offsets = row * stride + cols
    x = tl.load(in_ptr + offsets, mask=mask, other=-float("inf"))
    x = x - tl.max(x, axis=0)
    num = tl.exp(x)
    den = tl.sum(num, axis=0)
    tl.store(out_ptr + offsets, num / den, mask=mask)

When debugging a custom op, start by confirming the grid covers the problem, loads coalesce, and you are not synchronizing every token. Only then dive into bank conflicts and instruction-level tricks.

Two NVIDIA tools matter after PyTorch's profiler:

• Nsight Systems — timeline across CPU, kernels, copies, streams. Use it to see gaps between kernels (launch overhead, sync, allocator stalls).
• Nsight Compute — per-kernel deep dive: achieved bandwidth, compute throughput, occupancy, warp stalls.

For serving engines, also measure at the token metrics layer. A beautiful kernel that does not change TTFT under load may not be the bottleneck.

Common mistakes

• Inventing a matmul. cuBLAS and CUTLASS already exist.
• Writing Triton before profiling. Most decode pain is launch overhead or HBM, not a missing custom softmax.
• Judging a kernel by Nsight Compute alone. If TTFT did not move, you optimized the wrong thing.

Check yourself

1. You need a fused elementwise chain that PyTorch eager launches as 40 tiny kernels. Which tool first? Answer: Triton (or torch.compile / engine fusion), not a new cuBLAS.
2. When do you reach for raw CUDA over Triton? Answer: when you need control or intrinsics Triton cannot express cleanly.

Key takeaways

• Prefer FlashAttention, FlashInfer and cuBLAS. Use Triton for fusions. Raw CUDA is last-mile.
• Triton lets you write tiles in Python. The compiler handles a lot of the memory choreography.
• Profile PyTorch → Nsight Systems → Nsight Compute, and check token metrics before celebrating a kernel.

Further learning

Triton tutorials. Back to the CUDA programming model.