100 DAYS OF GPU Course preview
Day 86Module 9
draft

Writing kernels in Triton

Running pip install triton does not check whether the GPU can run its kernels. Triton's README lists support for "NVIDIA GPUs (Compute Capability 8.0+)" (https://github.com/triton-lang/triton/blob/main/README.md , checked 2026-09-01). TRITON_INTERPRET=1 runs kernels on the CPU for correctness checks but does not provide GPU timings.

This lesson implements a row softmax and an FP32 matmul as programs over blocks of data, beside the CUDA versions from day 47 and day 44, then compares what each language requires you to specify.

The unit is a block of data

Earlier CUDA kernels in this course define one thread's work and start by computing its index. A Triton kernel defines one block of data. tl.program_id(0) selects a block of the problem, and later operations act on arrays with that block's shape; the source contains no threadIdx.

Day 47's softmax gives one 256-thread block one row of 1024 columns. The block walks the row in a strided loop to find the maximum, reduces 256 partial maxima through shared memory in a halving tree, walks the row again for the sum, reduces again, then walks it a third time to divide. Five __syncthreads() in forty lines, two of them inside halving loops, and every one is a place the kernel can be wrong.

The Triton version loads the row as one value. tl.max(x, axis=0) and tl.sum(num, axis=0) are single calls on that value. No shared array is declared, no barrier is written, and the halving tree does not appear.

It still happens. The compiler picks the warp count from num_warps, allocates the shared memory, chooses the swizzle that keeps the reduction off a bank conflict, and emits the barriers. You provide two numbers instead: the block shape and the number of warps.

Diagram: where the barriers went. Two stacked bands over the same 1024-column softmax row, with a vertical line separating what you write on the left from what runs on the right. Band 1, day 47 in CUDA: 256 thread boxes, a strided-loop arrow across the row, then two shared-memory halving trees drawn in full on both sides of the line. Caption "40 lines, 5 barriers, 2 trees, all written by you." Band 2, day 86 in Triton: one program box holding the whole row as a single 1024-wide value, with tl.max and tl.sum as one node each on the left of the line. Caption "11 lines, no barrier in the source." Band 2b, greyed, on the right of the line under band 2: the same lanes, the same shared memory, the same two trees, with BLOCK_COLS and num_warps drawn as the only two inputs feeding them. Caption "the same machinery, sized by two compile-time numbers." Alt text: "The same 1024-column softmax row twice. Forty lines of CUDA spend five barriers and two shared-memory trees. Eleven lines of Triton spend none, and the compiler emits the same machinery from two numbers."

Triton's tuning parameters

Optimising a CUDA kernel can mean moving loads around: give a thread eight outputs instead of one, widen four scalar loads into a float4, pad a tile to dodge a bank conflict. Day 44 is four kernels of exactly that, and each one changes the inner loop.

In Triton, the inner loop is one tl.dot. A small set of compile-time values replaces the manual load schedule: the block shape, num_warps, and num_stages, which is how many iterations of the K loop the compiler keeps in flight at once. That last one is the software pipeline day 74 builds by hand out of cp.async and a barrier, written here as an integer.

Because they are compile-time constants, one kernel source is a family of kernels, and picking the member is a search. triton.autotune runs that search for you: hand it a list of triton.Config objects and a key, and the first launch at each new key times every configuration and caches the winner. The exercise below turns this page's hand-written sweep into one.

Day 44 explains where the autotuner's candidate shapes come from. Its 128 by 128 block tile with an 8 by 8 micro-tile maps to BLOCK_M, BLOCK_N and the register budget behind them, and the reason its occupancy fell by half while it got faster is the reason a bigger BLOCK_M here is often the right answer. You still choose the same values, but the compiler implements them.

Two kernels, five constants and one string

Full program in triton_softmax_matmul.py, under code/day86-triton/. The environment selects one of two modes. Under TRITON_INTERPRET=1, every kernel runs on the CPU through Triton's interpreter and the program prints no timings.

GPU mode requires compute capability 8.0 and exits with status 1 below that version.

The softmax is the whole kernel:

@triton.jit
def softmax_kernel(in_ptr, out_ptr, row_stride, n_cols,
                   BLOCK_COLS: tl.constexpr):
    # One program owns one row. Day 47's CUDA version hands one 256-thread
    # block the same row, walks it three times, and reduces it twice
    # through shared memory behind five written barriers. Here the row is
    # one BLOCK_COLS-wide value and both reductions are one call each.
    #
    # BLOCK_COLS is a power of two at or above n_cols, so the tail lanes are
    # masked. They load -inf, which loses the max and contributes exp(-inf)
    # = 0 to the sum. A 0.0 pad would win the max on an all-negative row.
    row = tl.program_id(0)
    cols = tl.arange(0, BLOCK_COLS)
    mask = cols < n_cols
    x = tl.load(in_ptr + row * row_stride + cols, mask=mask,
                other=-float("inf"))
    num = tl.exp(x - tl.max(x, axis=0))
    tl.store(out_ptr + row * row_stride + cols, num / tl.sum(num, axis=0),
             mask=mask)

The matmul's signature carries the whole tuning surface. Five arguments are tl.constexpr, so five numbers decide which kernel gets compiled:

@triton.jit
def matmul_kernel(a_ptr, b_ptr, c_ptr, m, n, k, stride_am, stride_ak,
                  stride_bk, stride_bn, stride_cm, stride_cn,
                  BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr,
                  BLOCK_K: tl.constexpr, GROUP_M: tl.constexpr,
                  IEEE: tl.constexpr):

GROUP_M controls the block order. The compiler owns the shared-memory layout, but not the order in which blocks run. The lines below walk GROUP_M tile rows before moving right, so the blocks resident at the same moment share more of A and B in L2 than a plain row-major walk would:

    pid = tl.program_id(0)
    tiles_m = tl.cdiv(m, BLOCK_M)
    tiles_n = tl.cdiv(n, BLOCK_N)
    per_group = GROUP_M * tiles_n
    first_m = (pid // per_group) * GROUP_M
    rows_here = min(tiles_m - first_m, GROUP_M)
    pid_m = first_m + ((pid % per_group) % rows_here)
    pid_n = (pid % per_group) // rows_here

And the K loop, which is where the string in the heading lives:

    # The accumulator is a value, not a buffer: no __shared__, no barrier,
    # no double buffer. IEEE is a compile-time constant, so exactly one of
    # these two branches survives into the compiled kernel.
    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for step in range(0, tl.cdiv(k, BLOCK_K)):
        left = k - step * BLOCK_K
        a = tl.load(a_ptrs, mask=offs_k[None, :] < left, other=0.0)
        b = tl.load(b_ptrs, mask=offs_k[:, None] < left, other=0.0)
        if IEEE:
            acc = tl.dot(a, b, acc, input_precision="ieee")
        else:
            acc = tl.dot(a, b, acc, input_precision="tf32")
        a_ptrs += BLOCK_K * stride_ak
        b_ptrs += BLOCK_K * stride_bk

input_precision decides "How to exercise the Tensor Cores for f32 x f32", its allowed NVIDIA values are "tf32", "tf32x3", "ieee", "bf16x3" and "bf16x6", and the default is "tf32" (https://triton-lang.org/main/python-api/generated/triton.language.dot.html , checked 2026-09-01). So a Triton FP32 matmul left alone uses tensor cores, and day 44's cuBLAS call deliberately does not: it pins CUBLAS_COMPUTE_32F under CUBLAS_DEFAULT_MATH. Compare the two as they come and the comparison measures that enum rather than the language.

The program runs both modes explicitly.

Counted rather than guessed, with the command in the README, these are non-blank non-comment lines from the signature to the end of the body:

Kernel CUDA Triton
row softmax 40, day 47's softmaxRows 11
FP32 matmul, register tiled 52, day 44's matmulRegTile2d 33
the same, with float4 loads 75, day 44's matmulRegTile2dVec4 33

The softmax has about one quarter as many lines. The matmul does not, because Triton still requires the pointer arithmetic, the masks and the block ordering are the same work in either language.

Results

Interpreter and T4 gate verified; the GPU path remains blocked. The interpreter ran with Triton 3.7.1 and torch 2.13.0+cu126, passed softmax and both matmul modes, and exited 0.

GPU mode identified the Tesla T4 at compute capability 7.5, printed the documented 8.0 floor and exited 1 before launching a kernel, as designed. No compute-capability-8 card has run the timing sweep.

Transcript What it settles Result
interpret-2026-09-02.txt, CPU both kernels against the float64 reference pass, exit 0
t4-gate-2026-09-02.txt, T4 node the 8.0 refusal, exit status 1 pass
run-*.txt, sm_80 or newer the same checks, then the timing sweep blocked

The blocked environment used the same free Colab T4 at compute capability 7.5 as day 79 and Compiler Explorer. Compute capability is the versioned hardware feature set used for this gate. The access snapshot in FACT-SHEET.md section 4 also records paid Colab L4 and A100 options and an isolated five-second Modal T4 run at about $0.011, with an sm_80 card priced above it.

Four things the first sm_80 run can falsify

  1. Open until an sm_80 run. The prediction is that the tf32 matmul is measurably less accurate than the ieee one, by at least a factor of a hundred on the normalised error the program prints, because tf32 keeps ten explicit mantissa bits against f32's twenty-three. The program gates on the direction and not the factor. If the two errors come back equal, one of the runs did not use the requested multiplier.
  2. Open until an sm_80 run. The prediction is that the tf32 matmul is at least twice as fast as the ieee one at 2048 cubed on any card at compute capability 8.0 or newer. This is the one speed claim made anywhere on the page, it compares a kernel against itself with one argument changed, and it is gated.
  3. Partly held. The interpreter passed softmax, and its ieee and tf32 normalized errors were identical at 1.86e-07, a ratio of 1. That confirms the precision gap collapses in the interpreter. Agreement with a real GPU softmax remains untested.
  4. Open until an sm_80 run. The prediction is that the best of the three shapes is not the smallest. 128, 128, 32 moves fewer bytes per output than 64, 64, 32 for the same reason day 44's 128 by 128 tile beat its 16 by 16 one, and it should win at 2048 cubed even though it asks for more registers and leaves fewer warps resident, which is the occupancy trade day 45 put a number on.

This page does not compare Triton and CUDA timings from different GPUs. Day 44's results used a T4 at compute capability 7.5, which cannot run this Triton GPU path. A ratio across different machines would not isolate the language.

Run it yourself

The GPU path needs compute capability 8.0 or newer. Triton supports Linux; use WSL2 on Windows. The setup options are listed at /setup/learn-cuda-without-a-gpu.

The correctness half costs nothing anywhere:

TRITON_INTERPRET=1 python3 triton_softmax_matmul.py

Exercise

Replace the three hand-written shapes in CONFIGS with a triton.autotune decorator over a wider space, at least eight triton.Config objects, keyed on the problem size. Report which shape it picks at 2048 cubed and whether it beats the best row the shipped sweep prints.

Time: 30 to 45 minutes. Submit: autotuned.py, the chosen config, the two times, and one sentence on why the first launch was slow.

Check: the harness runs the correctness half unchanged, both kernels against the float64 reference with the length-dependent tolerance from day 66 printed beside its own arithmetic, so a config that breaks the masks fails before it posts a time. On the measurement side it wants the autotuned time within five percent of the best hand-written row or better; a submission that is slower than a three-shape sweep has a bug in its key or is timing the search. On a card below compute capability 8.0 it prints the card and exits 1 without judging your kernel.

Hint 1

The three shapes in the file form a manual search. The autotuner must know when its cached answer no longer applies. Identify which problem dimensions can change the winning shape.

Hint 2

key=["m", "n", "k"]. The search runs once per distinct key and caches the winner, so the first launch at a new size pays for every config in your list. Look at where do_bench sits relative to that first launch.

Solution

The decorator goes above @triton.jit, takes configs=[...] and key=["m", "n", "k"], and the constexpr arguments it owns come out of the launch call. The first launch at each new key times every config, so it is slow by roughly the size of your list, and a timing loop that includes it reports the search rather than the kernel. Warm up once at the size you intend to measure.

Expect a ratio near one against the best hand-written row. Eight shapes may not improve much on three shapes chosen using day 44's constraints. The autotuner repeats the search when the problem-size key changes, while the fixed list does not.

Pitfalls

Installation succeeded, then the kernel failed. Triton's supported hardware is "NVIDIA GPUs (Compute Capability 8.0+)". pip does not check the GPU, so the failure occurs when the kernel runs.

Use TRITON_INTERPRET=1 for the correctness half, or move to a card at 8.0 or newer.

pip install torch triton==3.8.0 does not resolve. torch 2.13.0's own metadata reads triton==3.7.1; platform_system == "Linux" and python_version < "3.15" (https://pypi.org/pypi/torch/json , checked 2026-09-01), and the project pins triton==3.8.0. Take torch's pin or install torch first and upgrade Triton with --no-deps; the README gives both lines and asks the run to record which one it used.

Your Triton matmul beat your CUDA matmul under different precision modes. tl.dot's input_precision defaults to "tf32", so an FP32 Triton matmul reaches for tensor cores while day 44's cuBLAS comparison deliberately does not.

Pin the string, or say in the caption which multiplier each side used. This is day 44's own CUBLAS_COMPUTE_32F_FAST_TF32 pitfall arriving from the other direction.

You padded the softmax tail with zero. tl.load's other fills the masked lanes, and a 0.0 pad wins the maximum on a row whose real values are all negative, which silently changes every output on that row. The neutral pad for a maximum is negative infinity, and it stays neutral in the sum below it because exp takes it to zero.

You timed something under TRITON_INTERPRET=1. The setting "causes all Triton kernels to bypass compilation and be simulated by the interpreter using numpy equivalents of Triton operations", and "the interpreter processes each Triton program instance sequentially" (https://triton-lang.org/main/programming-guide/chapter-3/debugging.html , checked 2026-09-01). A millisecond figure from that describes Python.

The shipped program refuses to print one.

Your bfloat16 kernel works on the GPU but fails in the interpreter. Same page: "It does not support operations on bfloat16 numeric types. To perform operations on bfloat16 tensors, use tl.cast(tensor) to convert the tensor to float32."

The interpreter supports fewer cases than the GPU compiler.

Go deeper

Next

Day 87 wraps a kernel as a PyTorch operator with a backward pass. Day 89 compares Triton, cuTile, and hand-written CUDA without using unverified benchmark claims.

The kernel fusion argument from day 48 is one body here rather than one launch fewer, which is why an online softmax folded into a GEMM is the shape most Triton in the wild takes. That kernel is day 97, and reading the PTX it compiles to is how you inspect the compiler's choices.