Custom GPU kernels & Triton
A GPU can do far more arithmetic than it can feed itself data. The kernels that win are the ones that move fewer bytes — not the ones that do less math. This lesson builds the roofline model from first principles, shows why a fused softmax beats three library calls, walks a Triton kernel line by line, and reasons about occupancy — all modeled offline in pure Python, with the real GPU code labeled and left to read.
Learning objectives
- Classify any kernel as memory-bound or compute-bound using arithmetic intensity and the roofline ridge point.
- Explain why elementwise ops (softmax, LayerNorm, activations) are memory-bound and why kernel fusion is a data-movement win.
- Read a Triton kernel and map its
tl.load/tl.store/tl.dotblock operations to the memory hierarchy. - Reason about occupancy — how block size and register/SRAM pressure cap the warps an SM can keep resident.
- Decide, with a checklist, when to hand-write a kernel versus reach for cuBLAS/cuDNN/FlashAttention.
1 · Why kernels matter: the memory wall
A modern accelerator is a bandwidth-starved compute monster. Peak arithmetic throughput has grown far faster than memory bandwidth for two decades. The practical consequence: for a huge fraction of real kernels, the arithmetic units sit idle, waiting for data to arrive from HBM. When you make such a kernel faster, you almost never do it by doing less math — you do it by touching HBM fewer times.
The single number that decides which world you live in is arithmetic intensity (AI): FLOPs performed per byte moved from HBM. Compare it to the machine's ridge point — peak FLOP/s ÷ bandwidth. Below the ridge you are memory-bound (bandwidth caps you); above it you are compute-bound (the ALUs cap you). This is the roofline model, and it is the first thing to compute before optimizing anything.
2 · Lab · roofline & arithmetic intensity, offline
Let's make it a number. Given a kernel's FLOPs and its HBM byte traffic, this classifies it against a chosen GPU and reports the attainable throughput. Pure stdlib — it runs on a laptop and just does the roofline arithmetic.
roofline.pydef roofline(flops, bytes_moved, peak_flops, bw_bytes_per_s):
"""Arithmetic intensity = FLOPs per byte. Below the ridge (peak/bw) => memory-bound."""
ai = flops / bytes_moved # FLOP per byte for the kernel
ridge = peak_flops / bw_bytes_per_s # FLOP/byte where the GPU flips
attainable = min(peak_flops, ai * bw_bytes_per_s)
bound = "memory-bound" if ai < ridge else "compute-bound"
return ai, ridge, attainable, bound
# An H100-class GPU: ~990 TFLOP/s (BF16) and ~3.35 TB/s of HBM bandwidth.
PEAK, BW = 990e12, 3.35e12
N = 4096
for name, flops, byts in [
# softmax over one row: ~3 FLOPs/elem (max, exp, div), read+write the row
("softmax row (elementwise)", 3 * N, 2 * N * 4),
# dense GEMM N^3: 2*N^3 FLOPs, read A,B + write C (all N*N, fp32)
("4096-cube dense GEMM", 2 * N**3, 3 * N * N * 4),
]:
ai, ridge, att, bound = roofline(flops, byts, PEAK, BW)
print(f"{name:28} AI={ai:8.1f} FLOP/B ridge={ridge:6.1f} -> {bound}")
print(f"{'':28} attainable={att/1e12:8.1f} TFLOP/s ({att/PEAK*100:.1f}% of peak)")
softmax row (elementwise) AI= 0.5 FLOP/B ridge= 295.5 -> memory-bound
attainable= 1.7 TFLOP/s (0.2% of peak)
4096-cube dense GEMM AI= 1365.3 FLOP/B ridge= 295.5 -> compute-bound
attainable= 990.0 TFLOP/s (100.0% of peak)The softmax row lives far below the ridge and can only reach a fraction of a percent of peak — it starves for bandwidth. The big GEMM saturates the ALUs. This is exactly why vendors ship a hand-tuned GEMM (cuBLAS) but you fuse the elementwise glue around it: the GEMM is already at the compute roof, and the elementwise ops only cost you HBM trips.
3 · Kernel fusion: the bandwidth win
Every kernel launch reads its inputs from HBM and writes its outputs back. A chain of elementwise ops — say x → bias → GELU → dropout → residual — run as separate kernels round-trips the whole activation tensor through HBM once per op. Fusing them into one kernel keeps the intermediate in registers/SRAM and pays the HBM cost once. For a memory-bound chain the speedup is essentially the traffic-reduction ratio. FlashAttention is the famous instance; the same lever applies to LayerNorm, softmax, and activation stacks.
fusion.pydef fusion_traffic_gib(tensor_elems, k_ops, bytes_per=2):
"""Unfused: each of k elementwise ops does one read + one write.
Fused: one read + one write total for the whole chain."""
per_pass = tensor_elems * bytes_per
unfused = 2 * k_ops * per_pass # k ops x (read + write)
fused = 2 * per_pass # one read + one write
return unfused / 1024**3, fused / 1024**3
u, f = fusion_traffic_gib(tensor_elems=8192 * 8192, k_ops=5)
cut = (u - f) / u * 100
print(f"unfused (5 ops): {u:.2f} GiB HBM traffic")
print(f"fused (1 op) : {f:.2f} GiB HBM traffic")
print(f"traffic cut: {cut:.0f}% -> ~{u/f:.0f}x for a bandwidth-bound chain")
unfused (5 ops): 2.50 GiB HBM traffic
fused (1 op) : 0.50 GiB HBM traffic
traffic cut: 80% -> ~5x for a bandwidth-bound chain4 · A Triton kernel, walked
Writing a fused, tiled kernel traditionally meant CUDA C++. Triton (from OpenAI) lets you write GPU kernels in Python: you express the computation over blocks (tiles), and Triton's compiler handles thread scheduling, SRAM allocation, and memory coalescing. Much of the modern fused-kernel ecosystem — including FlashAttention variants — is authored this way. The block below is real GPU code, here to read, not to run in this offline lesson: a fused softmax over each row of a matrix.
triton_softmax.py# >>> NEEDS A GPU + Triton/PyTorch. Does NOT run in this offline lesson. <<<
# pip install triton torch # and a CUDA GPU
import triton
import triton.language as tl
@triton.jit
def softmax_row_kernel(X, Y, stride, n_cols, BLOCK: tl.constexpr):
row = tl.program_id(0) # one program per row
cols = tl.arange(0, BLOCK)
ptr = X + row * stride + cols
mask = cols < n_cols
x = tl.load(ptr, mask=mask, other=-float("inf")) # row tile -> SRAM (one read)
x = x - tl.max(x, axis=0) # numerically stable softmax, on chip
e = tl.exp(x)
y = e / tl.sum(e, axis=0)
tl.store(Y + row * stride + cols, y, mask=mask) # write once
# Host launch (also needs a GPU):
# import torch
# x = torch.randn(1823, 781, device="cuda")
# y = torch.empty_like(x)
# BLOCK = triton.next_power_of_2(x.shape[1])
# softmax_row_kernel[(x.shape[0],)](x, y, x.stride(0), x.shape[1], BLOCK=BLOCK)
tl.load pulls one row-tile into SRAM once, all the softmax math (max, exp, sum, divide) happens on-chip, and tl.store writes the result once. That is fusion made explicit: one read, one write, everything in between fused. The offline model in the next lab decides the BLOCK size the kernel uses.5 · Occupancy: block size vs resident warps
A kernel's BLOCK size is not free to pick. Each SM has a fixed budget of registers and shared memory (SRAM). The more each block consumes, the fewer blocks (and warps) the SM can keep resident at once. Occupancy is the ratio of resident warps to the hardware maximum. Low occupancy means the SM can't hide memory latency by switching to other warps while one stalls on HBM — so a bigger tile that spills registers can be slower than a smaller one that keeps the SM full.
Occupancy is not a goal in itself — a compute-bound kernel with enough instruction-level parallelism can run fast at modest occupancy. But for the memory-bound kernels that dominate LLM inference, keeping enough warps resident to hide HBM latency is usually the game. This lab models the tradeoff offline.
occupancy.pyimport math
def occupancy(block_threads, regs_per_thread, sram_per_block,
regs_per_sm=65536, sram_per_sm=101376, max_warps_per_sm=64,
max_blocks_per_sm=32):
"""Resident blocks are limited by whichever runs out first: registers,
shared memory, or the hard block cap. Occupancy = resident warps / max warps."""
warps_per_block = math.ceil(block_threads / 32)
regs_per_block = block_threads * regs_per_thread
by_regs = regs_per_sm // regs_per_block
by_sram = sram_per_sm // sram_per_block if sram_per_block else max_blocks_per_sm
blocks = min(by_regs, by_sram, max_blocks_per_sm)
warps = blocks * warps_per_block
return blocks, warps, warps / max_warps_per_sm
best = None
for block in (64, 256, 1024):
blocks, warps, occ = occupancy(block, regs_per_thread=32, sram_per_block=8192)
print(f"BLOCK={block:>4}: {warps:>2} warps/SM resident -> occupancy {occ*100:4.0f}%")
if occ >= 0.5:
best = block
print(f"sweet spot: largest block that stays >=50% occupancy = {best}")
BLOCK= 64: 64 warps/SM resident -> occupancy 100%
BLOCK= 256: 32 warps/SM resident -> occupancy 50%
BLOCK=1024: 8 warps/SM resident -> occupancy 12%
sweet spot: largest block that stays >=50% occupancy = 2566 · When to write a kernel vs use a library
Custom kernels are a sharp tool and a maintenance liability. The default is: use the library. cuBLAS/cuDNN own dense GEMM and convolutions; FlashAttention owns attention; torch.compile and the fused optimizers cover most elementwise glue. You write a kernel only when you have a specific memory-bound pattern the libraries don't fuse and you've proven with a roofline that bytes — not FLOPs — are the wall.
| Situation | Reach for |
|---|---|
| Dense matmul / conv | cuBLAS / cuDNN — already at the compute roof |
| Attention (any length) | FlashAttention / fused-attn library |
| Standard elementwise chain | torch.compile / fused kernels — auto-fused |
| Novel memory-bound fusion the libs miss | Write it (Triton first, CUDA if needed) |
| Compute-bound and vendor lib exists | Do not hand-write — you won't beat it |
A colleague benchmarks a LayerNorm and finds it slow, so they propose rewriting the mean/variance math with fewer operations. Why is this likely the wrong lever, and what is the right one?
Show answer
You increase a Triton kernel's BLOCK size expecting it to go faster (more work per program), but it slows down. Give the most likely cause.
Show answer
✓ Checkpoint — you can move on when you can…
- Compute a kernel's arithmetic intensity and classify it against the roofline ridge point.
- Explain why elementwise ops are memory-bound and why fusion is a data-movement win, not a math one.
- Read a Triton kernel and map its load/compute/store to the SRAM/HBM hierarchy.
- Reason about occupancy: how block size trades against resident warps and latency hiding.
- Apply the checklist to decide when to hand-write a kernel versus use a library.
🪜 Practice ladder beginner → industry
Six graded exercises, easy to real-world. Try each before opening its solution.
Context: Every kernel optimization starts by asking which wall you are hitting — bandwidth or ALUs. That answer decides whether you cut bytes or cut FLOPs.
Your task: Write a helper that takes an op's FLOPs and bytes moved and classifies it against a machine's FLOP-per-byte balance.
Requirements:
- Arithmetic intensity =
FLOPs / bytes_moved - Machine ratio (ridge) =
peak_FLOP/s / HBM_bytes/s - Compute-bound when intensity exceeds the ridge, else memory-bound
- Show a large matmul is compute-bound and an elementwise add is memory-bound
💡 Hint: Softmax, LayerNorm, and activations are the memory-bound ops fusion targets.
Show solution
Arithmetic intensity vs the machine's FLOP:byte ratio decides the wall. Runnable:
def bound(flops, bytes_moved, tflops=990.0, hbm_tb_s=3.35):
intensity = flops / bytes_moved # FLOP per byte
ridge = (tflops * 1e12) / (hbm_tb_s * 1e12) # machine FLOP:byte balance
kind = "compute-bound" if intensity > ridge else "memory-bound"
return f"AI={intensity:8.1f} ridge={ridge:6.1f} -> {kind}"
print("4096-cube matmul", bound(flops=2*4096**3, bytes_moved=3*4096**2*4))
print("elementwise add ", bound(flops=4096**2, bytes_moved=3*4096**2*4))
# matmul is compute-bound; the add is memory-bound (why fusion helps)The matmul does thousands of FLOPs per byte and saturates the ALUs; the add does a fraction of a FLOP per byte and starves for bandwidth. Softmax and LayerNorm behave like the add — which is exactly what kernel fusion targets.
Context: The roofline caps achievable FLOP/s at min(peak, intensity × bandwidth). Plotting an op against it shows how much performance you are leaving on the table.
Your task: Implement the roofline cap and report attainable TFLOP/s for a low- and a high-intensity op on a chosen GPU.
Requirements:
- Attainable =
min(peak_flops, intensity × bandwidth) - Report the ridge-point intensity (
peak / bandwidth) - Show the low-intensity op is stuck far below peak
- Show the high-intensity op is near peak
💡 Hint: Below the ridge you are bandwidth-limited; the attainable line is intensity times bandwidth.
Show solution
The roofline caps attainable throughput at min(peak, intensity×bandwidth). Runnable:
def attainable_tflops(intensity, peak_tflops=990.0, bw_tb_s=3.35):
bw_bytes = bw_tb_s * 1e12
mem_bound = intensity * bw_bytes / 1e12 # TFLOP/s allowed by bandwidth
return min(peak_tflops, mem_bound)
ridge = (990.0 * 1e12) / (3.35e12) # intensity at the corner
print(f"ridge point intensity = {ridge:.1f} FLOP/byte")
for name, I in [("softmax (low AI)", 0.5), ("big GEMM (high AI)", 1365.0)]:
a = attainable_tflops(I)
print(f"{name:20}: attainable = {a:7.1f} TFLOP/s ({a/990*100:.1f}% of peak)")
# low-AI op is stuck near 0% of peak; high-AI op saturates the ALUsBelow the ridge, attainable throughput is intensity × bandwidth — so a low-intensity
op is pinned far below peak no matter how fast the ALUs are. Fusion's whole job is to raise intensity and
slide the op rightward toward the compute roof.
Context: For a memory-bound elementwise chain, fusion's speedup is roughly the traffic-reduction ratio: K ops that each read+write collapse into one read and one write.
Your task: Model the HBM traffic saved by fusing K elementwise ops into a single kernel, and express it as a percentage cut.
Requirements:
- Unfused traffic =
2 × K × tensor_bytes - Fused traffic =
2 × tensor_bytes - Report unfused, fused, and the saving in GiB
- Express the saving as a percentage traffic cut
💡 Hint: No new math — just fewer trips to HBM; the percentage cut is the expected speedup.
Show solution
Fusion collapses K read+write passes into one. Runnable:
def fusion_saving_gib(tensor_elems, k_ops, bytes_per=2):
per_pass = tensor_elems * bytes_per
unfused = 2 * k_ops * per_pass # k ops x (read + write)
fused = 2 * per_pass # one read + one write
saved = (unfused - fused) / 1024**3
return unfused/1024**3, fused/1024**3, saved
u, f, s = fusion_saving_gib(tensor_elems=8192*8192, k_ops=5)
print(f"unfused={u:.2f} GiB fused={f:.2f} GiB saved={s:.2f} GiB")
print(f"traffic cut: {s/u*100:.0f}% (~{u/f:.0f}x for a bandwidth-bound chain)")Fusing 5 memory-bound ops cuts HBM traffic ~80% — a ~5× speedup for a bandwidth-bound chain, with no change to the arithmetic. That is the entire point of fusion: fewer trips to HBM.
Context: A block that spills registers starves the SM of resident warps, and a memory-bound kernel then cannot hide HBM latency. Occupancy quantifies that.
Your task: Given per-SM register and shared-memory budgets and a per-block cost, compute how many blocks fit and the resulting occupancy across a sweep of block sizes.
Requirements:
- Resident blocks =
min(regs_per_sm // regs_per_block, sram_per_sm // sram_per_block) - Warps per block scale with block size (32 threads per warp)
- Occupancy = resident warps / hardware max warps per SM
- Report the largest block that stays at or above 50% occupancy
💡 Hint: The binding constraint is whichever of registers or SRAM runs out first.
Show solution
Occupancy is limited by whichever of registers/SRAM runs out first. Runnable:
import math
def occupancy(block, regs_per_thread=32, sram_per_block=8192,
regs_per_sm=65536, sram_per_sm=101376, max_warps=64, max_blocks=32):
warps_per_block = math.ceil(block / 32)
by_regs = regs_per_sm // (block * regs_per_thread)
by_sram = sram_per_sm // sram_per_block
blocks = min(by_regs, by_sram, max_blocks)
warps = blocks * warps_per_block
return warps, warps / max_warps
best = None
for block in (64, 128, 256, 512, 1024):
warps, occ = occupancy(block)
tag = " <-- register pressure" if occ < 0.5 else ""
print(f"BLOCK={block:>4}: {warps:>2} warps -> {occ*100:4.0f}% occupancy{tag}")
if occ >= 0.5:
best = block
print(f"largest block at >=50% occupancy: {best}")As the block grows, per-block register use climbs and the SM keeps fewer blocks resident, so occupancy falls. For a memory-bound kernel that needs resident warps to hide HBM latency, the sweet spot is the largest block that stays above the occupancy floor — not the biggest block you can compile.
Context: Before touching Triton, the fused-softmax kernel is a tiling decision: pick a power-of-two BLOCK that fits SRAM, count tiles per row, and only then encode it as GPU code.
Your task: Write an offline planner that chooses the largest power-of-two BLOCK fitting in SRAM and the row, and reports tiles per row; keep the real Triton kernel labeled needs-GPU.
Requirements:
- Max elements per tile =
SRAM_bytes / bytes_per_element - BLOCK is the largest power of two ≤ both the SRAM limit and the row length
- tiles per row =
ceil(N / BLOCK) - Keep the Triton kernel as reference code labeled needs-GPU
💡 Hint: The kernel only encodes the tiling you decide here; the planner is pure arithmetic.
Show solution
Decide the tiling offline; the Triton kernel only encodes it. Runnable planner:
def tile_plan(n_cols, sram_kib=100, bytes_per=4):
max_elems = sram_kib * 1024 // bytes_per
block = 1
while block * 2 <= min(max_elems, n_cols):
block *= 2 # largest power-of-two block that fits
tiles = -(-n_cols // block) # ceil division
return block, tiles
block, tiles = tile_plan(n_cols=8192)
print(f"BLOCK={block}, tiles/row={tiles}")# needs GPU + Triton -- the SAME tiling as a kernel (documented Triton API):
# import triton, triton.language as tl
# @triton.jit
# def softmax_row(X, Y, stride, n, BLOCK: tl.constexpr):
# row = tl.program_id(0)
# cols = tl.arange(0, BLOCK)
# x = tl.load(X + row*stride + cols, mask=cols < n, other=-float("inf"))
# x = x - tl.max(x, axis=0)
# e = tl.exp(x); y = e / tl.sum(e, axis=0)
# tl.store(Y + row*stride + cols, y, mask=cols < n)The arithmetic that matters — how big a tile fits SRAM and how many tiles cover a row — is decided on a laptop. The kernel is a faithful transcription, which is why you can reason about the whole design without a GPU in front of you.
Context: A team hands you five slow ops and a deadline. Blindly hand-writing kernels burns weeks; the roofline plus a library-coverage check tells you which (if any) are worth it.
Your task: Write a triage function that, given each op's arithmetic intensity, whether a vendor library covers it, and whether it is on the hot path, recommends library / fuse / hand-write / leave-it.
Requirements:
- Compute-bound with a vendor lib → use the library (won't beat it)
- Memory-bound but covered by an auto-fuser (torch.compile) → let it fuse
- Memory-bound, on hot path, no lib fuses this exact pattern → hand-write (Triton)
- Not on the hot path → leave it and move on
- Return a per-op recommendation with a one-line reason
💡 Hint: The decision is a small truth table over (bound, lib-covered, hot-path) — encode it directly.
Show solution
Triage is a small truth table over (bound, lib-covered, hot-path). Runnable:
def triage(name, intensity, ridge, lib_covers, auto_fuses, hot_path):
memory_bound = intensity < ridge
if not hot_path:
return name, "LEAVE IT", "not on the hot path -- not worth the effort"
if not memory_bound and lib_covers:
return name, "USE LIBRARY", "compute-bound + vendor lib -- you won't beat cuBLAS/cuDNN"
if memory_bound and auto_fuses:
return name, "AUTO-FUSE", "memory-bound but torch.compile fuses this -- let it"
if memory_bound and not lib_covers:
return name, "HAND-WRITE", "memory-bound, hot, no lib fuses it -- Triton kernel"
return name, "USE LIBRARY", "covered by a vendor library"
RIDGE = 295.0
ops = [
("dense GEMM", 1365.0, True, False, True),
("layernorm chain", 0.6, False, True, True),
("novel fused mask", 0.4, False, False, True),
("startup reshape", 0.2, False, True, False),
]
for name, ai, lib, fuse, hot in ops:
n, rec, why = triage(name, ai, RIDGE, lib, fuse, hot)
print(f"{n:18} -> {rec:12} ({why})")Only the novel memory-bound op with no library coverage earns a hand-written kernel. The GEMM goes to cuBLAS, the standard chain to the auto-fuser, and the off-hot-path op is left alone. That triage is worth more than raw kernel skill: most of the wins are choosing not to write a kernel.