The bottleneck that survives FlashAttention
Every long-context memory story begins with the attention matrix, and every modern implementation has already killed it. FlashAttention never materialises QK^T; it tiles the computation and carries a running softmax state, so the [N, N] tensor is a compute artefact rather than a stored one. Problem solved — and then you raise the context length again and hit a second wall.
That wall is the FFN: a two-layer MLP that projects up to 4d and back down, holding two [N, 4d] tensors alive across the nonlinearity. Audit one layer:
Activations in one pre-LN layer (N tokens, dim d, FFN 4d):
residual stream [N, d] → N·d
Q, K, V 3×[N, d] → 3N·d
attn out + proj 2×[N, d] → 2N·d
logits QK^T [h, N, N] → h·N^2 <- Flash deletes
FFN hidden 2×[N, 4d] → 8N·d <- BPT deletes
After Flash: 8N·d (FFN) vs 6N·d (all the rest)
Once the h·N^2 term is gone, the FFN hidden states are not a rounding error — they are the largest line item, bigger than the residual stream, Q, K, V and the attention output combined.
The structural fact that makes fusion legal
Why can you fuse the FFN into the attention loop at all? Because the FFN is token-wise. Attention mixes information across positions: output row i needs keys and values from every position j the mask allows. The FFN does not. It applies the same W_1, nonlinearity and W_2 independently to each row, so row i of its output depends on row i alone.
Layer normalisation shares this property — it normalises within a token, across the feature axis — and so do the residual adds. So the moment a block of query rows has its attention output, everything remaining in the layer for those rows can be finished immediately. There is no reason to wait for the rest of the sequence. That is the whole insight; the rest is bookkeeping.
The fused blockwise algorithm
BPT keeps the two-level loop of memory-efficient attention — an outer loop over query blocks, an inner loop over key/value blocks with an online softmax — and appends the projection, residual and FFN to the body of the outer loop:
for i in query blocks: # outer — independent, parallel
x_i = X[i] # [b, d]
q_i = LN(x_i) W_Q
m, l, o = -∞, 0, 0 # online-softmax state
for j in key/value blocks: # inner — sequential reduction
k_j, v_j = LN(X[j]) W_K, LN(X[j]) W_V
s = q_i k_j^T / √d_k # [b, b] transient
m' = max(m, rowmax(s)); α = exp(m - m')
l = α·l + Σ exp(s - m')
o = α·o + exp(s - m') v_j
y_i = x_i + (o / l) W_O # [b, d]
h_i = act(LN(y_i) W_1) # [b, 4d] transient
Z[i] = y_i + h_i W_2 # [b, d], then free h_i
In the unfused version the outer loop finishes for all N tokens and writes a full [N, d] attention output, and only then does a separate kernel sweep the sequence to build a full [N, 4d] hidden state. Same arithmetic, bit-identical result — but the peak live set collapses from sequence-shaped to block-shaped.
Memory accounting: what actually shrinks
Be precise about which term BPT removes, because it is easy to overclaim. The residual stream stays: the layer’s input and output [N, d] tensors are genuinely sequence-shaped and nothing compresses them away. What becomes block-shaped is every intermediate — the score tile [b, b], the attention accumulator, and above all the FFN hidden [b, 4d].
Per layer the peak therefore falls from roughly 6N·d + 8N·d to 2N·d + O(b·d). The saving scales with N / b, exactly the number of outer-loop iterations, so it grows as you push the context out at fixed block size. That is why the original paper reports fitting sequences several times longer than memory-efficient attention alone allows, and an order of magnitude longer than a vanilla implementation.
A worked example at 32k context
Numbers make the trade concrete. Take d = 4096, 32 heads, a 4d FFN and bf16 activations at a 32768-token context:
N = 32768, d = 4096, h = 32, bf16, one layer:
N·d = 134.2M elem = 268 MB
FFN hidden, 8N·d = 2147 MB
logits h·N^2 = 68.7 GB (Flash removes)
BPT with b = 512: 8·b·d = 16.8M elem = 33.6 MB
reduction = N / b = 64×
32 layers: 68.7 GB of FFN activations → 1.1 GB
The 64× reduction is not a constant-factor tweak — it is the difference between a configuration that fits and one that does not. And after fusion the residual stream dominates, which tells you where to look next: checkpointing across layers, or sharding the sequence across devices. BPT moves the bottleneck, and knowing where it moved to is the point of doing the arithmetic.
Why , &amp;#x27;&amp;amp;amp;rsquo;&amp;#x27;: parallel&amp;amp;amp;rsquo; is in the name
The two loops have opposite characters. The inner loop over key/value blocks is sequential by construction: (m, l, o) is a running reduction, each step rescaling the last by α = exp(m - m'). The outer loop is not. Query block i never reads any other query block’s state; each independently produces its slice of the layer output.
So the outer loop is embarrassingly parallel — across thread blocks, CPU cores, or devices. This is also why the memory argument survives parallel execution: P concurrent query blocks cost P × O(b·d) of transient memory, and you tune P and b to the hardware instead of being hostage to N.
FLOPs are unchanged — so where is the win?
BPT is a scheduling change, not a cheaper algorithm. Per layer the arithmetic is identical to the unfused version:
attention = 4 N^2 d (QK^T and AV)
QKVO projections = 8 N d^2
FFN (4d) = 16 N d^2
attention dominates when 4N^2 d > 24 N d^2, i.e. N > 6d
For d = 4096 that crossover sits near N ≈ 24k: below it the FFN and projections are the bulk of the FLOPs, even though attention gets all the attention. The win from BPT is therefore not fewer operations but a smaller working set — which buys longer contexts, larger batches, and better locality, since the FFN hidden state is produced and consumed in fast memory instead of round-tripping through DRAM.
Choosing the block size: arithmetic intensity
Block size is the one real tuning knob, and it has a clean closed form. Each outer-loop iteration re-reads W_1 and W_2 — 8d^2 parameters, 16d^2 bytes in bf16 — and does 16·b·d^2 FLOPs with them:
intensity = 16 b d^2 FLOPs / 16 d^2 bytes = b FLOPs per byte
The block size is the arithmetic intensity of the block FFN. Compare it to machine balance — peak FLOP/s over memory bandwidth. A CPU near 10–30 FLOP/byte needs only b ≈ 64–128; an accelerator at 200–400 wants b ≥ 256–512. Go below and you become weight-bandwidth bound, re-streaming the same weights N / b times for no extra work; go far above and the [b, 4d] tile stops fitting in cache, recreating the original problem at smaller scale.
The backward pass
Training is where memory pressure actually bites, and BPT composes naturally with gradient checkpointing at block granularity. The forward pass stores only what is cheap and sequence-shaped: the layer input, the layer output, and the per-row softmax statistics (m, l).
Backward, each block recomputes its own s, attention output and FFN hidden from those statistics, then immediately consumes them to accumulate gradients into W_1, W_2, W_O and the QKV projections. The saved (m, l) make that recomputation exact and single-pass — you never re-derive the softmax normalisation. The price is the usual checkpointing tax of about one extra forward pass, paid to make a context length possible at all.