Attention & GPU kernels

1 · A mental model of the GPU essential

Start with a bill that shocks people the first time they see it. A team buys the biggest GPU they can afford — a card advertised at nearly a petaflop of matrix math — runs their long-context model on it, and profiles the hot loop. The GPU is 80% idle. Not broken, not misconfigured: idle. The arithmetic units are sitting there waiting, hands in their pockets, for data to arrive from memory. That gap between "how fast the chip can compute" and "how fast it can be fed" is the single most important thing to understand about modern accelerators, and it's why this section starts not with math but with a map of where data lives.

Here's the analogy that makes it click. Think of a chef (the compute units) working at a tiny cutting board (on-chip SRAM) next to a giant walk-in fridge down the hall (HBM). The chef is superhumanly fast with a knife — but every ingredient has to be walked from the fridge to the board and back. If the recipe makes the chef fetch one carrot, slice it, walk it back, fetch one onion, slice it, walk it back, the kitchen's throughput is set by the hallway, not the knife. A good kitchen loads a whole tray onto the board once and does all the chopping there. That is the entire discipline of GPU kernel design in one image: minimize trips down the hall.

Now the precise version. A GPU is not one big processor — it's a few dozen to ~130 Streaming Multiprocessors (SMs), each an independent cluster of arithmetic units with its own sliver of extremely fast on-chip memory. That fast memory comes in two tiers: registers (per-thread scratch, tens of KB) and SRAM / shared memory (a few hundred KB per SM). Surrounding the whole chip is HBM (High-Bandwidth Memory) — the gigabytes where your weights and activations actually live. HBM sounds fast, and in absolute terms it is (~2–3 TB/s), but it is roughly ten times slower than SRAM and it is physically far from the compute. Beyond HBM is host RAM over PCIe or NVLink, an order of magnitude slower again — a place you never want to visit on the hot path.

The common mistake that this model exists to kill: assuming "faster GPU" or "faster matmul library" automatically means "faster model." It often doesn't, because the bottleneck usually isn't the math — it's the bytes. Read the table below not as trivia but as a hierarchy of cost: every step down the list is bigger but slower and further away, and a fast kernel's whole job is to stay as high up this list as it possibly can.

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

Here's a mystery you may have hit yourself. You learned that attention is O(n²), so you reasonably assume the way to make it faster is to make the arithmetic faster — swap in a snappier matmul, use tensor cores harder, tune the GEMM. You do all that, and the kernel barely moves. Meanwhile someone hands you "FlashAttention," which does the exact same math, and it's several times faster on long sequences. If the FLOPs are identical, where did the speed come from? The answer is the whole point of this section: naive attention doesn't spend its time computing — it spends its time moving bytes.

To see why, follow the intermediate results, not the formulas. Attention is three steps: compute scores S = QKᵀ, turn them into probabilities P = softmax(S), then weight the values O = P·V. The trap is that a textbook implementation runs each step as a separate kernel launch, and a separate launch can only communicate with the next one through HBM — the slow fridge down the hall. So the giant S matrix gets written all the way out to HBM, then read all the way back for softmax, written again, and read a third time for the final matmul. The compute in each step is quick; the walking-down-the-hall between steps is where the clock burns.

And that intermediate is genuinely giant. S is N×N — it has one entry for every pair of tokens. At N=4096 that's ~16.7 million numbers, tens of megabytes at fp16, and it gets round-tripped through the slowest memory on the chip several times. Crucially, its size grows with the square of the sequence length: double the context and you quadruple this traffic. The actual math done per byte, by contrast, is small — each score is one dot product, each softmax entry one exponential. That lopsided ratio (few FLOPs, many bytes) has a name we'll formalize in §6: low arithmetic intensity, and it is the signature of a memory-bound kernel.

The reframe to carry out of this section: attention's cost, on real silicon, is not the O(n²) multiply-adds — it's the O(n²) bytes of the score matrix crossing the HBM boundary again and again. Once you believe that, the fix becomes obvious in shape even before you see it: stop writing the N×N matrix to HBM at all. The diagram traces exactly which trips cost you.

3 · Lab · an HBM-traffic estimator intermediate

Claims about performance are cheap; a number you computed yourself is not. We just argued that naive attention is slow because of bytes, not math — but "slow because of bytes" is the kind of sentence you can nod along to and still not truly believe. So in this lab we stop hand-waving and count. No GPU, no benchmark harness, no timing noise — just arithmetic on how many bytes each approach shoves across the HBM boundary. When you see the ratio grow from 17× to 257× as the sequence lengthens, the phrase "memory-bound" stops being jargon and becomes something you can defend in a design review.

The mental model for reading the estimator is a ledger with two columns. Both approaches pay the same unavoidable cost: read Q, K and V in once, and write the output out once — that's the O(N·d) traffic, linear in sequence length. The difference is entirely in the middle. The naive path adds two more line items — write-then-read the N×N scores, write-then-read the N×N probabilities — and because those terms are N×N they scale as O(N²). FlashAttention simply has no such line items: the N×N matrix is computed inside SRAM, tile by tile, and never lands in HBM at all. So the model is really "linear traffic that both pay" plus "quadratic traffic that only the naive path pays."

Why model bytes instead of timing the code? Because for a memory-bound kernel, time is basically bytes ÷ bandwidth — the arithmetic hides under the memory transfers, so bytes moved is a faithful proxy for wall-clock time, and it's deterministic (no warm-up, no thermal throttling, no scheduler jitter to average out). That is the quiet superpower of these offline labs: they let you reason about GPU performance on a laptop, because the thing that dominates the runtime is countable without the hardware. Read the code as a byte ledger, then watch the ratio in the output explode with N.

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

Everything so far has been diagnosis; this is the cure. We proved naive attention wastes its life writing and re-reading a giant N×N matrix through slow HBM. The obvious wish is: what if we just… didn't? What if the score matrix never had to exist in HBM at all? FlashAttention is the clean answer to exactly that wish, and the beautiful part is that it is not an approximation — it computes bit-for-bit the same attention, just without ever parking the big intermediate in slow memory. It won a place in essentially every serious inference stack for one reason: on modern hardware, not touching HBM is the speedup.

The core idea is tiling, and the everyday analogy is grading a stack of exams too tall to fit on your desk. You don't need all the papers spread out at once to compute the class average — you take a manageable stack (a tile), work through it, keep a running total and a running count, set it aside, and grab the next stack. At the end your running numbers give the exact average as if every paper had been laid out simultaneously. FlashAttention does the same with attention: load a block of queries and a block of keys/values into SRAM, compute that block's partial scores on chip, fold the result into a running output, discard the block, and move on. The N×N scores are still computed — but one small tile at a time, in fast on-chip memory, and never written to HBM. Only the O(N·d) inputs and output ever touch device memory, so peak memory drops from O(N²) to O(N) and the HBM round-trips collapse.

But tiling hides a real danger, and it's worth naming because it's the thing that makes FlashAttention non-trivial. Softmax is not a per-element operation you can just chop up — it needs the whole row, because it subtracts the row's maximum (for numerical stability) and divides by the sum of all the exponentials. If you naively softmax each tile in isolation, you get the wrong answer: each tile normalized against its own local max and sum, not the global ones. So the question becomes: how do you compute a row-wide softmax when you're only ever holding one slice of the row at a time? That is the puzzle online softmax solves, and it's the mathematical heart of the whole method — introduced just below and proven exact in the next lab.

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 lab where you stop trusting and start verifying. A skeptic reading the last section has every right to push back: "You're telling me I can process softmax in blocks, never see the full row, and somehow get the identical answer — not close, identical? Prove it." Fair. Because if online softmax were even slightly off, tiling would silently corrupt every attention output on long sequences, and no one would trust FlashAttention with a production model. So we don't ask you to take it on faith — we compute the same thing two ways and check they match to the last decimal.

Here's the trick, in plain terms before you meet the code. The whole difficulty is that a later tile might contain a bigger score than anything seen so far — and softmax's stability step subtracts the maximum, so a new maximum means everything computed earlier was scaled against a stale, too-small max. Online softmax handles this with three running numbers: m, the largest score seen so far; l, the running sum of exponentials; and acc, the running weighted sum of values. When a new tile reveals a bigger max m_new, it multiplies the old l and acc by a correction factor exp(m_old − m_new). That factor retroactively re-bases all the earlier work as if the new max had been used from the very beginning. It's the exam-grading analogy made exact: when you discover a paper scored higher than you assumed, you rescale your running total rather than starting over.

Why does that rescaling produce an exactly equal result rather than a good approximation? Because it's pure algebra, not estimation. Softmax is invariant to subtracting any constant from every score — softmax(s) = softmax(s − c) for any c — so which max you subtract can never change the final probabilities; it only changes intermediate magnitudes to keep them from overflowing. The correction factor is precisely the constant that converts "scaled by m_old" into "scaled by m_new," so the streaming path reconstructs the same normalization the batch path computes in one shot. The lab below runs both and prints their difference; watch for a difference of exactly 0.00e+00 — that zero is your license to tile.

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

By now you might over-learn the lesson and reach for FlashAttention-style tricks everywhere. That's a real failure mode: an engineer internalizes "fewer HBM trips = faster," starts fusing and tiling every kernel in sight, and burns a week making a compute-bound kernel move less memory — for zero speedup, because memory was never its bottleneck. FlashAttention worked because attention is memory-bound. The skill this section builds is knowing, before you optimize, whether a given kernel is limited by bandwidth or by arithmetic — so you turn the knob that actually moves the clock.

First, generalize the FlashAttention insight into its real name: kernel fusion. Every separate kernel launch has a fixed tax — it reads its inputs from HBM and writes its outputs back. If you run five elementwise ops in a row as five kernels, you pay that read-write tax five times on the same tensor. Fuse them into one kernel and the intermediate values stay in registers/SRAM between steps; you read once and write once. FlashAttention is just the marquee example of this: fusing the score matmul, the softmax, and the value matmul into a single kernel so the N×N intermediate never round-trips. The mental takeaway: fusion is almost always a memory-traffic optimization, not a compute one — it removes trips to the fridge, it doesn't make the chef chop faster.

Now the tool that tells you when fusion (or any bandwidth trick) will help: the roofline model. Picture a rooftop profile. Along the bottom axis is a kernel's arithmetic intensity — FLOPs performed per byte moved from HBM. Low intensity means you do very little math per byte you fetch (you starve for data); high intensity means you reuse each fetched byte for lots of math. The roof has two parts: a slanted section on the left (bandwidth limits you — you're memory-bound) and a flat ceiling on the right (the ALUs limit you — you're compute-bound). Where they meet is the ridge point, and it's a pure property of the machine: peak_FLOPs ÷ bandwidth. Plot your kernel's intensity: fall left of the ridge and your fix is move fewer bytes (fusion, FlashAttention); fall right of it and your fix is do less math or use lower precision (quantization). The lab computes exactly this classification for two very different kernels.

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

There's a moment in every systems engineer's life where the abstraction runs out. You've profiled, you've rooflined, you've concluded the win is a fused kernel that keeps the intermediate in SRAM — and now you have to actually write it. Historically that meant CUDA C++: manual thread indexing, shared-memory bank conflicts, warp-level shuffles, memory coalescing rules. Weeks of work and a skill set most ML engineers never acquire. That barrier is exactly why for a decade "just write a custom kernel" was advice, not a plan. Triton is the tool that lowered the barrier enough that this section can exist at all.

Triton (from OpenAI) lets you write GPU kernels in Python. The key shift is the unit of thought: instead of reasoning one scalar thread at a time — the CUDA mental model — you write the computation over blocks (tiles) of data, and Triton's compiler handles the fiddly parts underneath: assigning work to SMs, allocating SRAM, coalescing memory accesses, managing the warps. You say "load this tile, do this matmul on it, keep this running state, store the result," and Triton turns that block-level description into a well-scheduled GPU kernel. That's why much of the modern fused-kernel ecosystem — including production FlashAttention variants — is now authored in Triton: it gives you Python-level control with performance close to hand-tuned CUDA.

Here's the payoff that ties this whole lesson together, and the reason the snippet below is here to read rather than run. Look closely and you'll see the kernel carries the exact same m, l, acc running-state update you proved correct on your laptop in §5 — the online-softmax rescale, now expressed in Triton's block primitives. That's the deep lesson of the offline labs: the algorithm's correctness is separable from the silicon. You verified the math where it's cheap and reproducible (pure Python, no GPU), and the GPU kernel is a faithful transcription of that verified math into hardware-aware blocks. So you don't need a datacenter to understand — or even to trust — the kernel; you need it only to run it fast. Read the shape below and recognize the running state; that recognition is the whole point.

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

Picture the meeting where this knowledge earns its keep. An engineer proposes buying a fleet of faster GPUs to cut inference latency; another wants to spend a sprint hand-optimizing a matmul; product is asking whether you can double the context window without doubling the cost. Everyone's looking at you. You don't need to have written a single Triton kernel to answer well — you need the one mental model this lesson has been building: on modern accelerators, performance is a data-movement problem first and an arithmetic problem second. That single reframe turns most of these debates into a quick calculation instead of an argument.

Watch how the model resolves each question cleanly. "Faster GPUs will fix latency" — only if the kernel is compute-bound; roofline it first, because if it's memory-bound you're paying for FLOPs you can't use. "Let's hand-optimize the matmul" — for a memory-bound step like attention's softmax, that optimizes the part that isn't the bottleneck; the real lever is fusion and fewer HBM trips. "Can we double the context?" — with FlashAttention, attention itself is O(N) in memory, so the score matrix no longer explodes; the true limit becomes the KV-cache (which grows linearly with context and batch), tying this straight into the serving math of IC3. Notice that in every case the answer flows from the same source: where do the bytes go, and how many times?

The trap to help your team avoid is treating "it's O(n²)" as the end of the analysis rather than the start. Big-O counts operations; it says nothing about which memory tier those operations touch, and on a GPU that tier is the performance. A tech lead's job here is less about writing the fast kernel and more about asking the right diagnostic question before anyone optimizes — roofline it, find the bound, then spend effort only on the side of the ridge you're actually on. Adopt a fused-attention library for long context (the HBM savings are almost always the win); reach for lower precision when you're genuinely compute-bound; and always connect the kernel-level story back up to the system-level cost. The rest of this section distills that into a single sentence worth memorizing.

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)