28. Online Softmax Attention: Streaming vs Materialized Memory

Problem

Standard self-attention computes scores = Q @ K^T, which is a (T, T) matrix — for a batch of long sequences this materialized matrix dominates activation memory and is exactly the term that makes naive attention quadratic in memory as well as compute. FlashAttention's core trick is to never build that (T, T) matrix at all: process K/V in blocks, and maintain a running max and running (unnormalized) sum so the softmax can be finished incrementally as each block streams by, using only O(T * block_size) working memory instead of O(T^2). This is the "online softmax" algorithm, and it's the same running-statistics idea behind streaming log-sum-exp.

Implement the online-softmax version of causal self-attention: process the key/value sequence in fixed-size blocks, and after each block update a running max, running normalizer, and running weighted output using the standard online-softmax rescaling identity, so the final result is numerically identical to a full materialized-softmax attention but the (T, T) score matrix is never formed in full.

Source: src/28_online_softmax_attention.py

def attention_online(
    q: torch.Tensor,  # (B, H, T, D)
    k: torch.Tensor,  # (B, H, T, D)
    v: torch.Tensor,  # (B, H, T, D)
    block_size: int,
) -> torch.Tensor:  # (B, H, T, D), causal self-attention output
    ...

>>> q = k = v = torch.randn(1, 1, 8, 4)
>>> out = attention_online(q, k, v, block_size=3)
>>> out.shape
torch.Size([1, 1, 8, 4])

Step-by-Step Approach

1. Initialize three running accumulators over all query positions at once: m (running max logit, start at -inf), l (running softmax denominator, start at 0), and acc (running unnormalized weighted output, start at 0).
2. Loop over the key/value sequence in chunks of block_size. For each block, compute only that block's scores, q @ k_block^T — a (T, block_size) tensor, never the full (T, T) matrix.
3. Apply the causal mask within this block: a query at position i may only see a key at position j if j <= i; mask everything else to -inf before taking the block's max.
4. Compute the new running max new_m = max(m, block_max), then rescale the old accumulator into the new max's frame with alpha = exp(m - new_m) before adding this block's contribution — this rescaling is what lets the softmax be finished incrementally without ever seeing all the scores at once.
5. Update acc = acc * alpha + p @ v_block and l = l * alpha + p.sum(-1), where p = exp(scores - new_m) is this block's (still block-local) softmax numerators, then set m = new_m and move to the next block.
6. Guard the arithmetic for rows where a query hasn't seen any valid key yet (both m and new_m are -inf) — a naive exp(-inf - (-inf)) produces nan instead of the correct 0.
7. After the last block, the true softmax output is simply acc / l — the division that finally normalizes the running sum into real probabilities.

The key insight is the rescaling identity alpha = exp(m_old - m_new): it lets you retroactively correct an accumulator that was built using a stale (too-low) max, which is exactly what makes it possible to finish a softmax over data you're still streaming in rather than needing it all resident at once.

Reference solution

import math

import torch


def _safe_exp_diff(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    """exp(a - b), but forces 0 wherever a == -inf (instead of nan when b is
    also -inf) -- happens for queries that haven't seen any valid key yet."""
    raw = a - b
    safe = torch.where(torch.isneginf(a), torch.full_like(raw, float("-inf")), raw)
    return torch.exp(safe)


def attention_online(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    block_size: int,
) -> torch.Tensor:
    B, H, T, D = q.shape
    scale = 1.0 / math.sqrt(D)

    # Running statistics per query position -- this is the only state carried
    # across blocks; the full (T, T) score matrix is never materialized, only
    # one (T, block_size) chunk at a time.
    m = torch.full((B, H, T), float("-inf"), dtype=q.dtype)       # running max logit
    l = torch.zeros((B, H, T), dtype=q.dtype)                      # running softmax denominator
    acc = torch.zeros((B, H, T, D), dtype=q.dtype)                 # running unnormalized output

    query_pos = torch.arange(T)

    for start in range(0, T, block_size):
        end = min(start + block_size, T)
        k_blk = k[:, :, start:end, :]  # (B, H, Bk, D)
        v_blk = v[:, :, start:end, :]  # (B, H, Bk, D)

        scores = torch.matmul(q, k_blk.transpose(-1, -2)) * scale  # (B, H, T, Bk) -- only this block, not (T, T)

        # causal mask: key position j is only valid for query i if j <= i
        key_pos = torch.arange(start, end)
        invalid = key_pos.unsqueeze(0) > query_pos.unsqueeze(1)  # (T, Bk)
        scores = scores.masked_fill(invalid, float("-inf"))

        block_max = scores.max(dim=-1).values  # (B, H, T) -- may be -inf if this block is entirely future for a query
        new_m = torch.maximum(m, block_max)

        # rescale the accumulator built from previous blocks into the new max's frame
        alpha = _safe_exp_diff(m, new_m)          # (B, H, T)
        p = _safe_exp_diff(scores, new_m.unsqueeze(-1))  # (B, H, T, Bk) -- softmax numerators for this block only

        acc = acc * alpha.unsqueeze(-1) + torch.matmul(p, v_blk)
        l = l * alpha + p.sum(dim=-1)
        m = new_m

    return acc / l.unsqueeze(-1)


TEST_CASES = [
    {
        "name": "T divisible by block_size",
        "B": 1,
        "H": 2,
        "T": 16,
        "D": 8,
        "block_size": 4,
    },
    {
        "name": "T not divisible by block_size (ragged last block)",
        "B": 2,
        "H": 1,
        "T": 10,
        "D": 6,
        "block_size": 4,
    },
    {
        "name": "block_size larger than T (single block, degenerate case)",
        "B": 1,
        "H": 1,
        "T": 5,
        "D": 4,
        "block_size": 32,
    },
]


def main():
    import torch.nn.functional as F

    torch.manual_seed(0)

    for i, case in enumerate(TEST_CASES):
        B, H, T, D, block_size = case["B"], case["H"], case["T"], case["D"], case["block_size"]
        print(f"Test {i}: {case['name']} (B={B}, H={H}, T={T}, D={D}, block_size={block_size})")
        q = torch.randn(B, H, T, D)
        k = torch.randn(B, H, T, D)
        v = torch.randn(B, H, T, D)

        n_blocks = math.ceil(T / block_size)
        print(f"  streaming {n_blocks} K/V block(s) of shape up to (B={B}, H={H}, {block_size}, D={D})"
              f" -- peak score tensor (B, H, T, block_size) = {(B, H, T, min(block_size, T))}, vs"
              f" materialized (B, H, T, T) = {(B, H, T, T)}")

        out = attention_online(q, k, v, block_size)
        assert out.shape == (B, H, T, D)

        expected = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        torch.testing.assert_close(out, expected, atol=1e-4, rtol=1e-4)
        print(f"PASSED: out.shape={tuple(out.shape)}")


if __name__ == "__main__":
    main()

Key Functions & Tricks

• torch.maximum(m, block_max) — elementwise running max update, the core of the online algorithm's numerical-stability guarantee.
• torch.where(cond, a, b) — used to special-case the -inf vs -inf subtraction so masked/unseen rows evaluate to 0 instead of nan.
• masked_fill(mask, float('-inf')) — applies the per-block causal mask before computing the block's local max, matching how a full causal mask would behave.
• torch.matmul(p, v_blk) — block-local weighted sum of values, accumulated into acc across iterations instead of computed once over the full sequence.
• F.scaled_dot_product_attention(..., is_causal=True) — used as the oracle in the test to confirm the streaming implementation is numerically identical to PyTorch's fused causal attention.
• Rescaling accumulator by exp(m_old - m_new) — the FlashAttention-style trick that reconciles softmax statistics computed under different running maxima.

How to Recognize This Pattern