Attention Variants

Chapter 1: Dot-Product Attention Mechanics

Before we can fix attention, we need to understand exactly what it computes. Not the high-level "tokens look at other tokens" hand-wave — the actual matrix operations, from input to output, with every intermediate number visible.

The mechanism is called scaled dot-product attention, and it was introduced in the "Attention Is All You Need" paper (Vaswani et al., 2017). It has three inputs — queries, keys, and values — and produces a weighted combination of the values, where the weights come from comparing queries to keys.

Think of it like a library lookup. You walk in with a question (your query). Each book on the shelf has a label (its key). You compare your question to every label, figure out which books are most relevant (the attention weights), then read those books (the values) in proportion to their relevance. The output is a weighted blend of the relevant books' contents.

The Three Projections: Q, K, V

The input to attention is a sequence of token embeddings — one vector per token. For a sequence of n tokens with embedding dimension d_model, the input is a matrix X of shape (n, d_model).

Attention doesn't work on the raw embeddings directly. Instead, it projects them through three learned weight matrices to create three separate representations:

The projection matrices W_Q, W_K, and W_V are learned during training. This is crucial — the model doesn't just compare raw token embeddings. It learns what to look for (Q), how to be found (K), and what to send when found (V). These can be entirely different functions of the same input.

The Four Steps of Attention

Once we have Q, K, and V, attention proceeds in four steps:

Let's trace this with real numbers.

Hand Calculation: 4 Tokens, d_k = 3

Q matrix (4 × 3) — what each token is looking for:

K matrix (4 × 3) — what each token advertises:

V matrix (4 × 3) — the information payload:

Step 1: Compute Q · K^T

Each entry (i, j) is the dot product between query i and key j. Let's compute row 0 (token 0's query against all keys):

• q₀ · k₀ = 1.0×1.0 + 0.0×0.0 + 1.0×0.0 = 1.0
• q₀ · k₁ = 1.0×0.0 + 0.0×1.0 + 1.0×1.0 = 1.0
• q₀ · k₂ = 1.0×1.0 + 0.0×1.0 + 1.0×0.0 = 1.0
• q₀ · k₃ = 1.0×0.0 + 0.0×0.0 + 1.0×1.0 = 1.0

Token 0's query matches all keys equally! Now row 1 (token 1's query):

• q₁ · k₀ = 0.0×1.0 + 1.0×0.0 + 0.0×0.0 = 0.0
• q₁ · k₁ = 0.0×0.0 + 1.0×1.0 + 0.0×1.0 = 1.0
• q₁ · k₂ = 0.0×1.0 + 1.0×1.0 + 0.0×0.0 = 1.0
• q₁ · k₃ = 0.0×0.0 + 1.0×0.0 + 0.0×1.0 = 0.0

Token 1 strongly matches keys 1 and 2, ignores keys 0 and 3. Continuing for all rows, the full score matrix is:

Note row 2: token 2's query [1,1,0] dots strongly with key 2 [1,1,0] giving a score of 2.0. Token 2 attends most to itself — a common pattern in self-attention.

Step 2: Scale by √d_k

We divide every score by √3 ≈ 1.732:

Step 3: Softmax (row-wise)

Softmax converts each row to a probability distribution. For row 0, all values are equal (0.577), so softmax gives uniform weights: 0.25 each. For row 2, the 1.155 entry gets boosted:

Row 2: softmax([0.577, 0.577, 1.155, 0.000])

• Exponents: [e^0.577, e^0.577, e^1.155, e^0.0] = [1.781, 1.781, 3.174, 1.000]
• Sum = 1.781 + 1.781 + 3.174 + 1.000 = 7.736
• Weights: [0.230, 0.230, 0.410, 0.129]

Token 2 puts 41% of its attention on itself (key 2), 23% each on keys 0 and 1, and only 13% on key 3. The full attention weight matrix:

Each row sums to 1.0 — it's a probability distribution over which tokens to attend to. This is the "attention pattern." Token 0 attends uniformly. Token 1 focuses on tokens 1 and 2. Token 2 focuses on itself.

Step 4: Weighted sum of V

Each token's output is the weighted average of all value vectors, using the attention weights. For token 2:

• output₂ = 0.230 × V₀ + 0.230 × V₁ + 0.410 × V₂ + 0.129 × V₃
• dim 0: 0.230×0.5 + 0.230×1.0 + 0.410×0.0 + 0.129×0.5 = 0.115 + 0.230 + 0.0 + 0.065 = 0.410
• dim 1: 0.230×1.0 + 0.230×0.0 + 0.410×0.5 + 0.129×0.5 = 0.230 + 0.0 + 0.205 + 0.065 = 0.500
• dim 2: 0.230×0.0 + 0.230×0.5 + 0.410×1.0 + 0.129×0.5 = 0.0 + 0.115 + 0.410 + 0.065 = 0.590

Token 2's output is [0.410, 0.500, 0.590]. It's a blend of all value vectors, but skewed toward V₂ (its own value) because it had the highest attention weight on itself. The output "carries information" from the tokens that scored highest.

Why the √d_k Scaling Matters

Without the √d_k scaling, the dot products grow with the dimension of the key vectors. Here's why this is a problem.

If the entries of Q and K are roughly standard normal (mean 0, variance 1), then each dot product q·k is the sum of d_k products of independent standard normals. Each product has mean 0 and variance 1. By the central limit theorem, the sum has mean 0 and variance d_k, so the standard deviation is √d_k.

For d_k = 64 (a typical head size), that means dot products have standard deviation √64 = 8. For d_k = 512, it's √512 ≈ 22.6. Values of magnitude 20+ push softmax into extreme saturation — one entry gets weight ≈ 1.0, everything else gets ≈ 0.0. The output becomes nearly a hard lookup of a single token instead of a soft blend.

Worse: softmax gradients vanish in the saturated regime. The model can't learn to adjust attention patterns if the gradients are dead. Dividing by √d_k normalizes the dot products back to unit variance, keeping softmax in its sensitive region where gradients flow.

See It: Interactive Attention

Click any token (the "query") to see its attention pattern. Arrows from that token to all others show the attention weights — thicker arrows mean higher weight. The heatmap below shows the full attention matrix.

From Scratch: Attention in Code

First, the pure NumPy implementation — every step visible:

python
import numpy as np

def attention(Q, K, V):
    """Scaled dot-product attention from scratch."""
    d_k = Q.shape[-1]

    # Step 1: raw scores
    scores = Q @ K.T               # (n, n)

    # Step 2: scale
    scores = scores / np.sqrt(d_k) # prevents softmax saturation

    # Step 3: softmax (row-wise)
    exp_scores = np.exp(scores - scores.max(axis=-1, keepdims=True))
    weights = exp_scores / exp_scores.sum(axis=-1, keepdims=True)

    # Step 4: weighted sum of values
    output = weights @ V            # (n, d_v)
    return output, weights

# Test with our hand-calculated example
Q = np.array([[1,0,1],[0,1,0],[1,1,0],[0,0,1]], dtype=np.float32)
K = np.array([[1,0,0],[0,1,1],[1,1,0],[0,0,1]], dtype=np.float32)
V = np.array([[.5,1,0],[1,0,.5],[0,.5,1],[.5,.5,.5]], dtype=np.float32)

out, wts = attention(Q, K, V)
print("Attention weights:\n", wts.round(3))
print("Output:\n", out.round(3))

Now the PyTorch one-liner equivalent:

python
import torch
import torch.nn.functional as F

Q = torch.tensor([[1,0,1],[0,1,0],[1,1,0],[0,0,1]], dtype=torch.float32)
K = torch.tensor([[1,0,0],[0,1,1],[1,1,0],[0,0,1]], dtype=torch.float32)
V = torch.tensor([[.5,1,0],[1,0,.5],[0,.5,1],[.5,.5,.5]], dtype=torch.float32)

# PyTorch's built-in (handles batches, masking, dropout too)
out = F.scaled_dot_product_attention(
    Q.unsqueeze(0), K.unsqueeze(0), V.unsqueeze(0)
).squeeze(0)
print(out.round(decimals=3))  # matches our hand calculation

Chapter 2: Multi-Head Attention

One attention pattern isn't enough. Consider the sentence "The cat sat on the mat because it was tired." The word "it" needs to simultaneously attend to "cat" (to resolve the pronoun) and "sat" (to understand what action "it" performed). A single set of Q, K, V projections can only produce one attention pattern — one way of routing information.

Multi-head attention solves this by running multiple attention operations in parallel, each with its own learned Q, K, V projections. Each head can specialize: one head might focus on local syntax (nearby words), another on long-range semantic connections (pronouns to antecedents), another on positional patterns (sentence boundaries). The outputs are concatenated and projected back to the original dimension.

How the Split Works

Here's the key insight that makes multi-head attention efficient: we don't add more parameters for more heads. We split the existing dimension across heads.

If d_model = 512 and we want H = 8 heads, each head gets:

Each head operates on 64-dimensional queries, keys, and values instead of 512-dimensional ones. Since the total dimension across all heads is still 8 × 64 = 512, the total computation is the same as a single head with d_k = 512.

The Full Data Flow

Let's trace the exact tensor shapes through multi-head attention for a concrete example: batch size B = 1, sequence length n = 6, d_model = 512, H = 8 heads.

Step 1: Project Q, K, V

• Input X: shape (1, 6, 512)
• W_Q: shape (512, 512) — projects to all heads at once
• Q = X · W_Q: shape (1, 6, 512)
• Same for K and V: both (1, 6, 512)

Step 2: Reshape into heads

• Q reshaped: (1, 6, 512) → (1, 6, 8, 64) → transpose → (1, 8, 6, 64)
• Now Q[batch, head, token, dim] — each head has its own 6×64 slice
• Same for K: (1, 8, 6, 64)
• Same for V: (1, 8, 6, 64)

Step 3: Attention per head

• Scores = Q · K^T: (1, 8, 6, 64) · (1, 8, 64, 6) = (1, 8, 6, 6)
• Each head produces its own 6×6 attention matrix
• Scale by √64 = 8, apply softmax: still (1, 8, 6, 6)
• Output per head = weights · V: (1, 8, 6, 6) · (1, 8, 6, 64) = (1, 8, 6, 64)

Step 4: Concatenate heads

• Transpose back: (1, 8, 6, 64) → (1, 6, 8, 64)
• Reshape (concatenate): (1, 6, 8, 64) → (1, 6, 512)
• All 8 heads' 64-dimensional outputs are now stacked into a single 512-dimensional vector per token

Step 5: Output projection

• W_O: shape (512, 512)
• Output = concat · W_O: (1, 6, 512)
• Same shape as input — attention is a shape-preserving operation

Hand Calculation: Parameter Count

Each projection matrix maps d_model to d_model (we project all heads at once, then reshape):

• W_Q: 512 × 512 = 262,144 parameters
• W_K: 512 × 512 = 262,144 parameters
• W_V: 512 × 512 = 262,144 parameters
• W_O: 512 × 512 = 262,144 parameters

Total: 4 × 262,144 = 1,048,576 parameters (about 1 million).

Now the crucial comparison. What if we used a single head with d_k = 512?

• W_Q: 512 × 512 = 262,144
• W_K: 512 × 512 = 262,144
• W_V: 512 × 512 = 262,144
• W_O: 512 × 512 = 262,144
• Total: exactly 1,048,576 parameters — identical!

The parameter count is the same because we just split the projection into smaller pieces per head. The compute cost (FLOPs) is also the same — we're doing the same matrix multiplications, just organized differently. Multi-head attention is a free lunch: more expressive attention patterns at zero extra cost.

What Each Head Learns

In practice, different heads learn strikingly different attention patterns. Research on trained transformers (Voita et al., 2019; Clark et al., 2019) has identified several head types:

With only 1 head, the model must compromise — one attention pattern for all these functions. With 8+ heads, each can specialize. This is why multi-head attention is universal in modern transformers.

See It: Multi-Head Patterns

The simulation below shows attention patterns for different numbers of heads. With 1 head, you see a single blurred pattern trying to capture everything. Increase the head count and watch specialization emerge — each head focuses on a different relationship type.

From Scratch: Full MHA Implementation

python
import torch
import torch.nn as nn
import torch.nn.functional as F

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        assert d_model % n_heads == 0
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads  # 512 // 8 = 64

        # Project all heads at once for efficiency
        self.W_q = nn.Linear(d_model, d_model)  # (512 → 512)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)

    def forward(self, x):
        B, N, D = x.shape  # (batch, seq_len, d_model)

        # Step 1: Project
        Q = self.W_q(x)  # (B, N, 512)
        K = self.W_k(x)
        V = self.W_v(x)

        # Step 2: Reshape into heads
        # (B, N, 512) → (B, N, 8, 64) → (B, 8, N, 64)
        Q = Q.view(B, N, self.n_heads, self.d_k).transpose(1, 2)
        K = K.view(B, N, self.n_heads, self.d_k).transpose(1, 2)
        V = V.view(B, N, self.n_heads, self.d_k).transpose(1, 2)

        # Step 3: Scaled dot-product attention (per head)
        scores = Q @ K.transpose(-2, -1) / (self.d_k ** 0.5)
        weights = F.softmax(scores, dim=-1)  # (B, 8, N, N)
        attn_out = weights @ V                   # (B, 8, N, 64)

        # Step 4: Concatenate heads
        # (B, 8, N, 64) → (B, N, 8, 64) → (B, N, 512)
        attn_out = attn_out.transpose(1, 2).contiguous().view(B, N, D)

        # Step 5: Output projection
        return self.W_o(attn_out)  # (B, N, 512)

# Verify shapes
mha = MultiHeadAttention(d_model=512, n_heads=8)
x = torch.randn(1, 6, 512)  # (batch=1, seq=6, dim=512)
out = mha(x)
print(out.shape)  # torch.Size([1, 6, 512]) — same as input
print(f"Total params: {sum(p.numel() for p in mha.parameters()):,}")
# Total params: 1,050,624 (≈1M + biases)

Chapter 3: The KV Cache Bottleneck

During training, all tokens are processed in parallel — attention sees the full sequence at once. This is efficient because GPUs excel at large batch matrix multiplies. But during generation, we produce ONE token at a time. Each new token needs to attend to ALL previous tokens.

Without any tricks, generating token number 1000 would require recomputing the key and value vectors for all 999 previous tokens — even though those tokens haven't changed since the last step. That's a colossal waste.

The solution is the KV cache: store every previous token's key and value vectors in memory, and reuse them on each generation step. This turns attention from O(n²) recomputation into O(n) new work per step — but at a steep memory cost.

How Autoregressive Generation Works

Let's trace exactly what happens when an LLM generates a 5-token response to a 3-token prompt. We'll call the prompt tokens [A, B, C] and the generated tokens [D, E, F, G, H].

Step 1 (prefill): Process prompt [A, B, C]

• Compute K_A, K_B, K_C and V_A, V_B, V_C
• Run attention, predict next token → D
• Save K and V for [A, B, C] in the cache

Step 2: Generate token D

• WITHOUT cache: recompute K, V for [A, B, C, D] = 4 sets of K, V
• WITH cache: only compute K_D, V_D. Append to cached [K_A, K_B, K_C]
• Compute Q_D, attend to all 4 cached keys, predict → E

Step 3: Generate token E

• WITHOUT cache: recompute K, V for [A, B, C, D, E] = 5 sets
• WITH cache: only compute K_E, V_E. Cache now has 5 entries.
• Compute Q_E, attend to all 5 keys, predict → F

Without caching, by step N we recompute K and V for N tokens. The total K/V computations across all steps: 3 + 4 + 5 + 6 + 7 = 25. With caching: 3 + 1 + 1 + 1 + 1 = 7. That's a 3.6× savings for just 5 generated tokens. For 1000 generated tokens with a 1000-token prompt, the savings are over 500×.

Hand Calculation: KV Cache Memory

For each token in the sequence, we store one key vector and one value vector per layer, per head. The cache size per token:

The factor of 2 is for K and V (we cache both). Let's compute this for Llama 2 7B:

Bytes per token:

• = 2 (K+V) × 32 (layers) × 32 (heads) × 128 (d_head) × 2 (FP16 bytes)
• = 2 × 32 × 32 × 128 × 2
• = 2 × 32 × 8,192
• = 2 × 262,144
• = 524,288 bytes per token (512 KB)

At different sequence lengths:

• 512 tokens: 524,288 × 512 = 268 MB
• 2,048 tokens: 524,288 × 2,048 = 1.07 GB
• 4,096 tokens: 524,288 × 4,096 = 2.15 GB
• 8,192 tokens: 524,288 × 8,192 = 4.29 GB

Now here's where it gets alarming. The model weights themselves for Llama 2 7B in FP16 are about 14 GB. That's fixed — load once, serve forever. But the KV cache is per user, per request.

For a batch of 32 concurrent users at 4,096 context:

• KV cache = 2.15 GB × 32 = 68.7 GB
• Model weights = 14 GB
• Total = 82.7 GB
• H100 GPU memory = 80 GB
• It doesn't fit.

The KV cache alone needs 5× more memory than the model weights. This is why you can't "just add more users" to a GPU — each concurrent user adds ~2 GB of KV cache, and GPU memory is finite.

The Full Picture: Memory Breakdown

For the 70B model, the weights alone don't fit on a single H100. Even with tensor parallelism across 2 GPUs (160 GB total), after loading the 140 GB model, you have only 20 GB left — enough for just 2 concurrent users at 4K context. This is why KV cache optimization (the topic of the next chapters) is the most active area of LLM systems engineering.

See It: KV Cache Growth

The simulation below shows the KV cache growing in real time as tokens are generated. Click "Generate" to produce one token at a time and watch the cache expand. The memory counter shows exactly how many bytes are consumed.

From Scratch: KV Cache Calculator

python
def kv_cache_size_gb(
    n_layers,
    n_heads,
    d_head,
    seq_len,
    n_users=1,
    bytes_per_param=2,  # FP16
):
    """Calculate KV cache memory in GB."""
    bytes_per_token = 2 * n_layers * n_heads * d_head * bytes_per_param
    total_bytes = bytes_per_token * seq_len * n_users
    return total_bytes / (1024 ** 3)

# Llama 2 family
models = {
    "7B":  (32, 32, 128),   # (layers, heads, d_head)
    "13B": (40, 40, 128),
    "70B": (80, 64, 128),
}

print(f"{'Model':>6} {'SeqLen':>7} {'Users':>6} {'KV Cache':>10} {'Weights':>10} {'Total':>10}")
print("-" * 55)

for name, (nl, nh, dh) in models.items():
    weight_gb = {"7B": 14, "13B": 26, "70B": 140}[name]
    for seq in [2048, 4096, 8192]:
        for users in [1, 32]:
            kv = kv_cache_size_gb(nl, nh, dh, seq, users)
            total = kv + weight_gb
            print(f"{name:>6} {seq:>7} {users:>6} {kv:>9.2f}G {weight_gb:>9}G {total:>9.1f}G")

# Key output lines:
#     7B    4096     32      68.7G        14G      82.7G  ← doesn't fit H100!
#    70B    4096      1      10.0G       140G     150.0G  ← needs 2+ GPUs
#    70B    4096     32     320.0G       140G     460.0G  ← needs 6+ H100s

Why This Changes Everything

The KV cache bottleneck explains almost every engineering decision in modern LLM deployment:

• Quantization (FP16 → INT8 → INT4): Halving the bytes per parameter halves the KV cache. INT4 KV cache gives 4× more users.
• Multi-Query Attention (Ch 4): Share K and V across heads. Instead of 32 heads × 128d, use 1 K/V set. KV cache shrinks 32×.
• Grouped-Query Attention (Ch 5): Compromise — share K/V across groups of heads. 8 KV groups instead of 32: 4× smaller cache.
• Sliding window attention: Only cache the last W tokens instead of the full sequence. Caps KV cache at a fixed size.
• Paged attention (vLLM): Manage KV cache like virtual memory pages — avoid fragmentation, share cache across similar requests.

The rest of this lesson dives into each of these techniques. Every one of them exists because someone looked at the KV cache memory equation and found a term to shrink.

Chapter 4: Multi-Query Attention

Chapter 3 showed us the KV cache is the memory bottleneck. Every attention head stores its own K and V vectors for every token generated so far. For a model with 32 heads, that's 32 separate copies of keys and 32 separate copies of values, all sitting in GPU memory, all growing with every new token.

Here's the first clever fix: what if all 32 attention heads shared a single set of K and V vectors? Each head still asks its own question (different Q projections), but they all search through the same keys and retrieve from the same values. That's Multi-Query Attention (MQA), proposed by Noam Shazeer in 2019.

The idea sounds radical — surely different heads need different keys? But remember: Q is the "question," and K is the "index." If every head asks a different question, they'll get different attention patterns even against the same set of keys. The queries are still fully expressive; we're just sharing the lookup table.

The Architecture Change

In standard Multi-Head Attention (MHA), the input X gets projected into Q, K, and V per head. If there are H heads and each head has dimension d_h, the K and V projection matrices are each of size [d_model × (H · d_h)]. That's one d_h-dimensional key vector and one d_h-dimensional value vector per head, per token.

In MQA, the K and V projection matrices shrink to [d_model × d_h] — just one d_h-dimensional key vector and one d_h-dimensional value vector, shared across all H heads. The Q projection stays the same: each head still gets its own query.

Hand Calculation: Memory Savings

Let's get concrete with Llama 2 7B numbers. The model has 32 attention heads, each with d_head = 128. Weights are stored in float16 (2 bytes per value). The KV cache stores K and V for every token generated so far.

MHA KV cache per token:

• Keys: 32 heads × 128 dims × 2 bytes = 8,192 bytes
• Values: 32 heads × 128 dims × 2 bytes = 8,192 bytes
• Total K+V per token: 8,192 + 8,192 = 16,384 bytes (16 KB)

MQA KV cache per token:

• Keys: 1 shared × 128 dims × 2 bytes = 256 bytes
• Values: 1 shared × 128 dims × 2 bytes = 256 bytes
• Total K+V per token: 256 + 256 = 512 bytes

That's a 32× reduction — exactly equal to the number of heads.

Now scale it up. With a context window of 4,096 tokens:

• MHA: 16,384 bytes × 4,096 tokens = 67,108,864 bytes = 64 MB per sequence
• MQA: 512 bytes × 4,096 tokens = 2,097,152 bytes = 2 MB per sequence

Now consider serving 64 concurrent users (a modest production workload):

• MHA: 64 MB × 64 users = 4,096 MB ≈ 4 GB just for the KV cache
• MQA: 2 MB × 64 users = 128 MB — barely a rounding error on an H100

For larger models the savings are even more dramatic. PaLM 540B with 48 heads? MQA shrinks the KV cache by 48×. At scale, this is the difference between needing a cluster and fitting on a single GPU.

The Quality Concern

The obvious question: if all heads share the same K and V, don't they all attend to the same things? Don't we lose the diversity that makes multi-head attention powerful?

Surprisingly, no — or at least, not much. Each head still has its own Q projection, which means each head asks a different question of the shared keys. Head 1 might project its query to emphasize syntactic features; head 12 might emphasize positional features. They'll compute different attention weights (softmax(Q_hK^T/√d)) because the Q_h differs, even though K is the same for all.

Empirically, MQA drops quality by roughly 0.5–1% on standard benchmarks compared to MHA. That's a tiny cost for a 32× memory reduction. The tradeoff is overwhelmingly worthwhile for production inference.

Code: MQA Implementation

Here's how MQA modifies the standard MHA forward pass. The key change is in the projection dimensions: K and V are projected to d_head (not H × d_head), then broadcast across all heads during the attention computation.

python
import torch
import torch.nn as nn

class MultiQueryAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.n_heads = n_heads
        self.d_head = d_model // n_heads

        # Q: per-head projections (standard)
        self.W_q = nn.Linear(d_model, d_model)  # [d, H*d_h]

        # K, V: SINGLE head projection (the MQA trick)
        self.W_k = nn.Linear(d_model, self.d_head)  # [d, d_h]
        self.W_v = nn.Linear(d_model, self.d_head)  # [d, d_h]

        self.W_o = nn.Linear(d_model, d_model)

    def forward(self, x):
        B, T, D = x.shape
        H, d_h = self.n_heads, self.d_head

        # Project Q, K, V
        q = self.W_q(x).view(B, T, H, d_h).transpose(1, 2)  # [B, H, T, d_h]
        k = self.W_k(x).unsqueeze(1)                             # [B, 1, T, d_h]
        v = self.W_v(x).unsqueeze(1)                             # [B, 1, T, d_h]

        # k and v broadcast across H heads automatically
        scores = (q @ k.transpose(-2, -1)) / (d_h ** 0.5)  # [B, H, T, T]
        attn = scores.softmax(dim=-1)
        out = (attn @ v).transpose(1, 2).reshape(B, T, D)     # [B, T, D]
        return self.W_o(out)

The critical lines are the W_k and W_v projections: they output d_head dimensions (one head's worth), then unsqueeze(1) adds the head dimension so PyTorch's broadcasting does the rest. Each of the H heads computes Q_hK^T using the same K, without any copying.

Chapter 5: Grouped-Query Attention

MQA is aggressive. All heads share a single set of keys and values. For most large models, the quality drop is negligible. But some teams found it too steep — especially for models in the 7B–13B range where every fraction of a percent on benchmarks matters for competitive positioning.

What if we could dial the tradeoff between MHA and MQA? Instead of giving every head its own KV (MHA) or sharing one KV across all heads (MQA), what if we put heads into groups and shared one KV per group?

That's Grouped-Query Attention (GQA), introduced by Ainslie et al. at Google in 2023. Think of it as a knob: turn it all the way left and you get MQA (one group). Turn it all the way right and you get MHA (every head is its own group). GQA sits anywhere in between.

How Groups Work

Say your model has H = 32 query heads and you choose G = 8 KV groups. That means every group of 32/8 = 4 query heads shares a single K and V. The model has 8 independent K projection vectors and 8 independent V projection vectors, instead of 32 (MHA) or 1 (MQA).

During attention, heads 0–3 all use KV group 0. Heads 4–7 use KV group 1. And so on. Within each group, the mechanism is identical to MQA: different queries, shared keys and values.

Hand Calculation: Llama 2 70B

Llama 2 70B uses GQA with 64 query heads and 8 KV groups. Each head has d_head = 128. Let's compute the KV cache savings.

Heads per group: 64 query heads / 8 groups = 8 query heads per KV group.

GQA-8 KV cache per token:

• Keys: 8 groups × 128 dims × 2 bytes = 2,048 bytes
• Values: 8 groups × 128 dims × 2 bytes = 2,048 bytes
• Total: 2,048 + 2,048 = 4,096 bytes

MHA (for comparison) KV cache per token:

• Keys: 64 heads × 128 dims × 2 bytes = 16,384 bytes
• Values: 64 heads × 128 dims × 2 bytes = 16,384 bytes
• Total: 16,384 + 16,384 = 32,768 bytes

MQA (for comparison) KV cache per token:

• Total: 2 × 1 × 128 × 2 = 512 bytes

GQA-8 gives an 8× reduction over MHA. Not as extreme as MQA's 64×, but with much better quality retention. Here's the full comparison at 4,096 tokens context:

At 64 concurrent users, MHA eats 8 GB of VRAM just for the KV cache — half of a 16 GB card, before model weights even load. GQA-8 needs just 1 GB. MQA needs 128 MB. The sweet spot depends on your model size and serving constraints, but for most 70B-class models, GQA-8 is the winner.

The Conversion Trick

Here's a remarkable fact: you can convert an already-trained MHA model to GQA without retraining from scratch. The recipe is simple: take the K and V weight matrices from each group of heads and average them.

Say you have 64 heads and want 8 groups. Heads 0–7 form group 0. Take their 8 K weight matrices, average them element-by-element, and use that average as the single K matrix for the group. Do the same for V. Repeat for all 8 groups.

Then fine-tune the model for a few thousand steps to let it adapt to the shared parameters. Llama 2 70B was originally trained with MHA. Meta converted it to GQA-8 using exactly this procedure and recovered 97% of the original quality after a short fine-tuning run. The KV cache shrank 8× for a 3% quality cost that was further recovered during fine-tuning.

Code: GQA Implementation

The key operation in GQA is the group repeat: after projecting K and V to G groups, we need to "expand" them so each query head has matching K and V tensors. PyTorch handles this with repeat_interleave or manual reshaping.

python
import torch
import torch.nn as nn

class GroupedQueryAttention(nn.Module):
    def __init__(self, d_model, n_heads, n_kv_groups):
        super().__init__()
        self.n_heads = n_heads
        self.n_kv = n_kv_groups
        self.d_head = d_model // n_heads
        self.heads_per_group = n_heads // n_kv_groups

        # Q: full per-head projections
        self.W_q = nn.Linear(d_model, n_heads * self.d_head)
        # K, V: only n_kv_groups projections
        self.W_k = nn.Linear(d_model, n_kv_groups * self.d_head)
        self.W_v = nn.Linear(d_model, n_kv_groups * self.d_head)
        self.W_o = nn.Linear(d_model, d_model)

    def repeat_kv(self, x):
        # x: [B, G, T, d_h] -> [B, H, T, d_h]
        # Repeat each KV group heads_per_group times
        B, G, T, d = x.shape
        x = x[:, :, None, :, :].expand(B, G, self.heads_per_group, T, d)
        return x.reshape(B, self.n_heads, T, d)

    def forward(self, x):
        B, T, D = x.shape
        d_h = self.d_head

        q = self.W_q(x).view(B, T, self.n_heads, d_h).transpose(1, 2)
        k = self.W_k(x).view(B, T, self.n_kv, d_h).transpose(1, 2)
        v = self.W_v(x).view(B, T, self.n_kv, d_h).transpose(1, 2)

        # Expand K, V to match query head count
        k = self.repeat_kv(k)  # [B, H, T, d_h]
        v = self.repeat_kv(v)  # [B, H, T, d_h]

        scores = (q @ k.transpose(-2, -1)) / (d_h ** 0.5)
        attn = scores.softmax(dim=-1)
        out = (attn @ v).transpose(1, 2).reshape(B, T, D)
        return self.W_o(out)

The repeat_kv method is the core of GQA: it takes the G-grouped K or V tensor and repeats each group H/G times so the shapes align for the batched matrix multiply. In practice, most frameworks avoid actual memory copies by using expand (which creates a view, not a copy) — so the memory savings are real, not just theoretical.

Chapter 6: Sliding Window Attention

MQA and GQA attacked the KV cache size per token — shrinking how much memory each position needs. But there's a second dimension to the problem: the number of tokens each position attends to. In standard attention, every token looks at every previous token. That's the n×n attention matrix that makes the whole thing quadratic.

What if each token only looked at the W nearest tokens instead of all n? That's Sliding Window Attention — also called local attention. It caps the attention cost per token at W, making the total cost O(n · W) instead of O(n²). When the sequence is long and the window is moderate, the savings are enormous.

This isn't a new idea — local attention patterns appeared in research as early as 2019. But it became famous when Mistral 7B shipped with sliding window attention as a key architectural choice, showing that you could match or beat much larger models by being smart about what each token actually needs to see.

Why Local Context Is Usually Enough

Think about how language works. When you read the word "running" in "the dog is running through the park," you need the nearby tokens to understand it. You almost never need token #47 from 10,000 positions ago. Language is overwhelmingly local: subjects, verbs, and objects cluster together. Pronouns reference things within a few sentences. Semantic coherence is maintained over paragraphs, not pages.

There are exceptions — a reference to a character introduced chapters ago, or a variable defined hundreds of lines up in code. But these long-range dependencies are rare compared to local ones. Sliding window attention bets that most attention mass lands in a local neighborhood, and that bet pays off.

The Mask Pattern

Visualize the attention matrix as a grid. Rows are query positions (tokens asking questions). Columns are key positions (tokens being attended to). In standard causal attention, the lower-left triangle is filled — each token can attend to itself and all previous tokens.

In sliding window attention with window size W, each token can only attend to the W most recent tokens (including itself). The filled region shrinks from a triangle to a diagonal band of width W. Tokens outside the band get attention weight zero.

Hand Calculation: Computation Savings

Let's count the number of attention entries (cells in the matrix) for different configurations.

Full causal attention, n = 16,384:

The lower triangle has n · (n+1) / 2 entries. For n = 16,384:

• Entries = 16,384 × 16,385 / 2 = 134,225,920 (≈ 134M)

Sliding window, n = 16,384, W = 4,096:

Each of the first W tokens attends to all previous tokens (building up to the full window). After that, each token attends to exactly W tokens. A close upper bound:

• Entries ≤ n × W = 16,384 × 4,096 = 67,108,864 (≈ 67M)
• Savings: 134M / 67M = 2× fewer

At n = 32,768 with W = 4,096:

• Full: 32,768 × 32,769 / 2 = 536,887,296 (≈ 537M)
• Sliding: 32,768 × 4,096 = 134,217,728 (≈ 134M)
• Savings: 537M / 134M = 4× fewer

The longer the sequence relative to the window, the bigger the savings. At n = 131,072 with W = 4,096: full attention requires ∼8.6 billion entries, while sliding window needs ∼537 million. That's a 16× reduction. The savings scale as n / (2W) for large n.

The Receptive Field Trick

The obvious worry: if each token only sees W positions, doesn't the model lose access to information beyond the window? How can it handle a question like "Who was the main character introduced in chapter 1?" when you're 50,000 tokens later?

The answer is layer stacking. Each layer extends the effective reach by W positions. Think of it like a relay race:

• Layer 1: Token t can see tokens [t−W, t]. Information from position t−W is now "in" position t.
• Layer 2: Token t can see W positions back again, but those positions already contain information from their W-token windows. So token t effectively has information from [t−2W, t].
• Layer L: Token t has an effective receptive field of [t−L×W, t].

Mistral 7B has 32 layers and window size W = 4,096. Its effective receptive field is:

That's 131K tokens — far beyond the model's actual context window. The sliding window doesn't limit information flow; it just makes information travel through the network layer by layer rather than arriving all at once. Empirically, this works just as well for the vast majority of tasks.

The Rolling Buffer KV Cache

Sliding window attention enables another powerful optimization: a rolling buffer KV cache. In standard attention, the KV cache grows with every new token and never shrinks. At position 10,000, you're storing KV entries for all 10,000 previous positions.

With a sliding window, token 10,000 only attends to positions [5,905, 10,000] (for W = 4,096). Positions before 5,905 are outside every current and future token's window. They will never be attended to again. So we can evict them from the cache.

The result: the KV cache is bounded at W entries, regardless of how long the generation runs. Instead of KV cache growing linearly with sequence length, it stays constant. For long-form generation (writing a novel, processing a codebase), this is transformative.

Code: Sliding Window Mask

python
import torch

def sliding_window_mask(seq_len, window_size):
    # Create causal mask first (lower triangle)
    causal = torch.ones(seq_len, seq_len, dtype=torch.bool).tril()

    # Create window mask: token i attends to [i-W+1, i]
    rows = torch.arange(seq_len).unsqueeze(1)    # [n, 1]
    cols = torch.arange(seq_len).unsqueeze(0)    # [1, n]
    window = (rows - cols) < window_size        # within W positions

    # Combine: must be both causal AND within window
    mask = causal & window
    return mask  # True = attend, False = mask out

# Example: 8 tokens, window size 3
mask = sliding_window_mask(8, 3)
# Token 0 sees: [0]
# Token 1 sees: [0, 1]
# Token 2 sees: [0, 1, 2]
# Token 3 sees: [1, 2, 3]     <- window kicks in
# Token 7 sees: [5, 6, 7]

Chapter 7: Linear Attention

Every variant we've seen so far still computes the n×n attention matrix — they just make it sparser (sliding window) or share the KV cache (MQA, GQA). The matrix itself, that softmax(QK^T/√d) grid of attention weights, is still there. It's still quadratic.

What if we could avoid the n×n matrix entirely? Not approximate it with sparsity, not share components around it, but mathematically side-step it altogether?

That's what Linear Attention does. The key insight is startlingly simple: by changing the order of matrix multiplication, we can compute the same output without ever forming the n×n intermediate matrix. The cost drops from O(n²d) to O(nd²), and when d is much smaller than n (which it always is in practice), that's a massive win.

The Associativity Trick

Let's start with standard attention. For one head, given queries Q (n×d), keys K (n×d), and values V (n×d):

The problem is the QK^T term: (n×d) · (d×n) = n×n matrix. This requires O(n²d) operations and O(n²) memory.

Now, softmax is what forces us to compute QK^T first — we need every entry of that matrix before we can normalize rows. But what if we replaced softmax with a simpler function that doesn't require the full matrix?

Replace softmax with a kernel function φ applied element-wise to Q and K:

Read that carefully. The parentheses changed. Instead of computing (φ(Q) · φ(K)^T) · V, we compute φ(Q) · (φ(K)^T · V). Same letters, different grouping. This is the associativity of matrix multiplication.

Why does the grouping matter? Let's trace the shapes:

The intermediate matrix went from n×n (grows quadratically with sequence length) to d×d (fixed at d_head², regardless of sequence length). This is the core of linear attention.

Hand Calculation: Standard vs Linear Order

Let's work through a tiny example. n = 4 tokens, d = 3 dimensions. We'll skip the kernel φ (assume identity) and just compare the two multiplication orders.

Standard order: (QK^T)V

Step 1 — QK^T: multiply (4×3) by (3×4) = (4×4) matrix.

• Operations: each entry in the 4×4 result requires d = 3 multiply-adds.
• Total: 4 × 4 × 3 = 48 operations

Step 2 — (4×4) · V: multiply (4×4) by (4×3) = (4×3) output.

• Operations: 4 × 3 entries, each requiring 4 multiply-adds.
• Total: 4 × 3 × 4 = 48 operations

Grand total: 48 + 48 = 96 operations. Intermediate matrix: 4×4 = 16 elements.

Linear order: Q(K^TV)

Step 1 — K^TV: multiply (3×4) by (4×3) = (3×3) matrix.

• Operations: 3 × 3 entries, each requiring 4 multiply-adds.
• Total: 3 × 3 × 4 = 36 operations

Step 2 — Q · (3×3): multiply (4×3) by (3×3) = (4×3) output.

• Operations: 4 × 3 entries, each requiring 3 multiply-adds.
• Total: 4 × 3 × 3 = 36 operations

Grand total: 36 + 36 = 72 operations. Intermediate matrix: 3×3 = 9 elements.

At n = 4, the saving is modest: 96 vs 72 operations. But watch what happens as n grows. The standard approach scales as n² while the linear approach scales as n:

At n = 16,384 with d = 128 (typical for modern LLMs), linear attention is 128× faster than standard attention. The crossover where linear becomes cheaper is roughly n ≈ d — for any practical sequence length, linear wins.

The Kernel Function φ

We can't just remove softmax and call it a day. Softmax ensures attention weights are positive and sum to 1 — it's a proper probability distribution. If we want the associativity trick to work, we need a kernel function φ that preserves these properties (or close enough).

The original Linear Transformer (Katharopoulos et al., 2020) used a simple choice:

Where elu is the exponential linear unit: elu(x) = x if x > 0, e^x − 1 if x ≤ 0. Adding 1 ensures all values are positive (matching softmax's positivity). The output is then normalized by dividing by the sum of φ(Q)φ(K)^T along the key dimension.

Other kernels have been explored: random Fourier features (FAVOR+ in Performer), polynomial kernels, learned kernels. The choice of φ determines how well linear attention approximates the original softmax attention pattern.

The Catch

If linear attention is so much faster, why isn't everyone using it? The answer is quality. Softmax attention creates sharp, peaked distributions — it can put 90% of the attention weight on a single token when that token is the relevant one. This sharp retrieval is what makes attention so powerful for tasks like question answering, code completion, and factual recall.

The kernel approximation φ produces smoother distributions. Instead of a sharp peak on one token, it spreads weight more broadly. For most generation tasks (continuing text, summarizing, chat), this barely matters. But for "needle in a haystack" retrieval — finding a specific fact in a long context — linear attention is measurably weaker.

This is an active area of research. Models like RWKV, RetNet, and Mamba (which uses a related selective state space mechanism) are closing the gap. Each generation gets closer to softmax quality while keeping linear or near-linear scaling. The bet in the research community is that subquadratic attention will eventually match softmax — but we're not there yet for all tasks.

Code: Linear Attention in 15 Lines

python
import torch
import torch.nn.functional as F

def linear_attention(Q, K, V):
    # Q, K, V: [batch, heads, seq_len, d_head]

    # Apply kernel: phi(x) = elu(x) + 1 (ensures positivity)
    Q = F.elu(Q) + 1   # [B, H, n, d]
    K = F.elu(K) + 1   # [B, H, n, d]

    # THE TRICK: compute K^T V first (d×d, not n×n!)
    KV = K.transpose(-2, -1) @ V    # [B, H, d, d]

    # Then multiply Q by the d×d summary
    out = Q @ KV                        # [B, H, n, d]

    # Normalize (equivalent to softmax denominator)
    Z = Q @ K.transpose(-2, -1).sum(dim=-1, keepdim=True)
    # Simpler: Z = (Q * K.sum(dim=-2, keepdim=True)).sum(dim=-1, keepdim=True)

    return out / (Z + 1e-6)  # avoid division by zero

The entire trick is on one line: KV = K.transpose(-2, -1) @ V. This computes the d×d summary matrix. Then Q @ KV produces the output without ever materializing an n×n matrix. The normalization term Z ensures the output is properly scaled.

Chapter 8: FlashAttention

Every variant we've seen so far changes the attention algorithm — fewer heads, local windows, kernel tricks. FlashAttention takes a completely different approach: it computes exact standard attention, but reorganizes the computation to be friendly to the GPU's memory hierarchy. Same math, radically faster.

The insight is almost embarrassingly simple once you see it. Modern GPUs have two kinds of memory: a large, slow main memory (HBM — High Bandwidth Memory), and a tiny, blazing-fast on-chip cache (SRAM). Standard attention constantly shuffles enormous matrices between these two levels. FlashAttention restructures the work so that almost everything happens in the fast cache, and the huge intermediate matrix never gets written at all.

The GPU Memory Hierarchy

Think of HBM as a warehouse and SRAM as your workbench. The warehouse is huge (80 GB on an H100) but it takes time to haul things back and forth. Your workbench is tiny (about 50 MB) but everything on it is instantly accessible. Standard attention is like building furniture by carrying every plank to the warehouse after cutting it, then hauling it all back to assemble. FlashAttention keeps the planks on your workbench and only moves the finished piece once.

Concrete numbers for an NVIDIA H100:

The bottleneck in modern GPUs is rarely compute — it's memory bandwidth. Moving data to and from HBM is what makes attention slow. FlashAttention is an IO-aware algorithm: it minimizes the number of HBM reads and writes by doing as much work as possible while data sits in fast SRAM.

Standard Attention: The HBM Nightmare

Let's trace what standard attention does for a single head. Inputs: Q, K, V ∈ ℝ^{n × d}, all sitting in HBM.

Count the HBM traffic: we read/write that n × n matrix three times (write S, read for softmax, read P for multiply). For n = 8192 and d = 128 in FP16 (2 bytes per value):

With 32 heads and 80 layers, that's 32 × 80 × 128 MB = 320 GB of HBM traffic just for the attention scores — in a single forward pass. And HBM bandwidth is the bottleneck. This is where all the time goes.

FlashAttention: The Tiling Strategy

FlashAttention's core idea: never materialize the full n × n matrix. Instead, split Q into row blocks of size B_r and K, V into column blocks of size B_c. Load one Q-block and one K,V-block into SRAM at a time. Compute that tile of the attention matrix, immediately use it, and accumulate the result. The tile lives only in SRAM — it never touches HBM.

Hand Calculation: How Many Tiles?

SRAM must hold: one Q-block (B_r × d), one K-block (B_c × d), one V-block (B_c × d), and the partial output block (B_r × d). That's approximately 4 × B × d entries (if B_r = B_c = B).

For an A100 with 192 KB of SRAM per SM, d = 128, FP16 (2 bytes):

For a sequence of n = 8192 tokens:

Each tile computes a 187 × 187 sub-matrix = 34,969 entries. That tiny block lives entirely in SRAM. Standard attention would need all 67 million entries in HBM simultaneously. FlashAttention needs only ~35,000 at any moment — a 1,920× reduction in peak memory.

The Online Softmax Trick

There's a catch. Softmax requires the maximum over all scores in a row for numerical stability. The formula is:

This seems to require seeing ALL n scores before computing any output. If we're processing one tile at a time, we don't know the global max yet. How can we compute softmax?

The answer is the online softmax algorithm (Milakov & Gimelshein, 2018). It processes one block at a time, maintaining a running maximum and a running denominator. When a new tile arrives with a larger maximum, it rescales all previous results:

python
# Online softmax — the trick that makes FlashAttention possible
# Process tiles of scores one at a time

m_prev = -inf          # running maximum
l_prev = 0.0            # running sum of exp(scores - max)
O_prev = zeros(B_r, d)  # running output accumulator

for j in range(num_kv_blocks):
    # Load one K,V block from HBM into SRAM
    K_j = K[j * B_c : (j+1) * B_c]   # shape: (B_c, d)
    V_j = V[j * B_c : (j+1) * B_c]   # shape: (B_c, d)

    # Compute local attention scores for this tile
    S_ij = Q_block @ K_j.T             # shape: (B_r, B_c)

    # Update running max
    m_new = max(m_prev, rowmax(S_ij))

    # Rescale previous accumulations to new max
    correction = exp(m_prev - m_new)
    l_new = correction * l_prev + rowsum(exp(S_ij - m_new))

    # Update output: rescale old + add new contribution
    P_ij = exp(S_ij - m_new)          # local attention weights
    O_new = correction * O_prev + P_ij @ V_j

    m_prev, l_prev, O_prev = m_new, l_new, O_new

# Final output: normalize by total denominator
O = O_prev / l_prev                    # shape: (B_r, d)

The key insight: the correction = exp(m_prev - m_new) term rescales all prior work when the max changes. If the new tile has a larger maximum, previous exponentials were too large — multiply them down. This produces bit-for-bit identical results to standard softmax. No approximation whatsoever.

The Full Algorithm

Putting it all together:

HBM traffic for FlashAttention: read Q once (n × d), read K,V once per Q-block pass (n × d each, but reused), write O once (n × d). The n × n attention matrix? Never written to HBM. Never read from HBM. It exists only briefly in SRAM, one tile at a time.

The Simulation

Watch the two approaches side by side. Standard attention bounces data back and forth to HBM constantly. FlashAttention loads tiles into SRAM and keeps them there. Click Step to advance through each stage of the computation.

Performance Numbers

Here are benchmarks from the FlashAttention-2 paper (Dao, 2023) on an A100-80GB GPU, with head dimension d = 128, 12 heads:

The pattern: the longer the sequence, the bigger the win. At n = 16384, standard attention runs out of memory entirely while FlashAttention still fits comfortably. This is why every modern LLM uses FlashAttention — it's free performance with zero quality loss.

The Evolution

FlashAttention-3 on H100 can reach 740 TFLOPS for attention — 75% of the theoretical peak. Standard attention typically achieves less than 35% utilization. Same math, same answer, just dramatically better use of the hardware.

Chapter 9: The Attention Arena

You've learned seven different approaches to attention — from full multi-head attention down to linear approximations, from multi-query compression to FlashAttention's memory-hierarchy trick. But which one should you actually use? It depends on your sequence length, model dimension, and quality requirements.

Let's race them.

This arena simulates the throughput, memory usage, and quality tradeoffs of every attention variant at once. Adjust the sliders to match your deployment scenario: short prompts or long context? Small model or frontier-scale? Then hit Race! and watch the variants compete.

What the Arena Reveals

At short sequences (512–1024): All variants perform similarly. The n² cost is small, so clever tricks don't save much. Flash-MHA gives the best quality at competitive speed. This is why BERT-era models used full MHA — sequences were short enough that it didn't matter.

At medium sequences (2048–4096): The pack starts to separate. GQA and MQA pull ahead on throughput because they have smaller KV caches (less memory traffic). FlashAttention gives MHA a lifeline. Sliding window is great if your task is local.

At long sequences (8192+): The gap is dramatic. Standard MHA without FlashAttention may OOM entirely. Linear attention dominates raw throughput with its O(n) scaling. Sliding window excels if the task doesn't need global context. GQA + FlashAttention is the modern sweet spot — near-MHA quality with 4–8× less KV memory.

On quality: MHA and FlashAttention are lossless (they compute the same thing). GQA-8 is near-lossless (<0.1% degradation in most benchmarks). MQA has slight degradation (0.5–1%). Linear attention has noticeable degradation (1–3%) on tasks requiring precise long-range recall. Sliding window has zero loss for local patterns but misses distant information.

Throughput Scaling Behavior

(Throughput relative to MHA at n = 512. "OOM" means out of memory on a single A100-80GB.)

Memory Comparison

KV cache memory per token per layer (d_model = 4096, d_head = 128, H = 32, FP16):

Chapter 10: Cheat Sheet & Connections

You've mastered the complete attention variant toolkit — from standard multi-head attention through all its descendants, down to the hardware tricks that make it fast. This chapter is your reference card. No new concepts. Just the decision guide, the comparison table, and where to go next.

The Decision Flowchart

Follow the path that matches your situation:

Comprehensive Comparison

Symbol Reference

Summary of Everything

Connections

Attention variants don't exist in isolation. Here's where to go next:

• Attention & Transformers — The full Transformer architecture that wraps attention. See where MHA fits in the encoder-decoder stack.
• CS229S: FlashAttention Deep Dive — The Stanford lecture on IO-aware algorithms and the hardware-software interface.
• SSMs & Mamba — The alternative to attention entirely. O(n) by design, no approximation. The hybrid Mamba-Transformer models use both.
• GPT Architecture — See how these attention variants slot into a real model. GPT-2 used MHA; GPT-4 likely uses GQA + FlashAttention.
• Normalization — The layer that wraps attention. Pre-LN + RMSNorm placement interacts with attention stability.
• Optimizers — How the gradients flowing back through attention get turned into weight updates.

Key Papers