Efficient Attention — Flash, Sparse, Linear

Why this session exists

Materializing all attention scores costs quadratic memory in sequence length. That is an implementation choice, not a requirement of exact attention. FlashAttention changes the memory schedule; sparse and kernelized approaches change the attention operator; state-space models use a different sequence mechanism. Picking between them requires knowing which resource and which dependency pattern matter.

We derive streaming softmax, compare exact CPU outputs and gradients, measure a small CPU fixture and spell out a separate GPU-backend verification procedure. No CUDA kernel, pretrained model or serving benchmark is executed here.

Outcomes: count score/cache bytes; explain IO-aware tiling and backward recomputation; distinguish sparse, low-rank, kernel and state-space approaches; configure and verify a supported attention backend; and select candidates for long-context, streaming or edge workloads with quality checks.

Prerequisites

Use session 108's Q/K/V math, 109's head/cache shapes and 112's masks. Distinguish host RAM, accelerator HBM, on-chip shared memory and registers. An advertised memory capacity or bandwidth is hardware-specific, not a universal GPU constant.

(a) Intuition · 5 min

A meeting where everyone consults everyone else has quadratically many directed interactions. Attention includes self-pairs and direction, so its dense square has N² cells, not N(N−1)/2 undirected conversations. You need not retain a written transcript of every pair to compute a weighted summary. Streaming a correct summary is the core insight.

Timeline of the attention-efficiency race

The 2017 Transformer uses dense attention. Longformer, BigBird, Linformer, Performer and Reformer explored different long-sequence structures around 2020. FlashAttention (2022) is IO-aware exact attention; FlashAttention-2 (2023) improves work partitioning; FlashAttention-3 (2024) targets Hopper execution. State-space approaches such as Mamba are not new names for softmax attention. Long-context models existed before FlashAttention, so “all models were capped at 2K until 2022” is false.

The GPU memory hierarchy — the whole story lives here

A naive implementation writes scores to HBM, reads them for softmax, writes probabilities and reads them for the value multiplication. Capacity and traffic are related but different: a matrix may fit yet cost substantial bandwidth. Some shapes are compute-bound, others bandwidth- or launch-bound. “Arithmetic is free” is a misleading exaggeration; dense long-context arithmetic remains expensive.

For Q/K head dimension d_h and value width d_v, QKᵀ costs about 2BHN²d_h FLOPs, and AV about 2BHN²d_v. Explicit score storage costs BHN²b bytes at b bytes per entry, before probability arrays, gradients or temporary buffers. Width, heads and batch must be stated beside an O(N²) claim.

Flash-Attention's trick in one picture

For one row, let previously processed scores have maximum m, denominator l=Σexp(s−m) and numerator u=Σexp(s−m)v. For a new tile with scores t and values V_t, define

Finally output u'/l'. Every old contribution changes reference maximum by the same factor, so the recurrence exactly reconstructs softmax attention in real arithmetic. Floating-point operation ordering changes rounding: not bit-for-bit identical. Masks contribute no exponential mass; empty tiles can be skipped. An entirely masked query row needs a defined policy—this teaching code rejects it.

In backward, recompute local score/probability tiles from saved inputs and normalizers instead of storing the full probability matrix. This trades additional arithmetic for less global-memory traffic/storage [1]. A naive Python autograd loop may retain every tile's graph and does not inherit a custom Flash backward's memory bound. Our CPU loop proves the recurrence and gradients, not FlashAttention's GPU memory behavior.

The three families side by side

• Exact execution optimizations: FlashAttention tiles/recomputes without changing the dense softmax operator, up to floating-point differences. PagedAttention concerns KV-cache organization and allocation; it is not the same algorithm as online softmax. These mechanisms can compose.
• Sparse visibility: Longformer combines local windows with selected global positions; BigBird adds random links to local/global structure [2,3]. A windowed operator is exact for its sparse mask, but is not equal to full dense attention. Some token pairs require multiple layers to communicate; global tokens can shorten paths. Decoder architectures can also use sliding windows.
• Kernelized linear attention: replace the exponential dot-product kernel by features φ so $\phi(q)^T\phi(k)$ can be reassociated. Performer uses random-feature techniques to approximate softmax attention [4]. The simpler positive-feature fixture below is a different kernel, not a Performer implementation.

Low-rank sequence projection (Linformer), locality-sensitive hashing (Reformer), recurrent weighted-key/value architectures and selective state-space models are distinct approaches. Linformer reduces the sequence axis through a low-rank approximation [5]. Mamba uses input-dependent state-space dynamics [6]. “Linear in sequence length” does not make these mathematically interchangeable or guarantee a small fixed quality penalty.

Linear attention and streaming state

For positive features with width r,

The accumulated matrix has r×d_v entries and the denominator vector r. With fixed widths, processing N positions is linear in N; the width-dependent cost still matters. A causal version updates these summaries as keys arrive. Compression can lose distinctions that full stored K/V retains; unlike sparse masks, a kernelized model may let every pair interact through an altered similarity rather than dropping a specific pair. A state-space model's streaming state can be independent of sequence length, while training activations still grow with sequence length.

Which to reach for

First identify the bottleneck: score intermediates during training/prefill, persistent K/V during decode, weights, communication or latency. Use a supported exact kernel when its contract matches the model, then measure. GQA changes K/V parameter sharing and normally requires an appropriate checkpoint or adaptation; it is not a universally lossless no-retraining switch. Paging changes allocation, not context understanding. Gradient checkpointing reduces saved training activations, not inference cache size.

For arbitrary distant references, retain an exact-attention baseline and test it at target lengths. For local dependencies, evaluate a trained sparse/windowed model, including long-range negative controls. For streaming/edge constraints, compare recurrent/kernel/SSM models on both quality and state memory. Retrieval over a short context can control generator cost but indexing/search still depend on corpus size and infrastructure; retrieval is not constant-cost in every sense.

What the numbers actually look like

Use binary lengths here: N=32768, H=32, d_h=128, B=1, b=2 bytes. A single dense score tensor is 64 GiB. One output tensor (B,H,N,d_h) is 256 MiB; Q, K, V and other allocations are additional. This is not the entire memory cost of FlashAttention or a 7B model and does not prove it fits a particular device.

An explicitly materialized N×4096 windowed score tensor would be 8 GiB, eight times smaller than the dense score tensor—not eight times smaller than Flash's 256 MiB output. A fused sparse kernel can avoid materializing that matrix too. At N=8192 and FP32, one head's square scores are 256 MiB; 32 heads require 8 GiB for that one tensor.

KV-cache bytes are $2BLNH_{kv}d_hb$ for L layers, not query-head count when K/V are shared. For B=1,L=32,N=32768,H_kv=8,d_h=128,b=2, this is 4 GiB. Halving score storage does not halve that persistent cache. Units and tensor inventories prevent misleading “fits” conclusions.

(c) Hands-on · 25 min

The CPU program compares dense, streamed and SDPA outputs/gradients, implements a sparse-mask negative control and verifies the positive-feature reassociation against an independent pairwise kernel. A small timed fixture reports only local CPU timings; it is not an A100/H100 speedup claim.

import math,time,statistics
import torch
from torch.nn import functional as F
torch.set_num_threads(1);torch.manual_seed(115)
 
def dense(q,k,v,allowed):
    if not allowed.any(-1).all():raise ValueError('all-masked row')
    return (q@k.T/math.sqrt(q.shape[-1])).masked_fill(~allowed,-torch.inf).softmax(-1)@v
 
def online(q,k,v,allowed,tile=3):
    if tile<=0 or not allowed.any(-1).all():raise ValueError('invalid tile or all-masked row')
    rows=[]
    for i in range(len(q)):
        m=q.new_tensor(-torch.inf);den=q.new_tensor(0.);num=v.new_zeros(v.shape[1])
        for start in range(0,len(k),tile):
            legal=allowed[i,start:start+tile]
            if not legal.any():continue
            kk=k[start:start+tile][legal];vv=v[start:start+tile][legal]
            scores=q[i]@kk.T/math.sqrt(q.shape[-1]);newm=torch.maximum(m,scores.max())
            oldscale=torch.exp(m-newm);weights=torch.exp(scores-newm)
            den=oldscale*den+weights.sum();num=oldscale*num+weights@vv;m=newm
        rows.append(num/den)
    return torch.stack(rows)
q=torch.randn(5,4,dtype=torch.float64,requires_grad=True)
k=torch.randn(7,4,dtype=torch.float64,requires_grad=True)
v=torch.randn(7,3,dtype=torch.float64,requires_grad=True)
allowed=torch.arange(7)[None,:]<=torch.arange(5)[:,None]+2
ref=dense(q,k,v,allowed);stream=online(q,k,v,allowed)
sdpa=F.scaled_dot_product_attention(q[None,None],k[None,None],v[None,None],attn_mask=allowed,dropout_p=0.0)[0,0]
assert torch.allclose(ref,stream,atol=1e-12) and torch.allclose(ref,sdpa,atol=1e-12)
probe=torch.randn_like(ref)
g1=torch.autograd.grad((ref*probe).sum(),(q,k,v),retain_graph=True)
g2=torch.autograd.grad((stream*probe).sum(),(q,k,v))
g3=torch.autograd.grad((sdpa*probe).sum(),(q,k,v))
assert all(torch.allclose(a,b,atol=1e-12) for a,b in zip(g1,g2))
assert all(torch.allclose(a,b,atol=1e-12) for a,b in zip(g1,g3))
# Different tile sizes preserve the operator.
for t in (1,2,7,20):assert torch.allclose(ref,online(q,k,v,allowed,t),atol=1e-12)
# Negative controls: different sparse visibility and entirely illegal row.
sparse=allowed & (torch.arange(7)[None,:]>=torch.arange(5)[:,None])
assert not torch.allclose(ref,dense(q,k,v,sparse))
try:online(q,k,v,torch.zeros_like(allowed))
except ValueError:pass
else:raise AssertionError('all-masked row accepted')
# Positive feature kernel: reassociation equals its pairwise oracle, NOT softmax.
phi=lambda x:F.elu(x)+1
qq,kk=phi(q),phi(k)
pair=qq@kk.T
kernel_ref=(pair@v)/pair.sum(-1,keepdim=True)
summary=kk.T@v;norm=kk.sum(0)
linear=(qq@summary)/(qq@norm)[:,None]
assert torch.allclose(kernel_ref,linear,atol=1e-12)
assert not torch.allclose(linear,dense(q,k,v,torch.ones_like(allowed)))
# Math inventory; no large tensors allocated.
for n in (2048,4096,8192,16384,32768):
    score=32*n*n*2;output=32*n*128*2;cache=2*32*n*8*128*2
    print('N',n,'score GiB',score/2**30,'one output MiB',output/2**20,'cache GiB',cache/2**30)
assert 32*32768**2*2/2**30==64
assert 32*32768*128*2/2**20==256
assert 2*32*32768*8*128*2/2**30==4
# CPU timing, explicitly not Flash performance or memory measurement.
with torch.no_grad():
    a,b,c=torch.randn(1,2,64,16),torch.randn(1,2,64,16),torch.randn(1,2,64,16)
    for name,fn in [('dense',lambda:(a@b.transpose(-1,-2)/4).softmax(-1)@c),
                    ('CPU SDPA',lambda:F.scaled_dot_product_attention(a,b,c,dropout_p=0.0))]:
        for _ in range(3):fn()
        timings=[]
        for _ in range(9):
            start=time.perf_counter();fn();timings.append(time.perf_counter()-start)
        print('CPU fixture median seconds',name,statistics.median(timings))
print('streamed output/gradient parity, kernel reassociation, bytes and negative controls passed')

What each block does

online carries max/denominator/numerator per query and accepts a boolean allowed mask. dense is an independently written softmax reference. The gradient comparison uses the same nonconstant weighted output for dense, streamed and SDPA paths, checking all three Q/K/V gradients so a trivial reduction cannot conceal errors. The sparse example changes visibility while keeping Q/K/V fixed, demonstrating that different operators should not be labeled equal merely because their shapes agree.

The feature-kernel comparison computes φ(Q)φ(K)ᵀ explicitly only in the tiny oracle, and separately reassociates K/V. Matching those two does not establish matching softmax. The all-masked case is rejected. Temporary tile graphs in this educational autograd program are not a production memory-efficient backward.

GPU backend verification recipe — not executed here

For PyTorch 2.9, functional SDPA chooses among available implementations subject to device, dtype, shape, mask and other constraints [7]. The context manager torch.nn.attention.sdpa_kernel can restrict selection, for example SDPBackend.FLASH_ATTENTION. A successful CPU SDPA call is not proof of CUDA FlashAttention-2. Explicitly disabling fallback and observing a successful supported CUDA call, together with profiler kernel evidence and version/device records, establishes the selected path; a failure should preserve the actual diagnostic rather than be renamed “OOM.”

A careful GPU benchmark would:

1. Prepare small legal Q/K/V tensors on an existing supported device; warm each backend separately and check output/gradient tolerance before timing.
2. Synchronize before and after the timed region or use device events. Retain repeated measurements, dtype, shapes, causal/padding policy and software versions.
3. Record allocation baseline, reset peak stats, run the operation, synchronize, then report both peak allocated bytes and incremental peak above baseline. Reserved memory is a separate quantity; empty_cache does not free live tensors.
4. Keep unavailable/OOM/backend failures separate. If the math baseline failed, do not compute a speed ratio from an undefined timing variable.
5. Increase sequence length only within a known memory budget, preserving successful/failed cases. Do not intentionally exhaust a production device to draw a curve.

For a trusted local Hugging Face checkpoint, its supported attn_implementation choices can include eager, sdpa or flash_attention_2; the latter requires compatible installed dependencies and hardware [8]. It is not automatically faster than SDPA. Load one model at a time, use identical tokens and generation settings, distinguish prefill from decode, and do not allow a loader to download weights. No package installation, CUDA run or model load is part of this audit.

Practice: the memory wall without allocating it

Calculate score/output/cache bytes for N=2048,4096,8192,16384,32768 in the program. Doubling N quadruples explicit score bytes and doubles output/cache bytes under fixed dimensions. This arithmetic is exact for those tensor shapes, not measured peak memory. In a later GPU experiment, compare measured curves with this inventory and explain allocator/workspace differences instead of claiming every deviation is an error.

(d) Production reality · 15 min

There is no universal crossover at N=1024 and no monotonically increasing speedup guarantee. Short sequences can be launch-bound; head dimension, precision, masks, batch size and hardware can all change the winner. Benchmark the actual workload and diagnose fallback kernels before blaming the algorithm.

Ring Attention distributes sequence blocks across devices; each device keeps query work while key/value blocks circulate and contribute to blockwise attention/online normalization [9]. It can overlap communication and computation, but remains dense pairwise arithmetic overall. Network bandwidth, synchronization and per-device storage matter. It composes conceptually with tiled local kernels; this is not evidence that Gemini or Claude uses the cited implementation.

Long context is not purely a kernel race. Position handling, long-input training, cache/weight memory, evaluation and data quality matter. Sparse methods have not become mathematically obsolete because one exact implementation got faster. RAG may still be useful for permissions, freshness, corpus scale and evidence selection, even when a long context is available.

Where this shows up next

Retrieval decides what enters the context; serving decides how it is processed and cached. Multimodal inputs can increase token count, and agents can accumulate long histories, but neither implies a fixed hop limit without FlashAttention. An attention-only speedup reduces end-to-end cost according to the fraction of runtime it accelerates: if attention is 40% of time and doubles in speed, total time becomes .6+.4/2=.8, a 1.25× speedup, not 2×.

(e) Recall + stretch · 10 min

1. Quadratic memory? A naive implementation materializes BHN² scores/probabilities. Exact attention need not store them all.
2. Flash change? IO-aware tiling plus online normalization and backward recomputation reduce global intermediates/traffic while retaining dense attention mathematics.
3. Approximation? Not an approximation of the attention operator, but not bit-identical floating-point evaluation. Reduced-precision implementations add their own numerical considerations.
4. Online softmax? Rescale previous denominator and weighted numerator when a larger maximum appears; maintain all three accumulators, not just a sum.
5. Sparse workload? Long documents with measured local/global dependency structure are candidates. A legal/medical corpus is not permission to assume a quality loss is acceptable.
6. Modern linear-cost family? Selective SSMs such as Mamba are a distinct option; kernel attention, RWKV and low-rank methods must be described separately.
7. Enable Flash? Request/restrict a supported PyTorch backend, check diagnostics and profile the actual device path. Calling SDPA alone is not proof.
8. Slower Flash? Investigate shape, dtype, masking, backend selection and timing before asserting a universal short-context threshold.
9. KV caching? Reuse prior K/V; new queries still attend over legal cached keys. Cache storage remains linear in retained context.
10. Ring? Rotate K/V blocks across devices while accumulating each device's query outputs. It trades distributed memory/communication, not dense attention for a linear operator.

Explain-out-loud test

Derive the online recurrence, distinguish 64 GiB dense scores from a 256 MiB output, explain why GQA is not a free checkpoint switch, and design a backend proof that cannot accidentally report a fallback as Flash.

What comes next

Prompting operates above the inference machinery. Learning-plan hub.

Primary references

1. Dao et al., FlashAttention, 2022.
2. Longformer, 2020.
3. BigBird, 2020.
4. Performer, 2020 preprint.
5. Linformer, 2020.
6. Mamba, 2023 preprint.
7. PyTorch 2.9 SDPA.
8. Transformers 4.57.1 AttentionInterface.
9. Ring Attention, 2023.