Blockwise Parallel Transformer — BPT

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 , '’': parallel’ 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.

BPT and Ring Attention

BPT and Ring Attention come from the same line of work and solve orthogonal halves of one problem, which is why they are usually described together. BPT is the per-device kernel: given a shard of queries, compute their full layer in blocks without materialising anything sequence-shaped. Ring Attention is the distribution strategy: shard the sequence across devices, pin each device’s queries, and rotate key/value blocks around a ring until every query has met every key.

They compose because BPT’s inner loop already consumes K/V one block at a time. Ring Attention just changes where the next block comes from — a neighbour’s buffer, not local memory — and overlaps that transfer with compute. Stacked, the maximum context grows with device count, with no per-device tensor scaling in N.

On CPUs and small language models

On a CPU the argument is sharper still. CPU inference and fine-tuning are almost always memory-bandwidth bound, and the cache hierarchy is small enough that a sequence-shaped [N, 4d] tensor is guaranteed to live in DRAM and be streamed twice. For a 1B-parameter SLM with d = 2048 and a 4096-token prompt, that hidden state is 4096 × 8192 elements — about 67 MB in bf16, against an L3 of perhaps 16–32 MB.

Fuse it with b = 128 and the tile is 2 MB: L2-resident, produced and consumed without touching DRAM. You still stream the FFN weights once per block, which is precisely why the intensity rule matters — on CPU, pick the largest b whose [b, 4d] tile still fits the cache level you are targeting.

Pitfalls and misconceptions

Three traps recur. The name collision: BPT is not ‘blockwise parallel decoding’, an older speculative-decoding technique that predicts several future tokens at once. BPT changes nothing about what the model outputs; it is an exact reformulation of the same layer.

It is not a KV-cache optimisation. The win is proportional to N / b, the number of query blocks, so it pays during training and prefill, when you process many tokens at once, and pays essentially nothing during single-token decoding, where the query block is one row and there is no wide hidden state to shrink. For decode, reach for paged KV caches and GQA.

Remember causal masking. Key/value block j > i is entirely masked for query block i and should be skipped outright, not computed and discarded — forgetting that roughly doubles attention FLOPs, and it is the usual reason a hand-rolled blockwise loop underperforms.