3.3 A deep dive into Self-Attention

Self-Attention is the heart of the Transformer, and the module where modern LLMs concentrate the most compute and the richest set of optimizations. Whether you want to understand the IO ideas behind FlashAttention, or why Attention variants like GQA and MLA cut inference cost, you need a solid grip on Self-Attention. This article starts from the history of Attention, walks through every math step of Scaled Dot-Product Attention, writes a PyTorch implementation by hand, looks at compute bottlenecks, and then goes into FlashAttention and the main Attention variants, so you can build a full chain from intuition to formulas to engineering.

📑 Table of contents

• 1. The history of Attention
• 2. Understanding QKV through information retrieval
• 3. Scaled Dot-Product Attention in detail
• 4. Manual PyTorch implementation
• 5. Step-by-step dimension walkthrough
• 6. Why divide by $\sqrt{d_k}$: a variance-stability derivation
• 7. Numerical stability of Softmax
• 8. Multi-Head Attention
• 9. MHA vs MQA vs GQA vs MLA
• 10. Causal Mask: role and implementation
• 11. Bottleneck analysis for Self-Attention
• 12. The core idea of FlashAttention
• Self-check list
• References

---

1. The history of Attention

Before the Transformer landed in 2017, Attention had already been brewing for a few years in sequence modeling. Following that history helps you see what problem Self-Attention actually solves, and why it looks the way it does today.

1.1 Bahdanau Attention (2014): the starting point

Classic Encoder-Decoder setups (for example RNNs for machine translation) had a hard failure mode: the Encoder squeezed the whole input sequence into one fixed-length vector, then handed that to the Decoder. When the input got long, that fixed vector could not hold everything, and translation quality collapsed.

Bahdanau et al. proposed a real fix in 2014: stop forcing the compression. When the Decoder generates each word, let it look back at all Encoder hidden states and take what it needs. Concretely, the Decoder’s current hidden state acts as a “query”, is matched against every Encoder hidden state (a small feed-forward net scores the match), and those scores weight the Encoder states into a “context vector” that becomes extra input for the current decode step.

That is the early form of Attention: teach the model to “attend” to different parts of the input. Bahdanau Attention still used an extra feed-forward net as the match function (additive attention), so it was not cheap to compute.

1.2 Luong Attention (2015): cheaper matching

In 2015 Luong simplified Bahdanau Attention in two ways: a cheaper match function (dot product instead of a feed-forward net), and several alignment styles (global vs local). The dot-product Attention idea (score similarity by a plain inner product) later became the base of the Transformer.

Dot products are much cheaper than a feed-forward match: they are matrix multiplication, which GPUs parallelize well. That simple change opened the door for Self-Attention at scale.

1.3 Self-Attention (2017): Attention Is All You Need

The 2017 Transformer paper made a bold call: drop RNNs entirely and model sequences with Attention alone.

Earlier Attention was cross-sequence: the Decoder attended to the Encoder. Self-Attention is within one sequence: every position attends to every position in the same sequence (including itself). That gives two big wins:

• Parallelism: RNNs must run step by step (step $t$ depends on $t-1$). In Self-Attention every position can be computed at once, which fits GPUs.
• Long-range dependency: In an RNN, distant tokens only interact through many steps, and the signal decays. In Self-Attention any two positions interact in one step, so there is no distance decay in theory.

The cost? Self-Attention is $O(N^2)$: every pair of positions needs a match score, while RNNs are $O(N)$. That quadratic cost is where a huge amount of later optimization work starts.

---

2. Understanding QKV through information retrieval

The three core Attention ideas are Query (Q), Key (K), and Value (V). The names come straight from information retrieval. A everyday library analogy helps build intuition.

2.1 A library search analogy

You walk into a library with a fuzzy need: “I want material on parallel computing.” That need is your Query.

Every book spine has a label: “Operating Systems”, “Computer Networks”, “Intro to Parallel Computing”, “Data Structures”, and so on. Those labels are each book’s Key.

You compare your Query with every Key. “Intro to Parallel Computing” matches highly; “Operating Systems” is somewhat related; “Data Structures” is barely related.

After you have match scores, you do not take the spine labels (Keys). You take each book’s actual content (Value), weighted by the scores. Parallel-computing content dominates, OS contributes a bit, data structures almost nothing, and you end up with a blended summary centered on parallel computing with a little OS on the side.

2.2 Formal definition of the triple

Back to Self-Attention. Given input sequence $X$ with shape $(N, d_{model})$, three linear maps send $X$ into different spaces:

• $Q = X \cdot W_Q$: each token builds its own “query vector”: “what information do I need”
• $K = X \cdot W_K$: each token builds its own “index vector”: “what clues I offer about the information I have”
• $V = X \cdot W_V$: each token builds its own “content vector”: “the information I actually carry”

Q, K, and V come from three different linear maps of the same $X$, rather than from $X$ itself, because the roles need to be decoupled. What a token needs and what it can provide are often different. In a sentence, a verb may need to attend to its subject and object (Query direction), while as something others attend to it mainly supplies action semantics (Key/Value direction). Three independent projection matrices give the model room to learn those different maps.

---

3. Scaled Dot-Product Attention in detail

With Q, K, and V in hand, Attention breaks into five clear steps.

3.1 Step 1: linear projection

Input $X$ has shape $(N, d_{model})$. The three weight matrices $W_Q$, $W_K$, $W_V$ each have shape $(d_{model}, d_{model})$ (single-head case):

$$\begin{aligned} Q &= X \cdot W_Q \quad & (N, d_{model}) \times (d_{model}, d_{model}) = (N, d_{model}) \\ \end{aligned}$$ $$\begin{aligned} K &= X \cdot W_K \quad & (N, d_{model}) \times (d_{model}, d_{model}) = (N, d_{model}) \\ \end{aligned}$$ $$\begin{aligned} V &= X \cdot W_V \quad & (N, d_{model}) \times (d_{model}, d_{model}) = (N, d_{model}) \end{aligned}$$

These three matmuls are three independent GEMM (General Matrix Multiply) ops. In practice, to raise GPU utilization, people often fuse $W_Q$, $W_K$, $W_V$ into one big $W_{QKV}$ of shape $(d_{model}, 3 \times d_{model})$, run one GEMM, then split. That uses the GPU better.

3.2 Step 2: raw attention scores

Multiply $Q$ by $K^\top$ to score every token pair:

$$S = Q K^\top \quad (N, d_{model}) \times (d_{model}, N) = (N, N)$$

Each entry $S[i][j]$ is the dot product of token $i$‘s query and token $j$‘s key: how much token $i$ attends to token $j$ before scaling.

This step builds an $N \times N$ matrix. That is the root of Self-Attention’s quadratic cost.

3.3 Step 3: scaling

Divide the raw scores by $\sqrt{d_k}$ ($d_k$ is the Key dimension):

$$S_{\text{scaled}} = \frac{S}{\sqrt{d_k}}$$

Why scale? Briefly: when $d_k$ is large, dot products get huge and Softmax gradients go to zero. Full math is in section 6.

3.4 Step 4: Softmax normalization

Softmax each row of the scaled score matrix into a distribution:

$$A = \text{softmax}(S_{\text{scaled}}) \quad \in \mathbb{R}^{N \times N}, \text{ each row sums to 1}$$

Softmax does two jobs: map arbitrary reals into $(0, 1)$ so they can act as weights, and make each row sum to 1 so you get a valid distribution.

3.5 Step 5: weighted sum

Weight the Value matrix with the attention weights:

$$\text{Output} = A \cdot V \quad (N, N) \times (N, d_{model}) = (N, d_{model})$$

Each token position ends up with a $d_{model}$-dim vector that mixes the tokens it “attended” to, with weights from $A$.

3.6 Full formula

Put together, Self-Attention is one line:

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

Then apply an output projection $W_O$ of shape $(d_{model}, d_{model})$ to map back into the model’s hidden space. That projection matters even more in Multi-Head Attention: it remixes the concatenated heads.

---

4. Manual PyTorch implementation

Theory is cheaper than writing the code. Below are Single-Head and Multi-Head Self-Attention by hand.

4.1 Single-Head Self-Attention

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class SingleHeadSelfAttention(nn.Module):
    """Manual single-head Self-Attention, without PyTorch's built-in MHA"""
    def __init__(self, d_model: int):
        super().__init__()
        self.d_model = d_model
        # Three linear projections: in dim d_model, out dim d_model
        self.W_Q = nn.Linear(d_model, d_model, bias=False)
        self.W_K = nn.Linear(d_model, d_model, bias=False)
        self.W_V = nn.Linear(d_model, d_model, bias=False)
        # Output projection
        self.W_O = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x: torch.Tensor, mask: torch.Tensor = None):
        # x: (batch, seq_len, d_model)
        # Step 1: linear projections for Q, K, V
        Q = self.W_Q(x)  # (batch, seq_len, d_model)
        K = self.W_K(x)  # (batch, seq_len, d_model)
        V = self.W_V(x)  # (batch, seq_len, d_model)

        # Step 2: attention scores = Q @ K^T
        scores = torch.matmul(Q, K.transpose(-2, -1))  # (batch, seq_len, seq_len)

        # Step 3: scale so large dots do not kill softmax gradients
        scores = scores / math.sqrt(self.d_model)

        # Step 4: if a mask is given, set masked positions to -inf
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))

        # Step 5: Softmax to get attention weights
        attn_weights = F.softmax(scores, dim=-1)  # (batch, seq_len, seq_len)

        # Step 6: weighted sum over V
        context = torch.matmul(attn_weights, V)  # (batch, seq_len, d_model)

        # Step 7: output projection
        output = self.W_O(context)  # (batch, seq_len, d_model)
        return output

4.2 Multi-Head Self-Attention

class MultiHeadSelfAttention(nn.Module):
    """Manual multi-head Self-Attention: split d_model into num_heads subspaces"""
    def __init__(self, d_model: int, num_heads: int):
        super().__init__()
        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
        self.d_model = d_model
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads  # dim per head

        # QKV projections: in d_model, out d_model (all heads packed inside)
        self.W_Q = nn.Linear(d_model, d_model, bias=False)
        self.W_K = nn.Linear(d_model, d_model, bias=False)
        self.W_V = nn.Linear(d_model, d_model, bias=False)
        self.W_O = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x: torch.Tensor, mask: torch.Tensor = None):
        batch_size, seq_len, _ = x.shape

        # Step 1: linear projections
        Q = self.W_Q(x)  # (batch, seq_len, d_model)
        K = self.W_K(x)
        V = self.W_V(x)

        # Step 2: reshape to multi-head, transpose so head sits before seq_len
        # (batch, seq_len, d_model) -> (batch, seq_len, num_heads, head_dim)
        # -> (batch, num_heads, seq_len, head_dim)
        Q = Q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        K = K.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        V = V.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)

        # Step 3: attention scores per head
        # Q @ K^T: (batch, num_heads, seq_len, head_dim)
        #        @ (batch, num_heads, head_dim, seq_len)
        #        = (batch, num_heads, seq_len, seq_len)
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.head_dim)

        # Step 4: optional causal mask
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))

        # Step 5: Softmax + weighted sum
        attn_weights = F.softmax(scores, dim=-1)
        # (batch, num_heads, seq_len, seq_len) @ (batch, num_heads, seq_len, head_dim)
        # = (batch, num_heads, seq_len, head_dim)
        context = torch.matmul(attn_weights, V)

        # Step 6: concat heads: transpose back then merge last two dims
        context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)

        # Step 7: output projection remixes the concatenated heads
        output = self.W_O(context)  # (batch, seq_len, d_model)
        return output

The code walks every step from input to output. Note that .transpose(1, 2) turns (batch, seq_len, num_heads, head_dim) into (batch, num_heads, seq_len, head_dim), so later matmuls run independently inside each head and finish all heads in one batch-parallel pass.

---

5. Step-by-step dimension walkthrough

Abstract shapes get confusing. Below we track every tensor shape in Multi-Head Self-Attention with concrete numbers.

5.1 Setup

seq_len   = 6        # sequence length (6 tokens)
d_model   = 512      # model hidden dim
num_heads = 8        # number of attention heads
head_dim  = 512 / 8 = 64   # dim per head
batch     = 1        # simplify to batch=1

5.2 Step tracking

Input

X: (1, 6, 512)    # 1 sample, 6 tokens, each 512-dim

Linear projection

W_Q: (512, 512)    # nn.Linear weight matrix
Q = X @ W_Q^T:  (1, 6, 512) @ (512, 512) = (1, 6, 512)
K = X @ W_K^T:  (1, 6, 512) @ (512, 512) = (1, 6, 512)
V = X @ W_V^T:  (1, 6, 512) @ (512, 512) = (1, 6, 512)

Split heads + transpose

Q.view(1, 6, 8, 64):         (1, 6, 8, 64)
Q.transpose(1, 2):           (1, 8, 6, 64)    # 8 heads, each sees 6 tokens of 64-dim

K reshape likewise:          (1, 8, 6, 64)
V reshape likewise:          (1, 8, 6, 64)

Attention scores

K.transpose(-2, -1):         (1, 8, 64, 6)    # transpose last two dims
Q @ K^T:                     (1, 8, 6, 64) @ (1, 8, 64, 6) = (1, 8, 6, 6)
                             # each head produces a 6x6 score matrix

Scaling

scores / sqrt(64):           (1, 8, 6, 6)    # sqrt(64) = 8, divide each score by 8

Softmax

softmax(scores, dim=-1):     (1, 8, 6, 6)    # normalize each row (last dim) to a distribution
                             # each 6-dim row sums to 1

Weighted sum

attn_weights @ V:            (1, 8, 6, 6) @ (1, 8, 6, 64) = (1, 8, 6, 64)
                             # each head outputs 6 tokens of 64-dim

Concat heads

transpose(1, 2):             (1, 6, 8, 64)    # move head dim back
view(1, 6, 512):             (1, 6, 512)      # 8 * 64 = 512, concat to d_model

Output projection

W_O: (512, 512)
context @ W_O^T:             (1, 6, 512) @ (512, 512) = (1, 6, 512)

Final output

Output: (1, 6, 512)    # same shape as input X

End to end, input and output stay $(batch, seq\_len, d_{model})$, so Transformer Blocks can stack like Lego.

---

6. Why divide by $\sqrt{d_k}$: a variance-stability derivation

Dividing by $\sqrt{d_k}$ is easy to dismiss as a “tuning trick”, but there is a hard math reason.

6.1 The problem

Assume every component of $Q$ and $K$ is an independent random variable with mean 0 and variance 1. Compute the variance of the dot product $q \cdot k$.

Let $q = (q_1, q_2, \ldots, q_{d_k})$, $k = (k_1, k_2, \ldots, k_{d_k})$, with each $q_i$ and $k_j$ i.i.d., $E[q_i] = 0$, $\text{Var}(q_i) = 1$.

The dot product is:

$$q \cdot k = \sum_{i=1}^{d_k} q_i \cdot k_i$$

6.2 Derivation

Expectation and variance of one component product:

$$E[q_i \cdot k_i] = E[q_i] \cdot E[k_i] = 0 \times 0 = 0$$ $$\text{Var}(q_i \cdot k_i) = E[q_i^2 \cdot k_i^2] - (E[q_i \cdot k_i])^2$$

Because $q_i$ and $k_i$ are independent:

$$E[q_i^2 \cdot k_i^2] = E[q_i^2] \cdot E[k_i^2] = \text{Var}(q_i) \cdot \text{Var}(k_i) = 1 \times 1 = 1$$

So:

$$\text{Var}(q_i \cdot k_i) = 1 - 0 = 1$$

The dot product is a sum of $d_k$ independent terms. By additivity of variance:

$$\text{Var}(q \cdot k) = \sum_{i=1}^{d_k} \text{Var}(q_i \cdot k_i) = d_k$$

6.3 Conclusion

Variance of $q \cdot k$ is $d_k$. So when $d_k = 64$, the std of the dot product is about 8; when $d_k = 128$, about 11.3. Larger dimensions mean larger absolute dots and a more spread-out distribution.

Softmax $\text{softmax}(z_i) = \exp(z_i) / \sum \exp(z_j)$ is very sensitive to scale. When gaps get huge (say 50 vs -10), Softmax collapses toward a one-hot, gradients go near zero, and training stalls.

After dividing by $\sqrt{d_k}$:

$$\text{Var}\left(\frac{q \cdot k}{\sqrt{d_k}}\right) = \frac{\text{Var}(q \cdot k)}{d_k} = \frac{d_k}{d_k} = 1$$

Variance returns to 1. Whatever $d_k$ is, Softmax inputs stay in a sane range and gradients stay healthy. That is where the “Scaled” in Scaled Dot-Product Attention comes from.

---

7. Numerical stability of Softmax

Softmax looks like a simple formula, but GPU implementations hide an easy overflow trap.

7.1 The overflow problem

Standard Softmax:

$$\text{softmax}(z_i) = \frac{e^{z_i}}{\sum_{j=1}^{N} e^{z_j}}$$

$\exp$ grows fast. When $z_i$ is large (say $z_i = 1000$), $\exp(1000)$ overflows float32 and even float16 into inf. If numerator and denominator both overflow, inf / inf becomes NaN and the whole compute is toast.

7.2 Subtract the max

The fix uses Softmax shift invariance: subtract the max from every input, and the result is unchanged.

$$\text{softmax}(z_i) = \frac{e^{z_i - \max(z)}}{\sum_{j=1}^{N} e^{z_j - \max(z)}}$$

The proof is short: multiply numerator and denominator by $\exp(-\max(z))$, which is the same as dividing both by $\exp(\max(z))$. After the subtract, the largest exponent input is 0, so $\exp$ is at most 1 and cannot overflow.

That costs an efficiency hit: a naive implementation needs three passes over the data:

1. Pass 1: scan all elements for the max $m$
2. Pass 2: scan all elements, compute $\exp(z_i - m)$ and their sum
3. Pass 3: scan all elements, divide each $\exp(z_i - m)$ by the sum

Each pass reads from HBM (high-bandwidth memory). Three passes means three HBM reads. When $N$ is large (for example seq length 128K), that IO cost hurts.

7.3 Online Softmax

Milakov and Gimelshein (2018) proposed Online Softmax, which folds three passes into one pass. The idea is to maintain the running max and the normalization denominator while you scan.

While visiting element $i$, keep two variables:

• $m_i$: max of the first $i$ elements
• $d_i$: sum of exponentials of the first $i$ elements, relative to $m_i$

Recurrence:

$$\begin{aligned} m_i &= \max(m_{i-1}, z_i) \end{aligned}$$ $$\begin{aligned} d_i &= d_{i-1} \cdot e^{m_{i-1} - m_i} + e^{z_i - m_i} \end{aligned}$$

The key is the second line: when the max moves from $m_{i-1}$ to $m_i$, rescale the previous sum $d_{i-1}$ by $\exp(m_{i-1} - m_i)$ to change base. One pass gets you max and exp-sum; one more pass finishes normalization. Two passes total, one fewer than the naive version.

That idea is critical for FlashAttention: FlashAttention tiles the Attention matrix and updates Softmax intermediates per tile. Without online Softmax updates you would have to write the full $N \times N$ matrix to HBM and Softmax globally, which kills the point of tiling. Online Softmax makes tiling and exact Softmax compatible.

---

8. Multi-Head Attention

8.1 Why multiple heads

Single-head Attention has one QKV projection set, so the model learns one “attention pattern”. Token relations in language are multi-dimensional: the same word can hold syntactic ties (subject-verb agreement), semantic ties (synonymy), positional patterns (local neighbors), and more at once.

Analogy: in a project review, one reviewer can only push the angles they are good at. A review board (architecture, code quality, test coverage, docs) scores independently and merges into one report, and coverage is much broader.

Multi-Head Attention is that review board. Each head has its own $W_Q$, $W_K$, $W_V$ and catches different relation types in different subspaces.

8.2 Math

Assume $d_{model} = 512$, $num\_heads = 8$, so $head\_dim = 64$.

Full Multi-Head Attention:

$$\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h) \cdot W_O$$ $$\text{head}_i = \text{Attention}(Q \cdot W_Q^i, K \cdot W_K^i, V \cdot W_V^i)$$

Each head’s $W_Q^i$, $W_K^i$, $W_V^i$ has shape $(d_{model}, head\_dim) = (512, 64)$. In practice you do not keep $h$ small matrices. You project once with a big $W_Q$ of shape $(512, 512)$, then reshape to $(seq\_len, 8, 64)$. Same math: the big matrix is eight small ones stacked along the output dim.

8.3 Parameter count

With $d_{model} = 512$:

General rule: MHA params are $4 \times d_{model}^2$ (no bias). Head count does not change total params; it only changes the split. Add a head and each head gets thinner, so the product (total projected dim) stays $d_{model}$.

8.4 Engineering angle

Multi-head structure parallelizes naturally. Eight independent heads can:

• Parallelize inside a GPU: use the batch dim and process all heads in one CUDA kernel launch
• Tensor Parallelism across GPUs: put different heads on different GPUs. For example 8 heads on 4 GPUs means 2 heads each. One AllReduce merges the output projection results

---

9. MHA vs MQA vs GQA vs MLA

As LLMs moved into an inference-efficiency era, Attention head layouts kept evolving. The driver is simple: KV Cache is too big.

9.1 MHA (Multi-Head Attention)

In standard MHA every head has its own Q, K, V. With $h$ heads you cache $h$ K sets and $h$ V sets.

• $h$ Q heads, $h$ K heads, $h$ V heads
• Each Q/K/V set is independent; nothing shared

9.2 MQA (Multi-Query Attention)

Shazeer (2019) proposed: all Q heads share one K and one V. Still $h$ Q heads, but only 1 K head and 1 V head.

• $h$ Q heads, 1 K head, 1 V head
• All Q heads share the same K and V

KV Cache shrinks to $1/h$, so inference gets much faster. Expressiveness drops: every head is forced to pull from the same KV, so diversity is limited.

9.3 GQA (Grouped-Query Attention)

GQA is the MHA/MQA middle ground (Ainslie et al., 2023). Split $h$ Q heads into $g$ groups ($h/g$ Q heads each); each group shares one K and one V.

• $h$ Q heads, $g$ K heads, $g$ V heads
• Every $h/g$ Q heads share one K/V set
• $g = h$ falls back to MHA; $g = 1$ falls back to MQA

GQA balances quality and inference cost better. LLaMA-2-70B, Mistral-7B, and many other mainstream models use GQA.

9.4 Comparison table

With $d_{model} = 4096$, num_heads = 32, head_dim = 128, sequence length $N = 4096$, FP16:

Note: Cache size = 2 × B × L × num_heads × head_dim × d_type_bytes (2 is K and V, B is batch_size, L is generated sequence length so far) From MHA to GQA($g=8$), KV Cache shrinks to 1/4 and params drop, while quality barely moves. That is why GQA became the default choice.

9.5 MLA (Multi-Latent Attention)

DeepSeek-V2 proposed MLA (Multi-Latent Attention), another way to compress KV Cache.

The idea: do not cache full K and V. Compress them into a low-dimensional latent space and cache only that. At inference, decompress back to full K and V when needed.

• Standard MHA: $X \to W_K \to K$ (cache $K$, dim = num_kv_heads $\times$ head_dim)
• MLA: $X \to W_{DKV} \to c$ (cache $c$, much smaller than $K$), then at inference $c \to W_{UK} \to K$ on demand

MLA’s low-rank compression shrinks KV Cache a lot (DeepSeek-V2 reports about 93.3% compression) while carefully designed up-projections keep quality. Unlike GQA (fewer KV heads), MLA shrinks the dim of each representation. Same goal: make inference KV Cache as small as possible.

---

10. Causal Mask: role and implementation

10.1 Why mask

In autoregressive LMs (GPT family, LLaMA family), the training target is “predict the next token from the past”. During Attention, each position may only see itself and earlier tokens, never the future. Otherwise training peeks at the answer and never learns real prediction.

10.2 Implementation

A Causal Mask is an upper-triangular cover over the score matrix that sets future scores to -inf:

# Build a causal mask
def create_causal_mask(seq_len: int) -> torch.Tensor:
    """
    Build a lower-triangular matrix (including the diagonal); upper triangle is False.
    True positions are kept; False positions are set to -inf in Attention scores.
    """
    mask = torch.tril(torch.ones(seq_len, seq_len, dtype=torch.bool))
    return mask

# Example: seq_len = 4
# mask:
# [[True,  False, False, False],
#  [True,  True,  False, False],
#  [True,  True,  True,  False],
#  [True,  True,  True,  True ]]
#
# Masked Attention scores (upper triangle = -inf):
# [[s00,  -inf, -inf, -inf],
#  [s10,  s11,  -inf, -inf],
#  [s20,  s21,  s22,  -inf],
#  [s30,  s31,  s32,  s33 ]]
#
# After Softmax, -inf positions become weight 0; future info is fully blocked

Another engineering angle: only the lower triangle is live, so in theory you can skip the upper triangle. FlashAttention-2 does exactly that: skip upper-triangle tiles and cut nearly half the work.

---

11. Bottleneck analysis for Self-Attention

Self-Attention performance has two different regimes.

11.1 FLOPs breakdown

Four main matmuls sit inside Self-Attention:

Total: $8Nd^2 + 4N^2d$.

When $N$ is small (for example $N < d$), $8Nd^2$ dominates: bottleneck is QKV projection GEMM. When $N$ is large (for example $N > d$), $4N^2d$ dominates: bottleneck is the Attention matrix.

11.2 Compute Bound vs Memory Bound

GPU performance is limited by two numbers:

• Compute: how many FLOPS per second
• Bandwidth (Memory Bandwidth): how many GB/s you can read/write from HBM

Whether an op is Compute Bound or Memory Bound depends on Arithmetic Intensity = FLOPs / Bytes.

• If intensity > the GPU’s compute/bandwidth ratio, the op is Compute Bound: the ALU stays busy
• If intensity < that ratio, the op is Memory Bound: the ALU waits on data movement

A100 example: ~312 TFLOPS FP16, ~2 TB/s HBM, ratio about 156 FLOP/Byte.

Prefill (full prompt):

QKV projections are large matmuls with a large batch dim (seq_len), high intensity, usually Compute Bound. $Q \cdot K^T$ is also a large matmul, also Compute Bound.

Decode (token by token):

Each step has one token. QKV projections become GEMV, intensity crashes, Memory Bound. Attention becomes a vector against the whole KV Cache, also Memory Bound. Most time goes into hauling weights and KV Cache from HBM.

11.3 The $O(N^2)$ memory bottleneck

Naive Attention materializes the full $N \times N$ score matrix. With seq_len = 128$\text{K}$, num_heads = 32:

$$32 \times 128\text{K} \times 128\text{K} \times 2\text{B} = 32 \times 16{,}384\text{M} \times 2 = 1{,}048{,}576\text{ MB} = 1\text{ TB}$$

That is far beyond any single GPU. Even at “only” 8K, the Attention matrix needs $32 \times 8\text{K} \times 8\text{K} \times 2\text{B} = 4\text{ GB}$, which is already serious.

That $O(N^2)$ memory wall is exactly what FlashAttention targets.

---

12. The core idea of FlashAttention

FlashAttention (Dao et al., 2022) is one of the most influential Attention optimizations in recent years. It does not change the math result (exact, not approximate). It reorders the work so HBM traffic drops a lot.

12.1 GPU memory hierarchy refresher

To understand FlashAttention you need the GPU memory stack:

• HBM (High Bandwidth Memory): main GPU memory, large capacity (for example A100 80 GB) but relatively slow (~2 TB/s)
• SRAM (on-chip cache): shared memory and registers on each SM (Streaming Multiprocessor), tiny capacity (A100 ~192 KB shared memory per SM, ~20 MB across all SMs) but very fast (~19 TB/s)

Naive Attention: finish $Q \cdot K^T$ in HBM and write the full $N \times N$ matrix back; read it for Softmax and write back; read Softmax output and $V$ for the matmul. Every step reads/writes that huge $N \times N$ matrix in HBM. That wastes bandwidth.

12.2 Tiling

FlashAttention’s first key technique is Tiling: do not build the full $N \times N$ Attention matrix at once. Split Q, K, V into tiles. Load one Q tile and one K/V tile into SRAM, finish that tile’s Attention there, write results back to HBM.

Standard:
  Load all Q, K into HBM → full N x N matrix → write HBM → read Softmax → write HBM → ...

FlashAttention:
  Split Q into T_r tiles, K/V into T_c tiles
  For each Q_i tile:
    For each K_j, V_j tile:
      Load Q_i, K_j, V_j from HBM into SRAM (tiles are small enough to fit)
      In SRAM: Q_i @ K_j^T → local Softmax → multiply by V_j
      Accumulate into the output for this Q_i
    Write the final Q_i output back to HBM

Each step only holds one tile of the Attention matrix in SRAM. The full $N \times N$ matrix never lands in HBM.

12.3 Why Online Softmax matters

Tiling creates a hard problem: Softmax needs a full-row max and sum, but each tile only sees part of a row. How do you get global Softmax from local pieces?

That is where Online Softmax comes in. Recall section 7.3: when a new block arrives you can update the global max and denominator and correct earlier block results.

Concretely, after Attention of $Q_i$ against $K_1$ you have a local output $O_1$ and Softmax stats (local max $m_1$, local denominator $l_1$). Against $K_2$ you get $m_2$ and $l_2$. Online Softmax correction:

$$\begin{aligned} m_{\text{new}} &= \max(m_1, m_2) \end{aligned}$$ $$\begin{aligned} l_{\text{new}} &= l_1 \cdot e^{m_1 - m_{\text{new}}} + l_2 \cdot e^{m_2 - m_{\text{new}}} \end{aligned}$$ $$\begin{aligned} O_{\text{new}} &= \frac{O_1 \cdot l_1 \cdot e^{m_1 - m_{\text{new}}} + O_2^{\text{local}} \cdot l_2 \cdot e^{m_2 - m_{\text{new}}}}{l_{\text{new}}} \end{aligned}$$

Keep accumulating until every $K$ tile is done. The final result matches standard global Softmax exactly.

12.4 Why HBM traffic drops

Standard Attention HBM traffic:

• Write $S$ ($N \times N$): $O(N^2)$ writes
• Read $S$ for Softmax: $O(N^2)$ reads
• Write $P = \text{Softmax}(S)$: $O(N^2)$ writes
• Read $P$ for $P \cdot V$: $O(N^2)$ reads
• Total HBM traffic: $O(N^2)$ (plus $O(Nd)$ reads of Q, K, V themselves)

FlashAttention HBM traffic:

• Read Q, K, V: $O(Nd)$
• Write final output O: $O(Nd)$
• Intermediate Attention tiles live only in SRAM and are discarded; the full $N \times N$ matrix never hits HBM
• But Q, K, V, O tiles are reread: for V1, the outer loop walks $T_c = O(Nd/M)$ K/V tiles and each time sweeps full Q and O ($O(Nd)$)
• Total HBM traffic: $\Theta(N^2 d^2 / M)$ ($M$ is SRAM size in elements)

Because $d^2$ is much smaller than $M$ (for example $d = 128$ gives $d^2 = 16384$, while 192KB SRAM holds about 98304 FP16 elements), HBM traffic is several times below the standard $\Theta(N^2)$. It still grows with $N^2$; the constant is just much smaller. What truly drops to $O(N)$ is extra memory (no more stored $N \times N$ matrix). FLOPs stay $O(N^2d)$; only the IO pattern changes. FlashAttention speeds up by cutting slow HBM traffic and keeping work in fast SRAM.

12.5 FlashAttention-2 and FlashAttention-3

FlashAttention-2 (Dao, 2023) improves parallelism and warp scheduling on top of FlashAttention:

• Reorders inner/outer loops (outer over Q tiles, inner over K/V), cutting shared-memory traffic
• Better warp-level work assignment, higher GPU occupancy
• Skips all-zero Causal Mask tiles, cutting more compute

FlashAttention-3 (Dao et al., 2024) targets Hopper (H100), using TMA (Tensor Memory Accelerator) and wgmma (Warp Group MMA) to raise hardware utilization further.

FlashAttention is now standard in major inference and training stacks. PyTorch 2.0+ ships torch.nn.functional.scaled_dot_product_attention, which defaults to a FlashAttention kernel underneath.

---

🎯 Self-check list

After this article, test how deep your understanding goes:

• Trace Attention from Bahdanau to Self-Attention and say what each step fixed
• Write Scaled Dot-Product Attention from memory and explain what Q, K, V each do
• Derive from scratch why we divide by $\sqrt{d_k}$: assume $q_i, k_i$ i.i.d. standard normal and show dot-product variance is $d_k$
• Hand-write Single-Head and Multi-Head Self-Attention in PyTorch and track tensor shapes at each step
• Explain why Softmax subtracts the max, and the Online Softmax recurrence idea
• Given concrete seq_len, $d_{model}$, num_heads, track every tensor shape end to end
• Sketch MHA / MQA / GQA structure differences and estimate each KV Cache size
• Explain Causal Mask (upper triangle to -inf, Softmax zeros those weights) and how it is implemented
• Tell Compute Bound from Memory Bound, and which bottleneck Prefill vs Decode hits
• Explain FlashAttention: how Tiling + Online Softmax drops extra memory from $O(N^2)$ to $O(N)$ and HBM traffic from $\Theta(N^2)$ to $\Theta(N^2 d^2 / M)$

📚 References

Papers

• Neural Machine Translation by Jointly Learning to Align and Translate (Bahdanau et al., 2014): https://arxiv.org/abs/1409.0473 — Bahdanau Attention, the starting point of Attention
• Effective Approaches to Attention-based Neural Machine Translation (Luong et al., 2015): https://arxiv.org/abs/1508.04025 — Luong Attention, introduced dot-product Attention
• Attention Is All You Need (Vaswani et al., 2017): https://arxiv.org/abs/1706.03762 — original Transformer paper
• Fast Transformer Decoding: One Write-Head is All You Need (Shazeer, 2019): https://arxiv.org/abs/1911.02150 — Multi-Query Attention
• GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints (Ainslie et al., 2023): https://arxiv.org/abs/2305.13245 — Grouped-Query Attention
• DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model (DeepSeek-AI, 2024): https://arxiv.org/abs/2405.04434 — Multi-Latent Attention
• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (Dao et al., 2022): https://arxiv.org/abs/2205.14135 — FlashAttention v1
• FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning (Dao, 2023): https://arxiv.org/abs/2307.08691 — FlashAttention v2
• Online Normalizer Calculation for Softmax (Milakov & Gimelshein, 2018): https://arxiv.org/abs/1805.02867 — Online Softmax algorithm

Tutorials and blogs

• The Illustrated Transformer (Jay Alammar): https://jalammar.github.io/illustrated-transformer/ — illustrated Transformer intro
• The Annotated Transformer (Harvard NLP): https://nlp.seas.harvard.edu/annotated-transformer/ — paper lined up with a PyTorch implementation
• Andrej Karpathy: Let’s build GPT from scratch: https://www.youtube.com/watch?v=kCc8FmEb1nY — GPT from scratch
• ELI5: FlashAttention (Aleksa Gordic): https://gordicaleksa.medium.com/eli5-flash-attention-5c44017022ad — approachable FlashAttention walkthrough