Attention & GPU kernels

1 · A mental model of the GPU essential

Before kernels make sense you need a picture of the machine. A GPU is not one big processor — it's a few dozen to ~130 Streaming Multiprocessors (SMs), each a small independent core cluster with its own tiny, blisteringly fast on-chip memory. Around them sits a large pool of HBM (High-Bandwidth Memory) — gigabytes, but far away and comparatively slow to reach.

Two facts drive everything below. First, SRAM is tiny but ~10× the bandwidth of HBM and sits right next to the compute. Second, every trip to HBM is expensive — modern GPUs can do far more arithmetic per second than they can feed themselves data. So the game a kernel plays is: get data into SRAM once, do as much math on it as possible, and touch HBM as little as you can.

2 · Naive attention, byte by byte essential

Recall the attention math: scores S = QKᵀ (an N×N matrix for a sequence of length N), then P = softmax(S), then output O = P·V. A textbook implementation runs these as three separate kernels — and that's the problem. The N×N score matrix is materialized in HBM: written out by the matmul, read back for softmax, written again, read again for the second matmul. For N=4096 that intermediate is 4096×4096 — tens of MB — round-tripped through slow memory several times, while the actual arithmetic per byte is small.

3 · Lab · an HBM-traffic estimator intermediate

Let's make the traffic concrete. This models the bytes each approach moves through HBM: naive materializes and round-trips the N×N scores and probabilities; FlashAttention (next section) streams Q/K/V once and writes the output once, never materializing N×N. Pure stdlib — it runs offline and just counts bytes.

hbm_traffic.pydef bytes_moved(seq_len, d_head, dtype_bytes=2, tile=128):
    N, d = seq_len, d_head
    # Naive: read Q,K,V once, WRITE the N*N scores to HBM, READ them back for
    # softmax, WRITE softmax probs, READ them for the V-matmul, WRITE output.
    qkv = 3 * N * d * dtype_bytes
    scores_rw = 2 * (N * N * dtype_bytes)   # write + read scores
    probs_rw  = 2 * (N * N * dtype_bytes)   # write + read softmax probs
    out_w     = N * d * dtype_bytes
    naive = qkv + scores_rw + probs_rw + out_w
    # FlashAttention: tile K/V; Q,K,V each streamed once, output written once.
    # The N*N matrix is NEVER written to HBM -- it lives in on-chip SRAM per tile.
    flash = 3 * N * d * dtype_bytes + N * d * dtype_bytes
    return naive, flash

for N in (1024, 4096, 16384):
    naive, flash = bytes_moved(N, d_head=64)
    ratio = naive / flash
    print(f"N={N:>6}: naive={naive/1e6:8.2f} MB  flash={flash/1e6:6.2f} MB  "
          f"naive/flash={ratio:5.1f}x")

N=  1024: naive=    8.91 MB  flash=  0.52 MB  naive/flash= 17.0x
N=  4096: naive=  136.31 MB  flash=  2.10 MB  naive/flash= 65.0x
N= 16384: naive= 2155.87 MB  flash=  8.39 MB  naive/flash=257.0x

The gap widens with N because naive traffic grows as O(N²) (the score/prob matrices) while flash traffic grows as O(N·d). At a 16K context the naive path moves ~250× more bytes — and since attention is memory-bound, roughly that much more time on HBM.

4 · FlashAttention: tile + online softmax advanced

FlashAttention's insight: you never need the whole N×N score matrix at once. Process attention in tiles — load a block of queries and a block of keys/values into SRAM, compute that block's partial scores on chip, fold them into a running output, then move to the next block. The N×N matrix is computed but never written to HBM; only the O(N·d) inputs and output touch HBM. Memory drops from O(N²) to O(N), and HBM round-trips collapse — that's the speedup.

The trick that makes tiling correct is online softmax. Softmax normally needs the whole row (to subtract the max and divide by the sum). Online softmax keeps a running max m, a running denominator l, and a running weighted-value accumulator — and rescales the accumulators whenever a new tile reveals a larger max. The final result is bit-for-bit the same as computing softmax over the full row at once. The next lab proves that.

5 · Lab · online softmax equals batch softmax advanced

This is the mathematical heart of FlashAttention, in pure Python. We compute a softmax-weighted sum of values two ways — the full "materialize everything" batch softmax, and the streaming/tiled online softmax — and show they agree exactly. If this equality holds, tiling is safe.

online_softmax.pyimport math

def batch_softmax_weighted_sum(scores, values):
    """The 'materialize everything' path: full softmax over all scores, then dot V."""
    m = max(scores)
    exps = [math.exp(s - m) for s in scores]
    Z = sum(exps)
    probs = [e / Z for e in exps]
    return sum(p * v for p, v in zip(probs, values))

def online_softmax_weighted_sum(scores, values, tile=4):
    """FlashAttention's trick: stream tiles, keep a running max m, running
    denominator l, and running output acc -- rescaling on the fly. Never stores
    the full score row."""
    m = -math.inf   # running max
    l = 0.0         # running sum of exp(scores - m)
    acc = 0.0       # running weighted sum of values
    for i in range(0, len(scores), tile):
        s_tile = scores[i:i+tile]
        v_tile = values[i:i+tile]
        m_new = max(m, max(s_tile))
        # rescale the old accumulators to the new max
        correction = math.exp(m - m_new) if m != -math.inf else 0.0
        l = l * correction
        acc = acc * correction
        for s, v in zip(s_tile, v_tile):
            p = math.exp(s - m_new)
            l += p
            acc += p * v
        m = m_new
    return acc / l

scores = [3.0, 1.0, 2.5, 0.2, 4.1, -1.0, 2.2, 0.7, 3.3, 1.9]
values = [10, 20, 30, 40, 50, 60, 70, 80, 90, 100]
batch  = batch_softmax_weighted_sum(scores, values)
online = online_softmax_weighted_sum(scores, values, tile=4)
print(f"batch  softmax\u00b7V = {batch:.10f}")
print(f"online softmax\u00b7V = {online:.10f}")
print(f"max abs diff     = {abs(batch - online):.2e}")
print("MATCH" if abs(batch - online) < 1e-9 else "MISMATCH")

batch  softmax·V = 53.6863940430
online softmax·V = 53.6863940430
max abs diff     = 0.00e+00
MATCH

Exactly equal. That equality is the license to tile: because the streaming update reconstructs the same normalization, FlashAttention can process attention block-by-block in SRAM and skip the N×N HBM matrix entirely, with no loss of numerical fidelity.

6 · Kernel fusion & the roofline model expert

FlashAttention is one instance of a general lever: kernel fusion. Every separate kernel launch reads its inputs from HBM and writes its outputs back. Fusing several operations into one kernel keeps the intermediate in SRAM/registers and pays the HBM cost only once — the same reason softmax(x)·V fused beats three separate passes. Fusion is almost always a memory-traffic optimization, not a compute one.

The roofline model tells you when that matters. Plot a kernel by its arithmetic intensity (FLOPs performed per byte moved from HBM). The machine has a ridge point = peak-FLOPs ÷ bandwidth. Below the ridge you're memory-bound (bandwidth caps you; move fewer bytes); above it you're compute-bound (the ALUs cap you; do less math or use lower precision).

roofline.pydef roofline(flops, bytes_moved, peak_flops, bw_bytes_per_s):
    """Arithmetic intensity = FLOPs per byte moved. Compare to the machine's
    ridge point (peak_flops / bandwidth). Below the ridge => memory-bound."""
    ai = flops / bytes_moved                      # FLOP/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 = 990e12
BW   = 3.35e12

for name, flops, byts in [
    ("naive attention row",  2 * 4096 * 64,        4096 * 2),   # ~memory-bound
    ("big fused GEMM",        2 * 4096**3,          3 * 4096**2 * 2),
]:
    ai, ridge, att, bound = roofline(flops, byts, PEAK, BW)
    print(f"{name:22} AI={ai:8.1f} FLOP/B  ridge={ridge:6.1f}  -> {bound}")
    print(f"{'':22} attainable={att/1e12:7.1f} TFLOP/s")

naive attention row    AI=    64.0 FLOP/B  ridge= 295.5  -> memory-bound
                       attainable=  214.4 TFLOP/s
big fused GEMM         AI=  1365.3 FLOP/B  ridge= 295.5  -> compute-bound
                       attainable=  990.0 TFLOP/s

The low-intensity attention step sits well below the ridge — it can only reach ~214 of 990 TFLOP/s because it starves for bandwidth. The big GEMM is above the ridge and saturates the ALUs. Roofline is how you decide which knob to turn: for the memory-bound case the answer is fusion + fewer HBM trips (exactly FlashAttention); for the compute-bound case it's lower precision or fewer FLOPs.

7 · Triton at a glance expert

Writing a fused, tiled kernel like FlashAttention traditionally meant hand-written CUDA C++. Triton (from OpenAI) lets you write GPU kernels in Python: you express the computation over blocks (tiles), and Triton's compiler handles the SM scheduling, SRAM allocation, and memory coalescing. It's how much of the modern fused-kernel ecosystem (including FlashAttention variants) is now authored — Python-level control with near-CUDA performance.

The snippet below is real GPU code and is here to read, not to run in this course's offline environment. It sketches the shape of a Triton attention kernel: grab a tile, loop over key/value blocks, keep the online-softmax running state you just proved correct, and write the output once.

triton_flash.py# >>> NEEDS A GPU + Triton/PyTorch. Does NOT run in this offline lesson. <<<
# Illustrative shape of a fused, tiled attention kernel (the FlashAttention idea).
# pip install triton torch  # and a CUDA GPU
import triton
import triton.language as tl

@triton.jit
def flash_attn_kernel(Q, K, V, Out, N, d, BLOCK: tl.constexpr):
    pid = tl.program_id(0)                    # this program handles one Q tile
    q = tl.load(Q + pid * BLOCK)              # load a query tile into SRAM
    m = tl.full((BLOCK,), -float("inf"), tl.float32)   # running max
    l = tl.zeros((BLOCK,), tl.float32)                 # running denominator
    acc = tl.zeros((BLOCK, d), tl.float32)             # running output

    for start in range(0, N, BLOCK):         # stream K/V tiles through SRAM
        k = tl.load(K + start)
        v = tl.load(V + start)
        s = tl.dot(q, tl.trans(k))           # partial scores, on chip
        m_new = tl.maximum(m, tl.max(s, axis=1))
        p = tl.exp(s - m_new[:, None])       # online-softmax rescale...
        corr = tl.exp(m - m_new)
        l = l * corr + tl.sum(p, axis=1)
        acc = acc * corr[:, None] + tl.dot(p, v)
        m = m_new

    tl.store(Out + pid * BLOCK, acc / l[:, None])   # write output once

8 · Tech-lead — reasoning about the memory wall tech-lead

A tech lead doesn't hand-write kernels day to day — but they make the decisions that hinge on this model. Is this workload memory- or compute-bound? (Roofline it before optimizing.) Should we adopt a fused-attention library or a Triton kernel? (Almost always yes for long context — the HBM savings are the win.) Does a longer context window blow up memory or just compute? (With FlashAttention it's O(N) memory, so context length is bounded by KV-cache, not the score matrix — connect to IC3.) The through-line: on modern accelerators, performance is a data-movement problem first and an arithmetic problem second.

def bound(flops, bytes_moved, tflops=312.0, hbm_gib_s=2000.0):
    intensity = flops / bytes_moved                 # FLOPs per byte
    machine_ratio = (tflops * 1e12) / (hbm_gib_s * 1024**3)
    kind = "compute-bound" if intensity > machine_ratio else "memory-bound"
    return f"intensity={intensity:6.1f}  machine={machine_ratio:5.1f}  -> {kind}"

print("big matmul     ", bound(flops=2e11, bytes_moved=6e8))
print("elementwise add", bound(flops=1e8,  bytes_moved=8e8))
# matmul is compute-bound; elementwise is memory-bound (why fusion helps)

def naive_hbm_gib(N, bytes_per=2, passes=3):
    # QK^T write, softmax read+write, times bytes -- ~passes over the NxN matrix
    matrix_bytes = N * N * bytes_per
    return passes * matrix_bytes / (1024 ** 3)

for N in (1024, 4096, 8192, 16384):
    print(f"N={N:>5}: naive attention HBM traffic ~= {naive_hbm_gib(N):7.2f} GiB")
# quadruples when N doubles -- O(N^2) memory traffic is the wall

import math

def batch_softmax_dot(scores, values):
    m = max(scores)
    exps = [math.exp(s - m) for s in scores]
    Z = sum(exps)
    return sum(e / Z * v for e, v in zip(exps, values))

def online_softmax_dot(scores, values, tile=2):
    m, l, acc = -math.inf, 0.0, 0.0
    for i in range(0, len(scores), tile):
        s_t, v_t = scores[i:i+tile], values[i:i+tile]
        m_new = max(m, max(s_t))
        scale = math.exp(m - m_new) if m != -math.inf else 0.0
        l = l * scale + sum(math.exp(s - m_new) for s in s_t)
        acc = acc * scale + sum(math.exp(s - m_new) * v for s, v in zip(s_t, v_t))
        m = m_new
    return acc / l

sc = [3.1, -1.0, 2.2, 0.5, 4.0, -2.0]
va = [1.0,  2.0, 3.0, 4.0, 5.0,  6.0]
b = batch_softmax_dot(sc, va)
o = online_softmax_dot(sc, va)
print(f"batch  softmax.V = {b:.10f}")
print(f"online softmax.V = {o:.10f}")
print("MATCH" if abs(b - o) < 1e-9 else "MISMATCH")

def roofline_tflops(intensity, peak_tflops=312.0, bw_gib_s=2000.0):
    bw_bytes = bw_gib_s * 1024**3
    mem_bound = intensity * bw_bytes / 1e12          # TFLOP/s from bandwidth
    return min(peak_tflops, mem_bound)

ridge = (312.0 * 1e12) / (2000.0 * 1024**3)          # intensity at the corner
print(f"ridge point intensity = {ridge:.1f} FLOP/byte")
for name, I in [("naive attn (low I)", 5), ("fused attn (high I)", 120)]:
    print(f"{name:>22}: achievable = {roofline_tflops(I):6.1f} TFLOP/s")
# low-intensity op is stuck far below peak; fusion raises intensity toward peak

def fusion_saving_gib(tensor_elems, K, bytes_per=2):
    per_pass = tensor_elems * bytes_per
    unfused = 2 * K * 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=5)
print(f"unfused={u:.2f} GiB  fused={f:.2f} GiB  saved={s:.2f} GiB")
print(f"traffic cut: {s/u*100:.0f}%")
# fusing 5 memory-bound ops cuts HBM traffic ~80% -- pure bandwidth win

def tile_plan(N, sram_kib=100, bytes_per=4):
    # how big a BLOCK fits in SRAM, and tiles per row of length N
    max_elems = sram_kib * 1024 // bytes_per
    block = 1
    while block * 2 <= min(max_elems, N):
        block *= 2                       # largest power-of-two block that fits
    tiles = -(-N // block)               # ceil division
    return block, tiles

block, tiles = tile_plan(N=8192)
print(f"BLOCK={block}, tiles/row={tiles} (one row streamed in {tiles} SRAM-resident tiles)")

# needs GPU + Triton -- the SAME tiling, expressed as a kernel (documented Triton API):
# import triton, triton.language as tl
# @triton.jit
# def softmax_row(x_ptr, y_ptr, n, BLOCK: tl.constexpr):
#     row = tl.program_id(0)
#     offs = tl.arange(0, BLOCK)
#     x = tl.load(x_ptr + row*n + offs, mask=offs < 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_ptr + row*n + offs, y, mask=offs < n)