Scaled Dot-Product Attention

The previous chapter scored stored embedding vectors with cosine and dot product. A retriever compares one query vector to many keys.

A transformer layer turns that same routing idea inward: each token supplies a query and a key, while a mask decides which positions it may read.

Take the line "The cache stopped growing because it was full." You know "it" is the cache. Attention can bring information from that earlier position into the representation of "it", so later layers can use it.

Before naming any matrices, predict the operation: each token should choose source positions, turn those scores into weights, and mix the chosen content. Scaled dot-product attention learns that routing with matrix multiplication.

The mask answers a separate question from relevance: which positions are legal to read? A causal mask blocks future-key logits with $-\infty$, giving those keys zero weight. Compute the unmasked two-token slice by hand first, then put the mask back.

Every active query still needs at least one allowed key. An entirely blocked row has no softmax probability distribution. The textbook max-subtraction formula produces NaN there; some kernels explicitly return zeros instead. Don't rely on a backend's invalid-row behavior to define your masking semantics.

Three views of one token

Self-attention computes a dynamic, weighted mix of token representations. Each query gets its weights from learned relevance scores against the keys.^{[1]Reference 1Attention Is All You Need.https://arxiv.org/abs/1706.03762}

Raw dot products $X X^T$ use the same features to ask for context and to offer it. Learned projections let the model change that compatibility function and choose separately which content to transfer. The projections transform the available token states; they don't create information about an unseen token.

One distinction matters: $X X^T$ is symmetric, but its row-softmax weights needn't be. Each row has its own denominator, and a causal mask can block one direction. Symmetric scores alone don't force "it" and "cache" to attend equally to each other.

The model uses three learned linear transformations with different jobs:

Each token hidden state $x_i$ is projected three ways. $W_Q$ maps into the interrogation subspace, $W_K$ maps into the advertisement subspace, and $W_V$ maps into the payload subspace:

$Q = XW_Q, \quad K = XW_K, \quad V = XW_V$

Here $W_Q, W_K \in \mathbb{R}^{d_{\text{model}} \times d_k}$ and $W_V \in \mathbb{R}^{d_{\text{model}} \times d_v}$. $d_{\text{model}}$ is the residual stream width; $d_k$ and $d_v$ are the query/key and value widths.^{[1]Reference 1Attention Is All You Need.https://arxiv.org/abs/1706.03762} The learned score is $x_i W_Q W_K^T x_j^T$. It can be asymmetric, but $W_Q \neq W_K$ doesn't guarantee unequal scores for every pair. For example, $W_Q=I$ and $W_K=2I$ still give symmetric raw scores.

Try to predict which statement survives this counterexample: symmetric raw scores, symmetric weights, or both?

import math

X = [[1.0, 0.0], [1.0, 1.0]]
scores = [[sum(a * b for a, b in zip(x, y)) for y in X] for x in X]

def softmax(row):
    exps = [math.exp(x - max(row)) for x in row]
    return [x / sum(exps) for x in exps]

weights = [softmax(row) for row in scores]
# W_Q = I and W_K = 2I are unequal, yet the raw score matrix stays symmetric.
unequal_projection_scores = [[2 * s for s in row] for row in scores]
print("raw scores symmetric:", scores[0][1] == scores[1][0])
print("opposite weights:", round(weights[0][1], 6), round(weights[1][0], 6))
print("unequal projections, symmetric scores:",
      unequal_projection_scores[0][1] == unequal_projection_scores[1][0])

raw scores symmetric: True
opposite weights: 0.5 0.268941
unequal projections, symmetric scores: True

For the running example, token cache uses these toy vectors after projection:

Predict the split before reading the diagram: Q and K meet to produce a routing score, while V travels past that score to remain content. The self-score for cache uses $q$ and $k$, not $v$.

Q and K determine where to look. Their dot products produce routing scores, and softmax converts those scores into weights.

V determines what information to extract. Separating routing (Q, K) from content (V) lets attention route information flexibly across representation subspaces.

The core formula

Read the formula as a lookup followed by a weighted read. Compatibility of queries with keys becomes weights; those weights mix values.^{[1]Reference 1Attention Is All You Need.https://arxiv.org/abs/1706.03762} Softmax turns each query row into non-negative weights that sum to 1:

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

When some key positions aren't allowed, add a mask matrix $M$ after scaling. An allowed location has $M_{ij}=0$; a blocked logit gets $M_{ij}=-\infty$:

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

These formulas assume attention-weight dropout is disabled. Before dropout, each valid row is a distribution and its context lies in the convex hull of the visible value vectors. Training-time dropout can change the row sum and that convex-hull property. The later output projection and residual addition also needn't preserve either property.

Step by step

The order is the contract: scores choose locations, softmax makes each choice a distribution, and only then do those weights touch V. Predict the shapes before following the picture: two queries compared with two keys make a $2 \times 2$ routing matrix, while the value width survives into the output.

The shape flow uses batch size $B$, number of heads $h$, query count $N_q$, key count $N_k$, and value width $d_v$. Query and key widths must match ($d_k$); value width needn't. Queries and keys build one routing matrix; values join only after softmax:

A trace with real numbers

Walk through cache and grew without a mask first. Each uses 2-dimensional toy projections:

Step 1: raw scores

Compute the dot product of every query with every key:

Step 2: scale

With $d_k = 2$, divide by $\sqrt{2} \approx 1.414$:

Step 3: softmax

Normalize each row so it sums to 1:

Step 4: weighted values

Multiply the weights by the value vectors:

• Context for "cache" $\approx 0.53 \cdot [2.0, 1.0] + 0.47 \cdot [1.0, 2.0] = [1.53, 1.47]$
• Context for "grew" $\approx 0.42 \cdot [2.0, 1.0] + 0.58 \cdot [1.0, 2.0] = [1.42, 1.58]$

Each head produces a blend of values, not a replacement for the entire token state. A transformer block projects the heads' outputs and combines them with the residual stream. Here cache pulls slightly more from its own value (0.53) than from grew (0.47); grew mixes both nearly evenly. All displayed decimals are rounded after computation.

The same arithmetic is easy to verify without a tensor library:

import math

Q = [[1.0, 0.5], [0.5, 1.0]]
K = [[0.8, 0.2], [0.3, 0.9]]
V = [[2.0, 1.0], [1.0, 2.0]]

def softmax(row: list[float]) -> list[float]:
    if not row or any(math.isnan(x) or x == math.inf for x in row):
        raise ValueError("softmax needs finite logits or masked -inf entries")
    shift = max(row)
    if shift == -math.inf:
        raise ValueError("each query needs at least one finite logit")
    exps = [math.exp(x - shift) for x in row]
    total = sum(exps)
    return [x / total for x in exps]

scores = [[sum(q_i * k_i for q_i, k_i in zip(q, k)) for k in K] for q in Q]
scaled = [[score / math.sqrt(2) for score in row] for row in scores]
weights = [softmax(row) for row in scaled]
outputs = [
    [sum(weight * value[col] for weight, value in zip(row, V)) for col in range(2)]
    for row in weights
]

print([[round(x, 2) for x in row] for row in weights])
print([[round(x, 2) for x in row] for row in outputs])
print([round(sum(row), 3) for row in weights])

[[0.53, 0.47], [0.42, 0.58]]
[[1.53, 1.47], [1.42, 1.58]]
[1.0, 1.0]

Add the causal mask

The mask changes who may be read, not how the allowed weights are computed. This function reuses Q, K, V, and softmax from the previous cell. With a lower-triangular mask, cache can only read itself, while grew can read both tokens, matching the opening figure.

causal = [[True, False], [True, True]]

def scaled_dot_product_attention(
    Q: list[list[float]],
    K: list[list[float]],
    V: list[list[float]],
    mask: list[list[bool]] | None = None,
) -> tuple[list[list[float]], list[list[float]]]:
    def shape(matrix: list[list[float]]) -> tuple[int, int]:
        if not matrix or not matrix[0]:
            raise ValueError("matrices must be nonempty")
        width = len(matrix[0])
        if any(len(row) != width for row in matrix):
            raise ValueError("matrices must be rectangular")
        if any(not math.isfinite(x) for row in matrix for x in row):
            raise ValueError("Q, K, and V must be finite")
        return len(matrix), width

    n_q, d_k = shape(Q)
    n_k, key_width = shape(K)
    n_v, d_v = shape(V)
    if d_k != key_width or n_k != n_v:
        raise ValueError("Q/K widths and K/V row counts must match")
    if mask is not None:
        if len(mask) != n_q or any(len(row) != n_k for row in mask):
            raise ValueError("mask must have shape (Nq, Nk)")
        if any(type(x) is not bool for row in mask for x in row):
            raise ValueError("this teaching mask must contain booleans")
        if not all(any(row) for row in mask):
            raise ValueError("each query row must have at least one visible key")
    scores = [
        [sum(q_i * k_i for q_i, k_i in zip(q, k, strict=True)) / math.sqrt(d_k) for k in K]
        for q in Q
    ]
    if any(not math.isfinite(x) for row in scores for x in row):
        raise ValueError("score arithmetic overflowed; finite inputs aren't sufficient")
    if mask is not None:
        scores = [
            [logit if visible else float("-inf") for logit, visible in zip(row, visible_row)]
            for row, visible_row in zip(scores, mask)
        ]
    weights = [softmax(row) for row in scores]
    outputs = [
        [sum(weight * value[col] for weight, value in zip(row, V)) for col in range(d_v)]
        for row in weights
    ]
    if any(not math.isfinite(x) for row in outputs for x in row):
        raise ValueError("value aggregation overflowed")
    return outputs, weights

output, weights = scaled_dot_product_attention(Q, K, V, causal)
pretty_weights = [[round(value, 3) for value in row] for row in weights]
print("causal weights:", pretty_weights)
print("cache future weight:", pretty_weights[0][1])
print("output rows:", len(output), "width:", len(output[0]))

causal weights: [[1.0, 0.0], [0.421, 0.579]]
cache future weight: 0.0
output rows: 2 width: 2

Two masking details matter in real code. First, softmax must run along the key axis so each query row sums to 1; normalizing along the query axis silently changes the operation.

Second, a row with no permitted key has no valid attention distribution. Our teaching function rejects it before softmax, and deliberately requires an exact mask shape rather than supporting broadcasting. It also rejects non-finite intermediate arithmetic: finite Q and K can still overflow when multiplied. Max-shifting helps with exponentiation; it can't repair an already corrupted score.

Replacing every logit in a fully blocked row with the same finite negative value hides the bug: softmax assigns weight to blocked keys because they all tie. Adding a finite negative mask instead can preserve differences between blocked scores, but a fully blocked row still gets total weight 1. In a valid row, very negative finite logits may underflow to zero; that numerical accident isn't a masking contract. Ensure each active query has an allowed key, or explicitly suppress outputs for padded query rows.

Predict the failure before running this case: with a finite fill value, a fully blocked row gives every blocked key the same score, so softmax invents a uniform distribution. The suppression step must handle that row explicitly.

This short failure case makes the second rule concrete:

import math

scores = [[0.8, 0.1], [0.4, -0.2]]
visible = [[True, False], [False, False]]

def softmax(row: list[float]) -> list[float]:
    shift = max(row)
    exps = [math.exp(x - shift) for x in row]
    total = sum(exps)
    return [x / total for x in exps]

filled = [
    [logit if allowed else -1e4 for logit, allowed in zip(row, mask_row)]
    for row, mask_row in zip(scores, visible)
]
finite_fill_weights = [softmax(row) for row in filled]
served_weights = [
    row if any(mask_row) else [0.0] * len(row)
    for row, mask_row in zip(finite_fill_weights, visible)
]

print("finite fill, invalid row:", [round(x, 4) for x in finite_fill_weights[1]])
print("after padded-query suppression:", served_weights[1])
print("valid row ignores blocked key:", [round(x, 4) for x in served_weights[0]])

finite fill, invalid row: [0.5, 0.5]
after padded-query suppression: [0.0, 0.0]
valid row ignores blocked key: [1.0, 0.0]

In PyTorch, F.scaled_dot_product_attention is an entry point for the same math, with backend dispatch that can select a fused implementation.

A fused kernel isn't guaranteed for every device, shape, and dtype. Two API traps matter: True in a boolean attn_mask means the position participates, unlike nn.MultiheadAttention's boolean key_padding_mask; and dropout_p is applied whenever it's greater than zero, so pass 0.0 at eval.^{[2]Reference 2torch.nn.functional.scaled_dot_product_attentionhttps://docs.pytorch.org/docs/2.14/generated/torch.nn.functional.scaled_dot_product_attention.html}

Common shape mistakes

When attention breaks, start with two questions: do the dimensions produce one score per query-key pair, and did the score scale stay stable? The two mistakes below produce distinct symptoms.

Forgetting to transpose K

With $Q$ shape $(B,h,N_q,d_k)$ and $K$ shape $(B,h,N_k,d_k)$, you need $(N_q \times d_k)(d_k \times N_k)$. In PyTorch that's Q @ K.transpose(-2, -1), not Q @ K. Usually the wrong version crashes. If $N_k=d_k$ by coincidence, it can run and compute the wrong operation. Test unequal dimensions too.

Forgetting to scale

If you skip attn_scores / math.sqrt(d_k), the code won't crash. Under the independent unit-variance setup below, a $d_k=64$ raw dot product has standard deviation 8 rather than 1. That larger spread can saturate softmax and reduce routing gradients.

Inspect score statistics and attention entropy when debugging, then restore the scale factor unless you're intentionally testing a different attention formulation.

Why scale by $\sqrt{d_k}$? The variance proof

The two-token trace hides a silent failure mode because its head dimension is only two numbers wide. Each attention score adds one feature product per dimension into a single scalar logit.

Under the initialization model below, increasing head width $d_k$ widens the spread of raw dot products. Larger differences between logits can make softmax concentrate nearly all weight on one key, reducing its derivatives. Dividing by $\sqrt{d_k}$ counteracts that width-dependent spread.^{[1]Reference 1Attention Is All You Need.https://arxiv.org/abs/1706.03762}

The derivation models all components of both vectors as mutually independent, with mean 0 and variance 1. Shared token states and learned projections can violate these assumptions, so this is a motivation for the scale, not a theorem about every learned attention row.^{[1]Reference 1Attention Is All You Need.https://arxiv.org/abs/1706.03762}

Deriving the raw dot-product variance

Assume $q, k \in \mathbb{R}^{d_k}$ satisfy that mutual-independence model, with $\mathbb{E}[q_i] = 0$, $\text{Var}(q_i) = 1$, $\mathbb{E}[k_i] = 0$, and $\text{Var}(k_i) = 1$. The dot product sums their elementwise products:

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

For each dimension $i$, independence gives an expected product of zero:

$\mathbb{E}[q_i k_i] = \mathbb{E}[q_i]\mathbb{E}[k_i] = 0 \cdot 0 = 0$

The variance of each individual product term is:

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

Because all $d_k$ terms are mutually independent across dimensions, the variance of their sum equals the sum of their variances:

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

The standard deviation of the raw dot product is therefore $\sigma = \sqrt{d_k}$. As head dimension grows, raw scores disperse much more widely across zero:

The softmax gradient saturation trap

What can wider scores do to routing gradients? For the unscaled case, softmax converts logits $z_j = q \cdot k_j$ into probability weights:

$a_j = \frac{e^{z_j}}{\sum_r e^{z_r}}$

The gaps matter, not the absolute logit values. If the largest logit exceeds another by 10, their unnormalized weights differ by $e^{10} \approx 22,026$. Whether the winning probability approaches one also depends on the number and scores of the other keys. A standard deviation alone doesn't determine the largest gap in a particular row.

Look at the Jacobian, the matrix of softmax derivatives with respect to its input logits. Here $\delta_{ij}$ is 1 when $i=j$ and 0 otherwise:

$\frac{\partial a_i}{\partial z_j} = a_i (\delta_{ij} - a_j)$

When softmax saturates:

• For the winning key ($i = j = \text{max}$): $\frac{\partial a_{\text{max}}}{\partial z_{\text{max}}} = a_{\text{max}}(1 - a_{\text{max}}) \approx 1.0 \cdot (1.0 - 1.0) = 0$.
• For any nonwinning row ($i \neq \text{max}$): the factor $a_i$ approaches zero, so $a_i(\delta_{ij}-a_j)$ approaches zero for every $j$.

As a row approaches a one-hot distribution, its softmax Jacobian approaches zero. During backpropagation, gradients into query and key projections pass through this Jacobian:

$\frac{\partial \ell}{\partial q} = \sum_j \frac{\partial \ell}{\partial z_j} k_j, \quad \text{where } \frac{\partial \ell}{\partial z_j} = \sum_i \frac{\partial \ell}{\partial a_i} \frac{\partial a_i}{\partial z_j}$

With bounded upstream gradients and bounded query/key norms, that concentration can make this row's routing gradients small. It doesn't permanently freeze the model: other rows and training steps can still update shared projections. The value path has a different derivative, derived below.

Would adding 1,000 to every logit cause the same problem as making one logit 20 larger than its neighbors?

import math

def softmax(logits):
    exps = [math.exp(x - max(logits)) for x in logits]
    return [x / sum(exps) for x in exps]

uniform = softmax([0.0, 0.0, 0.0])
shifted = softmax([1000.0, 1000.0, 1000.0])
gapped = softmax([20.0, 0.0, 0.0])
print("common offset leaves weights unchanged:", uniform == shifted)
print("uniform self-derivative:", round(uniform[0] * (1 - uniform[0]), 6))
print("gapped self-derivative:", f"{gapped[0] * (1 - gapped[0]):.2e}")

common offset leaves weights unchanged: True
uniform self-derivative: 0.222222
gapped self-derivative: 4.12e-09

Restoring unit variance

Dividing the dot product by $\sqrt{d_k}$ scales the logit variance by $(\frac{1}{\sqrt{d_k}})^2 = \frac{1}{d_k}$:

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

Under those assumptions, the scaled population standard deviation is 1 regardless of head width. This controls one source of score growth. It doesn't guarantee a nonzero gradient: a row with just one allowed key has constant weight 1, and identical values can make routing irrelevant even when weights aren't saturated.

Independence is doing real work. Let every coordinate be an independent random sign. If $q$ and $k$ are independent draws, the variance proof applies. If $k=q$, the coordinates of each vector still have mean 0 and variance 1, but $q\cdot k=d_k$ is constant. Enumerate the complete four-coordinate distribution to see the difference:

import itertools
import math
import statistics

width = 4
vectors = list(itertools.product((-1.0, 1.0), repeat=width))
independent = [sum(a * b for a, b in zip(q, k)) for q in vectors for k in vectors]
correlated = [sum(a * a for a in q) for q in vectors]  # k = q
print("independent raw variance:", statistics.pvariance(independent))
print("independent scaled variance:", statistics.pvariance([x / math.sqrt(width) for x in independent]))
print("k=q mean and variance:", statistics.mean(correlated), statistics.pvariance(correlated))

independent raw variance: 4.0
independent scaled variance: 1.0
k=q mean and variance: 4.0 0.0

Verify the prediction in code. This experiment samples 5,000 independent unit-variance vector pairs across three head widths to confirm that raw standard deviation tracks $\sqrt{d_k}$ while scaled standard deviation stays centered near 1.0:

import math
import random
import statistics

rng = random.Random(7)

def dot_products(width: int, samples: int = 5000) -> list[float]:
    return [
        sum(rng.gauss(0, 1) * rng.gauss(0, 1) for _ in range(width))
        for _ in range(samples)
    ]

for width in (16, 64, 512):
    raw = dot_products(width)
    raw_std = statistics.pstdev(raw)
    scaled_std = statistics.pstdev([x / math.sqrt(width) for x in raw])
    print(f"d_k={width:3d}: raw std={raw_std:5.2f}, scaled std={scaled_std:4.2f}")

d_k= 16: raw std= 4.07, scaled std=1.02
d_k= 64: raw std= 8.09, scaled std=1.01
d_k=512: raw std=22.39, scaled std=0.99

The finite-sample values differ from the population prediction. In this independent Gaussian experiment, raw spread tracks $\sqrt{d_k}$ and scaled spread stays near one; this isn't a measurement of a trained model's logits.

Three types of attention

The formula doesn't decide who may read whom. The mask and the sources of $Q$, $K$, and $V$ do.

Bidirectional attention lets every visible position read every other visible position. Causal attention allows only the current and earlier positions. Its training queries can be computed in parallel; generation adds tokens sequentially. Cross-attention lets a target sequence read from a separate source sequence, such as a decoder reading encoder states.

1. Bidirectional self-attention (encoder)

Every non-padding token may attend to every other non-padding token; there's no future-token mask. Encoder architectures such as BERT, and the Vision Transformer you'll meet next, use this pattern when the full input is already available:

"The cache grew quickly"
 Token "grew" attends to: [The, cache, grew, quickly]  (full context)

2. Causal self-attention (decoder)

Each token can only attend to itself and previous tokens. Future positions are masked with $-\infty$. Decoder architectures such as GPT-style and other autoregressive language models use this pattern for generation tasks where the model must predict the next token without seeing the future.

When processing a sequence step-by-step, the model progressively builds context but remains strictly blind to upcoming words:

"The cache grew quickly"
 Token "grew" attends to: [The, cache, grew]  (only past + self)
 Token "quickly" attends to: [The, cache, grew, quickly]  (full history)

The causal mask is a lower-triangular matrix that lets each position look only at itself and the positions before it. This minimal Python version takes the sequence length as input and outputs a boolean matrix where True indicates an allowed connection and False indicates a masked one.

def create_causal_mask(seq_len: int) -> list[list[bool]]:
    if type(seq_len) is not int or seq_len <= 0:
        raise ValueError("sequence length must be a positive integer")
    return [[key_pos <= query_pos for key_pos in range(seq_len)] for query_pos in range(seq_len)]

mask = create_causal_mask(4)
for row in mask:
    print(row)
print(mask[0] == [True, False, False, False])
print(mask[1] == [True, True, False, False])
print(mask[3] == [True, True, True, True])

[True, False, False, False]
[True, True, False, False]
[True, True, True, False]
[True, True, True, True]
True
True
True

Cached queries need absolute positions

Before switching sequences, there's one subtle causal case: cached decoding. A new query at absolute position 4 may read keys at positions 0 through 4. Its local row index is 0, but that doesn't make it the first token of the sequence. A naive 1 x 5 lower triangle would incorrectly allow only key 0.

Use absolute positions when building a mask for a cache or a chunk of new tokens. This example also covers a two-token chunk starting at position 3:

def causal_positions(query_positions, key_positions):
    queries, keys = list(query_positions), list(key_positions)
    if any(type(p) is not int or p < 0 for p in queries + keys):
        raise ValueError("absolute positions must be nonnegative integers")
    return [[k <= q for k in keys] for q in queries]

single = causal_positions([4], range(5))
chunk = causal_positions([3, 4], range(5))
wrong_local_triangle = [[k <= 0 for k in range(5)]]
assert single == [[True] * 5]
assert chunk[0] == [True, True, True, True, False]
assert chunk[1] == single[0]
assert causal_positions([3, 4], iter(range(5))) == chunk
print("new query at position 4:", single[0])
print("chunk query at position 3:", chunk[0])
print("wrong local triangle:", wrong_local_triangle[0])

new query at position 4: [True, True, True, True, True]
chunk query at position 3: [True, True, True, True, False]
wrong local triangle: [True, False, False, False, False]

PyTorch documents is_causal=True as upper-left aligned for non-square attention. For an offset cache, supply the appropriate explicit mask with is_causal=False, or use a supported causal-bias representation with the required alignment. A single new query with only past/current, non-padding keys needs no additional causal restriction.^{[2]Reference 2torch.nn.functional.scaled_dot_product_attentionhttps://docs.pytorch.org/docs/2.14/generated/torch.nn.functional.scaled_dot_product_attention.html}

Now combine two constraints: a query at absolute position 2 can read all three cached positions, but position 1 is padding. The correct allowed mask is [True, False, True]. Predict the context when the two visible keys have equal scores. Then compare SDPA with the inverse padding-mask convention in nn.MultiheadAttention.^{[3]Reference 3torch.nn.MultiheadAttentionhttps://docs.pytorch.org/docs/2.14/generated/torch.nn.MultiheadAttention.html}

import math
import torch
import torch.nn.functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel

# (batch, heads, queries/keys, width), already projected.
q = torch.tensor([[[[1.0, 0.0]]]], dtype=torch.float64)
k = torch.tensor([[[[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]]]], dtype=torch.float64)
v = torch.tensor([[[[1.0, 0.0], [0.0, 1.0], [2.0, 2.0]]]], dtype=torch.float64)
allowed = torch.tensor([[True, False, True]])

# Pin the CPU math backend for this correctness check, not a GPU benchmark.
with sdpa_kernel(backends=[SDPBackend.MATH]):
    cached = F.scaled_dot_product_attention(q, k, v, attn_mask=allowed, dropout_p=0.0)
    upper_left = F.scaled_dot_product_attention(q, k, v, is_causal=True, dropout_p=0.0)

logits = (q @ k.transpose(-2, -1)) / math.sqrt(2)
reference = torch.softmax(logits.masked_fill(~allowed, -torch.inf), dim=-1) @ v
assert torch.allclose(cached, reference, atol=1e-12, rtol=1e-12)

# Identity projections make MHA's projected Q/K/V equal to the tensors above.
mha = torch.nn.MultiheadAttention(2, 1, dropout=0.0, bias=False,
                                batch_first=True, dtype=torch.float64).eval()
with torch.no_grad():
    identity = torch.eye(2, dtype=torch.float64)
    mha.in_proj_weight.copy_(torch.cat([identity] * 3))
    mha.out_proj.weight.copy_(identity)
    mha_output, _ = mha(q[0], k[0], v[0], key_padding_mask=~allowed, need_weights=False)

agrees = torch.allclose(mha_output, cached[0], atol=1e-12, rtol=1e-12)
assert agrees
print("explicit cached mask:", cached.flatten().tolist())
print("upper-left causal mask:", upper_left.flatten().tolist())
print("SDPA/MHA inverse masks agree:", agrees)

explicit cached mask: [1.5, 1.0]
upper-left causal mask: [1.0, 0.0]
SDPA/MHA inverse masks agree: True

The upper-left mask reads only key 0 here. It doesn't encode the query's offset or the padding constraint. SDPA's boolean attn_mask uses True for allowed entries; MHA's boolean key_padding_mask and attn_mask use True for blocked entries. Don't pass both an explicit SDPA mask and is_causal=True; construct the combined allowed mask instead.^{[2]Reference 2torch.nn.functional.scaled_dot_product_attentionhttps://docs.pytorch.org/docs/2.14/generated/torch.nn.functional.scaled_dot_product_attention.html[3]Reference 3torch.nn.MultiheadAttentionhttps://docs.pytorch.org/docs/2.14/generated/torch.nn.MultiheadAttention.html}

Padding keys and padding queries are different

A key padding mask prevents every query from reading a padded source. It doesn't remove a padded query from the output. With left padding, a causal query at the first padded position may have no allowed key at all.

For a batch, combine visibility as causal AND key_valid AND query_valid. The key-valid tensor needs shape (B, 1, 1, Nk) to broadcast across heads and queries. The query-valid tensor needs (B, 1, Nq, 1). Passing a raw (B, Nk) mask can align the batch dimension with the query axis by accident.

Handle padded rows before an undefined softmax enters the computation. In this reference implementation, padded query rows temporarily receive zero logits, then their weights are explicitly replaced with zeros. Active rows still require a visible key. This avoids computing softmax([-inf, ...]) and hoping that multiplying its NaN output by zero will remove it.

For this padding test, change both real queries to [1, 0] and the real keys to [1, 0] and [1, 1]. Their scores now tie, isolating mask behavior from compatibility. Keep the original value vectors.

Predict the three output rows for the left-padded sequence [PAD, cache, grew]. The padded query should output zero, cache should read itself, and grew should mix the two real values equally:

import math
import torch
import torch.nn.functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel

q = torch.tensor([[[[1., 0.], [1., 0.], [1., 0.]]]],
                 dtype=torch.float64, requires_grad=True)
k = torch.tensor([[[[9., 9.], [1., 0.], [1., 1.]]]],
                 dtype=torch.float64, requires_grad=True)
v = torch.tensor([[[[999., 999.], [2., 1.], [1., 2.]]]],
                 dtype=torch.float64, requires_grad=True)
valid = torch.tensor([[False, True, True]])  # (B, N)
positions = torch.arange(3)
causal = positions[None, :] <= positions[:, None]  # (Nq, Nk)
active = valid[:, None, :, None]  # (B, 1, Nq, 1)
allowed = causal[None, None] & valid[:, None, None, :] & active
assert ((~active) | allowed.any(dim=-1, keepdim=True)).all()
assert all(torch.isfinite(t).all() for t in (q, k, v))

logits = (q @ k.transpose(-2, -1)) / math.sqrt(q.shape[-1])
masked = logits.masked_fill(~allowed, -torch.inf)
safe_logits = torch.where(active, masked, torch.zeros_like(masked))
weights = torch.softmax(safe_logits, dim=-1).masked_fill(~active, 0.)
context = weights @ v

# Compare only active queries, whose rows define valid distributions.
with sdpa_kernel(backends=[SDPBackend.MATH]):
    library = F.scaled_dot_product_attention(
        q[:, :, 1:], k, v, attn_mask=allowed[:, :, 1:], dropout_p=0.)
torch.testing.assert_close(context[:, :, 1:], library)
context[..., 0].sum().backward()
assert all(torch.isfinite(t.grad).all() for t in (q, k, v))
assert torch.equal(q.grad[:, :, 0], torch.zeros_like(q.grad[:, :, 0]))
assert torch.equal(k.grad[:, :, 0], torch.zeros_like(k.grad[:, :, 0]))
assert torch.equal(v.grad[:, :, 0], torch.zeros_like(v.grad[:, :, 0]))
print("weight row sums:", weights.sum(dim=-1).flatten().tolist())
print("context:", context.detach()[0, 0].tolist())
print("active rows match SDPA and padding gradients are zero")

weight row sums: [0.0, 1.0, 1.0]
context: [[0.0, 0.0], [2.0, 1.0], [1.5, 1.5]]
active rows match SDPA and padding gradients are zero

The padded row has sum zero because it is deliberately suppressed, not a probability distribution. Keep its loss excluded too. This reference materializes scores for clarity. Production implementations can skip padded queries instead. Keep all Q/K/V entries finite, including padding: masking scores doesn't repair NaN values, and 0 * NaN is NaN.

3. Cross-attention (encoder-decoder)

Queries come from one sequence, Keys and Values from another. The original Transformer uses this pattern in its decoder: encoder outputs provide keys and values, while current decoder states provide queries.^{[1]Reference 1Attention Is All You Need.https://arxiv.org/abs/1706.03762}

The resulting weights choose which source positions contribute to each target representation:

Encoder output (source): "cache memory pressure"   provides K, V
Decoder state (target):  "decode slows ___"        provides Q

Q from decoder times K from encoder gives attention weights
Weights times V from encoder give decoder context

In an offline encoder-decoder model, cross-attention usually applies a source padding mask so decoder tokens don't attend to padded encoder positions. It generally doesn't need a causal source mask, because the encoder has already seen the whole source. Streaming systems can impose additional availability constraints.

Full-sequence self-attention produces a square query-by-key matrix. Cached self-attention and cross-attention can both be rectangular. Two target queries reading three source positions produce a 2 x 3 routing matrix.

import math

decoder_queries = [[1.0, 0.0], [0.0, 1.0]]
encoder_keys = [[1.0, 0.0], [0.2, 0.8], [0.0, 1.0]]
encoder_values = [[1.0, 0.0], [0.5, 0.5], [0.0, 1.0]]

def softmax(row: list[float]) -> list[float]:
    exps = [math.exp(x - max(row)) for x in row]
    return [value / sum(exps) for value in exps]

logits = [
    [sum(q_i * k_i for q_i, k_i in zip(q, k)) / math.sqrt(2) for k in encoder_keys]
    for q in decoder_queries
]
weights = [softmax(row) for row in logits]
context = [
    [sum(weight * value[d] for weight, value in zip(row, encoder_values)) for d in range(2)]
    for row in weights
]

print(f"routing shape: {len(weights)} x {len(weights[0])}")
print("row sums:", [round(sum(row), 3) for row in weights])
print("context:", [[round(x, 3) for x in row] for row in context])

routing shape: 2 x 3
row sums: [1.0, 1.0]
context: [[0.623, 0.377], [0.393, 0.607]]

Multi-head attention

A single attention head constructs one distribution over keys for each query token. Multiple heads let the same token use several distributions at once. For example, one distribution might emphasize nearby positions while another emphasizes an earlier noun. These are possible learned patterns, rather than tasks assigned to particular head numbers.

Multi-Head Attention (MHA) uses $h$ separately parameterized sets of projections. A common choice is $d_k = d_v = d_{\text{model}} / h$.^{[1]Reference 1Attention Is All You Need.https://arxiv.org/abs/1706.03762} In self-attention, each head computes:

$\text{head}_i = \text{Attention}(XW_i^Q, XW_i^K, XW_i^V)$

Here $W_i^Q, W_i^K \in \mathbb{R}^{d_{\text{model}} \times d_k}$ and $W_i^V \in \mathbb{R}^{d_{\text{model}} \times d_v}$. The heads are trained jointly and their outputs can correlate; separate projections don't imply statistically independent or orthogonal features. Studies have found positional and syntactic patterns in some translation heads, and copy-like induction circuits in other tested models.^{[4]Reference 4Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Prunedhttps://arxiv.org/abs/1905.09418[5]Reference 5In-context Learning and Induction Heads.https://arxiv.org/abs/2209.11895} A heatmap showing a pronoun attending to a noun is a clue to investigate, not proof that the head performs coreference resolution.

The $h$ head outputs concatenate to width $h d_v$. The output projection has shape $W^O \in \mathbb{R}^{h d_v \times d_{\text{model}}}$; it is square in the common case $h d_v=d_{\text{model}}$:

$\text{MultiHead}(Q,K,V) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h) W^O$

$W^O$ is a learned linear mixture into the residual stream. It isn't constrained to be a rotation, preserve distances, or keep head contributions separate.

What do attention heads learn?

These interpretability results are evidence about specific trained models, not a promise about every transformer. Voita et al. found positional and syntactic patterns among heads in neural machine translation encoders.^{[4]Reference 4Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Prunedhttps://arxiv.org/abs/1905.09418}

Michel et al. found head redundancy in their tested WMT translation model and fine-tuned BERT.^{[6]Reference 6Are Sixteen Heads Really Better than One?https://arxiv.org/abs/1905.10650} Olsson et al. presented causal evidence for induction-head effects in small attention-only models and correlational evidence in larger models with MLPs. Their broader explanation of in-context learning remains a hypothesis.^{[5]Reference 5In-context Learning and Induction Heads.https://arxiv.org/abs/2209.11895}

Same asymptotic attention FLOPs

Multi-head attention doesn't increase the asymptotic cost of the attention core when you keep $d_{\text{model}}$ fixed. It restructures the work.

With $h d_k=h d_v=d_{\text{model}}$, single-head attention on width 512 uses the same leading-order matrix-product FLOPs (Floating Point Operations) as 8-head attention with $d_k=d_v=64$. Both query/key and value widths matter to that comparison.

The dense Q/K/V and output projections still cost $O(n d_{\text{model}}^2)$ either way.

Predict before running the arithmetic: with fixed $d_{\text{model}}$, one 512-wide head and eight 64-wide heads should do the same leading score-and-value work. Splitting a fixed width into more heads doesn't change the total number of score-and-value multiply-adds in the attention core:

seq_len = 2048
d_model = 512

for heads in (1, 8, 16):
    d_head = d_model // heads
    score_and_value_work = 2 * heads * seq_len**2 * d_head
    print(f"heads={heads:2d}, d_head={d_head:3d}, core MACs={score_and_value_work:,}")

heads= 1, d_head=512, core MACs=4,294,967,296
heads= 8, d_head= 64, core MACs=4,294,967,296
heads=16, d_head= 32, core MACs=4,294,967,296

The count above is multiply-accumulate operations (MACs) for the two matrix products; counting multiplication and addition separately gives roughly twice as many FLOPs. More heads still mean more softmax rows and, if materialized, more score matrices. Equal leading matmul work doesn't promise equal runtime or memory.

After each dense Q/K/V projection to width $D$, split that projected width into $h$ heads. This example uses distinct coordinate labels so a wrong reshape can't hide behind zeros. It tests only the split/merge layout, not the learned projection or attention itself:

B, N, D, n_heads = 2, 3, 8, 4
d_k = D // n_heads
assert D % n_heads == 0

projected = [[[100 * b + 10 * n + d for d in range(D)] for n in range(N)] for b in range(B)]
# (B, N, D) -> (B, h, N, d_k)
heads = [
    [[projected[b][n][h * d_k:(h + 1) * d_k] for n in range(N)] for h in range(n_heads)]
    for b in range(B)
]
# (B, h, N, d_k) -> (B, N, D)
concat = [
    [
        [coord for head in range(n_heads) for coord in heads[b][head][n]]
        for n in range(N)
    ]
    for b in range(B)
]
print("input shape:", (B, N, D))
print("head shape:", (len(heads), len(heads[0]), len(heads[0][0]), len(heads[0][0][0])))
assert concat == projected
assert heads[1][2][1] == [114, 115]
print("batch 1, head 2, token 1:", heads[1][2][1])
print("every coordinate restored:", concat == projected)

input shape: (2, 3, 8)
head shape: (2, 4, 3, 2)
batch 1, head 2, token 1: [114, 115]
every coordinate restored: True

Gradient flow through attention

For one query with attention-weight dropout disabled, write $o=\sum_j a_j v_j$, where $a=\mathrm{softmax}(z)$ and $z_j=q\cdot k_j/\sqrt{d_k}$. Substituting the softmax derivative $\partial a_i/\partial z_j=a_i(\delta_{ij}-a_j)$ into the weighted sum gives a useful interpretation:^{[7]Reference 7FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.https://arxiv.org/abs/2205.14135}

$\frac{\partial o}{\partial z_j}=a_j(v_j-o).$

Increasing a key's score pulls the output toward that key's value, away from the current mixture. If its value already equals the mixture, changing its weight does nothing.

For a scalar loss $\ell$ with upstream gradient $g=\partial\ell/\partial o$, let $r_j=a_j\,g^T(v_j-o)$. The three paths are:

For the grew row, take the toy loss $\ell=o_0$, the output's first coordinate. Cache has first value coordinate 2, grew has 1, and the mixture is about 1.421. Raising cache's score increases the loss; raising grew's score decreases it. This isn't a useful training objective by itself, just a derivative we can check exactly.

The next cell reuses the attention function and compares its analytic query gradient with central finite differences:

q = [0.5, 1.0]
out, weights = scaled_dot_product_attention([q], K, V)
a = weights[0]
o = out[0]
r = [a[j] * (V[j][0] - o[0]) for j in range(len(K))]
analytic = [sum(r[j] * K[j][d] for j in range(len(K))) / math.sqrt(2) for d in range(2)]

def loss(query):
    return scaled_dot_product_attention([query], K, V)[0][0][0]

epsilon = 1e-5
numeric = []
for d in range(2):
    plus, minus = q.copy(), q.copy()
    plus[d] += epsilon
    minus[d] -= epsilon
    numeric.append((loss(plus) - loss(minus)) / (2 * epsilon))
assert max(abs(x - y) for x, y in zip(analytic, numeric)) < 1e-9
assert abs(sum(r)) < 1e-12  # Shifting every score equally changes nothing.
print("output:", [round(x, 4) for x in o])
print("query gradient:", [round(x, 5) for x in analytic])
print("finite differences match:", all(abs(x - y) < 1e-9 for x, y in zip(analytic, numeric)))

output: [1.4211, 1.5789]
query gradient: [0.08619, -0.12066]
finite differences match: True

With only one allowed key, its weight is exactly 1 and its value equals the output: all routing gradients for that query are zero. The value path still receives $g$. A saturated row approaches this behavior; the scale factor helps avoid accidental saturation, but can't create a routing choice where the mask permits none.

Implement softmax backward without a Jacobian

The Jacobian formula is useful for derivation, but allocating an $N_k \times N_k$ derivative matrix for every row is unnecessary. Let $t_j=\partial\ell/\partial a_j$ be the upstream gradient into a softmax row. Multiplying it by the Jacobian simplifies to one row reduction:^{[7]Reference 7FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.https://arxiv.org/abs/2205.14135}

Each logit gradient compares its upstream gradient with the attention-weighted average. The gradients sum to zero because increasing all logits equally cannot change softmax. A blocked key with weight zero also gets logit gradient zero.

For attention, $t_j=g^T v_j$, so the average is $\sum_j a_j g^T v_j=g^T o$. This recovers $r_j=a_j g^T(v_j-o)$ from the previous derivation. Across all queries, value gradients add together as $\partial\ell/\partial V=A^T G$, where $G$ contains the output gradients.

This cell implements the row reduction and compares it with PyTorch autograd. The upstream gradient of 99 on the blocked third key deliberately tests that masking blocks its gradient too:

import torch

z = torch.tensor([[0., 0., 5.]], dtype=torch.float64, requires_grad=True)
allowed = torch.tensor([[True, True, False]])
a = torch.softmax(z.masked_fill(~allowed, -torch.inf), dim=-1)
upstream = torch.tensor([[2., 1., 99.]], dtype=torch.float64)
manual = a * (upstream - (a * upstream).sum(dim=-1, keepdim=True))
(a * upstream).sum().backward()
torch.testing.assert_close(manual, z.grad)
assert torch.allclose(manual.sum(dim=-1), torch.zeros(1, dtype=z.dtype))
assert z.grad[0, 2] == 0
print("weights:", a.detach().flatten().tolist())
print("logit gradients:", z.grad.flatten().tolist())
print("manual backward matches autograd")

weights: [0.5, 0.5, 0.0]
logit gradients: [0.25, -0.25, 0.0]
manual backward matches autograd

Don't substitute the cross-entropy shortcut $a-y$ here. That expression requires a particular loss on softmax probabilities. Attention generally receives gradients through its value mixture, so its upstream vector $t$ comes from $G V^T$.

Complexity analysis

There are two bottlenecks to keep separate: quadratic query-key pairs and width-heavy projections. Big-O notation describes how work or storage grows as sequence length $n$ and model width $d_{\text{model}}$ increase.

It hides constant factors, so use it for growth trends and use concrete byte/FLOP arithmetic for capacity planning.

Concrete memory example

Batch of 8, 32 heads, $n = 8192$, FP16 (16-bit floating point format):

$\text{Attention matrix} = 8 \times 32 \times 8192^2 \times 2 \text{ bytes} \approx 34.4 \text{ GB} \approx 32 \text{ GiB}$

How the numbers work

8 sequences in the batch x 32 attention heads x 8192² entries per attention map x 2 bytes per number (FP16) = about 34.4 GB of raw storage, or about 32 GiB, for one attention-score tensor before values needed for backpropagation, model weights, or optimizer state.

batch = 8
heads = 32
seq_len = 8192
bytes_per_fp16 = 2

bytes_total = batch * heads * seq_len**2 * bytes_per_fp16
gb = bytes_total / 1_000_000_000
gib = bytes_total / 1024**3

print(round(gb, 1), "GB")
print(round(gib, 1), "GiB")
print(round(gb, 1) == 34.4, round(gib, 1) == 32.0)

34.4 GB
32.0 GiB
True True

In the naive formulation, $O(n^2)$ temporary memory becomes a bottleneck for long sequences. Doubling sequence length makes one materialized score tensor four times larger. Fused kernels can avoid storing that full tensor, but exact dense attention still computes interactions across all query-key pairs.

Training vs. inference bottlenecks

Two attention workloads stress memory differently; model weights and other operations can also dominate:

• During training (or prompt prefill): All $n$ tokens are processed together. A naive implementation materializes the entire $n \times n$ attention score matrix in GPU high-bandwidth memory (HBM), creating an $O(n^2)$ memory wall.
• During autoregressive decoding: The model generates one token at a time. The query sequence length is just $N_q = 1$, so attention is a $1 \times N_k$ vector-matrix operation. The memory problem isn't score matrix storage; it's repeatedly loading the accumulated KV cache (all past keys and values) from HBM to on-chip SRAM for every single generated token.^{[8]Reference 8Fast Transformer Decoding: One Write-Head is All You Need.https://arxiv.org/abs/1911.02150}

Head architectures: MHA, MQA, and GQA

To attack this decode bandwidth bottleneck, modern transformers modify how query heads route to key and value heads:

1. Multi-Head Attention (MHA): $h_Q$ query heads pair with $h_{\text{KV}} = h_Q$ key and value heads. Each head has separate K/V projections; raw cache payload scales with $h_{\text{KV}}$.
2. Multi-Query Attention (MQA): $h_Q$ query heads share one key head and one value head ($h_{\text{KV}} = 1$). For fixed widths, precision, and cached lengths, raw cache payload shrinks by $h_Q\times$. Shazeer's tested translation models traded a small quality loss for faster decoding; that experiment doesn't establish a universal effect on reasoning or multilingual tasks.^{[8]Reference 8Fast Transformer Decoding: One Write-Head is All You Need.https://arxiv.org/abs/1911.02150}
3. Grouped-Query Attention (GQA): Each group of $h_Q/G$ query heads shares one K/V head, with $h_{\text{KV}}=G$ and $h_Q$ divisible by $G$. Intermediate group counts trade sharing against separate K/V projections; MQA and MHA are the endpoints. Ainslie et al. uptrained T5 checkpoints and found quality close to MHA on their tested tasks. Their results aren't a guarantee for every model.^{[9]Reference 9GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints.https://arxiv.org/abs/2305.13245} As a concrete architecture, Llama 3.1 70B uses 64 query heads and 8 KV heads.^{[10]Reference 10The Llama 3 Herd of Models.https://arxiv.org/abs/2407.21783}

Changing the head layout requires a checkpoint trained or uptrained for it. You can't turn an arbitrary MHA checkpoint into equivalent GQA by reshaping its cache. GQA also retains $h_Q$ query distributions; it doesn't divide dense query-key interactions or a materialized $h_Q n^2$ score tensor by the KV sharing ratio.

Count the raw KV payload first

For $L$ identical layers, $B$ equally long cached sequences, equal key/value width $d_k=d_v$, and $P_{\text{bytes}}$ bytes per stored coordinate, the raw dense payload is:

$\text{Memory}_{\text{KV}} = 2 \times L \times h_{\text{KV}} \times d_k \times N \times B \times P_{\text{bytes}}$

Breaking down every variable in the equation:

• $2$: Stores both Key and Value tensors.
• $L$: Number of transformer layers in the model.
• $h_{\text{KV}}$: Number of Key/Value heads ($h_Q$ in MHA, $1$ in MQA, $G$ in GQA).
• $d_k$: Key and value width in this equal-width calculation.
• $N$: Context sequence length (prompt tokens plus generated tokens).
• $B$: Batch size (number of concurrent requests being served).
• $P_{\text{bytes}}$: Bytes per coordinate ($2$ for FP16/BF16, $1$ for FP8). Packed INT4 has an ideal $0.5$ bytes per coordinate before scales, metadata, and padding.

If widths differ, replace $2d_k$ with $d_k+d_v$. If cached lengths differ by request or layer, sum their payloads rather than pretending every slot is occupied equally. Actual device memory also includes page or allocator slack, quantization metadata, possible replicated shards, model weights, and working space.

Use the Llama 3.1 70B dimensions ($L=80$, $h_Q=64$, $d_k=d_v=128$) for a hypothetical layout comparison, with $B=16$, $N=8,192$, and two-byte coordinates. The GQA row matches its head counts; the MHA and MQA rows are alternative shapes, not equivalent released checkpoints.^{[10]Reference 10The Llama 3 Herd of Models.https://arxiv.org/abs/2407.21783}

The decode memory-bandwidth bottleneck

Why can a smaller cache accelerate decode? Each query head performs two products: scoring against K and mixing V. Counting a multiplication and addition as two FLOPs gives approximately $2Nd_k+2Nd_v$ FLOPs. With $d_k=d_v=d$, that's $4Nd$, while reading one FP16 K/V head once transfers $4Nd$ bytes.

The arithmetic intensity is:

$\text{MHA core intensity} \approx \frac{4Nd \text{ FLOPs}}{4Nd \text{ bytes}} = 1 \text{ FLOP/byte}$

If a grouped kernel reuses each loaded K/V head across its query group, the idealized ratio is:

$\text{Grouped core intensity} \approx \frac{4Nh_Qd}{2Nh_{\text{KV}}dP_{\text{bytes}}} = \frac{2h_Q}{h_{\text{KV}}P_{\text{bytes}}}.$

For FP16, that is 1 FLOP/byte for MHA and 8 for an 8:1 GQA layout. This counts only the two attention products and one K/V read, excluding softmax, projections, other transfers, and repeated loads. A backend that expands shared K/V heads can lose the reuse benefit.

In a roofline estimate, achievable FLOPs/s is bounded by the smaller of compute capacity and bandwidth times arithmetic intensity. For illustrative capacities of 312 TFLOPs/s and 2 TB/s, the crossover is 156 FLOPs/byte. That suggests this simplified decode core is bandwidth constrained; it doesn't measure idle cycles or predict an 8x end-to-end speedup. Batch size, kernel reuse, other layers, and communication still matter.

Count raw payload and idealized core intensity together. This smaller fixture has one request, 32 layers, and 32 query heads; its GQA sharing ratio is 4:1:

layers = 32
tokens = 8192
head_dim = 128
bytes_per_value = 2  # FP16

def cache_gib(kv_heads: int) -> float:
    bytes_total = layers * tokens * kv_heads * head_dim * 2 * bytes_per_value
    return bytes_total / 1024**3

mha = cache_gib(32)
for label, kv_heads in [("MHA", 32), ("GQA", 8), ("MQA", 1)]:
    size = cache_gib(kv_heads)
    print(f"{label}: kv_heads={kv_heads:2d}, cache={size:.3f} GiB, reduction={mha / size:.0f}x")
    core_flops = 4 * tokens * 32 * head_dim  # scoring AND value mixing, one layer
    kv_bytes = 2 * tokens * kv_heads * head_dim * bytes_per_value
    print(f"  idealized core intensity={core_flops / kv_bytes:.0f} FLOPs/byte")

MHA: kv_heads=32, cache=4.000 GiB, reduction=1x
  idealized core intensity=1 FLOPs/byte
GQA: kv_heads= 8, cache=1.000 GiB, reduction=4x
  idealized core intensity=4 FLOPs/byte
MQA: kv_heads= 1, cache=0.125 GiB, reduction=32x
  idealized core intensity=32 FLOPs/byte

FlashAttention and MQA/GQA target different bottlenecks. FlashAttention cuts temporary attention I/O during prompt processing and training, while MQA/GQA cut persistent KV-cache size during decoding.

Softmax numerical stability

The two-token walkthrough already subtracted the row max before exp. That isn't a style choice.

Naive softmax overflows: $e^x$ grows so fast that a large positive logit becomes Inf, and Inf / Inf is NaN. Attention kernels need a stable form so large but valid scores don't corrupt routing weights.

The max-shift technique

Subtract the maximum value from the input vector before exponentiating:

$\text{softmax}(x_i) = \frac{e^{x_i - \max(\mathbf{x})}}{\sum_j e^{x_j - \max(\mathbf{x})}}$

Before running the example, predict what stays invariant: shifting $[1000, 1001, 999]$ by its maximum changes the numbers sent to exp, not the relative probabilities. The largest shifted logit should receive the largest weight, and the row should still sum to 1.

Subtracting $\max(\mathbf{x})$ shifts every logit by the same constant. That constant factors out of the exponentials and cancels between numerator and denominator, so the probabilities don't change.

For finite input logits, the largest exponent is $e^0 = 1$, which blocks positive overflow. Very small terms can still underflow to zero, losing their tiny probability and gradient contributions. Max-shifting doesn't repair non-finite inputs or prevent a genuinely concentrated distribution.^{[11]Reference 11Online normalizer calculation for softmax.https://arxiv.org/abs/1805.02867}

Online softmax (Milakov and Gimelshein, 2018) extends the same identity to a streaming pass. It tracks a running maximum and a running sum of exponentials, then rescales the sum when a larger maximum appears. That's the algebra FlashAttention uses to tile a query row without storing every score.^{[11]Reference 11Online normalizer calculation for softmax.https://arxiv.org/abs/1805.02867}

import math

def stable_softmax(logits: list[float]) -> list[float]:
    shift = max(logits)
    exps = [math.exp(x - shift) for x in logits]
    total = sum(exps)
    return [x / total for x in exps]

probs = stable_softmax([1000.0, 1001.0, 999.0])
print([round(p, 4) for p in probs])
print(round(sum(probs), 6))
print(probs[1] == max(probs))

[0.2447, 0.6652, 0.09]
1.0
True

Overflow can happen before softmax

Max-shifting assumes the logits are still usable numbers. A manually materialized FP16 score can overflow during Q @ K.T, even when Q and K are finite. Casting that damaged result to FP32 before softmax preserves the infinity. Recompute the product in a wider type instead.

This CPU example uses head width 1 so scaling cannot obscure the failure. The exact scores are [90000, 0]. Check which operation first becomes non-finite, then compare casting after the product with casting before it:

import torch

q = torch.tensor([[300.]], dtype=torch.float16)
k = torch.tensor([[300.], [0.]], dtype=torch.float16)
half_scores = q @ k.T
late_cast = torch.softmax(half_scores.float(), dim=-1)
wide_scores = q.float() @ k.float().T
wide_weights = torch.softmax(wide_scores, dim=-1)
assert torch.isfinite(q).all() and torch.isfinite(k).all()
assert not torch.isfinite(half_scores).all()
assert not torch.isfinite(late_cast).all()
assert torch.isfinite(wide_weights).all()
print("FP16 scores:", half_scores.flatten().tolist())
print("late FP32 cast repairs weights:", torch.isfinite(late_cast).all().item())
print("FP32 recomputation:", wide_weights.flatten().tolist())

FP16 scores: [inf, 0.0]
late FP32 cast repairs weights: False
FP32 recomputation: [1.0, 0.0]

The repaired row is still saturated: [1, 0] reflects the huge score gap. Finite weights and useful routing gradients are separate checks. PyTorch documents FP32 intermediates for its SDPA math backend with FP16/BF16 inputs, so this manual materialization example doesn't predict every fused kernel's behavior.^{[2]Reference 2torch.nn.functional.scaled_dot_product_attentionhttps://docs.pytorch.org/docs/2.14/generated/torch.nn.functional.scaled_dot_product_attention.html}

When debugging NaN attention, inspect Q/K/V finiteness, the scores before masking, and allowed-key counts before inspecting softmax. An invalid mask, score overflow, and a valid but concentrated distribution require different fixes.

FlashAttention: tiled exact attention

For long sequences, a materialized $n \times n$ score matrix can become a major memory cost. It's an intermediate on the way to $O=PV$, so an implementation can compute the final output without storing the full matrix.

FlashAttention (Dao et al., 2022) uses tiling and recomputation to avoid writing the full score or probability matrix to high-bandwidth memory (HBM). Additional row statistics scale linearly with sequence length; Q, K, V, and the output still occupy their own storage. The computation remains dense attention, with floating-point differences rather than a sparse or low-rank approximation.^{[7]Reference 7FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.https://arxiv.org/abs/2205.14135}

As of October 10, 2026, the official repository also provides FlashAttention-4 in CuTe DSL for Hopper and Blackwell. The identity below explains the shared foundation. A call to framework SDPA doesn't establish that it uses this particular implementation.^{[12]Reference 12FlashAttention-4: CuTe DSL Implementationhttps://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/cute/README.md}

The GPU memory hierarchy bottleneck

The useful distinction is where intermediate data lives:

• HBM: Large device memory holding tensors that many compute blocks can access. Repeatedly writing and reading a quadratic intermediate costs bandwidth.
• On-chip shared memory and registers: Much smaller storage used by active compute blocks. Tiles must fit the resources available to those blocks; adding up SRAM across the whole GPU doesn't make it one shared tile buffer. GPU cores perform the arithmetic on these operands.

Consider a materializing implementation with three logical stages. Actual kernel counts also depend on scaling, masking, dropout, and fusion:

1. Matrix multiplication 1: Reads $Q$ and $K$ from HBM, computes raw scores $S = QK^T / \sqrt{d_k}$, and writes the full $B \times h \times n \times n$ tensor back to HBM ($O(n^2)$ write).
2. Softmax kernel: Reads $S$ from HBM, computes $P = \text{softmax}(S)$, and writes $P$ back to HBM ($O(n^2)$ read and write).
3. Matrix multiplication 2: Reads $P$ and $V$ from HBM, computes $O = PV$, and writes $O$ to HBM ($O(n^2)$ read).

Those score and probability reads and writes can dominate time even when the matrix products are efficient. A percentage of stalled cycles requires profiling a specific workload; the logical operation count doesn't tell you that percentage.

SRAM tiling and the online softmax trick

The forward computation fuses the stages using two core ideas:^{[7]Reference 7FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.https://arxiv.org/abs/2205.14135[11]Reference 11Online normalizer calculation for softmax.https://arxiv.org/abs/1805.02867}

Here take $d_k=d_v=d$ and suppress batch/head axes to keep the tile shapes readable.

1. SRAM Tiling: Instead of processing full sequences, FlashAttention partitions $Q$ into blocks $Q_i \in \mathbb{R}^{B_r \times d}$ and streams blocks of $K_j, V_j \in \mathbb{R}^{B_c \times d}$ sized to fit entirely within on-chip SRAM.
2. Online Softmax Rescaling: Standard softmax requires computing the maximum over the entire sequence before evaluating exponentials. Online softmax (Milakov and Gimelshein, 2018) updates the row maximum and normalizer incrementally as blocks stream through SRAM.

To understand the identity, follow one query block while K/V blocks stream past it. Implementations can choose different loop orders and locations for saved state. The active block maintains:

• $m$: the running maximum of dot-product scores.
• $\ell$: the running sum of exponentials (the softmax normalizer).
• $u$: the running weighted-value output accumulator.

When a new key-value block arrives with local scores $\tilde{S}_{ij} = Q_i K_j^T / \sqrt{d_k}$ and local row maximum $\tilde{m} = \max(\tilde{S}_{ij})$:

$\text{New running max: } m_{\text{new}} = \max(m_{\text{old}}, \tilde{m})$

$\text{Rescaling factor: } \alpha = \exp(m_{\text{old}} - m_{\text{new}})$

$\text{Updated normalizer: } \ell_{\text{new}} = \alpha \cdot \ell_{\text{old}} + \sum_k \exp(\tilde{S}_{ij, k} - m_{\text{new}})$

$\text{Updated accumulator: } u_{\text{new}} = \alpha \cdot u_{\text{old}} + \sum_k \exp(\tilde{S}_{ij, k} - m_{\text{new}}) V_{j, k}$

When all key-value blocks have streamed through, the exact output vector is:

$O_i = \frac{u_{\text{final}}}{\ell_{\text{final}}}$

Initialize $m=-\infty$, $\ell=0$, and $u=0$. Skip an empty or entirely masked tile before applying the update: its local maximum would otherwise be $-\infty$, making $-\infty-(-\infty)$ undefined. A valid query must encounter at least one allowed key across all tiles.

Transient score and probability tiles are consumed on chip. The full $n \times n$ matrices are never materialized in HBM. This removes quadratic score storage, not the quadratic number of dense query-key interactions.

Backward pass recomputation

A conventional backward implementation can save the $n \times n$ probability matrix $P$ from the forward pass. That consumes quadratic memory; saving it isn't a mathematical requirement of backpropagation.

FlashAttention saves Q, K, V, the output, and linear row statistics needed for backward, rather than saving the full $P$. Depending on the implementation, statistics may be stored as a maximum and normalizer or their log-sum-exp equivalent. Backward reloads the saved tensors and recomputes score and probability tiles on chip.^{[7]Reference 7FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.https://arxiv.org/abs/2205.14135}

Recomputation adds arithmetic while reducing HBM traffic. The original paper's experiments show how doing extra arithmetic can reduce elapsed time by avoiding large memory transfers.^{[7]Reference 7FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.https://arxiv.org/abs/2205.14135} That benefit depends on the device, dimensions, precision, mask, dropout, and comparison implementation. For a local measurement, record those conditions, time forward and backward separately, and compare outputs and gradients with a reference before interpreting the speedup.

The streaming identity works identically on scalars or matrix tiles. In this example, the first tile contains score 1000 and value 2; the second contains scores 1001 and 999, with values 1 and 4:

import math

def online_attention(tiles):
    maximum, total, numerator = -math.inf, 0.0, 0.0
    for logits, values in tiles:
        logits, values = list(logits), list(values)
        if len(logits) != len(values):
            raise ValueError("each tile needs one value per logit")
        if any(math.isnan(x) or x == math.inf for x in logits):
            raise ValueError("logits must be finite or masked -inf")
        if any(not math.isfinite(v) for v in values):
            raise ValueError("values must be finite")
        if not logits or max(logits) == -math.inf:
            continue  # no contribution; avoids -inf - (-inf)
        updated_max = max(maximum, max(logits))
        correction = math.exp(maximum - updated_max)
        weights = [math.exp(x - updated_max) for x in logits]
        total = correction * total + math.fsum(weights)
        try:
            numerator = math.fsum([correction * numerator] +
                                  [w * v for w, v in zip(weights, values, strict=True)])
        except OverflowError as error:
            raise ValueError("weighted accumulator overflowed") from error
        if not math.isfinite(total) or not math.isfinite(numerator):
            raise ValueError("attention accumulator overflowed")
        maximum = updated_max
    if total == 0:
        raise ValueError("query needs at least one allowed key across all tiles")
    return numerator / total

scores, values = [1000.0, 1001.0, 999.0], [2.0, 1.0, 4.0]
exps = [math.exp(s - max(scores)) for s in scores]
full = sum(w * v for w, v in zip(exps, values)) / sum(exps)
split = online_attention([(scores[:1], values[:1]), (scores[1:], values[1:])])
reversed_tiles = online_attention([(scores[1:], values[1:]), (scores[:1], values[:1])])
assert abs(split - full) < 1e-12
assert abs(reversed_tiles - full) < 1e-12
with_blocked_tile = online_attention([([-math.inf], [20.0]), (scores, values)])
assert abs(with_blocked_tile - full) < 1e-12
print("full-row output:", round(full, 6))
print("two-tile output:", round(split, 6))
print("tile order agrees within tolerance:", abs(reversed_tiles - full) < 1e-12)
print("leading masked tile is skipped:", abs(with_blocked_tile - full) < 1e-12)

full-row output: 1.51482
two-tile output: 1.51482
tile order agrees within tolerance: True
leading masked tile is skipped: True

The cell accepts finite logits or masked $-\infty$, requires matching finite values, and skips empty or entirely blocked tiles. It rejects a query with no allowed key and arithmetic that overflows its scalar accumulator. Even when the normalized answer would be representable, an unnormalized accumulator can overflow; a production kernel's precision strategy still matters.

You can use this through framework primitives without writing a custom CUDA kernel. Check which backend actually runs before attributing a speedup to FlashAttention.^{[2]Reference 2torch.nn.functional.scaled_dot_product_attentionhttps://docs.pytorch.org/docs/2.14/generated/torch.nn.functional.scaled_dot_product_attention.html}

Mistakes that break attention

Going deeper

"Why not use additive attention instead of dot-product?"

Bahdanau et al. (2015)^{[13]Reference 13Neural Machine Translation by Jointly Learning to Align and Translate.https://arxiv.org/abs/1409.0473} used additive attention: $\text{score}(q,k)=u^T\tanh(W_q q+W_k k)$, where $u$ is a learned scoring vector, not the value vector. Both approaches learn compatibility. The Transformer paper favors dot products because they map efficiently to matrix multiplication, and introduces scaling to address the wider score distribution at large $d_k$.^{[1]Reference 1Attention Is All You Need.https://arxiv.org/abs/1706.03762}

"Can attention attend to nothing / everything equally?"

Before attention-weight dropout, equal allowed scores give a uniform distribution over allowed keys. If one dominates, context approaches that key's value. A valid softmax row can't assign zero weight to every key: weights sum to 1. Values can still cancel to a zero context vector, which is different from attending to nothing.

Practice drill

Keep attention-weight dropout disabled. Change both toy value rows to $[3,-1]$ while keeping Q and K different. Predict the output and query gradient before running the attention and finite-difference cells. Then restore distinct values but permit only one key. Which gradient path remains active in both cases?

The output is independent of routing when all allowed values are identical. Query/key gradients vanish for this isolated attention operation, but changing the values still changes its output. That distinction is more informative than an attention heatmap alone.