Lecture 06 · Part 2

Kernels, Triton, XLA

Percy turns Tatsu's high-level GPU picture into working code. The recipe is benchmark and profile first, then write Triton kernels, and the lecture builds four of them in increasing difficulty — elementwise GeLU, row-wise softmax, a row sum whose rows overflow a block, and a tiled matmul fused with ReLU. One idea carries all of them, namely load a tile into shared memory, do as much work as possible there, and write back to HBM once.

▶ Lecture video Executable lecture (.py) ↗

The GPU you’re programming, one more time

Tatsu gave the high-level overview on Monday; this lecture dives into the code — Triton kernels, benchmarking, profiling. Percy opens by re-drawing the machine. A GPU has on the order of 100–200 streaming multiprocessors (SMs), a count that has barely moved across generations; what grows is the memory:

	A100	H100	B200
SMs	108	132	148
Registers per SM	256 KB	256 KB	256 KB
L1 + shared memory per SM	192 KB	256 KB	256 KB
L2 cache (whole chip)	40 MB	50 MB	96–126 MB
HBM	80 GB	80 GB	192 GB
HBM bandwidth	2 TB/s	3.35 TB/s	8 TB/s

Bandwidth runs inversely to size: registers fastest (~447 TB/s on B200), then L1/shared memory, then L2, then HBM — “although eight terabytes a second is still not that slow in the grand scheme of things.” Big memory is far and slow; fast memory lives on the SM and is small.

The programming model has three levels. A thread executes code on a small piece of data; a grid is the whole launch; in between sits the level that matters:

Triton, where this lecture is headed, makes you think natively in thread blocks — “once you get the hang of thinking of thread blocks… it makes your life a lot easier.”

Where the abstraction leaks: five hardware considerations

The programming model is clean and, for correctness, all you need. But “performance is very sensitive to the hardware,” and the whole reason to write kernels is performance. Percy walks five places the hardware shows through.

Warps. Within a block, threads are grouped into warps of 32, and all threads in a warp execute the same instruction in lockstep. Control divergence — an if/else where threads take different branches — forces the A-threads and B-threads to run sequentially, which is why branching in kernels is bad. The redeeming feature: an SM keeps multiple warps resident and switches between them at zero cost, so when one warp blocks on a 100-cycle HBM read, another immediately runs. Latency gets hidden, not eliminated.

Warp occupancy. Each thread may use at most 255 registers, and the SM’s pool is fixed, so fatter threads mean fewer of them.

Worked example — Computing occupancy from register pressure

A block of 128 threads, each using 160 registers, on a B200 (65,536 registers per SM, at most 64 concurrent warps):

num_registers_per_block = 128 * 160          # 20,480
num_blocks = 65536 // num_registers_per_block  # 3 blocks fit
num_warps = num_blocks * 128 / 32            # 12 warps
occupancy = num_warps / 64                   # 0.1875

Only 12 of a possible 64 warps run — about 18% occupancy, because register use per thread is high. But low occupancy isn’t automatically bad: with thread coarsening, fewer, fatter threads each process several elements. “If your threads are very light, maybe you want to fatten them up a little bit.”

Bank conflicts (shared memory). Shared memory is divided into 32 banks, 4 bytes wide, and each bank serves at most one thread per cycle. If 32 threads all hit the first column of a row-major matrix, that’s a 32-way bank conflict — fully serialized. It’s unavoidable in general (a matmul must access rows of one matrix and columns of the other); the fix is swizzling, rearranging the shared-memory layout to spread accesses across banks.

Memory coalescing (HBM). When a warp’s 32 threads read HBM, accesses are combined into 128-byte cache-line transactions. Best case — full coalescing — all 32 threads hit the same cache line and one fetch serves everyone; striding down a column fetches lines you mostly throw away. It feels like bank conflicts but is a different constraint on a different memory.

Block occupancy. A B200 has 148 SMs. Launch 160 thread blocks and the first wave runs 148, then a second wave runs 12 while 136 SMs idle — the wave quantization problem. “Maybe it’s a good idea to make the number of thread blocks divide the number of SMs.”

Benchmark, change, benchmark again

Before writing any kernel: “you should always just measure what’s going on in your code and figure out what the bottlenecks are.” The recipe for success is a loop — benchmark and profile, make changes, benchmark and profile again.

Benchmarking measures end-to-end wall-clock time. torch.utils.benchmark exists, but the class rolls its own to expose the gotchas: warm up first (compilation shouldn’t count — steady state is what you’ll live in), time multiple trials for variance, use CUDA events rather than CPU timers, and call torch.cuda.synchronize() before reading the clock, because GPU execution is asynchronous:

start_event = torch.cuda.Event(enable_timing=True)
end_event = torch.cuda.Event(enable_timing=True)
start_event.record()
run()
end_event.record()
torch.cuda.synchronize()  # wait for CUDA threads to finish (important!)
times.append(start_event.elapsed_time(end_event))

Scaling a square matmul from dim 256 up to 8192 shows time roughly constant until around dim 2,000, then cubic — GPUs are shaped for large matmuls, and a small one can’t fill the machine.

Profiling tells you where the time goes — and, independent of time, what’s happening under the hood. Profile a + b and you find a CUDA kernel with a long name containing CUDAFunctor_add. Profile a @ b and you get something like cutlass3x_sm100_..._f32_..._64x64x16_...: cutlass is NVIDIA’s CUDA linear-algebra library, sm100 means a kernel built for Blackwell, f32 the dtype, 64x64x16 the tile shape. Change the matrix from 2048 to 128 and a different kernel with a different tile shape gets invoked. a @ b was never one thing.

GeLU three ways: the case for fusion

Take GeLU’s tanh approximation and implement it three ways:

def naive_gelu(x):
    return 0.5 * x * (1 + torch.tanh(0.79788456 * (x + 0.044715 * x * x * x)))

builtin_gelu = lambda x: torch.nn.functional.gelu(x, approximate="tanh")
compiled_gelu = torch.compile(naive_gelu)

All three agree numerically; benchmarked on a 16,384-dim input, the built-in and compiled versions are dramatically faster than the naive one. The profiler explains why. The naive version launches a separate kernel for every primitive in the computation graph — multiply, add, tanh, each with its own HBM round-trip, “because between kernel invocations, things have to go back to HBM.” The built-in is one hand-written GeluCUDAKernelImpl. And torch.compile inspects the computation graph and emits a single fused kernel — written in Triton.

Your first Triton kernel (and the PTX underneath)

In CUDA you specify what each thread does — fine-grained control, but you manage synchronization and shared memory yourself. In Triton (developed by OpenAI) you specify what each thread block does: load data into shared memory, operate on it, write back to global memory. It’s generally powerful enough, and it reads like Python.

Code walkthrough — Triton GeLU: launch, wake up, load, compute, store

The launch is ordinary PyTorch plus a grid. Triton isn’t functional — you allocate the output yourself and the kernel writes into it:

y = torch.empty_like(x)
num_elements = x.numel()           # 8,192
BLOCK_SIZE = 1024
num_blocks = triton.cdiv(num_elements, BLOCK_SIZE)   # 8 blocks
triton_gelu_kernel[(num_blocks,)](x, y, num_elements, BLOCK_SIZE=BLOCK_SIZE)

Inside, x_ptr and y_ptr are pointers — “you can think about these as just integers, they’re addresses.” Every kernel has the same shape: wake up and figure out who you are, compute your span, load, compute, store.

@triton.jit
def triton_gelu_kernel(x_ptr, y_ptr, num_elements, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(axis=0)      # which block am I?
    start = pid * BLOCK_SIZE
    offsets = start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < num_elements    # don't run past the end

    x = tl.load(x_ptr + offsets, mask=mask)      # read from HBM

    # tl.tanh doesn't exist: tanh(a) = (exp(2a) - 1) / (exp(2a) + 1)
    a = 0.79788456 * (x + 0.044715 * x * x * x)
    exp = tl.exp(2 * a)
    tanh = (exp - 1) / (exp + 1)
    y = 0.5 * x * (1 + tanh)

    tl.store(y_ptr + offsets, y, mask=mask)      # write back to HBM

The mask handles the case where the block size doesn’t divide the tensor — the last block reads and writes only up to the boundary.

Triton compiles down to PTX, an intermediate assembly language for GPUs. Reading it: ld.global and st.global are the HBM reads and writes; %ctaid.x is the block index and %tid.x the thread index (the same compiled code runs on every thread, distinguished only by its IDs); %f* are floating-point registers, %r* integer ones. And the compiler chose thread coarsening on its own: one thread processes eight elements.

Reductions and tiling: softmax, row sum, matmul

The remaining examples climb a ladder: reduction where a row fits in a block, reduction where it doesn’t, and full tiling.

Softmax exponentiates and normalizes each row. The naive PyTorch version — max per row for stability, subtract, exponentiate, sum, divide — is five kernels; counting traffic gives $5MN + M$ reads and $3MN + 2M$ writes for an $M \times N$ matrix, when in principle you need only $MN$ of each: a ~4x speedup on the table. The Triton kernel takes one block per row (block size = the column count rounded up to a power of 2, “for good luck”), loads the row with masked-out entries set to $-\infty$ (the softmax equivalent of zero), and the core computation is nearly verbatim naive code:

x_row = x_row - tl.max(x_row, axis=0)
numerator = tl.exp(x_row)
denominator = tl.sum(numerator, axis=0)
y_row = numerator / denominator

“If you can fit it into a block, you can just write normal PyTorch, almost.”

Row sum, when the row doesn’t fit. With 4,096 columns and block size 1,024, the block must iterate: break the row into tiles, sweep them into a per-thread accumulator, then reduce to a scalar at the end.

acc = tl.zeros([BLOCK_SIZE], dtype=tl.float32)
for start in range(0, N, BLOCK_SIZE):
    cols = start + tl.arange(0, BLOCK_SIZE)
    x = tl.load(x_ptr + row * N + cols, mask=cols < N, other=0.0)
    acc += x
result = tl.sum(acc, axis=0)

Don’t confuse this with the GeLU case: there the pieces of the vector were independent blocks; here the pieces are tiles, all processed by one block, because the sum couples them. This is where Triton stops looking like PyTorch — the data no longer fits in shared memory at once.

Matmul, the finale — with a ReLU fused on, since a linear layer plus activation is exactly this shape. The naive kernel computes one element of $C$ at a time, reading $A[m,k]$ and $B[k,n]$ from HBM for every $k$ : order $MKN$ reads against order $MKN$ FLOPs, arithmetic intensity $O(1)$ . Terrible. The idealized approach loads all of $A$ and $B$ into shared memory — $MK + KN$ reads, intensity $O(N)$ — but real matrices don’t fit in a few hundred KB. So:

The implementation (BLOCK_M, BLOCK_N, BLOCK_K = 64, 64, 32, a 2-D grid over output tiles) is the row-sum loop with two-dimensional pointer arithmetic — strides map the (row, column) index into linearized memory:

acc = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
for k in range(0, K, BLOCK_K):
    a = tl.load(a_ptrs, mask=..., other=0.0)
    b = tl.load(b_ptrs, mask=..., other=0.0)
    acc += tl.dot(a, b)              # tiles are on-chip: just matmul them
    a_ptrs += BLOCK_K * stride_ak    # next row tile of A
    b_ptrs += BLOCK_K * stride_bk    # next column tile of B

acc = tl.maximum(acc, 0.0)           # the fused ReLU, free before the store
tl.store(c_ptrs, acc, mask=...)

The fusion bonus is one line: the output tile is already in fast memory before the final store, so applying an elementwise activation there costs nothing — the same lesson torch.compile learned for GeLU. With these ingredients (elementwise, row-wise reduction, tiled reduction, tiled matmul) you have everything needed to implement FlashAttention in the assignment. Next lecture: more than one GPU.

Back of the chalk tray — what to remember

The programming model maps onto the memory hierarchy: grid ↔ HBM, thread block ↔ shared memory (one SM), thread ↔ registers. Correctness needs only the model; performance needs the hardware.
Warps execute 32 threads in lockstep (branching serializes), and SMs hide latency by switching warps at zero cost; occupancy, bank conflicts, memory coalescing, and wave quantization are where the hardware shows through.
Benchmark and profile before writing kernels: warm up, use CUDA events, synchronize; the profiler reveals which CUDA kernels actually run and that they change with tensor shape.
Naive PyTorch launches one kernel per primitive with an HBM round-trip between each; kernel fusion (built-in kernels or torch.compile, which emits Triton) reads once and writes once.
A Triton kernel specifies what one thread block does: identify yourself via program_id, compute masked offsets, tl.load, compute, tl.store.
Tiling is the whole game for matmul: sweep tile pairs through shared memory, accumulate on-chip, write once — arithmetic intensity O(tile size) instead of O(1), with elementwise fusion free at the end.

Check yourself

Hardware arithmetic and read/write accounting, on paper.

Problem 6.1 occupancy arithmetic

A kernel launches thread blocks of 256 threads, each thread using 96 registers, on a B200 SM (65,536 registers, at most 64 concurrent warps). Compute the warp occupancy. Is the constraint registers or the warp cap?

Show solution

Registers per block: $256 \cdot 96 = 24{,}576$ . Blocks per SM limited by registers: $\lfloor 65{,}536 / 24{,}576 \rfloor = 2$ . Warps: $2 \cdot 256 / 32 = 16$ , so occupancy $= 16/64 = 25\%$ . The register pool is the binding constraint — the warp cap would allow 64. A third block would need $73{,}728$ registers total, which doesn’t fit. Whether 25% is bad depends on what each thread does: if threads are coarsened and heavy, low occupancy can be fine.

Problem 6.2 wave quantization

You tile a $4096 \times 4096$ output matrix with $64 \times 64$ output tiles (one thread block each) on a B200 with 148 SMs. How many waves of thread blocks run, and how many SMs sit idle during the last wave?

Show solution

Blocks: $(4096/64)^2 = 64^2 = 4{,}096$ . Waves: $\lceil 4096 / 148 \rceil = 28$ . The first 27 waves occupy all 148 SMs ( $27 \cdot 148 = 3{,}996$ blocks); the last wave has $4{,}096 - 3{,}996 = 100$ blocks, leaving $148 - 100 = 48$ SMs idle. That last partial wave is the wave-quantization problem — the lecture’s suggestion is to pick block counts that divide the SM count when you can.

Problem 6.3 read/write accounting

For an $M \times N$ matrix, the naive PyTorch softmax does $5MN + M$ reads and $3MN + 2M$ writes. Do the same accounting for computing y = relu(a * x + b) elementwise in naive PyTorch (with a, b scalars and x of $n$ elements), and for the fused version. What’s the traffic ratio?

Show solution

Naive: three kernels. Multiply reads $x$ ( $n$ ) and writes $ax$ ( $n$ ); add reads $n$ and writes $n$ ; ReLU reads $n$ and writes $n$ . Total $3n$ reads, $3n$ writes, $6n$ element transfers. Fused: read $x$ once, do all three ops in registers, write $y$ once — $2n$ transfers. Ratio 3x. For a memory-bound elementwise chain, that’s roughly a 3x speedup — the whole argument for kernel fusion, and exactly the pattern torch.compile found for GeLU.

Problem 6.4 tiling intensity

In the tiled matmul with square tiles of size $T$ (take $M = N = K$ and $\text{BLOCK\_M} = \text{BLOCK\_N} = \text{BLOCK\_K} = T$ ), count the HBM reads needed to produce one $T \times T$ output tile, and show the arithmetic intensity is $O(T)$ . What stops you from just growing $T$ forever?

Show solution

One output tile sweeps $K/T$ tile pairs; each iteration loads a $T \times T$ tile of $A$ and of $B$ : $2T^2 \cdot K/T = 2TK$ reads (plus $T^2$ writes at the end). FLOPs for the tile: $2 T^2 K$ (each of $T^2$ outputs needs $K$ multiply-adds). Intensity $\approx 2T^2K / 2TK = T$ , so it grows linearly in the tile size. The ceiling is shared memory: both tiles plus the $T \times T$ accumulator must fit in the SM’s ~256 KB of L1/shared memory, and larger blocks also consume more registers per SM, cutting occupancy — the same trade-offs from problem 6.1.