Since its introduction in Vaswani et al. (2017), the Transformer architecture has replaced recurrence and convolution across modern natural language processing, computer vision, and multimodal foundation models. At the core of every Transformer lies the Multi-Head Self-Attention (MHSA) mechanism.
In this comprehensive engineering guide, we rigorously derive the mathematics of Scaled Dot-Product Attention, prove the probabilistic necessity of the 1 / √(d_k) variance scaling factor, evaluate FlashAttention-2 online softmax tiling, and implement an industrial, batched PyTorch module with causal masking and KV-cache mechanics.
The Query, Key, Value Formulation#
Attention can be conceptualized as a differentiable, soft information retrieval system operating over continuous vector representations:
- Query (Q): Represents the current token's active search vector—what syntactic or semantic attributes the token needs to resolve its context.
- Key (K): Represents the index vector of candidate tokens—what features or contextual properties each candidate token offers.
- Value (V): The actual informational payload extracted when a query matches a candidate key.
For a batch of input token sequence representations X ∈ ℝ^(B × L × d_model), where B is batch size, L is sequence length, and d_model is embedding dimension, we compute linear projections using learned weight matrices W_Q, W_K, W_V ∈ ℝ^(d_model × d_k):
Q = X · W_Q, K = X · W_K, V = X · W_V Token Embeddings X [B, L, d_model]
| | |
v v v
Linear Linear Linear
(W_Q) (W_K) (W_V)
| | |
v v v
Query (Q) Key (K) Value (V)
[B, L, d_k] [B, L, d_k] [B, L, d_v]Scaled Dot-Product Attention Mathematics#
The semantic compatibility between tokens is quantified via the inner product of their respective Query and Key vectors. The canonical scaled dot-product attention formulation is defined as:
Attention(Q, K, V) = softmax((Q · K^T) / √(d_k)) · V Query (Q) Key (K)
\ /
\ /
MatMul: Q · K^T --> Logits Shape [B, H, L, L]
|
v
Scale: / √(d_k)
|
v
Mask (Causal / Padding)
|
v
Softmax --> Probability Matrix A [B, H, L, L]
|
v
MatMul with V --> Context Output [B, H, L, d_k]Formal Mathematical Proof: Why Divide by √(d_k)?#
To understand why the scaling factor 1 / √(d_k) is mandatory for numerical stability during gradient backpropagation, consider two independent random vectors q, k ∈ ℝ^(d_k). Assume their individual components q_i and k_i are independent, identically distributed (i.i.d.) random variables with zero mean and unit variance:
E[q_i] = 0, Var(q_i) = 1
E[k_i] = 0, Var(k_i) = 1The unscaled dot product of q and k is defined as the sum of d_k component-wise products:
S = q · k = Σ_{i=1}^{d_k} (q_i · k_i)By the linearity of mathematical expectation:
E[S] = Σ_{i=1}^{d_k} E[q_i · k_i] = Σ_{i=1}^{d_k} (E[q_i] · E[k_i]) = Σ_{i=1}^{d_k} (0 · 0) = 0Because q_i and k_i are mutually independent, the variance of their product is:
Var(q_i · k_i) = E[(q_i · k_i)^2] - (E[q_i · k_i])^2
= E[q_i^2] · E[k_i^2] - 0
= Var(q_i) · Var(k_i) = 1 · 1 = 1The variance of the sum S is the sum of the variances:
Var(S) = Var(Σ_{i=1}^{d_k} q_i · k_i) = Σ_{i=1}^{d_k} Var(q_i · k_i) = d_k · 1 = d_kThe standard deviation of the unscaled dot product is therefore:
σ_S = √(Var(S)) = √(d_k)The Vanishing Gradient Problem: In modern LLMs where head dimension d_k = 128, the unscaled dot product has a standard deviation of √(128) ≈ 11.31. When inputs to the softmax function have values on the order of ±15 to ±30, the exponential term e^(s_i) explodes for the maximum value and approaches zero for all other components.
The derivative of the softmax function with respect to logit z_i is:
∂(softmax(z)_i) / ∂z_j = softmax(z)_i · (δ_{ij} - softmax(z)_j)When softmax saturates into an almost one-hot vector (e.g., [0.9999, 0.0001, ...]), this derivative evaluates to near zero (≈ 0). Consequently, during backpropagation, the error gradients vanish entirely, arresting parameter optimization in earlier layers. Dividing the dot product by √(d_k) rescales the variance strictly back to 1.0 (Var(S / √(d_k)) = Var(S) / d_k = d_k / d_k = 1.0), preserving healthy gradient flow throughout training.
Multi-Head Attention: Subspace Diversity#
Rather than performing a single attention function with d_model-dimensional queries, keys, and values, Multi-Head Attention projects Q, K, and V linearly h times with distinct learned projection matrices. This allows the model to jointly attend to information from different representation subspaces at different token positions:
MultiHead(Q, K, V) = Concat(head_1, head_2, ..., head_h) · W_O
where head_i = Attention(Q · W_i^Q, K · W_i^K, V · W_i^V)- Head Dimension: d_k = d_model / h
- Output Projection: W_O ∈ ℝ^(d_model × d_model)
Input Embeddings X
|
+---------------------+---------------------+
| | |
Head 1 Head 2 Head h
(Syntax & Grammar) (Entity Extraction) (Long-Range Context)
[B, L, d_k] [B, L, d_k] [B, L, d_k]
\ | /
+--------------------+--------------------+
|
Concatenation [B, L, d_model]
|
Linear (W_O)
|
Output TokensArchitectural Evolutions: MHA vs. MQA vs. GQA#
In large autoregressive language models, the memory bandwidth required to fetch Key and Value caches from GPU High Bandwidth Memory (HBM) to on-chip SRAM during inference becomes the primary latency bottleneck. To address this, modern architectures optimize head configurations:
- Multi-Head Attention (MHA): Every attention head maintains its own independent Query, Key, and Value projections. High expressiveness, but maximum KV-cache memory footprint.
- Multi-Query Attention (MQA): All h Query heads share a single common Key head and Value head. Reduces KV-cache memory bandwidth by factor h, but can slightly degrade reasoning fidelity.
- Grouped-Query Attention (GQA): A balanced compromise (used in Llama 3, Mistral). The h Query heads are partitioned into G groups, where each group shares 1 Key and 1 Value head. Offers near-MHA reasoning accuracy with 4× to 8× KV-cache compression.
Memory Complexity & FlashAttention Tiling#
During standard training, the intermediate attention matrix S = Q · K^T requires materializing an [L, L] tensor in GPU High Bandwidth Memory (HBM). For sequence length L = 131,072 (128K context), storing a single attention matrix per head requires:
Memory = L^2 · 2 bytes (FP16) = (131,072)^2 · 2 = 34,359,738,368 bytes = 34.36 GB per head!Multiplied across 32 heads and 64 transformer layers, standard attention demands terabytes of VRAM, making long context windows impossible without architectural changes.
FlashAttention-2: Online Softmax and SRAM Tiling#
FlashAttention (Dao et al., 2022, 2023) solves this memory wall by never materializing the [L, L] attention matrix in GPU HBM. Instead, it tiles queries, keys, and values into small blocks that fit directly inside fast on-chip SRAM (20 MB per Streaming Multiprocessor on NVIDIA H100):
+---------------------------------------------------------------+
| GPU High Bandwidth Memory (HBM - 80 GB, 3.35 TB/s) |
| Stores: Q, K, V, O in chunked tiles |
+---------------------------------------------------------------+
| ^
Block Read | | Write O
v |
+---------------------------------------------------------------+
| GPU On-Chip SRAM (Shared Memory - ~228 KB/SM, 33 TB/s) |
| Computes: Online Softmax over blocks Q_i, K_j, V_j |
| Fuses scaling, causal masking, and reduction without saving A|
+---------------------------------------------------------------+The mathematical core of FlashAttention relies on the Online Softmax Trick. Given a vector x split into two sub-blocks x = [x^(1), x^(2)], the softmax normalizer can be updated iteratively without re-reading past elements:
m_new = max(m_old, max(x^(2)))
d_new = d_old · e^(m_old - m_new) + Σ e^(x_i^(2) - m_new)By maintaining running maximums m and normalizers d in fast GPU registers, FlashAttention-2 achieves a 2× to 4× end-to-end wall-clock speedup and reduces memory complexity from O(L^2) to O(L).
Vectorized PyTorch Implementation#
Below is a production-grade, batched PyTorch module implementing Multi-Head Self-Attention with causal autoregressive masking, QKV fusion, and dropout:
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
class MultiHeadAttention(nn.Module):
"""
Production-grade batched Multi-Head Attention module.
Optimized with fused QKV linear projections and causal masking support.
"""
def __init__(self, d_model: int, num_heads: int, dropout: float = 0.1):
super().__init__()
assert d_model % num_heads == 0, f"d_model ({d_model}) must be divisible by num_heads ({num_heads})"
self.d_model = d_model
self.num_heads = num_heads
self.d_k = d_model // num_heads
# Fused linear projection: Computes Q, K, and V simultaneously in a single GEMM call
self.qkv_proj = nn.Linear(d_model, 3 * d_model, bias=False)
self.out_proj = nn.Linear(d_model, d_model, bias=False)
self.dropout = nn.Dropout(dropout)
# Precompute scale factor
self.scale = 1.0 / math.sqrt(self.d_k)
def forward(
self,
x: torch.Tensor,
is_causal: bool = False,
attn_mask: torch.Tensor = None
) -> torch.Tensor:
"""
Forward computation for Multi-Head Self-Attention.
Args:
x: Input tensor of shape [batch_size, seq_len, d_model]
is_causal: If True, applies an upper-triangular autoregressive mask (-inf)
attn_mask: Optional custom boolean or float attention mask
Returns:
Output tensor of shape [batch_size, seq_len, d_model]
"""
batch_size, seq_len, _ = x.shape
# Step 1: Project inputs to fused QKV [B, L, 3 * d_model]
qkv = self.qkv_proj(x)
# Step 2: Reshape and permute into multi-head format [3, B, num_heads, L, d_k]
qkv = qkv.reshape(batch_size, seq_len, 3, self.num_heads, self.d_k)
qkv = qkv.permute(2, 0, 3, 1, 4)
q, k, v = qkv[0], qkv[1], qkv[2] # Each has shape [B, num_heads, L, d_k]
# Step 3: Compute Scaled Dot Products [B, num_heads, L, L]
# (B, H, L, d_k) @ (B, H, d_k, L) -> (B, H, L, L)
scores = torch.matmul(q, k.transpose(-2, -1)) * self.scale
# Step 4: Apply Causal Masking if requested
if is_causal:
# Upper triangular mask where row < column is masked out with -inf
causal_mask = torch.triu(
torch.full((seq_len, seq_len), float("-inf"), device=x.device, dtype=scores.dtype),
diagonal=1
)
scores = scores + causal_mask.unsqueeze(0).unsqueeze(0)
if attn_mask is not None:
scores = scores.masked_fill(attn_mask == 0, float("-inf"))
# Step 5: Softmax normalization over key dimension
attn_weights = F.softmax(scores, dim=-1)
attn_weights = self.dropout(attn_weights)
# Step 6: Context aggregation [B, num_heads, L, d_k]
context = torch.matmul(attn_weights, v)
# Step 7: Concatenate heads and project output [B, L, d_model]
context = context.permute(0, 2, 1, 3).contiguous().reshape(batch_size, seq_len, self.d_model)
output = self.out_proj(context)
return outputCheckpoint Practice Problems#
Test your mastery of attention mechanisms, probabilistic derivations, and tensor complexity.
Problem 1: Softmax Scaling Factor Derivation#
Prompt: Suppose an engineering team experiments with an alternative head dimension d_k = 256.
- If the scaling factor 1 / √(d_k) is omitted, what is the expected standard deviation of the dot product between two standard normal random vectors?
- By what exact numerical factor does the variance change when the scaling factor 1 / √(256) is applied?
Solution:
- For two independent random vectors q, k ∈ ℝ^256 where each component has zero mean and unit variance:
- Var(S) = Σ_{i=1}^{256} Var(q_i · k_i) = 256 · (1 · 1) = 256
- Standard deviation σ = √(Var(S)) = √(256) = 16.0.
- Without scaling, the standard deviation is 16.0, pushing logits into extreme saturation.
- When applying the scale factor c = 1 / √(256) = 1 / 16:
- Var(c · S) = c^2 · Var(S) = (1 / 256) · 256 = 1.0.
- The variance is reduced by a factor of exactly 256, restoring unit variance.
Problem 2: Memory Footprint of KV-Cache under MHA vs. GQA#
Prompt: A language model has embedding dimension d_model = 4096, number of layers N = 32, number of attention heads h = 32 (head dim d_k = 128), and sequence length L = 8192 tokens. Calculate the KV-cache memory footprint per sequence in FP16 (2 bytes per parameter) under:
- Standard Multi-Head Attention (MHA).
- Grouped-Query Attention (GQA) with G = 4 KV head groups (each group has 8 query heads sharing 1 key and 1 value head).
Solution:
- For MHA:
- Keys and Values per layer = 2 × (L × d_model) = 2 × 8,192 × 4,096 = 67,108,864 elements.
- In FP16 bytes: 67,108,864 × 2 bytes = 134,217,728 bytes = 128.0 MB per layer.
- Across 32 layers: 128.0 MB × 32 = 4,096 MB = 4.00 GB per sequence.
- For GQA (G = 4 groups):
- Total KV heads = 4 (instead of 32).
- KV dimension per token = 4 heads × 128 dim = 512 parameters (compared to 4096 in MHA).
- Keys and Values per layer = 2 × (8,192 × 512) = 8,388,608 elements.
- In FP16 bytes: 8,388,608 × 2 bytes = 16,777,216 bytes = 16.0 MB per layer.
- Across 32 layers: 16.0 MB × 32 = 512 MB = 0.50 GB per sequence.
- GQA achieves an 8× memory reduction, enabling 8× larger batch sizes on serving GPUs.
Problem 3: FLOPs Complexity Comparison#
Prompt: Calculate the total floating-point operations (FLOPs) required to compute scaled dot product attention (Q · K^T and A · V) for a sequence of L = 4,096 tokens with head dimension d_k = 64 across h = 32 heads.
Solution:
- MatMul 1 (Q · K^T):
- Matrix Q has shape [L, d_k] = [4096, 64].
- Matrix K^T has shape [d_k, L] = [64, 4096].
- Multiplying [L, d_k] by [d_k, L] takes 2 · L · L · d_k FLOPs per head (1 multiply + 1 add per MAC).
- FLOPs per head = 2 · 4096 · 4096 · 64 = 2,147,483,648 FLOPs (≈ 2.15 GFLOPs).
- Across 32 heads: 32 · 2,147,483,648 = 68,719,476,736 FLOPs (68.72 GFLOPs).
- MatMul 2 (A · V):
- Matrix A has shape [L, L] = [4096, 4096].
- Matrix V has shape [L, d_k] = [4096, 64].
- Multiplying [L, L] by [L, d_k] also takes 2 · L · L · d_k FLOPs per head.
- Across 32 heads: 68.72 GFLOPs.
- Total Attention FLOPs:
- Total = 68.72 + 68.72 = 137.44 GFLOPs.
- Notice the quadratic dependence on L (O(L^2 · d_model)), which drives computational cost in long-context models.
Problem 4: Causal Masking Invariant#
Prompt: Explain why an autoregressive causal mask must set upper-triangular logits to -∞ rather than 0 prior to the softmax function.
Solution:
- If logits are set to 0 before the softmax function:
```
softmax([z_1, z_2, 0, 0])_j = e^0 / (e^(z_1) + e^(z_2) + e^0 + e^0) = 1 / (e^(z_1) + e^(z_2) + 2) > 0
```
- Because e^0 = 1, future tokens would receive non-zero probability mass, leaking information from the future and violating the autoregressive causal invariant.
- When logits are set to -∞:
```
e^(-∞) = lim_{x → -∞} e^x = 0
```
- Consequently, the numerator for future token indices becomes 0, ensuring that future token positions receive exactly 0 attention weight and cannot influence the current token prediction.