24. Tiled/Blockwise Attention (Flash-Attention-Style Online Softmax)

Problem

Naive attention materializes the full (seq_q, seq_k) score matrix before softmax can even start, because softmax needs every score in a row before it can normalize any of them. For long sequences that intermediate matrix is the memory bottleneck — O(seq_q · seq_k) memory that never even makes it into the final output, just gets read once for the softmax reduction and discarded.

FlashAttention's key trick avoids ever materializing that full matrix: sweep through the keys/values in blocks, and maintain a running (incremental) softmax normalizer and a running weighted output accumulator, updating both as each new block arrives and correcting the running totals for the fact that the true row-max wasn't known until later blocks were seen. Implement that "online softmax" update rule in plain PyTorch — without the fused-kernel/IO-avoidance machinery a real CUDA kernel provides — so you can reason about the numerics kernel engineers actually have to get right.

Source: src/24_tiled_blockwise_attention.py

def tiled_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, block_size: int) -> torch.Tensor: ...

>>> q = torch.randn(1, 5, 4)
>>> k = v = torch.randn(1, 5, 4)
>>> tiled_attention(q, k, v, block_size=2).shape
torch.Size([1, 5, 4])

Step-by-Step Approach

1. Initialize three running accumulators before the block loop: running_max = -inf per query (shape (batch, seq_q, 1)), running_sum = 0 (softmax denominator so far), running_out = 0 (unnormalized weighted value sum so far).
2. For each key/value block, compute only that block's scores: q @ k_block.T / sqrt(d), shape (batch, seq_q, block_size) — small and bounded regardless of total sequence length.
3. Find this block's local max, then new_max = max(running_max, block_max) — the softmax max can only grow as more blocks are seen.
4. Compute a correction factor exp(running_max - new_max) and rescale both running_sum and running_out by it — they were previously normalized against a (possibly) too-small max, so this brings them in line with the newly discovered max before adding anything new.
5. Add this block's contribution, normalized against the new max: block_exp = exp(scores_block - new_max); accumulate running_sum += block_exp.sum(-1) and running_out += block_exp @ v_block. Update running_max = new_max and move to the next block.
6. After the last block, divide: running_out / running_sum is the final, fully-normalized attention output — mathematically identical to softmax(full_scores) @ v computed in one shot, for any block size and any number of blocks (including a non-evenly-dividing last block).

The key insight is that softmax's normalizer can be computed incrementally as long as every partial sum is rescaled whenever a larger max is discovered — this is exactly what lets the block loop process key/value chunks one at a time, in bounded memory, and still produce output that is bit-for-bit equivalent (up to floating point rounding) to seeing the whole row of scores at once.

Reference solution

import math
import torch


def tiled_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, block_size: int) -> torch.Tensor:
    batch, seq_q, d = q.shape
    seq_k = k.shape[1]
    scale = 1.0 / math.sqrt(d)

    running_max = torch.full((batch, seq_q, 1), float("-inf"), device=q.device, dtype=q.dtype)
    running_sum = torch.zeros(batch, seq_q, 1, device=q.device, dtype=q.dtype)
    running_out = torch.zeros(batch, seq_q, d, device=q.device, dtype=q.dtype)

    for start in range(0, seq_k, block_size):
        end = min(start + block_size, seq_k)
        k_block = k[:, start:end, :]
        v_block = v[:, start:end, :]

        scores_block = q @ k_block.transpose(-1, -2) * scale  # (batch, seq_q, block)

        block_max = scores_block.max(dim=-1, keepdim=True).values
        new_max = torch.maximum(running_max, block_max)

        # rescale old totals -- they were normalized against a stale max
        correction = torch.exp(running_max - new_max)

        block_exp = torch.exp(scores_block - new_max)
        block_sum = block_exp.sum(dim=-1, keepdim=True)

        running_sum = running_sum * correction + block_sum
        running_out = running_out * correction + block_exp @ v_block
        running_max = new_max

    return running_out / running_sum

Key Functions & Tricks

• torch.maximum(running_max, block_max) — the running max only ever grows across blocks, which is what makes the correction-factor rescale valid (exponentiating a shrinking gap, never an expanding one, avoids overflow).
• torch.exp(running_max - new_max) as a rescale/correction factor — the crux of online (streaming) softmax: it retroactively fixes up totals computed against a stale normalizer without ever re-reading earlier blocks.
• Slicing key/value blocks with k[:, start:end, :] and min(start + block_size, seq_k) — handles a last block smaller than block_size without special-casing.
• Deferred normalization (running_out / running_sum only at the very end) — accumulating the unnormalized weighted sum and dividing once, rather than renormalizing every block, avoids compounding rounding error.
• Equivalence check against ordinary full-matrix softmax attention across multiple block sizes (1, an evenly-dividing size, an unevenly-dividing size, and the whole sequence as one block) — the real correctness bar, since a tiling bug often only shows up at specific block-size/sequence-length combinations.

How to Recognize This Pattern