Part 3 ended on a square. Two tensors in one decoder block carry T twice, [B, H, T, T]: a T × T square per head, H heads, B sequences. At the widget’s defaults they were 89% of the block’s memory. This post reads what fills them. It’s a weighted average. The average is ordinary; the weights are computed from the input, and that is the new part.
I wrote attention in four lines and checked it against F.scaled_dot_product_attention, torch’s own implementation. Then I read LlamaAttention.forward with the shapes in the margin, and priced it for Llama 3.1 8B at the context lengths people promise.
A dictionary that matches a little
A dictionary lookup is exact. The key matches or it doesn’t, and you get one value back. Attention is the same operation with the match made soft. Every key matches a little, and you get back the average of all the values, weighted by how well each key matched.
Those are attention weights, not parameters. They aren’t stored anywhere; the model computes them from the input on every call. Until the cost section, “weights” means these. Three learned matrices turn each input row into a query, a key, and a value. The query is compared against every key, and the scores become weights:
q, k, v = x @ Wq, x @ Wk, x @ Wv # [B, T, d] eachscores = q @ k.transpose(-1, -2) / sqrt(d) # [B, T, T]weights = scores.softmax(-1) # every row sums to 1out = weights @ v # [B, T, d]Read the shapes. x is Part 3’s [B, T, d]. The three projections keep it. scores is [B, T, T]: one number for every pair of positions, which is the square. The softmax turns each row into weights that sum to one, and the last line is the average. The four lines agree with F.scaled_dot_product_attention to 1e-5.
Two properties follow from “average”. Each output coordinate sits between the smallest and largest value at that coordinate; attention never leaves that range. And if you replace a softmax row with a one-hot row, a single 1 and zeros elsewhere, the output is exactly the value at that position. The soft dictionary contains the hard one as a limit.
Why the square root
The division by sqrt(d) looks like a detail. It’s the difference between a model that trains from the first step and one that starts saturated. The reason is arithmetic you can check in a shell.
A dot product of two random vectors of width d, with unit-variance entries, has variance d. I measured it: 4,104 at d = 4096. Feed scores that wide into a softmax and it saturates: one weight goes to nearly 1 and the rest to nearly 0. With 1,024 keys, the largest weight in a row averaged 0.976. The row had 0.09 bits of entropy out of a possible 10. That’s a hard lookup wearing a soft one’s clothes, at initialisation; trained models sharpen rows on purpose, from a start they can learn from. Divide by sqrt(d) and the variance is 1.0, the largest weight 0.017, the entropy 9.27 bits. The gradient reaching the scores is larger by one to two orders of magnitude: 180 times for a single weight, 20 through the average. A saturated softmax has nothing to learn from; its gradient is zero almost everywhere.
Llama writes it once: self.scaling = self.head_dim**-0.5, at modeling_llama.py:226. The width that matters is the head’s, not the model’s.
Heads are a view, the mask is a triangle
Multi-head attention is the four lines above run once per head, H times, on slices of d. The slicing is free. q.view(B, T, H, d // H).transpose(1, 2) gives [B, H, T, d/H] and shares memory with q. The head axis is a stride. Scores are then [B, H, T, T], one square per head. Head 0 of the batched call equals attention run on head 0 alone. What the heads buy is H different sets of weights over the same positions. Each lives in a slice of width d/H: 128 in Llama 3.1 8B, with 32 heads.
The mask is where a decoder stops being an average over everything. To predict position t, the model may read positions 0..t and nothing after, so the scores above the diagonal are set to -inf before the softmax and their weights come out as zero. I tested it the direct way: attention over eight positions, then rewrite the keys and values at positions 5 to 7. Positions 0 to 4 returned the same output to 1e-6. Without the mask, all eight changed.
T = 12: rows are queries, columns are keys, and a lit cell is a score the softmax will see. Pick a mask and read how many cells stay live.Query t reads keys 0..t. The scores above the diagonal are −inf before the softmax; their weights are zero.
Pick the three masks and read the count. The third chip, the window, is the last lever in the cost section. The triangle composes with Part 3’s padding mask by addition. Left-pad a sequence by two and its first two query rows are all -inf; a row of all -inf is 0/0 in the softmax, NaN. Same trap, one part later, in the four-line form. transformers fills its eager mask with the dtype’s minimum instead (masking_utils.py:602-604); those rows come out uniform. Quiet, and still meaningless.
The eager path computes the exponentials and their sum in fp32 whatever the model’s dtype, then casts the row back: dtype=torch.float32 at modeling_llama.py:208. A normaliser summed in bf16’s eight significant bits is the wrong place to save memory.
Reading LlamaAttention.forward
modeling_llama.py:243-281 in transformers 5.17.0 is the four lines with the heads, the mask, and a cache written out. Condensed, with shapes for the 8B model:
query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) # [B, 32, T, 128]key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) # [B, 8, T, 128]value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) # [B, 8, T, 128]query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)if past_key_values is not None: key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)attn_output, attn_weights = attention_interface(self, query_states, key_states, value_states, attention_mask, scaling=self.scaling)attn_output = self.o_proj(attn_output.reshape(*input_shape, -1))Three things in it aren’t in the four lines. past_key_values.update is the cache. At generation time the keys and values of every earlier token are kept, so each new token computes one query row and reads T stored keys.
The key and value have 8 heads where the query has 32: k_proj maps 4,096 to 1,024. That’s grouped-query attention (GQA), against one key-value head per query head (MHA). It exists for the cache. And attention_interface is a dispatch whose fallback is eager_attention_forward, modeling_llama.py:191-213: the name means the four lines run as written, plus repeat_kv, which copies the 8 key heads to 32 at runtime so the matmul lines up. Set _attn_implementation to sdpa, the scaled_dot_product_attention from the first section, or to flash_attention_2, and a different kernel (the routine that runs the arithmetic) computes the same maths. The rotary line is Part 7’s.
The cost is in the data, not the code
In software you read a function’s cost from its code. A loop over the input is linear; two nested loops are quadratic; you can tell before it runs. Attention is one function; its code is the four lines. There is no loop in them. The square is inside a single @, and nothing in the code’s shape says quadratic. Its cost is set by T in the request: positions squared, per head, per layer.
Take Llama 3.1 8B’s shape. The scores for one sequence in one layer, in bf16, are 268 MB at 2,048 tokens, 4.3 GB at 8,192, and 1.1 TB at the 131,072 the config declares; the eager path holds about three of those squares at its peak. The weights of the whole model are 16 GB. Nobody materialises that last number. The kernel is why.
[B, H, T, T]; the cache is every earlier token's keys and values across all layers. The weights are 16 GB, drawn for scale.Drag T to 131k with eager selected and the scores line passes the weights. Switch to sdpa and the number is struck out: never allocated.
I timed the last three lines on my laptop’s CPU with 8 heads of width 64. From 512 tokens to 8,192, sixteen times the input, the eager form went from 1.9 ms to 284 ms: 150 times. A linear layer over the same input went from 0.23 ms to 2.5 ms, eleven times. Doubling T from 4,096 to 8,192 multiplied the linear layer by 1.95 and attention by 4.3: the square, within timing noise.
Three levers move the number. The first is the kernel. F.scaled_dot_product_attention with its flash or memory-efficient backend computes the same numbers in blocks and never holds the T × T matrix; its math fallback still does. At 4,096 tokens, in fp32, the profiler charged 545 MB to the eager form’s matmuls, the 537 MB square plus its output, and 8.5 MB to sdpa, and it ran 2.4 times faster.
Flash attention is that idea on a GPU. The arithmetic is still quadratic. What changes is what must exist in memory at once. The line that picks it is one config field, _attn_implementation. In 5.17.0 the default at load is sdpa; eager is what you get by asking for it. sdpa still holds the mask itself, [B, 1, T, T] booleans, 17 MB at 4,096.
The other two change what’s computed. Grouped-query attention shrinks the cache. With 8 key-value heads the 8B model stores 131 kB per token across its 32 layers, 17 GB for a 128k context, and with 32 heads it would be 524 kB per token and 69 GB, more than the weights. A sliding window caps how far back a query reads. Llama 3.1 has none; Mistral and Gemma do. transformers writes it as kv_idx > q_idx - sliding_window, at masking_utils.py:100. At 32,768 tokens a 4,096 window leaves 126 million live cells where the causal triangle has 537 million. Below the window the cost is a triangle. Above it, a band of width W. A band grows with T, not its square.
Be fair about the square. Count operations rather than bytes, over the prompt pass, and the linear layers of the 8B model do more of them than attention until roughly 30,000 tokens, about half that if you count only the triangle a causal kernel computes. At 8,192, attention is about a fifth of the block’s floating-point operations, its FLOPs.
The square is a memory problem long before it’s a compute problem. The number to compute before promising a context length is the memory one.
What the ML engineer watches
Four things, and the first is a multiplication.
Price [B, H, T, T] and the cache before promising a context. The widget above is the multiplication: heads times T squared times two bytes, per layer. Add the cache per token times the longest sequence times the batch. Do it at the batch you’ll serve. The config’s max_position_embeddings is what the model accepts after its rotary scaling (131,072 here, from 8,192 trained), not what your card can hold.
Know which kernel runs. eager returns the weights and materialises the square; sdpa and flash don’t. The choice is _attn_implementation, and it decides whether the 545 MB above exists. A debugging tool that wants the weights can’t get them from sdpa: output_attentions=True returns an empty tuple there, and the fix, loading with attn_implementation="eager", materialises the square on every call.
Compose the masks on purpose. Causal, padding, and window are ANDed; added, in the -inf form. A query row with nothing left to read is NaN in that form and uniform in the library’s, and neither is a number you want. Part 3’s habit applies twice now: the last-token index comes from the mask, and the mask travels with the batch.
Measure latency at the longest prompt you’ll accept. The linear layers scale with T. Attention scales with its square. The cache read at each generated token scales with everything already in it. A p99 measured at 500 tokens says nothing about 8,000. The 1.9 ms to 284 ms above is the shape of the surprise.
None of this is modelling work either. A weighted average is four lines, and the weights come from three learned matrices. The square they live in is the first cost in this series that the input sets and the code doesn’t.