Custom GPU kernels & Triton

1 · Why kernels matter: the memory wall

Start with a benchmark that lied to you. You profile your model and find softmax eating a surprising slice of wall-clock time. So you do the obvious thing: you count the FLOPs. Softmax over a row is a handful of operations per element — a max, an exponential, a divide. Almost nothing. The GEMM next to it does billions of FLOPs and somehow finishes faster per byte touched. You stare at the numbers and they make no sense: the cheap op is the slow one. Nothing about "count the math" explains it, because you're measuring the wrong resource.

Here's the reframe that fixes it, in one sentence: a GPU can do far more arithmetic than it can feed itself data. Think of the ALUs as a huge kitchen full of chefs and HBM (the GPU's main memory) as a single narrow doorway they all fetch ingredients through. Adding chefs doesn't help if they're all queued at the door. For softmax, the "cooking" (the exp, the divide) is trivial, but every element still has to walk through that doorway to be read and walk back out to be written — so the doorway, not the chefs, sets the pace. That is what "memory wall" means: for a large class of real kernels the arithmetic units sit idle, waiting on bytes.

This inverts the beginner instinct. The naive optimizer asks "how do I do less math?" The kernel engineer asks "how do I touch HBM fewer times?" On memory-bound work, shaving FLOPs off an op that's already starving for bandwidth changes nothing — you optimized the chefs while the doorway stayed the same width. The whole discipline of custom kernels, and everything in this lesson — fusion, FlashAttention, tiling, mixed precision — is a set of tricks for moving fewer bytes, not doing less arithmetic. The common, expensive mistake is to profile, see a slow elementwise op, and rewrite its math; you can burn a week that way and gain nothing.

So before you touch a single line, you compute one number that tells you which world you're in. That number is arithmetic intensity, and the picture that reads it is the roofline — introduced next and used in every lab that follows.

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

You've heard "memory-bound" and "compute-bound" thrown around in code reviews as if everyone can eyeball which is which. They can't — and neither can you, reliably, until you turn the words into a division. The roofline model exists so the judgment stops being a vibe and becomes a calculation you can do on a napkin before you write any kernel.

Two numbers do all the work. The first is arithmetic intensity (AI): how many FLOPs the op performs for every byte it moves from HBM — FLOPs ÷ bytes. Picture it as fuel economy: miles per gallon, but "math per byte." A dense matmul reuses each byte it loads across many multiply-adds, so it has high intensity (a fuel-efficient op). Softmax reads a value, does almost nothing to it, writes it back — one or two FLOPs per byte, terrible intensity. The second number belongs to the machine, not the op: the ridge point = peak FLOP/s ÷ memory bandwidth. It's the break-even intensity where the hardware flips from being bottlenecked by bytes to being bottlenecked by arithmetic.

• AI below the ridge → memory-bound: the doorway is the limit; the cure is fewer bytes (fuse, tile, use lower precision).
• AI above the ridge → compute-bound: the ALUs are the limit; you're already near the roof and a vendor library probably owns this op.

The subtle part is that the ridge is a property of the GPU you happen to be running on, not the algorithm. On an H100 the ridge sits near ~295 FLOP/byte (as the lab computes), which is startlingly high — it means most elementwise ops, and even many "reasonable-looking" kernels, land well below it and are memory-bound. That single fact explains why so much kernel work is really data-movement work. The mistake to avoid: quoting an op as "memory-bound" in the abstract. Bound relative to what machine? A kernel that's memory-bound on an H100's fat bandwidth can be a different story on hardware with a different FLOP-to-bandwidth ratio. The lab below makes the whole thing concrete: feed it an op's FLOPs and bytes, pick a GPU, and it tells you the side of the ridge you're on and what fraction of peak you can even hope to reach.

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

Picture a five-step recipe where every step makes you walk to the pantry and back. You fetch the flour, do one thing, return the bowl to the pantry. Then walk back, fetch the same bowl, do the next thing, return it. Five steps, ten trips down the hall — even though the bowl never really needed to leave your hands. That is exactly what a chain of elementwise ops does when each runs as its own GPU kernel: bias, then GELU, then dropout, then residual — every one launches separately, reads the whole activation tensor from HBM, does its little bit of math, and writes the whole tensor back. The intermediate results are round-tripped through the slow doorway again and again for no reason.

Kernel fusion is keeping the bowl on the counter. You merge the chain into a single kernel: read the tensor from HBM once, do all the math while the data lives in fast on-chip registers/SRAM, and write the result back once. The arithmetic is identical — you didn't remove a single FLOP — but the trips to HBM collapse from "one read + one write per op" to "one read + one write total." For a memory-bound chain, where the doorway is the bottleneck, the speedup is essentially the traffic-reduction ratio: fuse five ops and you can cut HBM traffic roughly fivefold, and the kernel runs roughly that much faster.

This is not a niche trick — it's the shape of most of the famous wins. FlashAttention is fusion applied to the attention block (it never writes the giant intermediate score matrix to HBM at all). Fused LayerNorm, fused softmax, fused activation-plus-bias stacks are the same idea at smaller scale. Once you see it, you see it everywhere: the win is the byte you didn't move.

But there's a trap the roofline warns you about, and it catches people who over-apply fusion. Fusion only helps the ops that were memory-bound to begin with. If the dominant cost in your chain is a big matmul that's already compute-bound — already pinned at the ALU roof — then fusing the cheap elementwise glue around it barely moves the total, because those ops were never the wall. The discipline is always the same order: roofline first to find which part is bandwidth-starved, then fuse that part. The lab below turns the traffic math into a number so you can predict the payoff before writing the kernel.

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 chain

4 · A Triton kernel, walked

For years, "write a fused kernel" was a sentence that ended most conversations. It meant CUDA C++: managing threads and warps by hand, allocating shared memory yourself, worrying about which thread reads which byte so the hardware can coalesce the accesses, and debugging with tools that assume you already think in grids and blocks. The idea for the kernel was often ten lines; the CUDA to express it safely was two hundred. So most engineers who knew the fusion win existed still couldn't reach for it — the cost of authoring was the barrier, not the concept.

Triton (from OpenAI) lowers that barrier by letting you write GPU kernels in Python at the level of blocks — tiles of data — instead of individual threads. You say "each program instance handles one row" and "load this block into fast memory"; Triton's compiler decides how to map that onto threads and warps, how to lay out SRAM, and how to coalesce the memory accesses so neighboring lanes read neighboring addresses. You describe the tiling; it handles the plumbing. This is why so much of the modern fused-kernel ecosystem, including FlashAttention variants, is authored in Triton rather than raw CUDA.

A few vocabulary anchors before you read the code, because the mental model matters more than the syntax. A @triton.jit function is the kernel — the body that runs on the GPU, compiled just-in-time. tl.program_id(0) is "which instance am I?" — the launch fires many copies of the kernel in parallel, and each asks its ID to figure out which slice of data it owns (here, its row). tl.arange(0, BLOCK) builds the vector of column offsets that instance will touch; tl.load pulls that tile from HBM into on-chip memory once; the math runs on-chip; tl.store writes the result back once. The mask handles the ragged edge when a row's length isn't a clean multiple of the tile.

Watch for the fusion signature as you read: exactly one tl.load and one tl.store, with all the softmax work — max, subtract, exp, sum, divide — sandwiched between them and never spilling back to HBM. That is section 3's principle made literal in code: one read, one write, everything fused in the middle. The block below is real GPU code, here to read, not to run in this offline lesson — but you'll be able to reason about every line, and the next lab decides the BLOCK size it should use.

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)

5 · Occupancy: block size vs resident warps

Here's a result that feels like a bug the first time you hit it. You've got a working Triton kernel, and the obvious way to make it faster is to give each program more work — so you bump BLOCK from 256 to 1024, expecting fewer launches and better efficiency. You run it and it gets slower. You didn't change the math, you didn't add memory traffic, you just made each tile bigger, and the kernel regressed. Every kernel author trips over this once; the explanation is occupancy, and it's the reason "bigger tile" is not the same as "faster kernel."

The intuition first. A GPU hides the slowness of HBM not by making any single memory read fast, but by keeping many groups of threads (warps) resident on each SM at once — when one warp stalls waiting for bytes from the doorway, the SM instantly switches to another warp that's ready to compute. It's a call center: while one agent is on hold, the others keep taking calls, so the line never goes quiet. But each SM has a fixed budget of registers and shared memory, and every resident block spends from that budget. A bigger BLOCK uses more registers and SRAM per block, so fewer blocks — fewer warps — fit at once. Push the tile too big and you have a call center with three agents: the moment they're all on hold, everyone waits, and the memory-bound kernel stalls.

Occupancy is the name for that ratio: resident warps ÷ the hardware maximum. It's capped by whichever resource runs out first — registers, shared memory, or a hard per-SM block limit — so the binding constraint is a min over those three. Now the paradox dissolves: the big block regressed because it dropped occupancy below the point where the SM could hide HBM latency, and a memory-bound kernel lives or dies on latency hiding.

One honest caveat, so you don't over-index on the number: occupancy is a means, not the goal. A compute-bound kernel with lots of independent work per thread (high instruction-level parallelism) can run near the roof at modest occupancy — it doesn't need a deep bench of warps because it's rarely waiting. But for the memory-bound kernels that dominate LLM inference, keeping enough warps resident to cover the doorway latency is usually the whole game. The mistake is chasing 100% occupancy blindly, or its opposite — grabbing the largest block that compiles. The right move, which the lab models offline, is to pick the largest tile that still keeps occupancy above the floor where latency stays hidden.

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 = 256

6 · When to write a kernel vs use a library

The most expensive kernel is the one you didn't need to write. There's a seductive failure mode for engineers who've just learned all of the above: a profiler flags a slow op, and — armed with roofline and Triton — you spend a week hand-crafting a beautiful fused kernel, only to discover cuBLAS was already faster, or torch.compile would have fused it for free, or the op wasn't even on the hot path. The skill this section teaches is mostly the skill of restraint: knowing when the answer is "use the library and go home."

The reason restraint usually wins is that the vendor libraries represent thousands of engineer-hours of tuning against specific hardware, and they own the compute-bound territory outright. cuBLAS and cuDNN are at or near the compute roof for dense matmul and convolution — you will not beat them, and trying is a trap because the roofline already told you those ops are ALU-limited, not something a clever fusion can rescue. FlashAttention owns attention at any sequence length. torch.compile and the fused optimizers sweep up most of the standard elementwise glue automatically. Against all of that, a hand-written kernel is not just effort spent — it's a permanent liability: code that must be re-tuned for the next GPU generation, tested, and maintained by someone after you've moved on.

So the space where a custom kernel actually earns its keep is narrow and specific: a memory-bound pattern (roofline-proven — bytes, not FLOPs, are the wall), that no library fuses for you, sitting on the hot path where the win is worth the maintenance, and expressible as tiles that fit SRAM. Miss any one of those and the honest recommendation flips back to "use the library." The common mistake is to reason from "this op is slow" straight to "so I'll write a kernel," skipping the two questions in between — which wall, and does something already fuse it. The table and checklist below turn that judgment into a decision you can defend in a design review, and the final triage lab encodes it as a small truth table over (bound, library-covered, hot-path).

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.

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)

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 ALUs

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)")

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}")

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)

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})")