Scaled Dot-Product Attention

Attention is a lookup in which nothing is looked up. Every position asks a question, every position advertises an answer, and instead of picking the best match the model keeps all of them and takes a weighted average. That single design choice is why a transformer trains in parallel, why processing a prompt costs the square of its length, and why generating the next token after it does not.

Three vectors come out of each position, produced by three different learned matrices applied to the same input. The query is what this position is looking for. The key is what this position offers to anyone looking. The value is what actually gets copied if someone looks. The query and the key are separate projections of the same token, which is the part that surprises people: a token advertises something different from what it asks for, so attention from position 3 to position 1 is not the same number as attention from 1 to 3.

The operation itself is four steps, and Vaswani et al. wrote it on one line in 2017:

Attention(Q, K, V) = softmax( Q K^T / √d_k ) V

Read right to left. Q KT is one dot product for every (query, key) pair — a raw score saying how well position i's question matches position j's advertisement. Dividing by √dk, the square root of the key vector's length, keeps those scores in a range the softmax can work with. The softmax turns each row of scores into weights that are positive and sum to 1. Multiplying by V averages the value vectors using those weights. The output for position i is a blend of every value it was allowed to see.

Below is one attention head with real numbers in it. Eight words, each a real 8-dimensional vector; three real projection matrices; real dot products; a real softmax. Nothing is drawn from a formula — the heatmap is the arithmetic. Start by moving inspect query to position 2 and reading the panel underneath it, then turn the causal mask off.

With the defaults, row 2 — the query from sat — puts 0.77 of its weight on cat and 0.07 on the. The verb found its subject. It did not find it by searching; it computed three dot products, exponentiated them, and divided by their sum. That is the whole mechanism, and the reason it is called soft: the did not lose, it got 7%.

Now switch the causal mask off. Row 2 hands 0.23 of its weight to mat, a word three positions in its future. Row 0 and row 4 — both the token the — become numerically identical, because without the mask the two positions see exactly the same set of tokens and attention has no notion of order to tell them apart. Order enters a transformer through positional encoding, not through this operation.

What the causal mask actually masks

The mask sets a score to negative infinity before the softmax, not after it. That ordering is the whole trick. exp(−∞) is 0, so the masked position contributes nothing to the normalising sum, and the surviving weights still add to exactly 1. Zeroing weights after the softmax would leave a row summing to less than 1 and quietly rescale every output. The 2017 paper says it in one clause: they implement it "inside of scaled dot-product attention by masking out (setting to −∞) all values in the input of the softmax which correspond to illegal connections."

What that buys is not privacy and not speed. It is this: row i of a masked pass over n tokens is bit-for-bit the row you would get from a pass over only the first i + 1 tokens. The simulation checks that claim on every render — the truncate line in the arithmetic panel recomputes the inspected row against a sequence cut off at that position and compares the two. With the mask on they agree. Turn the mask off and the same line reports a mismatch and prints the largest disagreement.

That equality is why a transformer can be trained on a whole sentence in one forward pass and still be a next-token predictor, and it is why prefill — pushing an entire prompt through in a single pass — produces exactly the state that n sequential steps would have produced. Without it you would have to run the model once per token during training, and there would be no transformer.

Two things the mask does not do. It does not save arithmetic in a naive implementation: the score above the diagonal is computed, then overwritten with −∞, then exponentiated to zero. The readout labelled pairs computed then discarded counts exactly those. Put the causal mask back on — the paragraph above turned it off — and the sequence length to 8: it is 28 of 64 pairs, thrown away after being paid for. With the mask off the same readout is 0 of 64, because nothing is being discarded. Kernels that know the mask is triangular skip whole blocks above the diagonal and do get the saving; that is one of the things FlashAttention does, alongside never writing the full square of scores — one number for every ordered pair of positions — to memory at all.

And it does not apply during decode. Switch the phase control to decode: the mask toggle stops changing anything, and the log says why. The new token is at the end of the sequence, so every key that exists is a key it is allowed to see. There is nothing in its future to hide. This is the answer to a question people keep asking about generation code — the causal mask is a prefill and training construct, and a single-token decode step passes no mask at all.

Where the n² actually lives

Watch the hero readout as you drag sequence length from 3 to 8 in prefill: dot products go 9, 16, 25, 36, 49, 64. Now switch to decode and drag it again: 3, 4, 5, 6, 7, 8. Same model, same mask, same context. The quadratic term is a property of how many queries you run at once, not of the context length.

Write n for the number of tokens in the sequence. Prefill runs n queries against n keys, so it does n² dot products. A decode step runs one query against n keys, so it does n. The KV cache — the keys and values of every earlier position, kept in memory instead of recomputed — is what makes the second number possible. Positions 0 to n−2 were projected on earlier steps and their keys and values were retained, so the new token only has to produce its own.

Generating a whole response is still quadratic in total — the steps cost n, n+1, n+2, … and that sum grows with the square of the length. But it is a sum of linear steps, and that distinction is what people are measuring when they report that their decode looks linear. It does, because each step is.

The second thing people get wrong is how much of the bill attention is. The panel below counts floating-point operations for a real model, splitting them into the attention scores plus the weighted sum of values, and everything else — the query, key, value and output projections, the feed-forward network and the output layer, all of which are linear in n.

At a 4,096-token prompt on Qwen3-8B, attention is 7.4% of prefill arithmetic. Everything else — the projections and the feed-forward network — is the other 93%. Push the prompt to 32,768 and attention reaches 39%; at 131,072 it is 72%. The crossover, where the quadratic term finally equals everything else, sits at 51,324 tokens for a kernel that skips the masked blocks, and half that — 25,662 — for a naive one that computes them. Below that length, "attention is quadratic" is true and irrelevant.

The crossover is not a constant of nature. It is matmul_params ÷ (2 × layers × heads × head_dim), doubled when the kernel is triangular, so it moves with the shape of the model. Select GPT-2 small: the crossover falls to 13,404 tokens, a factor of 3.8, because a 124M-parameter model has very little "everything else" to hide the quadratic term behind. But GPT-2's maximum position count is 1,024, where attention is 7.1% of prefill. The quadratic term was in the architecture from the beginning and was simply never reachable.

The boundary. The reason none of this settles the performance question is that floating-point operations are not time. Put the model back to Qwen3-8B — the paragraph above left you on GPT-2 — set the prompt to 256 and generate 4,096 tokens: attention is 0.50% of prefill and 0.99% of a decode step, under 1% in both phases, and yet that request is dominated by decode, which is slow for a reason this panel cannot show — every step re-reads the whole weight set and the whole cache from memory. The KV cache lesson is where that argument lives. Attention arithmetic and attention bytes are different constraints, and the second one usually wins.

Why the square root, and what happens without it

Put the first simulation back to prefill, sequence length 6, inspect query 2, causal mask on — the sequence-length ladder above left it in decode at 8 — and then turn off divide scores by √d_k. Row 2's weight on cat goes from 0.77 to 0.97, and effective tokens attended falls from 1.98 to 1.17. The distribution collapsed toward putting everything on one token. In a 5-dimensional head that is a mild effect, because removing the divisor multiplies every score by only √5 ≈ 2.24.

In a real head it is not mild. The paper's argument is about variance: if the components of q and k are independent with mean 0 and variance 1, then q · k is a sum of d_k such products, so it has mean 0 and variance d_k — standard deviation √d_k. At the head_dim of 128 that Qwen3 and Llama both use, unscaled scores swing about ±34 across three standard deviations. The ratio between the largest and smallest exponential in that range is e⁶⁸, roughly 10²⁹. Softmax of that is one-hot to within rounding, its gradient is zero everywhere, and training stops. The paper's words for it are that the dot products "grow large in magnitude, pushing the softmax function into regions where it has extremely small gradients."

The panel below is that argument as arithmetic. It draws a query and n keys with independent unit-variance components from a fixed pseudo-random generator — the paper's assumption, literally — computes the real dot products, and softmaxes them with and without the divisor. Drag d_k and watch the two rows separate.

Over 16 keys, the unscaled softmax spreads across 6.68 effective keys at d_k = 4, 2.84 at 16, 1.44 at 128 and 1.19 at 512. It is collapsing onto one key, monotonically, as the dimension rises. The scaled row over the same scores gives 11.19, 10.96, 10.71 and 10.76 — it does not move. That flatness is the result worth taking away: dividing by √d_k makes the sharpness of attention independent of the head dimension, so the architect can widen a head without also changing how peaked its attention is. Nothing about the vectors changed across that sweep except how many components the dot product sums over.

The measured standard deviation of the scores tracks √d_k the whole way — 1.82 against 2.00 at d_k = 4, 10.72 against 11.31 at 128, 21.48 against 22.63 at 512. It runs a few percent low because it is estimated from only 16 samples per draw, not because the theory is off.

This also explains a failure mode that looks like a bug elsewhere. A head whose scores have drifted large produces attention that is effectively a hard argmax, so gradients vanish for every key but one and the head stops learning. The fix has never been to clip the weights; it is to keep the scores in range, which is why the divisor is inside the operation rather than a training trick bolted on outside it.

Checking it on a real system

The numbers in the cost panel come from four fields of a config.json, and you can compute them for your own model in a minute. Attention arithmetic per (query, key) pair is 4 × num_hidden_layers × num_attention_heads × head_dim. For Qwen3-8B that is 4 × 36 × 32 × 128 = 589,824 floating-point operations. Multiply by n² for a prefill and by n for one decode step. Note it is num_attention_heads here, the query heads — unlike the KV cache formula, which uses num_key_value_heads. Under grouped-query attention the keys are shared across query heads but every query head still does its own dot products, so grouped-query attention shrinks the cache without reducing this term at all.

To see the split in a running server rather than on paper, read time to first token against inter-token latency as you vary prompt length alone. If doubling the prompt from 2k to 4k roughly doubles time to first token, prefill is still in its linear regime and attention is a minority of the work. When the same doubling starts to cost closer to four times, you have crossed into the quadratic regime and are past the number the panel calls the crossover. In vLLM that transition is also where chunked prefill stops being free, because the chunk that gets folded into a decode step now carries real arithmetic.

In a PyTorch model, the concrete thing to check is which kernel you are getting. torch.nn.functional.scaled_dot_product_attention has three implementations behind one call — FlashAttention-2, a memory-efficient kernel, and the one the docs call "the PyTorch C++ implementation", which the backend enum names SDPBackend.MATH — and it "attempts to automatically select the most optimal implementation based on the inputs." Each fused kernel has input limitations, and when none of them qualifies you get the math one, which builds the whole n × n score matrix in memory. PyTorch does warn: the docs promise that "in the event that a fused implementation is not available, a warning will be raised with the reasons why the fused implementation cannot run." That warning is emitted once and is easy to lose in a training log, which is why it is worth reading deliberately.

Two specifics worth knowing before you write the call. Setting both attn_mask and is_causal=True raises an error rather than combining them, so a causal model passes one or the other. Passing your own triangular tensor where is_causal=True would have done is the common way to lose the fast path, and it costs the block-skipping saving as well, because a generic additive mask tells the kernel nothing about which blocks are empty. If you need a guarantee rather than a preference, wrap the call in torch.nn.attention.sdpa_kernel with the math backend disabled — the documentation's own recommendation — which converts a silent degradation into an exception.

If you want to look at the weights themselves you have to give up the fused kernels entirely, because they never build the matrix you are asking to see. In Hugging Face transformers that means loading with attn_implementation="eager". Understand what you are buying: the score matrix is batch × heads × n² × bytes per layer, which at a 4,096-token prompt on a 32-head model is 1 GiB per layer in bfloat16 for a single sequence — and the eager path upcasts the softmax to float32, so the peak is larger than that. Ask for attention weights on a long prompt and you will reproduce, exactly, the quadratic memory problem that FlashAttention exists to remove.

Next: the cache that makes a decode step linear, the KV cache; the kernel that computes this operation without ever writing the matrix down, FlashAttention; and the same head run many times over slices of one vector, multi-head attention. For a lookup that does the opposite of this one — hard top-k, everything else discarded — see vector retrieval. For the architectures that give up the n × n matrix entirely, linear attention.