Advanced Attention & Core Transformer Architecture¶

This document delves into the foundational mechanics of modern Transformer architectures, focusing strictly on computational complexity, memory bandwidth constraints, and hardware-aware optimizations.

1. Standard Multi-Head Attention: Complexity & Bottlenecks¶

Standard scaled dot-product attention computes the mapping of a query and a set of key-value pairs to an output. For a sequence length $N$, embedding dimension $d$, and number of heads $h$ (where $d = h \times d_k$), the operation is defined as:

\[ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V \]

Computational Complexity¶

$QK^T$ (Logits): Multiplying an $(N \times d)$ matrix by a $(d \times N)$ matrix takes $O(N^2 d)$ FLOPs.
Softmax: Row-wise softmax takes $O(N^2)$ FLOPs.
$AV$ (Output): Multiplying the $(N \times N)$ attention matrix by the $(N \times d)$ value matrix takes $O(N^2 d)$ FLOPs.
Total Compute: $O(N^2 d)$

Memory Complexity (Space)¶

The raw materialized attention matrix $A = QK^T$ requires $O(N^2)$ space per head.
During training, these $O(N^2)$ matrices must be materialized and stored in High Bandwidth Memory (HBM) for the backward pass, establishing the primary bottleneck for scaling context length.

The Autoregressive Decoding Bottleneck¶

During inference (autoregressive generation), the sequence length is effectively $N=1$ for the query. We must read the past Keys and Values (the KV cache) from memory to compute attention for each generated token. - Compute: $O(N_{past} \cdot d)$ - Memory IO: To perform this computation, we must load the entire KV cache from HBM to the GPU's streaming multiprocessor (SM) SRAM. For batch size $B$, the KV cache size is $O(B \cdot N_{past} \cdot d)$. The arithmetic intensity (FLOPs / Bytes transferred) is incredibly low ($O(1)$), making autoregressive generation fundamentally memory-bandwidth bound, not compute bound.

2. KV Cache Scaling: MQA and GQA¶

To alleviate the memory bandwidth bottleneck during decoding, structural modifications to the attention heads are employed.

Multi-Query Attention (MQA)¶

MQA utilizes a single Key and Value head across all Query heads. - Memory Bandwidth Reduction: The KV cache size is reduced by a factor of $h$. Instead of reading $B \cdot N_{past} \cdot d$ bytes, the hardware reads $B \cdot N_{past} \cdot d_k$ bytes. - Trade-off: This significantly increases arithmetic intensity during decoding but limits the expressivity of the attention mechanism, potentially leading to capacity degradation.

Grouped-Query Attention (GQA)¶

GQA interpolates between MHA and MQA. It divides the $h$ query heads into $G$ groups, where each group shares a single KV head. - Mechanics: If $h=32$ and $G=8$, every 4 query heads share 1 KV head. - Implications: GQA achieves inference speeds virtually identical to MQA while retaining the model quality of MHA. The KV cache footprint is reduced by a factor of $h/G$.

3. Hardware-Aware Attention: FlashAttention¶

Standard attention implementation performs multiple round trips to HBM (High Bandwidth Memory) to read/write intermediate tensors (like $QK^T$, $P = \text{softmax}(S)$). On modern NVIDIA GPUs (e.g., A100/H100), SRAM (L1/shared memory) is small (~20MB total, ~200KB per SM) but orders of magnitude faster (19TB/s) than HBM (~1.5TB/s - 3TB/s).

FlashAttention achieves exact attention with $O(N)$ memory complexity and massive wall-clock speedups by fusing the operation and keeping intermediates in SRAM.

FlashAttention-1: Tiling and Recomputation¶

Tiling: Divides the $N \times d$ $Q, K, V$ matrices into blocks that fit entirely within SRAM. It iteratively loads these blocks, computes a local $QK^T$, updates the running softmax statistics (using an online safe softmax trick), computes the local $AV$, and accumulates the result in HBM.
Recomputation: Instead of saving the $O(N^2)$ attention matrix for the backward pass, FA1 stores the $O(N)$ softmax normalization statistics (logsumexp). During the backward pass, it quickly recomputes the blocks of the attention matrix on-the-fly in SRAM.
Result: Reduces HBM reads/writes from $O(N^2)$ to $O(N)$, shifting the bottleneck from memory I/O back to FLOPs.

FlashAttention-2: Work Partitioning¶

FA2 optimizes the execution by reducing non-matmul FLOPs and improving SM occupancy. - Parallelism: FA1 parallelized over batch size and number of heads. FA2 adds parallelization over the sequence length dimension for the Query, allowing better scaling for long contexts when batch size is small. - Inner Loop Tweaks: FA2 flips the loop order to keep the $Q$ block in SRAM while iterating over $K$ and $V$, reducing the number of times SRAM needs to be flushed and reloaded.

FlashAttention-3: Asynchrony and FP8 on Hopper (H100)¶

FA3 is explicitly co-designed with NVIDIA Hopper microarchitecture constraints. - Tensor Memory Accelerator (TMA): Uses asynchronous DMA to fetch blocks from HBM to SRAM in the background while the Tensor Cores compute the previous block. - WGMMA (Warpgroup Matrix-Multiply-Accumulate): Uses Hopper's WGMMA instructions, bypassing shared memory for operands where possible (fetching directly from registers). - FP8 Support: Exploits Hopper's FP8 Tensor Cores for massive throughput increases, employing custom block-wise scaling to manage precision loss.

4. Rotary Positional Embeddings (RoPE)¶

Transformers are permutation invariant; position information must be injected. Absolute positional embeddings add a vector $P_i$ to $X_i$. RoPE, however, is a Relative Positional Encoding implemented via a rotation matrix.

Mechanics¶

RoPE treats pairs of features within the embedding vector $x \in \mathbb{R}^d$ as complex numbers. To encode position $m$, it rotates the vector in a 2D plane by an angle proportional to the position index. For a feature pair $(x^{(1)}, x^{(2)})$, the rotation is applied as:

\[ \begin{pmatrix} q_m^{(1)} \\ q_m^{(2)} \end{pmatrix} = \begin{pmatrix} \cos(m\theta) & -\sin(m\theta) \\ \sin(m\theta) & \cos(m\theta) \end{pmatrix} \begin{pmatrix} x^{(1)} \\ x^{(2)} \end{pmatrix} \]

Where $\theta_i = 10000^{-2i/d}$.

Why RoPE is Superior¶

Inner Product Invariance: The attention score is the dot product of $Q_m$ and $K_n$. Because a rotation preserves the vector norm, the dot product between a query at position $m$ and a key at position $n$ depends only on the relative distance $(m - n)$: $$ \langle R_m q, R_n k \rangle = q^T R_m^T R_n k = q^T R_{n-m} k $$
Long Context Extrapolation: Unlike absolute embeddings which fail when encountering positions unseen during training, RoPE naturally provides a mechanism for length extrapolation.
Decay Property: The dot product naturally decays as the relative distance $|m - n|$ increases, which provides a strong inductive bias that tokens closer in the sequence are generally more relevant than distant tokens.