Attention

Attention is the core operation that turns a stack of token vectors into a context-sensitive machine. An embedding by itself only says what a token looked like when it entered the model. Attention asks a more useful question: given where we are in the sequence right now, which earlier tokens matter, by how much, and what information should flow forward from them? In the running example What is the capital of France?, the model does not solve the whole sentence at once. It repeatedly builds relevance maps so the final position can look backward and discover that capital and France are especially informative for predicting the next token Paris.

If you want one sentence to keep in your head for the whole chapter, use this: attention is a learned lookup over prior context. It is not a hash table lookup with exact keys, and it is not symbolic reasoning in the classical sense. It is a soft, weighted, high-dimensional search. Every token asks what it is looking for, every earlier token advertises what it contains, and the model blends information from the best matches.

1. The intuition

Historically, attention replaced the older mental model of sequence processing dominated by recurrent neural networks. In an RNN, information about token one had to travel through token two, then token three, then token four, and so on. That serial chain made long-distance dependencies both hard to learn and expensive to compute. Attention breaks that chain. Token seven can talk directly to token two in a single layer. That direct edge is the reason transformers parallelise well during prefill and why they usually preserve distant context better than classical recurrent architectures.

The best analogy is the cocktail party. You stand in a loud room full of overlapping conversations. Most voices are background noise. But if somebody says your name, or mentions the topic you care about, your brain can lock onto that voice. Attention does something similar with vectors. The current token emits a query. Earlier tokens each present a key. Similarity between the query and the keys determines which voices rise above the noise. The winning voices then contribute their value vectors to the updated representation.

attention(i, j) ∝ sim(q_i, k_j)
output_i = Σ_j α_ij v_j

In plain English, the equation says: for the current token i, compute a weight for every candidate source token j, normalize those weights, and use them to blend source information into a new vector. The shapes underneath this intuition are already concrete: you begin with X[T, d_model], project to Q[T, d_head], K[T, d_head], and V[T, d_head] for each head, and end with one updated vector per token.

# Conceptual single-head view
for i in range(T):
    scores = []
    for j in range(T):
        if j <= i:               # causal model
            scores.append(dot(Q[i], K[j]))
        else:
            scores.append(-inf)
    weights = softmax(scores)
    Y[i] = sum(weights[j] * V[j] for j in range(T))

2. Queries, Keys, and Values

The three projections are the vocabulary of attention. A query answers the question what am I looking for right now? A key answers what kind of information do I contain that others might search for? A value answers if selected, what content should I actually contribute? All three come from the same input activations X, but they are projected through different learned matrices, so the model is free to separate matching behavior from payload behavior.

Q = X · Wq
K = X · Wk
V = X · Wv

The immediate translation is mechanical: take the token-state matrix X[T, d_model] and multiply it by three parameter matrices. For one head, Wq, Wk, and Wv usually have shape [d_model, d_head], producing Q[T, d_head], K[T, d_head], and V[T, d_head]. In fused implementations, all three are often computed in one large GEMM and then split, because one bigger matrix multiplication is usually better for hardware than three smaller ones.

For the running example, imagine the token states for What, is, the, capital, of, France, and ? already exist in X. The last position does not contain the answer by itself; it contains a state that says something like I am at the end of a question asking for a capital city. After projection, its query vector is shaped to match keys associated with geography and answer-bearing tokens. Meanwhile the key for France may advertise that it is a country entity, and the key for capital may advertise that the question is about a capital-city relation.

# Fused QKV projection for one layer
X        # [T, d_model]
Wqkv     # [d_model, H_q*d_head + 2*H_kv*d_head]
QKV = X @ Wqkv
Q, K, V = split(QKV, [H_q*d_head, H_kv*d_head, H_kv*d_head])
Q = reshape(Q, [T, H_q, d_head])
K = reshape(K, [T, H_kv, d_head])
V = reshape(V, [T, H_kv, d_head])

3. The dot product score

Once you have queries and keys, attention scoring is just pairwise similarity. For every destination token and every candidate source token, compute a dot product between their vectors. If the directions align strongly, the score is high. If they are orthogonal or opposed, the score is lower. Collecting these pairwise comparisons for all tokens yields the score matrix. In a single head with sequence length T, that matrix has shape [T, T].

S = Q · Kᵀ

The immediate explanation is: take every query row and compare it to every key row. If Q is [T, d_head] and K is [T, d_head], then Kᵀ is [d_head, T], so the product is [T, T]. In multi-head form the shape is [H, T, T]. The entry S[i, j] is the raw relevance score telling us how much token i wants to read from token j before masking and normalization.

Geometrically, the dot product combines direction and magnitude. If two vectors point in similar directions, their dot product is positive and can be large. If they point in unrelated directions, it is near zero. This is why attention can be thought of as a content-based routing mechanism: the current token does not need an explicit address for France; it only needs a query whose direction aligns well with keys representing the country and the requested relation.

Read the last row as the final question token asking which earlier token matters most for answering the prompt. In this toy example it scores France at 5, higher than What or capital. In a real model you would have all prompt positions in the matrix, not just three summarized ones, and every head would produce its own score matrix, often attending to quite different things.

# Score computation for one head
Q  # [T, d_head]
K  # [T, d_head]
S = Q @ transpose(K)   # [T, T]

4. Scaling by √d_k

A raw dot product grows with vector dimension. If query and key entries are roughly zero-mean with unit variance, then summing d_k products produces a value whose variance grows linearly with d_k. Large raw scores are a problem because the next step is softmax. If some scores become too large relative to others, softmax saturates. Saturated softmax makes learning harder during training and makes numeric behavior less pleasant during inference. The standard fix is simple: divide scores by √d_k.

scaled_scores = (Q · Kᵀ) / √d_k

If q_a, k_a ~ N(0, 1), then:
Var(q · k) = d_k
Var((q · k) / √d_k) = 1

The immediate explanation is that dividing by the square root of the head dimension keeps score variance roughly stable as head size changes. If d_k = 128, the raw standard deviation of scores is about 11.3. Passing numbers of that size into exponentials quickly yields extremely peaky distributions. After scaling, typical score magnitudes are pulled back toward a saner range where softmax still differentiates among candidates without instantly collapsing to almost one-hot outputs.

Without scaling, attention would become accidentally sharper as models used wider heads. That is not a desirable semantic effect; it is an artifact of arithmetic growth. The scaling factor makes attention behavior more invariant to representational width. Some newer variants instead normalize queries and keys directly or use cosine-style attention, but the baseline transformer formulation remains the scaled dot-product form.

scores = (Q @ K.transpose(-1, -2)) * (1.0 / sqrt(d_head))

5. Causal masking

Language-model inference is autoregressive. When predicting token i, the model must not peek at tokens i + 1, i + 2, or anything else from the future. During training this rule prevents cheating. During generation it preserves causality by ensuring that the output depends only on the prompt and previously generated tokens. The causal mask enforces this by overwriting illegal future positions with negative infinity before softmax.

masked_scores[i, j] = scores[i, j]   if j ≤ i
masked_scores[i, j] = -∞             if j > i

The shape stays the same as the score matrix: [T, T] per head in prefill, or [1, T] per head in single-token decode. The important structural property is lower triangular visibility. Token zero can only see token zero. Token five can see tokens zero through five. Nobody can see token six until token six actually exists. After the mask, softmax will assign exactly zero probability mass to future positions because exp(-∞) = 0.

# Lower-triangular causal mask for T = 4
[[0, -inf, -inf, -inf],
 [0,    0, -inf, -inf],
 [0,    0,    0, -inf],
 [0,    0,    0,    0]]

Prefill and decode handle masking differently in practice. In prefill, you often build or implicitly apply the full lower-triangular mask because many query positions are active at once. In decode, the query length is usually one, so there are no future positions relative to that new token. The effective mask is then mostly a length boundary on the existing cache rather than a dense triangular matrix you explicitly materialize.

6. Softmax

Raw scores are not yet weights. They can be negative, positive, or on incompatible scales across rows. Softmax converts each row of scores into a probability-like distribution: every weight is non-negative and the row sums to one. That normalization matters because attention output is supposed to be a convex blend of values, not an uncontrolled accumulation whose scale drifts with context length.

α_ij = exp(s_ij) / Σ_j exp(s_ij)

In English: exponentiate every score in the row, then divide by the sum of all exponentials in that row. Large scores become disproportionately influential because exponentials amplify differences. This is why softmax is often called a softmax or soft argmax: it does not choose exactly one source, but it tends to concentrate mass on the best few. For the example row [1, 4, 5], subtracting the row max gives [-4, -1, 0]; exponentiating yields roughly [0.018, 0.368, 1.0]; normalizing produces about [0.013, 0.265, 0.722].

The numerical stability trick is mandatory in real implementations: subtract the maximum score in each row before exponentiating. This does not change the final distribution, because adding or subtracting the same constant from every element cancels out in the normalization. But it prevents overflow when one score is very large. Fused kernels often compute the row max, row sum of exponentials, and normalized outputs in one pass or tile-aware sequence.

row_max = max(scores[i])
exps = exp(scores[i] - row_max)
weights = exps / sum(exps)

7. Weighted value accumulation

After softmax, attention finally performs the actual information transfer. The weight matrix tells each destination token how much of each source value vector to collect. This is the moment where context is aggregated. Up to now the model has only computed match scores. Here it uses those scores to rewrite token states based on what earlier tokens contributed.

Y = A · V

The shape story is simple and worth memorizing. If A is [T, T] and V is [T, d_head], then the output Y is [T, d_head]. In multi-head form you do this independently for every head. Using the example weights [0.013, 0.265, 0.722] and toy value vectors v_What = [0, 0], v_capital = [3, 0], and v_France = [0, 4], the resulting output is approximately [0.795, 2.888]. That vector is not literally the word Paris. It is a context-enriched latent feature vector that now strongly reflects the relevant source tokens.

This distinction between keys and values is important. Keys decide who wins the lookup; values decide what content is delivered once they win. That separation gives the model flexibility. A token can be easy to match on one feature yet contribute payload emphasizing some other feature. In hardware terms, this step is another dense multiply, but semantically it is the moment where information actually moves across positions.

# Blend values with attention weights
A  # [T, T]
V  # [T, d_head]
Y = A @ V   # [T, d_head]

8. Multi-head attention

A single attention head can learn one family of relationships, but language has many kinds of structure happening at once. Some dependencies are local syntax. Some are entity references. Some are long-range topic cues. Some are formatting patterns. Multi-head attention lets the model learn several distinct query-key-value subspaces in parallel. Rather than one head of width d_model, the model uses H heads of width d_head such that H · d_head = d_model in the standard case.

head_h = Attention(Q_h, K_h, V_h)
MultiHead(Q, K, V) = Concat(head_1, …, head_H) · Wo

The tensor shapes are the operational truth. Suppose d_model = 4096, H = 32, and therefore d_head = 128. Then queries are shaped like Q[32, T, 128]. Each head independently builds scores [T, T], applies masking and softmax, and produces an output [T, 128]. The outputs from all heads are concatenated into [T, 4096] and then multiplied by W_o[4096, 4096] to mix head results back into model space.

In the running example, one head might focus on entity identity and heavily attend from the final token to France. Another might focus on the relation implied by capital. Another might mostly care about punctuation and question structure. The important thing is not that every head corresponds cleanly to a human concept; often they do not. The important thing is that the architecture gives the model parallel channels for distinct matching patterns.

# Multi-head attention shapes
Q = reshape(X @ Wq, [T, H, d_head]).transpose(1, 0, 2)   # [H, T, d_head]
K = reshape(X @ Wk, [T, H, d_head]).transpose(1, 0, 2)   # [H, T, d_head]
V = reshape(X @ Wv, [T, H, d_head]).transpose(1, 0, 2)   # [H, T, d_head]

for h in range(H):
    Y[h] = attention(Q[h], K[h], V[h])                    # [T, d_head]

Y = transpose(Y, [1, 0, 2]).reshape([T, H * d_head])     # [T, d_model]
O = Y @ Wo                                                # [T, d_model]

9. Multi-Query Attention (MQA)

Classic multi-head attention gives every head its own queries, keys, and values. That is expressive but expensive at decode time because the KV cache grows with the number of K/V heads. Multi-query attention changes one design choice: every query head keeps its own Q, but all heads share the same K and V. In notation, H_q can be large while H_kv = 1.

Q ∈ [H_q, T, d_head]
K ∈ [1,   T, d_head]
V ∈ [1,   T, d_head]

The immediate consequence is dramatic memory reduction. If a model had 32 query heads and switched from full multi-head attention to MQA, its per-layer K/V cache would shrink by roughly 32× compared with storing distinct keys and values for all heads. The trade-off is that different query heads no longer get head-specific source representations. They can ask different questions, but they all read from the same key-value inventory.

Why is this attractive? Because decode performance is dominated less by arithmetic and more by reading the cache. Shared keys and values mean less memory to store, less memory to move, and better batch density across many concurrent requests. The quality loss is often modest enough to be worth it, especially in serving systems where long contexts and large batches make cache pressure painful.

# MQA decode sketch
Q_new = project_q(x_new)               # [H_q, 1, d_head]
K_new = project_k(x_new)               # [1,   1, d_head]
V_new = project_v(x_new)               # [1,   1, d_head]
append_to_cache(K_new, V_new)

for h in range(H_q):
    scores[h] = Q_new[h] @ K_cache[0].T
    out[h] = softmax(scores[h]) @ V_cache[0]

10. Grouped Query Attention (GQA)

Grouped query attention is the compromise between classic multi-head attention and MQA. Instead of giving every query head its own key and value, or forcing all query heads to share one global key and value set, GQA groups query heads so each group shares a K/V head. This keeps most of the cache savings of MQA while recovering much of the quality of full multi-head attention.

group_size = H_q / H_kv
kv_head = floor(q_head / group_size)

A concrete example makes this intuitive. Llama-2 70B uses 64 query heads but only 8 K/V heads. That means each K/V head serves a group of 8 query heads. The query tensor is Q[64, T, 128], while the key and value tensors are K[8, T, 128] and V[8, T, 128]. During attention, query head 17 might reuse key-value head 2 because 17 // 8 = 2.

The implementation detail that matters most is consistency of head mapping. You either physically store fewer K/V heads and broadcast them logically at use time, or you materialize an expanded view without copying. Good kernels do the former. The arithmetic behavior is almost the same as MHA from the query side, but the cache footprint scales with H_kv, not H_q. That is why GQA is now common in practical inference models.

# GQA mapping
group_size = H_q // H_kv
for qh in range(H_q):
    kvh = qh // group_size
    scores[qh] = Q[qh] @ K[kvh].T
    out[qh] = softmax(scores[qh]) @ V[kvh]

11. The KV cache

Now we reach the state object that dominates real autoregressive inference. During decode, only one new token arrives at a time. The query for that new token must be computed fresh, but the keys and values for all earlier tokens are unchanged. Recomputing the old K and V tensors on every step would be absurdly wasteful, so implementations store them in a per-layer cache. When token t + 1 arrives, the layer computes its new K and V, appends them to the cache, and attends over the entire history already stored there.

KV cache shape per layer:
K_cache[max_seq, H_kv, d_head]
V_cache[max_seq, H_kv, d_head]

Full model conceptual shape:
[num_layers, max_seq, num_kv_heads, d_head]

The memory math is brutal enough that every systems engineer should do it by hand at least once. Suppose a model has 32 layers, context length 8192, 8 K/V heads, head width 128, and uses bf16 storage at 2 bytes per number. Because keys and values are both stored, the cache size per sequence is:

32 layers × 8192 tokens × 8 heads × 128 dims × 2 tensors × 2 bytes
= 1,073,741,824 bytes
≈ 1 GiB per sequence

If the same model used full 64 K/V heads instead of 8, that single-sequence cache would jump to about 8 GiB. This is why MQA and GQA matter so much. At long context, the cache, not the weights, often becomes the per-request memory bottleneck. The arithmetic for one decode token may be modest, but the system still has to drag enormous history tensors through memory every step.

# Decode-time cache update
for layer in range(num_layers):
    q = project_q(x_new)                    # fresh every step
    k = project_k(x_new)                    # one new row
    v = project_v(x_new)                    # one new row
    K_cache[layer, pos] = k
    V_cache[layer, pos] = v
    y = attend(q, K_cache[layer, :pos+1], V_cache[layer, :pos+1])

Cache management strategies differ by serving stack. A simple implementation allocates a contiguous buffer per request per layer. Sliding-window models can use ring buffers because old tokens eventually fall out of visibility. High-throughput engines often use paged or block-based allocators so many variable-length sequences can share device memory efficiently. No matter the policy, correctness depends on associating every cached row with the exact layer, logical position, head group, and positional encoding state that produced it.

12. Prefill vs decode

People often say transformer inference as if it were one thing, but two very different regimes exist inside that phrase. Prefill processes the initial prompt, often many tokens at once, and populates the KV cache. Decode then processes one new token at a time, attending against the full cached history. The same layer equations apply in both modes, yet the performance profile changes dramatically because the tensor shapes change.

Prefill attention cost per layer ≈ O(T² · d_head · H_q)
Decode attention cost per new token per layer ≈ O(T · d_head · H_q)

The asymptotics already hint at the difference. Prefill is full-matrix work. Large prompt chunks mean large GEMMs, which accelerators handle well. Decode is a narrow query against a long history. Arithmetic per step is much smaller, but the system must read the entire visible cache for every layer and every generated token. That is why prefill tokens per second and decode tokens per second behave so differently on the same hardware. One is a throughput problem dominated by dense compute; the other is a latency and bandwidth problem dominated by repeatedly rereading history.

This distinction affects batching strategy, scheduler design, and kernel choice. Prefill wants big contiguous work and benefits from chunking large prompts together. Decode wants careful continuous batching across many requests so the hardware stays occupied while each request only contributes a tiny query. It is common for a serving engine to treat these modes almost like separate workloads sharing the same weights.

# Prefill
X_prompt = embed(prompt_tokens)
for layer in layers:
    X_prompt, K_cache[layer], V_cache[layer] = attention_prefill(X_prompt)

# Decode
while not stop:
    x_new = embed(last_token)
    for layer in layers:
        x_new = attention_decode(x_new, K_cache[layer], V_cache[layer])
    last_token = sample(x_new)

13. FlashAttention

Naive attention has an ugly memory habit: it often materializes the full score matrix [T, T], sometimes the masked matrix, and sometimes the post-softmax weight matrix too. At long context, these intermediates are huge and expensive to move to and from high-bandwidth memory. FlashAttention attacks this problem not by changing the mathematical result, but by changing the execution schedule. It tiles the computation so the full attention matrix never has to exist in device memory at once.

Instead of:
S = QKᵀ
A = softmax(S)
Y = AV

FlashAttention computes tiled blocks and maintains
running row max m_i and running row sum l_i
to produce the same Y without materializing S or A globally.

The crucial trick is online softmax. Softmax usually needs the whole row to know the row max and denominator. FlashAttention processes the row in chunks but keeps enough running statistics to combine chunks exactly. For each tile it updates the running maximum and running normalization term, rescales partial outputs when a new tile changes the maximum, and accumulates the final output directly. The result is mathematically equivalent to standard attention up to floating-point details, yet dramatically lower in memory traffic.

Why can it be faster even if the arithmetic count is similar or slightly higher? Because modern accelerators are often bottlenecked by memory movement rather than raw multiply-add throughput. Moving a giant score matrix out to HBM and then reading it back is expensive. Recomputing or rescaling inside on-chip SRAM can be cheaper than storing and reloading massive intermediates. That is the essence of the IO complexity story.

# Highly simplified FlashAttention sketch
for q_tile in tiles(Q):
    m = -inf
    l = 0
    o = 0
    for k_tile, v_tile in zip(tiles(K), tiles(V)):
        s = q_tile @ k_tile.T
        s = apply_mask_and_scale(s)
        m_new = max(m, rowmax(s))
        l = exp(m - m_new) * l + sum(exp(s - m_new))
        o = exp(m - m_new) * o + exp(s - m_new) @ v_tile
        m = m_new
    output_tile = o / l

FlashAttention matters most in prefill, where the score matrix is widest and densest. Decode already has query length one, so the specific benefit of not materializing a huge square matrix is smaller, though related flash-decoding kernels still matter. Variants such as FlashAttention-2 improve work partitioning and GPU utilization. The conceptual lesson is broader than the specific kernel family: in transformer inference, the schedule of memory movement can matter more than the nominal FLOP count.

14. Sliding window attention

Not every task needs every token to see the entire past. Many dependencies are local, and recent context often matters more than ancient context. Sliding window attention exploits that by limiting visibility to the most recent W tokens. Instead of attending over all prior positions 0..i, token i attends only over max(0, i - W + 1)..i. Mistral-style models use this idea successfully.

visible(i) = { j | max(0, i - W + 1) ≤ j ≤ i }

The shape implication is straightforward. Instead of a score matrix [T, T] per head in prefill, the active neighborhood is conceptually [T, W]. During decode, instead of reading the entire cache of length T, the layer only needs the most recent W rows. If the implementation truly evicts or overwrites older cache entries, memory usage becomes O(W) rather than O(T) for the active state.

Why does this often work? Because many language tasks rely heavily on recent syntax, nearby discourse, and the latest instructions. The model can still keep long-range information indirectly if that information is propagated forward layer by layer or summarized into recent tokens. The trade-off is obvious: you save memory and improve throughput, but hard long-distance retrieval becomes less reliable unless the architecture adds some other mechanism such as periodic global tokens or retrieval augmentation.

# Sliding-window decode read
start = max(0, pos - window + 1)
K_view = K_cache[start:pos+1]
V_view = V_cache[start:pos+1]
scores = q @ K_view.T
weights = softmax(scores)
out = weights @ V_view

15. Paged attention

Once many requests of different lengths share one accelerator, cache management becomes an allocator problem. Simple contiguous per-sequence buffers waste memory because real requests finish at different times and grow at different rates. Paged attention, popularized by systems such as vLLM, treats the KV cache more like virtual memory. Instead of requiring one long contiguous region per request, it stores K/V rows in fixed-size blocks and keeps a mapping from each sequence's logical token positions to physical blocks in device memory.

physical_blocks[num_blocks, block_size, H_kv, d_head]
page_table[sequence_id][logical_block] -> physical_block_id

The operating-system analogy is good and worth keeping. Logical position 5000 in a sequence does not have to live next to logical position 4999 in memory if both positions can be located through an indirection table. This means the serving engine can allocate, free, and reuse fixed-size cache pages as requests arrive and complete. Fragmentation drops, continuous batching becomes easier, and the accelerator can stay busier because more variable-length requests fit into the same memory budget.

There is, of course, an extra level of address calculation. The engine must translate logical positions into physical blocks before reading K/V. But that indirection cost is usually far smaller than the memory waste avoided by abandoning contiguous reservation. Paged attention pairs especially well with continuous batching, speculative execution, and other serving strategies where requests are frequently entering, leaving, or being reordered.

# Conceptual paged cache lookup
logical_block = pos // block_size
offset = pos % block_size
physical_block = page_table[seq_id][logical_block]
k = K_blocks[physical_block, offset]
v = V_blocks[physical_block, offset]

16. Putting it all together

By now the pieces should feel less mystical. Attention starts with token states, projects them into queries, keys, and values, scores every allowed source token against every destination token, scales, masks, normalizes, blends values, and projects the multi-head result back into model space. The mathematics is compact; the systems consequences are not. Every design choice around heads, cache layout, masking, batching, and kernel fusion changes real serving behavior.

# One complete causal self-attention layer, GQA form
# Inputs:
#   X           [T_q, d_model]              prefill: T_q = T, decode: T_q = 1
#   K_cache     [T_k, H_kv, d_head]
#   V_cache     [T_k, H_kv, d_head]
# Parameters:
#   Wq          [d_model, H_q * d_head]
#   Wk          [d_model, H_kv * d_head]
#   Wv          [d_model, H_kv * d_head]
#   Wo          [H_q * d_head, d_model]

Q = reshape(X @ Wq, [T_q, H_q, d_head])
K = reshape(X @ Wk, [T_q, H_kv, d_head])
V = reshape(X @ Wv, [T_q, H_kv, d_head])

Q = apply_rope(Q, positions_q)
K = apply_rope(K, positions_q)

append(K_cache, K)
append(V_cache, V)

K_all = view(K_cache, visible_positions)              # [T_k_total, H_kv, d_head]
V_all = view(V_cache, visible_positions)              # [T_k_total, H_kv, d_head]

for qh in range(H_q):
    kvh = qh // (H_q // H_kv)
    scores = Q[:, qh, :] @ transpose(K_all[:, kvh, :])    # [T_q, T_k_total]
    scores *= 1 / sqrt(d_head)
    scores = apply_causal_or_window_mask(scores)
    weights = softmax(scores)
    context[:, qh, :] = weights @ V_all[:, kvh, :]         # [T_q, d_head]

Y = reshape(context, [T_q, H_q * d_head])
O = Y @ Wo                                                 # [T_q, d_model]
return O, K_cache, V_cache

Notice how many implementation details are hidden inside seemingly small verbs. append means managing persistent cache memory. view might mean following a page table or ring-buffer offset. apply mask may encode full causal visibility, sliding windows, or paged lookup boundaries. softmax may be a fused online kernel. The pseudocode is short because the abstraction is clean, not because the engineering is trivial.

X[T_q, d_model]
  ↓ linear projections
Q[T_q, H_q, d_head]   K_new[T_q, H_kv, d_head]   V_new[T_q, H_kv, d_head]
  ↓ RoPE               ↓ RoPE                     ↓
Q_rot                  K_rot                      V_new
  ↓                    ↓ append to K cache        ↓ append to V cache
  └───────────────visible K/V view from cache───────────────┘
                    ↓
scores[H_q, T_q, T_visible]
  ↓ scale + mask
masked scores
  ↓ softmax
weights[H_q, T_q, T_visible]
  ↓ weighted sum with V
context[H_q, T_q, d_head]
  ↓ transpose + concat
[T_q, d_model]
  ↓ Wo
O[T_q, d_model]

That total is for prefill at a long prompt length. Decode looks very different. For one new token at sequence length 4096, the Q projection is small, the K/V projections are tiny, and the QK plus AV work per layer is only on the order of tens of millions of floating-point operations. Yet the system still must read thousands of cached rows for each layer, and that memory traffic often dominates latency. This is why real transformer inference is not just about FLOPs. It is about the shape of the state you move through memory step after step.

When a model predicts Paris after the prompt What is the capital of France?, attention is the mechanism that made the prompt coherent at the decision point. Some heads learned to route from the final position toward France. Some learned to route toward capital. Some maintained the broader question structure. Together they transformed the last token's hidden state into one that strongly supports the right output token in vocabulary space. The prediction feels holistic, but the execution is local, repeated, and exact: project, compare, normalize, blend, cache, repeat.