Attention — All Variants & Tricks

Updated July 2026 with 2025–2026 SOTA additions — new entries marked ★. Algorithm names link to their papers (arXiv / project page).

July 2026 · Updated Edition

Foundations: Scaled Dot-Product Attention
Position Encodings (Attention's Constant Companion)
KV-Cache Compression (Inference Memory)
Sparse / Local / Global Attention Patterns
Linear / Sub-Quadratic Attention
Memory-Efficient Exact Attention: FlashAttention
Distributed Attention: Sequence and Context Parallelism
Attention Patterns for Vision
Cross-Attention Patterns
Diffusion-Specific Attention Tricks
Attention in MM-DiT and Modern Diffusion Backbones
Speculative Decoding and Tree Attention
Attention Biases, Masks, and Manipulation
Hardware-Aware and Kernel-Level Tricks
Diagnostic and Interpretability Patterns
Putting It All Together: A Production Stack

Appendix A: Twenty Things to Memorize

1. Foundations: Scaled Dot-Product Attention

1.1 The single equation

Inputs: query \(Q \in \mathbb{R}^{n \times d_k}\), key \(K \in \mathbb{R}^{m \times d_k}\), value \(V \in \mathbb{R}^{m \times d_v}\).

\[\mathrm{Attn}(Q,K,V) = \mathrm{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}} + M\right) V \in \mathbb{R}^{n \times d_v}.\]

Mask \(M \in \{0, -\infty\}^{n \times m}\) adds \(-\infty\) to forbidden positions.

1.2 Why √dₖ?

If \(q_i, k_j\) are zero-mean, unit-variance with i.i.d. entries, then \(\mathrm{Var}(q^\top k) = d_k\). Logits proportional to \(\sqrt{d_k}\) would saturate the softmax, killing gradients. Dividing by \(\sqrt{d_k}\) keeps logits at unit scale.

1.3 Multi-head attention (MHA)

\(H\) heads, head dim \(d_h = d/H\):

\[Q_i = XW_Q^{(i)}, \; K_i = XW_K^{(i)}, \; V_i = XW_V^{(i)}\]

\[\mathrm{MHA}(X) = \mathrm{Concat}(h_1, \ldots, h_H) W_O, \quad h_i = \mathrm{Attn}(Q_i, K_i, V_i).\]

Different heads learn different relations (syntax, coreference, induction patterns, spatial).

1.4 Self vs cross attention

Self: \(Q, K, V\) all from the same \(X\). Cross: \(Q\) from decoder/target, \(K, V\) from encoder/context. Same equation; only the source differs.

1.5 Computational complexity

Time: \(O(n \cdot m \cdot d_h)\) per head for the \(QK^\top\) matmul plus the same for \(\mathrm{softmax} \cdot V\). Total per layer: \(O(H \cdot n^2 \cdot d_h) = O(n^2 d)\) for self-attention with \(n = m\).

Memory: \(O(n^2)\) for the attention matrix. This is the bottleneck FlashAttention removes (§6).

1.6 Numerical stability: stable softmax

\[\mathrm{softmax}(z)_i = \frac{e^{z_i - \max_j z_j}}{\sum_j e^{z_j - \max_j z_j}}.\]

Subtract row-max before exponentiating to avoid overflow. The shift cancels in the ratio.

1.7 Online softmax (the FlashAttention key trick)

Streaming computation across K-blocks:

\[m^{\mathrm{new}} = \max(m, \tilde{m}), \qquad \ell^{\mathrm{new}} = e^{m - m^{\mathrm{new}}} \ell + e^{\tilde{m} - m^{\mathrm{new}}} \tilde{\ell},\]

\[O^{\mathrm{new}} = \frac{e^{m - m^{\mathrm{new}}} \ell\, O + e^{\tilde{m} - m^{\mathrm{new}}} \tilde{\ell}\, \tilde{O}}{\ell^{\mathrm{new}}}.\]

After processing all blocks, \(O^{\mathrm{new}}/\ell^{\mathrm{new}}\) equals the exact softmax-weighted output.

1.8 Pre-norm Transformer block

\[z' = z + \mathrm{MHA}(\mathrm{LN}(z)), \quad z'' = z' + \mathrm{MLP}(\mathrm{LN}(z')).\]

Residuals + LN before the sub-layer. Standard for all modern stacks; post-norm is harder to train at depth.

Key

The four equations to memorize: (1) scaled dot-product, (2) multi-head concat, (3) stable softmax, (4) online-softmax recurrence. Everything in this document is a variant or optimization of these four.

2. Position Encodings (Attention's Constant Companion)

2.1 Why needed?

Attention is permutation-equivariant: \(\mathrm{Attn}(\pi X) = \pi\, \mathrm{Attn}(X)\). Without position info, "cat on couch" = "couch on cat."

2.2 Sinusoidal absolute (Vaswani 2017)

\[\mathrm{PE}(p, 2i) = \sin(p/10000^{2i/d}), \quad \mathrm{PE}(p, 2i+1) = \cos(p/10000^{2i/d}).\]

Property: \(\mathrm{PE}(p+\Delta)\) is a linear function of \(\mathrm{PE}(p)\) for any fixed \(\Delta\), encouraging the model to learn relative offsets implicitly.

2.3 Learned absolute

Lookup table of size \(L \times d\). Hard to extrapolate beyond training length \(L\).

2.4 Relative position bias (T5, Swin)

Add bias \(b_{i-j}\) to logit:

\[A_{ij} = \frac{q_i^\top k_j}{\sqrt{d_k}} + b_{i-j}.\]

T5: log-bucketed bias with shared head; Swin: 2D version \(b_{\Delta x, \Delta y}\) over a window.

2.5 Rotary Position Embedding (RoPE)

Pair adjacent feature dims \((2i, 2i+1)\), rotate by angle \(p\theta_i\), \(\theta_i = 10000^{-2i/d}\):

\[R_p = \mathrm{blockdiag}\!\left( \begin{pmatrix} \cos p\theta_i & -\sin p\theta_i \\ \sin p\theta_i & \cos p\theta_i \end{pmatrix} \right)_{i=0}^{d/2-1}.\]

\[q'_p = R_p q_p, \; k'_{p'} = R_{p'} k_{p'}. \; \text{Then} \; \langle q'_p, k'_{p'} \rangle = \langle q_p, R_{p'-p} k_{p'} \rangle,\]

which depends only on relative position \(p' - p\). Standard in modern LLMs and VLMs.

2.5.1 2D-RoPE (for VLMs)

Split feature dim into row and column halves; apply 1D RoPE to each indexed by patch row \(r\) and column \(c\).

Inner product depends only on \((r - r', c - c')\). Used in Qwen2-VL, InternVL3, FLUX.

2.5.2 Long-context RoPE: NTK-aware, YaRN

Naive RoPE extrapolation degrades. NTK-aware: scale base \(\theta\) such that high-frequency components stay intact.

YaRN: piecewise rescaling by frequency band; supports \(10\text{–}32\times\) context extension. Position Interpolation: linearly interpolate position indices to fit longer sequences; cheap but loses high-freq detail.

2.6 ALiBi (Attention with Linear Biases)

Add a head-specific linear bias on the logit:

\[A_{ij} = \frac{q_i^\top k_j}{\sqrt{d_k}} - m_h \cdot |p_i - p_j|,\]

with \(m_h\) a fixed per-head slope (geometric sequence \(2^{-8/H}, 2^{-16/H}, \ldots\)). No learnable position; extrapolates to longer contexts well.

2.7 Stable Diffusion 3 / FLUX 2D-RoPE for image tokens

For each patch at \((r, c)\), use 2D-RoPE on the patch dim. Plus a separate 2D learned PE for text tokens; text-image cross-modal positions handled via offset.

Key

By 2026 the consensus stack is RoPE for LLMs/VLMs, 2D-RoPE for image tokens in MM-DiT, and ALiBi for some long-context use cases. Sinusoidal absolute is now legacy.

3. KV-Cache Compression (Inference Memory)

3.1 The KV-cache size problem

Each generated token writes \(K, V \in \mathbb{R}^{H \cdot d_h}\) for every of \(L\) layers. Total cache:

\[\text{KV-bytes per token} = 2 \cdot L \cdot H \cdot d_h \cdot \text{bytes}_{\text{dtype}}.\]

For Llama-2 70B (\(L=80\), \(H=64\), \(d_h=128\)) in FP16: \(2 \cdot 80 \cdot 64 \cdot 128 \cdot 2 = 2.6\) MB per token. A 100K-token context needs 260 GB just for KV cache. Compression is mandatory.

3.2 MHA → MQA → GQA → MLA

MHA (original): \(H\) separate KV heads. Largest cache, best quality.
Multi-Query Attention (MQA): 1 KV head shared by all \(H\) Q heads. \(H\times\) smaller cache. Slight quality loss.
Grouped-Query Attention (GQA): \(H_{kv} < H\) KV heads, each shared by \(H/H_{kv}\) Q heads. \(\mathrm{Quality} \approx \mathrm{MHA}\), cache \(H/H_{kv}\times\) smaller. Standard in Llama 2/3, Mistral.
Multi-Latent Attention (MLA) (DeepSeek-V2/V3): low-rank latent KV cache.

3.2.1 MLA in detail

Project to a low-rank latent \(c_t = W_{DKV} h_t \in \mathbb{R}^{d_c}\), \(d_c \ll d\). Cache only \(c_t\). At attention time:

\[K_t = W_{UK} c_t, \quad V_t = W_{UV} c_t.\]

The up-projection is absorbed into \(W_Q\) and \(W_O\) at inference (no extra matmul):

\[q^\top k = (W_Q^{(i)} h_q)^\top (W_{UK} c_t) = h_q^\top \big(W_Q^{(i)\top} W_{UK}\big) c_t.\]

DeepSeek-V3: \(d_c = 512\), \(d = 7168\), \(H = 128\). KV cache reduced \(\sim 14\times\) vs MHA without quality loss.

3.3 Paged KV cache (vLLM)

Break KV cache into fixed-size blocks (e.g. 16 tokens). Per-request page table maps logical tokens \(\to\) physical blocks (like virtual memory). Benefits:

No fragmentation: arbitrary request lengths fit.
Prefix sharing: multiple requests with same system prompt share blocks (CoW for divergence).
Higher utilization: more concurrent requests fit in same memory.

Combined with continuous batching: per-token, not per-request, batching at the kernel level.

3.4 Sliding window + sink tokens

Mistral, StreamingLLM: keep the first \(k\) "sink" tokens (always attended) plus a sliding window of \(w\) recent tokens. Effective infinite context with \(O(k + w)\) cache.

Why sinks work: the first few tokens act as bias absorbers — the softmax wants something to put low-attention mass on, and dropping them causes catastrophic divergence.

3.5 KV cache quantization

Per-channel INT8 or INT4 cache (KV-Cache-INT4, KIVI). Asymmetric quantization with per-token scales. Quality preserved up to INT4 with careful per-channel handling of outlier dimensions.

4. Sparse / Local / Global Attention Patterns

4.1 Causal mask

\[M_{ij} = \begin{cases} 0 & j \le i \\ -\infty & j > i \end{cases}.\]

Standard for autoregressive generation. Lower-triangular.

4.2 Padding mask

\(-\infty\) on padding positions; ensures gradient flows only on real tokens.

4.3 Prefix-LM mask

Bidirectional within prefix, causal after. Used in T5, XGLM. Combines understanding and generation.

4.4 Sliding window attention (Longformer, Mistral)

Each token attends only to a window of \(\pm w\) tokens. Complexity \(O(nw)\). Mistral uses \(w = 4096\) at training; longer sequences stack windows recurrently.

4.5 Dilated attention (Longformer, BigBird)

Strided attention pattern with gap \(d\) between attended positions. Combined with local windows, captures long-range info at sub-quadratic cost.

4.6 Global tokens (Longformer, BigBird)

Special tokens (e.g. [CLS]) attend to / are attended by every position. \(O(g\, n)\) extra. Often combined with sliding window.

4.7 Block-sparse / block-diagonal attention

Sequences packed into one batch with block-diagonal mask (each block one example). Avoids padding waste; standard in modern training.

4.8 LSH attention (Reformer)

Use locality-sensitive hashing to bucket similar \(Q, K\) together; attend only within buckets. \(O(n \log n)\). Practical for very long sequences but with quality trade-off.

4.9 Routing attention (Routing Transformer)

Cluster tokens with online k-means; attend only within clusters. \(O(n\sqrt{n})\).

4.10 Native Sparse Attention (DeepSeek NSA)

Hardware-friendly hybrid sparse: compressed (down-sampled global), selected (top-k blocks via importance), sliding (local window). Trained end-to-end with sparsity from scratch. Strong long-context performance with \(\sim 10\times\) speedup vs full attention.

★ 2026 SOTA update — MoBA mixture-of-block attention

MoBA: MoE-style block-sparse attention (Moonshot/Kimi). Splits keys into blocks; a router gates each query to its top-k blocks (plus current block), so the model learns where to attend rather than using fixed windows. Switches losslessly between full and sparse; deployed in Kimi long-context.

★ 2026 SOTA update — DeepSeek Sparse Attention

DeepSeek Sparse Attention (DSA): fine-grained sparse attention in DeepSeek-V3.2-Exp. A lightweight FP8 'lightning indexer' scores query-key relevance, top-k tokens are selected, then attention runs only over them. Turns \(O(L^2)\) toward linear, ~2-3x faster long-context inference at near-parity quality.

★ 2026 SOTA update — XAttention antidiagonal block scoring

XAttention: training-free block-sparse attention (MIT). Uses the sum of antidiagonal entries of each attention block as a cheap importance proxy to prune non-essential blocks, giving high sparsity. Up to ~13.5x attention speedup on RULER/LongBench/VideoMME at full-attention accuracy.

5. Linear / Sub-Quadratic Attention

5.1 The kernel-feature trick

Standard attention: \(O = \mathrm{softmax}(QK^\top)V\). Replace softmax similarity by a non-negative kernel \(\phi(\cdot)\):

\[O_i = \frac{\sum_j \phi(q_i)^\top \phi(k_j)\, v_j}{\sum_j \phi(q_i)^\top \phi(k_j)} = \frac{\phi(q_i)^\top \sum_j \phi(k_j) v_j^\top}{\phi(q_i)^\top \sum_j \phi(k_j)}.\]

Compute \(S = \sum_j \phi(k_j) v_j^\top\) and \(z = \sum_j \phi(k_j)\) once per layer. Per-token cost \(O(d^2)\).

5.1.1 Kernel choices

Linear Attention (Katharopoulos): \(\phi(x) = \mathrm{elu}(x) + 1\).
Performer (FAVOR+): \(\phi(x) = \exp(\omega^\top x)/\sqrt{m}\) with random features \(\omega \sim \mathcal{N}(0, I)\). Unbiased estimator of softmax kernel.
cosFormer: \(\phi(x) = \mathrm{ReLU}(x)\) with a position-dependent re-weighting.

5.2 Linformer

Project \(K, V\) along the sequence dim to a fixed length \(k\):

\[K' = EK \in \mathbb{R}^{k \times d_k}, \quad V' = FV \in \mathbb{R}^{k \times d_v},\]

then standard attention \(\mathrm{softmax}(Q(K')^\top/\sqrt{d_k}) V'\). \(O(nk)\) time and memory.

5.3 Mamba / state-space models

Recurrence:

\[h_t = \bar{A} h_{t-1} + \bar{B} x_t, \quad y_t = C h_t + D x_t.\]

Selective SSM (Mamba): matrices depend on input via a selection mechanism \(\Delta_t = f(x_t)\), \(\bar{A} = \exp(\Delta A)\).

Hardware-aware parallel scan in \(O(n \log n)\).

Vision Mamba / VMamba: scan in 2D (multi-direction); competitive with ViT at long sequences.

5.4 RWKV

Linear-attention recurrence with token-shift mixing:

\[\mathrm{wkv}_t = \frac{\sum_{i \le t} \exp(-(t-i-1)w + k_i)\, v_i}{\sum_{i \le t} \exp(-(t-i-1)w + k_i)}.\]

Trainable in parallel as Transformer, runnable as RNN at inference.

5.5 RetNet

Retention mechanism with parallel and recurrent dual form:

\[S_t = \gamma S_{t-1} + K_t^\top v_t,\]

\[O_t = Q_t S_t.\]

\(\gamma\) is a fixed decay per head. Linear inference, parallel training.

Watch out

Linear attention variants generally trade quality for speed at long sequences. As of 2026, exact attention with FlashAttention 3 + GQA + ring attention dominates production deployments. SSMs and RWKV remain interesting for very long sequences and edge inference.

★ 2026 SOTA update — Kimi Linear hybrid attention

Kimi Linear: hybrid linear-attention architecture (Moonshot) that beats full attention in fair short/long/RL comparisons. Core is Kimi Delta Attention (KDA), a finer-gated Gated-DeltaNet with a DPLR chunkwise kernel; layers interleave KDA with periodic MLA full-attention.

★ 2026 SOTA update — MiniMax Lightning Attention

MiniMax Lightning Attention: I/O-aware linear-attention variant scaled to a 456B-param (45.9B active) MoE. Hybrid schedule uses one softmax block after every seven lightning-attention blocks; trains at 1M-token context and extrapolates to ~4M at inference.

6. Memory-Efficient Exact Attention: FlashAttention

6.1 The I/O bottleneck

Standard attention reads and writes \(S = QK^\top \in \mathbb{R}^{n \times n}\) to/from HBM, plus \(P = \mathrm{softmax}(S)\), plus \(O = PV\).

Total HBM traffic \(O(n^2)\) even though arithmetic is \(O(n^2 d)\). For long \(n\), attention is memory-bound, not compute-bound.

6.2 FlashAttention 1

Tiling: split \(Q, K, V\) into blocks \(B_r \times d\) and \(B_c \times d\). Stream blocks of \(K, V\) into SRAM; compute partial attention; aggregate with online softmax.

Recompute backward: don't store \(S\) or \(P\); recompute them on the backward pass from the saved \((Q, K, V, O, \ell, m)\).

Trades extra compute for \(n\)-vs-\(n^2\) memory.

HBM traffic: \(O(n^2 d^2 / M)\) where \(M\) is SRAM size (typically 100–200 KB). Empirically \(5\text{–}10\times\) speedup on long sequences.

6.3 FlashAttention 2

Improvements over FA1:

Better work partitioning: parallelize over heads and over the query sequence dim, not just over batch.
Reduce non-matmul FLOPs (avoid unnecessary divisions).
Better SM occupancy on Ampere (A100).

Achieves \(\sim 50\%\text{–}73\%\) of theoretical peak FLOPs.

6.4 FlashAttention 3

Hopper-specific (H100) with:

Asynchronous TMA for memory copies, overlapped with WGMMA matmul.
Producer-consumer pattern across warpgroups.
FP8 support with block quantization (per-128-element scaling).

Reaches \(\sim 75\%\) of H100 BF16 peak; \(\sim 1.2\) PFLOPs FP8 attention.

6.5 Attention backward in FlashAttention

For backward, recompute \(S = QK^\top\) and \(P = \mathrm{softmax}(S)\) tile-by-tile; compute

\[dV = P^\top dO, \quad dP = dO\, V^\top, \quad dS = P \odot (dP - \rho_i \mathbf{1}^\top),\]

where \(\rho_i = \sum_j P_{ij} dP_{ij}\) is the per-row dot product. Then \(dQ = dS\, K\), \(dK = dS^\top Q\).

★ 2026 SOTA update — FlashAttention-4 for Blackwell

FlashAttention-4: Blackwell (B200/GB200) kernel targeting asymmetric hardware scaling (tensor-core throughput doubled, other units not). Adds software-emulated exp, conditional softmax rescaling, and 2-CTA tensor-core use; ~1.3x over cuDNN 9.13 and ~2.7x over Triton in BF16, written in CuTe-DSL.

7. Distributed Attention: Sequence and Context Parallelism

7.1 Tensor parallelism on heads

Split \(H\) heads across TP devices. Each device computes attention for its head subset; output \(W_O\) all-reduces.

Requires high-bandwidth interconnect (NVLink) within a TP group.

7.2 Sequence parallelism

For LayerNorm and dropout (parts not affected by TP split), shard activations along sequence dim. Complements TP; reduces activation memory.

7.3 Ring attention

For very long sequences, split sequence across \(P\) devices. Each device holds its own \(Q\) shard plus initially its \(K, V\) shards. Compute attention block-by-block while passing \(K, V\) around the ring:

Device \(i\) computes attention with its own \(K_i, V_i\).
Send \(K_i, V_i\) to device \(i + 1\) (ring).
Receive \(K_{i-1}, V_{i-1}\) from device \(i - 1\).
Accumulate using online softmax.
Repeat for \(P\) steps until each \(Q\) has seen all \(K, V\).

Total comm: \(O(n\, d/P) \times P\) rounds vs \(O(n^2)\) activations. Enables million-token contexts on \(\sim 32\) H100s.

7.4 Striped attention

Variant of ring that stripes \(K, V\) blocks differently to improve causal-mask load balancing (devices late in the ring have less work in causal setting; striping evens it out).

7.5 Context parallelism (NVIDIA Megatron)

Sequence parallelism specifically for the attention compute, with all-gather of KV at attention time. Different decomposition trade-off: more memory, less comm than ring.

8. Attention Patterns for Vision

8.1 Vanilla ViT attention

\(N = HW/p^2\) patch tokens + 1 [CLS] token. Full \(O(N^2)\) self-attention. Standard ViT-B: \(N = 196\), totally manageable. ViT-L at 1024 resolution: \(N = 4096\), becomes the bottleneck.

8.2 Window attention (Swin)

Partition image into non-overlapping \(M \times M\) windows (\(M = 7\) default). Attention only within each window.

Complexity \(O(N M^2)\) — linear in image size.

8.2.1 Shifted window

In alternating layers, shift windows by \(\lfloor M/2 \rfloor\) to enable cross-window communication. Cyclic shift + masked attention to handle wrap-around efficiently.

8.2.2 Swin V2 improvements

Log-spaced relative position bias for resolution generalization.
Post-norm with residual + LN at end (stable for huge models).
Scaled cosine attention: \(\cos(q, k)/\tau\) instead of \(q^\top k\).

8.3 Pooled attention (MViTv2)

Pool \(K, V\) spatially before attention: \(K' = \mathrm{Pool}(K)\), \(V' = \mathrm{Pool}(V)\). Reduces \(m\) in \(O(n\, m\, d)\). Stride-2 pool gives \(4\times\) reduction in self-attention cost.

8.4 Block + grid attention (MaxViT)

Two passes per layer:

Block attention: within local \(7 \times 7\) blocks (local).
Grid attention: dilated across the image grid (global).

Combines local detail with global context at \(O(n)\) cost.

8.5 Deformable attention (Deformable DETR, BEVFormer)

Each query attends to a small set of \(K\) predicted offsets per feature level:

\[\mathrm{DeformAttn}(q, p, x) = \sum_{m=1}^{M} W_m \sum_{k=1}^{K} A_{mk} \cdot W'_m\, x(p + \Delta p_{mk})\]

with \(\Delta p_{mk}\) and \(A_{mk}\) predicted from \(q\). Sparse, content-aware sampling.

8.6 3D / spatiotemporal attention for video

Three factorizations:

Joint space-time: attention over all \(T\, H\, W\) tokens. Most expressive, most expensive.
Factorized space-then-time (TimeSformer): alternate 2D spatial and 1D temporal attention layers.
Tube attention: 3D tubelets as patches; attention within a tubelet, cross-tubelet via standard self-attention.

8.7 Hiera

Pure ViT with simple pooling between stages; no fancy local attention. Argues that mask-MAE pretraining recovers any locality bias on its own.

8.8 Register tokens (DINOv2 / V-JEPA)

Append \(K\) extra learnable tokens to the patch sequence. They absorb "high-norm" artifacts that otherwise dump on random patch tokens, cleaning up attention maps and improving dense-task quality.

9. Cross-Attention Patterns

9.1 Encoder-decoder cross-attention (original Transformer)

\[Q = X_{\mathrm{dec}} W_Q, \quad K = X_{\mathrm{enc}} W_K, \quad V = X_{\mathrm{enc}} W_V.\]

Standard machine-translation pattern.

9.2 Q-Former (BLIP-2)

\(n_q \sim 32\) learnable query tokens cross-attend over frozen image encoder features:

\[Q' = \mathrm{CrossAttn}(Q_{\mathrm{learn}}, f_\phi(x), f_\phi(x)),\]

output is a fixed-size \(n_q\) token block fed to LLM. Trade-off: token-budget control vs lossy for dense reasoning.

9.3 Gated cross-attention (Flamingo)

Insert cross-attention layers into a frozen LLM at intervals; gate with \(\tanh(\alpha) \cdot \mathrm{CrossAttn}(\cdot)\), \(\alpha\) initialized to 0 so the LM is undisturbed at start of training. Train only the cross-attn modules + tanh gates.

9.4 Cross-attention in U-Net (Stable Diffusion 1/2/XL)

Each U-Net resblock alternates self-attention and cross-attention to text:

\[h' = h + \mathrm{CrossAttn}(\mathrm{LN}(h), \mathrm{TextEmb}(c)).\]

Text embeddings from CLIP/T5 frozen during diffusion training.

9.5 Joint attention in MM-DiT (SD3, FLUX)

Text and image tokens concatenated into one sequence; joint self-attention computes attention over both, with separate \(W_Q, W_K, W_V, W_O\) per modality:

\[Q = [Q_{\mathrm{txt}}; Q_{\mathrm{img}}], \quad K = [K_{\mathrm{txt}}; K_{\mathrm{img}}], \quad V = [V_{\mathrm{txt}}; V_{\mathrm{img}}],\]

\[\to \mathrm{softmax}\!\left(\frac{[Q_{\mathrm{txt}}; Q_{\mathrm{img}}][K_{\mathrm{txt}}; K_{\mathrm{img}}]^\top}{\sqrt{d_k}}\right) [V_{\mathrm{txt}}; V_{\mathrm{img}}].\]

Cleaner cross-modal coupling than U-Net cross-attention; standard for new diffusion architectures.

9.6 IP-Adapter (image conditioning for SD)

Decouple text and image cross-attention paths:

\[h' = h + \mathrm{CrossAttn}(h, c_{\mathrm{text}}) + \alpha \cdot \mathrm{CrossAttn}(h, c_{\mathrm{image}}),\]

where the image cross-attn has its own \(W_K, W_V\). Personalization without fine-tuning.

9.7 ReferenceNet (AnimateAnyone, Magic Animate)

A trainable copy of the U-Net runs on the reference image and produces feature banks; the main U-Net's self-attention concatenates these as additional KV:

\[K = \mathrm{Concat}(K_{\mathrm{self}}, K_{\mathrm{ref}}), \quad V = \mathrm{Concat}(V_{\mathrm{self}}, V_{\mathrm{ref}}).\]

Identity preservation across video frames.

9.8 ControlNet's zero-init connections

Encoder of diffusion model cloned into a trainable copy that takes a control input. Connections from the copy (\(1 \times 1\) convs) back to the base model use zero-initialized convs, so at start the base output is unchanged. During training, gradients flow through; the copy's contribution grows organically.

10. Diffusion-Specific Attention Tricks

10.1 Token Merging for Stable Diffusion (ToMe SD)

Merge highly similar tokens in self-attention layers via bipartite soft matching:

Split tokens into two sets \(A, B\).
For each \(a \in A\), match to most similar \(b \in B\).
Take top-r matches; merge each \(a\) into its \(b\) (weighted average).
Run attention on the reduced set.
Unmerge for the residual addition.

Token count drops by \(r\) per layer; SD inference \(\sim 2\times\) speedup at minimal quality loss.

10.2 DeepCache (block caching)

Diffusion U-Nets compute similar features at adjacent timesteps. Cache the deep block outputs and reuse for \(k\) steps before refreshing. Inference \(\sim 2\text{–}4\times\) speedup.

10.3 TGATE (cross-attention gating)

Cross-attention contributes most early in sampling, then becomes redundant. After step \(\tau^*\), freeze cross-attention output and reuse. \(\sim 1.5\times\) speedup with no quality loss.

10.4 Self-Attention Guidance (SAG)

Given an attention map \(A\), blur the regions of high attention in the input, denoise the blurred input, and use the difference as a guidance direction:

\[\tilde{\epsilon} = \epsilon_\theta(x_t) + s \cdot (\epsilon_\theta(x_t) - \epsilon_\theta(\hat{x}_t^{\mathrm{blurred}})).\]

Improves sample quality without retraining.

10.5 Perturbed-Attention Guidance (PAG)

Replace the attention matrix with the identity in a chosen layer; use the difference between perturbed and original predictions as guidance:

\[\tilde{\epsilon} = \epsilon_\theta(x_t, c) + s \cdot (\epsilon_\theta(x_t, c) - \epsilon_\theta^{\mathrm{PAG}}(x_t, c)).\]

Acts as a form of self-guidance; works without classifier-free guidance.

10.6 Prompt-to-Prompt (P2P)

Manipulate cross-attention maps to do edits:

Word swap: replace cross-attn of one word with another, keep the rest.
Add token: inject cross-attn for new token.
Re-weight: scale cross-attn map of a target token by \(c\).

Localized edits without changing global structure.

10.7 Attend-and-Excite

At each step, ensure each subject token has at least one "excited" attention peak:

\[\mathcal{L} = \max_n \Big(1 - \max_{i,j} A_{ij}^{(n)}\Big)\]

where \(A^{(n)}\) is the cross-attention map for the \(n\)-th subject. Backprop through \(\mathcal{L}\) to nudge \(x_t\) at each step. Fixes "missing object" failures.

10.8 Region-based / spatial attention control

Mask cross-attention to restrict each token's influence to a region:

\[A_{ij} = \mathrm{softmax}\!\left(\frac{q_i^\top k_j}{\sqrt{d_k}} + B_{ij}\right), \quad B_{ij} = -\infty \; \text{if} \; i \notin \mathrm{Region}(j).\]

Used in MultiDiffusion, Attention Refocusing, GLIGEN.

10.9 Cross-attention map manipulation for editing

Null-text inversion fixes DDIM inversion drift by per-step optimization of the unconditional embedding so that DDIM forward + reverse exactly reconstructs. Combined with P2P for high-quality real-image editing.

10.10 Decoupled cross-attention for image animation

AnimateDiff inserts a temporal attention module after each spatial attention block:

\[\mathrm{TemporalAttn} : \mathbb{R}^{B \cdot HW \times T \times d} \to \mathbb{R}^{B \cdot HW \times T \times d},\]

that attends across the time dim per spatial location. Layer is initialized to zero so spatial generation is undisturbed at start.

11. Attention in MM-DiT and Modern Diffusion Backbones

11.1 adaLN conditioning (DiT)

Scale and shift LayerNorm with timestep + class condition:

\[\mathrm{adaLN}(x, c) = \gamma(c) \cdot \frac{x - \mu}{\sigma} + \beta(c),\]

\(\gamma, \beta\) predicted by an MLP from condition embedding \(c\). Optionally a third gate \(\alpha(c)\) scales the residual:

\[y = x + \alpha(c) \cdot f(\mathrm{adaLN}(x, c)).\]

adaLN-zero: initialize \(\alpha\) to zero so the block is a no-op at init.

Two parallel streams (text and image), interacting only through the joint self-attention (not separate cross-attn):

Each stream has its own \(W_Q, W_K, W_V, W_O\) and MLP.
Concatenated \(Q, K, V\) flow into one big attention.
adaLN-zero on each stream for timestep conditioning.

11.3 Pixart-α / Σ

Standard DiT with cross-attention to T5 text. Less unified than MM-DiT but cheaper.

11.4 FLUX, Lumina-T2I, Hunyuan-DiT

Variants of MM-DiT with various tweaks: shared QKV between modalities, 2D-RoPE, more / fewer joint attention layers, token compression.

11.5 Block-cache (FLUX)

Cache attention outputs across consecutive denoising steps; recompute every \(k\) steps. Same idea as DeepCache adapted to MM-DiT.

12. Speculative Decoding and Tree Attention

12.1 Speculative decoding setup

Draft model \(q\) proposes \(k\) tokens \(\hat{y}_{1:k}\); target \(p\) runs one forward pass over the prompt \(+\, \hat{y}\) returning logits at each position. For each token:

Accept with probability \(\min\!\big(1, p(\hat{y}_t | \cdot)/q(\hat{y}_t | \cdot)\big)\).
On reject: resample from \(\mathrm{normalize}(\max(0, p - q))\) at that position; discard rest.

Expected speedup \(\approx \mathbb{E}[\text{accepted}] + 1\) tokens per target call.

12.2 Tree attention (Medusa, EAGLE)

Instead of one draft chain, propose a tree of candidate continuations. Single target forward pass evaluates the entire tree via a custom causal mask:

\[M_{ij} = \begin{cases} 0 & j \text{ is an ancestor of } i \text{ in the tree} \\ -\infty & \text{otherwise}. \end{cases}\]

Largest accepted prefix is committed. EAGLE-2/3 use a learned tree-structured draft head that conditions on target's hidden state.

12.3 Lookahead decoding

No draft model: maintain a verification window of \(W\) tokens generated \(W\) steps ago and a n-gram pool. At each step, propose tokens from the pool, verify against target, advance.

12.4 Multimodal speculative decoding

Draft is a smaller VLM with same vision encoder but smaller LLM. Same accept/reject logic; vision tokens shared (no extra encode cost).

13. Attention Biases, Masks, and Manipulation

13.1 Attention sinks

First few tokens absorb "unwanted" attention mass. Removing them in long-context inference causes catastrophic perplexity spikes (StreamingLLM observation). Solution: keep first \(k \sim 4\) tokens always in cache.

13.2 Massive attention spikes

ViTs trained at scale develop tokens with extremely large norms in middle layers (Darcet et al.). These act like attention sinks for vision. Fix: register tokens (§8.10), or train with gradient clipping on activation norms.

13.3 Sliding-window mask + sink

\[M_{ij} = 0 \; \text{if} \; (i - w \le j \le i) \; \text{or} \; (j < k_{\mathrm{sink}}), \quad -\infty \; \text{else}.\]

Used in Mistral, StreamingLLM.

13.4 Block-diagonal mask for packed sequences

When packing multiple short sequences into one row to avoid padding, use a block-diagonal mask so each example only attends to itself.

13.5 Attention rolling

For very long generation: periodically discard early tokens beyond the window, but keep the sink tokens. Effectively rolling-window long-context.

13.6 Cross-attention masking for editing

In Prompt-to-Prompt edits, mask the cross-attention of unwanted tokens (e.g. words from the source prompt in a word-swap). Localizes the edit.

14. Hardware-Aware and Kernel-Level Tricks

14.1 FlashAttention 3 specifics on H100

Asynchronous TMA (Tensor Memory Accelerator) for \(K, V\) tile loads.
WGMMA (warpgroup matmul) instructions; producer/consumer split across warpgroups.
FP8 (E4M3) attention with per-128 block scaling. Matmul accumulator in FP32.
Reaches \(\sim 75\%\) of H100 BF16 peak; \(\sim 1.2\) PFLOPs FP8.

14.2 ThunderKittens / Liger / xFormers

Modern fused attention kernels with template metaprogramming. Often \(1.1\text{–}1.3\times\) over FA2 in specific regimes (small head dim, MQA/GQA shapes, weird sequence lengths).

14.3 2:4 sparsity attention

On Hopper / Blackwell, weights with 2 of every 4 nonzero entries enable \(2\times\) matmul throughput. Some attention projections (\(W_Q, W_K, W_V, W_O\)) are 2:4-sparsified post-training.

14.4 INT8 / FP8 attention

Per-tensor or per-block scaling. Quantization-aware training (or careful calibration) needed because softmax saturates badly with quantization noise.

14.5 Continuous batching kernel design

Tokens from different requests packed into one matmul; per-token KV indices passed via a separate page table tensor; output scattered back to per-request tensors.

★ 2026 SOTA update — SageAttention3 FP4 attention

SageAttention3: microscaling FP4 attention on Blackwell FP4 tensor cores (1x16 block quant, two-level attention-map quantization). ~1038 TOPS on RTX5090, ~5x over the fastest FlashAttention; also pioneers 8-bit attention for training (lossless in fine-tuning).

15. Diagnostic and Interpretability Patterns

15.1 Attention rollout

Aggregate attention across layers by matrix multiplication of (per-layer attention + identity):

\[\bar{A} = \prod_{\ell=L}^{1} \big(A_\ell + I\big).\]

Approximates "how much each input position influences each output."

15.2 Attention as a graph

View self-attention as a weighted graph where nodes are tokens and edge weight is attention probability. Useful for visualization (BertViz) and for graph-based pruning of unimportant edges.

15.3 Induction heads (Anthropic mechanistic interpretability)

In LLMs, certain heads in layer \(\sim L/2\) implement a copy-and-shift pattern that enables in-context learning.

Identifiable by their attention pattern: token \(t\) attends to the position of an earlier occurrence of the same n-gram.

15.4 Probing attention heads for tasks

Some heads specialize: syntactic heads (attend to dependency parents), coreference heads, positional heads.

Probing reveals these via diagnostic classifiers on attention vectors.

16. Putting It All Together: A Production Stack

16.1 Modern LLM training stack

RoPE position encoding (NTK-scaled or YaRN for long ctx).
GQA (4–8 KV heads) for inference cache.
Pre-norm with RMSNorm.
SwiGLU MLP.
FlashAttention 3 (or FA2 on A100).
Sliding-window + sink for streaming inference.
Sequence parallelism + ring attention for long-context training.

16.2 Modern VLM training stack

Same as LLM, plus:
2D-RoPE for image patch positions.
Native dynamic resolution via patch-position packing.
Pixel unshuffle (\(r = 2\)) for token compression.
Either MLP projector (LLaVA) or unified joint attention (Chameleon-style native multimodal).

16.3 Modern diffusion stack (image)

MM-DiT with adaLN-zero for timestep conditioning.
2D-RoPE on image tokens.
Joint self-attention over text + image (no separate cross-attn).
Flow matching loss instead of \(\epsilon\) DDPM.
Block-cache (e.g. TGATE) for inference acceleration.
Distillation to 1–4 step (LCM, DMD2, Hyper-SD).

16.4 Modern diffusion stack (video)

Causal 3D VAE for tokenization.
Spatiotemporal MM-DiT.
3D-RoPE on \((T, H, W)\) tokens.
Joint image + video training (\(T=1\) special case).
Variable resolution / aspect / fps via native packing.
Audio joint generation via additional token stream (MovieGen, Veo 3).