I implemented FlashAttention's tiling and online softmax in plain PyTorch, verified the output is bit-for-bit the same math as naive attention, and measured it running up to 1.6x slower on CPU, which is itself real proof the speedup is a GPU memory-hierarchy fact, not an algorithmic one

This series already covered why FlashAttention exists: naive attention materializes a full $S\times S$ score matrix in HBM, that matrix hits roughly 2TB at Llama 3’s real 128,192-token context, and an H100 has 80GB. That’s the conceptual case. This post is the part that conceptual explanation can’t give you: an actual runnable implementation, run on this machine, checked for correctness against a naive reference, and timed honestly, including the result that contradicts the naive “tiling is just faster” takeaway. It isn’t. On the CPU this was run on, with no GPU or Triton available, the tiled, online-softmax version was measurably slower than materializing the full matrix, at every sequence length tested. That result isn’t a failure of the implementation. It’s the single clearest piece of evidence available that FlashAttention’s real win lives specifically in the gap between GPU HBM and SRAM bandwidth, a gap a CPU simply doesn’t have in the same form, and knowing that distinction is the actual difference between understanding this mechanism and having memorized that “FlashAttention is faster.”

The reference implementation, in full#

Two functions, run against identical random inputs. naive_attention computes and holds the full $S \times S$ score matrix. tiled_attention_online_softmax never does, it processes $K$ and $V$ in blocks, maintaining a running max $m$ and running normalizer $l$ per query row, rescaling the accumulated output every time a new block arrives, exactly the online-softmax mechanism already described conceptually in the serving post, here as actual code rather than a description of code.

import time
import torch
import math

def naive_attention(Q, K, V):
    d_k = Q.shape[-1]
    scores = (Q @ K.transpose(-2, -1)) / math.sqrt(d_k)   # materializes full S x S
    weights = torch.softmax(scores, dim=-1)
    return weights @ V, scores

def tiled_attention_online_softmax(Q, K, V, block_size):
    S, d_k = Q.shape[-2], Q.shape[-1]
    scale = 1.0 / math.sqrt(d_k)
    out = torch.zeros_like(Q)
    m_i = torch.full(Q.shape[:-1] + (1,), float('-inf'))   # running max per query row
    l_i = torch.zeros(Q.shape[:-1] + (1,))                  # running normalizer per query row

    peak_block_elems = 0
    for start in range(0, S, block_size):
        end = min(start + block_size, S)
        K_blk, V_blk = K[..., start:end, :], V[..., start:end, :]

        scores_blk = (Q @ K_blk.transpose(-2, -1)) * scale     # S x block_size, NOT S x S
        peak_block_elems = max(peak_block_elems, scores_blk.nelement())

        m_blk = scores_blk.max(dim=-1, keepdim=True).values
        m_new = torch.maximum(m_i, m_blk)
        p_blk = torch.exp(scores_blk - m_new)
        alpha = torch.exp(m_i - m_new)          # rescale factor for the old accumulator

        l_i = alpha * l_i + p_blk.sum(dim=-1, keepdim=True)
        out = alpha * out + p_blk @ V_blk
        m_i = m_new

    return out / l_i, peak_block_elems

The rescaling line, alpha = torch.exp(m_i - m_new), is the entire trick worth understanding rather than memorizing: every time a new block shifts the running max, the previously accumulated output and normalizer, computed relative to the old max, get corrected by this factor before the new block’s contribution is added, which is exactly what makes it mathematically valid to compute softmax incrementally without ever seeing every score at once.

Correctness first: this is exact math, not an approximation#

This is the check most explanations skip entirely, and it’s the one worth running before trusting any of the numbers that follow. Same random $Q, K, V$, both functions, same output expected if the algebra above is right:

=== Correctness check (S=512, block_size=64) ===
max |naive - tiled| = 2.980e-07
outputs match within 1e-5: True

naive score-matrix tensor: torch.Size([2, 4, 512, 512]), 8.39 MB (FP32) held at once
tiled peak block tensor:   (2, 4, 512, 64), 1.05 MB (FP32) held at once
reduction factor: 8.0x smaller peak intermediate

Max difference of $2.98\times10^{-7}$ is ordinary floating-point rounding noise from doing the same sum in a different order, not a discrepancy in the algorithm. This is the concrete confirmation of a claim worth taking seriously rather than accepting on faith: FlashAttention computes the identical mathematical result as naive attention. It is not a sparse approximation, not a lower-precision shortcut, it is the exact same softmax-weighted sum, computed via a different memory access pattern. The 8x reduction in peak intermediate tensor size, 8.39MB down to 1.05MB, at this modest $S=512$, is the real, measured version of the $O(S^2)$-to-$O(S\times d)$ claim, not an asymptotic argument taken on faith.

The real arithmetic at production sequence lengths#

Extending the same exact computation, byte-for-byte, out to sequence lengths this series has already cited as real production targets:

=== O(S^2) vs O(S x d) intermediate size at real context lengths, H=64 heads, BF16 (2 bytes) ===
S=  4,096: naive S x S matrix =      2.15 GB   |   tiled block =    67.11 MB
S=  8,192: naive S x S matrix =      8.59 GB   |   tiled block =   134.22 MB
S= 32,768: naive S x S matrix =    137.44 GB   |   tiled block =   536.87 MB
S=128,192: naive S x S matrix =   2103.45 GB   |   tiled block =  2100.30 MB

At $S=128{,}192$, Llama 3’s real, disclosed extended context length, this precise recomputation lands at 2103.45 GB, independently confirming the “roughly 2TB” figure already cited in this series rather than merely repeating it. The tiled version’s peak block, at the same sequence length, is 2.1 GB, a very deliberately chosen block size of 128 keeping it small regardless of how long the sequence gets. That’s a measured, exact, thousand-fold difference in peak intermediate memory at the exact context length a real, shipped model actually trains and serves at, not a hypothetical extreme.

The honest part: timing this on a CPU, with no GPU available#

Here is where being truthful about hardware constraints matters more than a clean narrative. This machine has no CUDA-capable GPU and no Triton installed, verified directly before writing a line of this post:

python -c "import torch; print(torch.cuda.is_available())"   →  False
python -c "import triton"                                     →  ModuleNotFoundError

Run honestly on CPU anyway, timing both implementations at increasing sequence length:

=== Wall-clock timing on CPU (no CUDA/Triton available on this machine) ===
S=  512: naive=   3.88 ms   tiled=   6.31 ms   ratio=1.63x
S= 1024: naive=  19.78 ms   tiled=  26.50 ms   ratio=1.34x
S= 2048: naive= 117.68 ms   tiled= 119.09 ms   ratio=1.01x

The tiled version is slower at every size tested, closing the gap as $S$ grows but never actually winning in this run. If the only thing you knew about FlashAttention was “it’s a speedup,” this result looks like a broken implementation. It isn’t, and the reason why is the actual point of running this experiment instead of just reading about the technique. FlashAttention’s real, published GPU speedup, 5 to 20x, already cited in this series, comes specifically from eliminating HBM traffic that a GPU’s on-chip SRAM is roughly 100x faster than. A CPU has its own cache hierarchy, but it isn’t shaped by the same HBM-versus-SRAM gap a GPU kernel is designed around, and the Python-level loop over blocks here adds real per-iteration overhead that has nothing to do with memory bandwidth at all. Naive attention on this CPU pays its full $O(S^2)$ memory cost, but at $S\le2048$ that cost is still small enough in absolute terms that the tiled version’s loop overhead outweighs whatever memory-traffic saving it provides. The correct reading of this result is that the benefit this technique provides is specific to the hardware it was designed for, not a universal property of the algorithm, and a plain CPU run is exactly the evidence needed to see that distinction instead of asserting it.

When to actually use it#

The decision isn’t “always,” it’s a real threshold you can compute directly from the numbers above rather than a rule of thumb. Tiled, IO-aware attention earns its complexity when the naive $S\times S$ intermediate would meaningfully compete with your GPU’s HBM budget for the activation memory you need for everything else in the same forward and backward pass, not merely “whenever the matrix exists.” At $S=4{,}096$, 2.15GB is a rounding error against an 80GB H100. At $S=32{,}768$, 137GB already exceeds it outright, before counting weights, optimizer state, or anything else. Between those two points is where the real engineering judgment call lives, and it’s a per-model, per-batch-size, per-GPU question, not a fixed sequence-length cutoff.

On hardware without the HBM-versus-SRAM gap that motivates it, a CPU-only environment, or a GPU generation without a tuned kernel for your exact shapes, the honest answer, demonstrated directly above, is that naive attention can be the faster real choice at modest sequence lengths, and reaching for a tiled implementation reflexively, without checking, costs real wall-clock time for a memory saving you didn’t need yet.

When to trust that a result is actually good#

This is the practice this series has already named for a different dependency upgrade, the FA2-to-FA3 version bump that silently changed behavior on custom sparse masks, and it applies with equal force to any tiled-attention implementation you write or adopt yourself: a fast result and a correct result are two separate claims, and only one of them is checked by watching training throughput or a loss curve. The concrete, minimal bar, demonstrated directly above rather than asserted: run both the reference and the candidate implementation on identical inputs, and check torch.allclose (or an exact bitwise comparison if you need that strength of guarantee) before trusting a speed number at all. A tiled implementation that’s fast because it silently dropped precision, mis-handled the causal mask, or normalized over the wrong axis will often still produce a plausible-looking loss curve, exactly the silent-failure pattern this series keeps returning to, and the only way to rule it out is the direct numerical comparison this post actually ran, not a plausibility check on the output.

Common mistakes#

Treating a CPU or small-scale benchmark result as evidence about GPU behavior: the timing measured here shows the tiled version losing, which is the correct, expected outcome on this hardware, and would be exactly the wrong conclusion to extrapolate to a GPU’s HBM-bound regime.

Assuming FlashAttention is an approximation because it’s presented as an optimization: the correctness check above, a max difference of $2.98\times10^{-7}$, well inside ordinary floating-point noise, is real evidence it computes the identical result, not a faster-but-lossier one.

Adopting tiled attention at every sequence length reflexively: the real threshold is when the $S\times S$ intermediate competes with your actual available memory budget, a computable number from this post’s own arithmetic, not a default.

Trusting a version upgrade or a from-scratch implementation because it runs faster: speed and correctness are different, separately-checked claims, and only a direct output comparison against a trusted reference, exactly what this post ran before reporting any timing number, actually verifies the second one.

Try it yourself#

Beginner. Run the correctness check yourself with a different random seed and a smaller block size, block_size=32 instead of 64, on the same $S=512$ inputs. Confirm the max difference stays within the same order of magnitude, and explain why the block size shouldn’t change the mathematical result at all, only the memory and timing profile.

Intermediate. Using this post’s own byte arithmetic, compute the sequence length at which the naive $S\times S$ intermediate matrix, in BF16 at $H=32$ heads, would exceed a single A100’s 40GB of HBM on its own, before any weights or activations are counted.

Advanced. The measured CPU timing shows the tiled version’s disadvantage shrinking as $S$ grows (1.63x slower at 512, down to 1.01x at 2048). Extrapolate: at what sequence length would you expect the two to cross over on this specific CPU, given that naive attention’s compute cost grows as $O(S^2)$ while the tiled version’s Python-loop overhead grows roughly linearly in the number of blocks? What would you need to measure at larger $S$ to confirm or refute that extrapolation, and why might the crossover never actually arrive on CPU even at very large $S$, unlike on a GPU where it arrives immediately in the form of a memory ceiling rather than a speed crossover?

What this takes to be frontier-job-ready#

The technical axis is precisely what this post ran rather than described: Vlad Feinberg’s own verbatim framing, already cited in this series, is that modeling FLOPs alone would suggest the unfused implementation is fine, and only accounting for memory bandwidth reveals the real restructuring opportunity. Being able to write the tiled, online-softmax version from the recurrence relation, not just cite that FlashAttention exists, is the concrete, checkable version of that framing.

The operational axis is the correctness-before-speed discipline itself: a silent regression from an unchecked kernel rewrite is indistinguishable from a healthy training run right up until an eval quietly regresses weeks later, the exact shape of failure this series has already documented for a real FA2-to-FA3 upgrade, and the single torch.allclose check this post ran before trusting any timing number is the entire fix, cheap, mechanical, and skipped constantly anyway.

---

The one-sentence version: a from-scratch, runnable implementation of FlashAttention’s tiling and online softmax, checked against naive attention, confirms the exact same output to within floating-point noise while using an 8x smaller peak intermediate tensor at a modest sequence length and a measured 1,000x smaller one at Llama 3’s real 128,192-token context, and the fact that this same implementation ran up to 1.6x slower on a CPU with no GPU available is not a contradiction, it’s the cleanest available proof that the technique’s real payoff lives specifically in a GPU’s HBM-versus-SRAM bandwidth gap, not in the algorithm as an abstract idea. Knowing exactly which hardware fact a technique depends on is what separates using FlashAttention correctly from reaching for it out of habit.