Long context and efficient architectures

Attention (see primer.ml.attention) lets every token look at every earlier token. That is where its power comes from, and it is also where its bill comes from. The bill has two lines:

• Compute. Every token is scored against every other token, so the work grows with the square of the context length n.
• Memory. During generation the model keeps every earlier token's keys and values (the KV cache, see primer.ml.inference), so memory grows with n, per layer, per conversation.

Everything in this lesson attacks one of those two lines, in one of three ways:

1. Look at fewer tokens: sliding-window and sparse attention.
2. Summarize the past into a fixed-size state: linear attention and state-space models (Mamba), which run like a recurrent network.
3. Store the past more compactly: fewer key/value heads, a small latent vector per token, fewer bits per number.

flowchart TB
  P["Long context is expensive<br/>compute grows with n², memory with n"] --> F["Look at fewer tokens"]
  P --> S["Keep a fixed-size summary"]
  P --> C["Store the past compactly"]
  F --> F1["sliding window"] & F2["sparse: local + global, strided"]
  S --> S1["linear attention"] & S2["state-space models, Mamba"]
  C --> C1["GQA / MQA"] & C2["latent KV"] & C3["quantized KV"]
  S2 --> H["hybrids: a few attention layers<br/>among many SSM layers"]
  F1 --> H

Everyday picture A dinner party where every new guest must shake hands with every guest already there. Ten guests make 45 handshakes; a thousand guests make half a million. On top of that, the coat check keeps one coat per guest on every floor of the building, so the coat racks grow with every arrival. Attention pays both bills: the handshakes are the query-key scores, and the coats are the KV cache.

Tiny worked example Take a Llama-3-8B-shaped model: 32 layers, 8 key/value heads, 128 numbers per head, 16-bit numbers (2 bytes). One token costs 2 × 32 × 8 × 128 × 2 = 131,072 bytes of cache, about 128 KB.

Going from 8k to 1M is 128 times more tokens, 16,384 times more pairs, and a cache that no longer fits on one 80 GB GPU.

flowchart LR
  T["new token n"] --> Q["its query"]
  Q -->|"scored against"| KC[("KV cache<br/>n − 1 keys and values<br/>per layer")]
  KC --> O["output for token n"]
  T -->|"its key and value<br/>are appended"| KC

L, H_kv, d_h, b = 32, 8, 128, 2
# bytes per token: a key and a value, per layer, per KV head
per_token = 2 * L * H_kv * d_h * b
per_token  # → 131072
contexts = [8_192, 131_072, 1_048_576]
# pairs per head per layer
[f"{n * n:,}" for n in contexts]  # → ['67,108,864', '17,179,869,184', '1,099,511,627,776']
# GB of cache for one conversation
[round(n * per_token / 1e9, 1) for n in contexts]  # → [1.1, 17.2, 137.4]

Why it matters in practice. The compute bill is paid once per prompt (prefill), and exact tricks such as FlashAttention make it faster without changing it. The memory bill is paid for as long as a conversation is open, and it decides how many users one GPU can serve. That is why most of the techniques below target the cache.

In code: context_cost returns both lines of the bill for any context length, using primer.ml.inference.kv_cache_bytes for the cache.

Everyday picture You read a long scroll through a letterbox slot that shows only the last w words. On your own you would lose the beginning. But suppose a row of readers sits one above the other, each reading the notes of the reader below through a slot of the same width. Each reader passes things along a little further, like a bucket brigade, so a stack of readers can carry a fact much further back than any one slot shows.

Tiny worked example Eight tokens, window w = 3: each token reads itself and the two tokens before it.

            key:  0 1 2 3 4 5 6 7
query 0           ✓ · · · · · · ·
query 1           ✓ ✓ · · · · · ·
query 2           ✓ ✓ ✓ · · · · ·
query 3           · ✓ ✓ ✓ · · · ·
query 4           · · ✓ ✓ ✓ · · ·
query 5           · · · ✓ ✓ ✓ · ·
query 6           · · · · ✓ ✓ ✓ ·
query 7           · · · · · ✓ ✓ ✓

That is 1 + 2 + 3 × 6 = 21 scores instead of the 36 a full causal mask needs. Now stack three such layers. Token 7 reads token 5 in the top layer; token 5 had read token 3 in the layer below; token 3 had read token 1 in the layer below that. So after three layers, token 1 can influence token 7, six positions back, but token 0 cannot.

flowchart RL
  subgraph L3["layer 3"]
    a7["token 7"]
  end
  subgraph L2["layer 2"]
    b5["token 5"]
  end
  subgraph L1["layer 1"]
    c3["token 3"]
  end
  subgraph IN["input"]
    d1["token 1"]
    d0["token 0: out of reach"]
  end
  a7 -->|"reads 2 back"| b5 -->|"reads 2 back"| c3 -->|"reads 2 back"| d1

n, w = 8, 3
# M_ij for the row of token 5
[int(0 <= 5 - j < w) for j in range(n)]  # → [0, 0, 0, 1, 1, 1, 0, 0]
# pairs scored: row i keeps min(i + 1, w) keys
sum(min(i + 1, w) for i in range(n))  # → 21
# full causal attention, for comparison
n * (n + 1) // 2  # → 36
# reach after L layers
L = 3
L * (w - 1)  # → 6
# Mistral 7B: 32 layers, a 4,096-token window
32 * (4096 - 1)  # → 131040

In code: sliding_window_mask builds the band, pairs_computed counts its cells, receptive_field applies the mask layer after layer to find who can reach whom, and reach is the L(w − 1) formula.

The rolling buffer: a cache that stops growing

Everyday picture A whiteboard with room for exactly w notes. When it is full, the next note goes over the oldest one. You never need a bigger board, however long the meeting runs.

Tiny worked example With w = 4, after token 9 the buffer holds the keys and values of tokens 6, 7, 8 and 9. Token 10 overwrites token 6's slot. On the running 8B example at 128k tokens, a 4,096-token window holds 4,096 × 128 KB = 0.54 GB instead of 17.2 GB: 32 times less, and the same at 1M tokens.

flowchart LR
  subgraph B["buffer of w = 4 slots, after token 9"]
    s0["slot 0: token 8"]
    s1["slot 1: token 9"]
    s2["slot 2: token 6 (oldest)"]
    s3["slot 3: token 7"]
  end
  N["token 10 arrives"] -->|"overwrites the oldest"| s2
  B --> A["token 10 attends to<br/>tokens 7, 8, 9, 10"]

Why it matters in practice. Mistral 7B combined the window with this rolling buffer. Several later model families interleave sliding-window layers with a few full-attention layers: the local layers are cheap, and the full layers keep long-range lookups exact. The weakness is plain from the diagram: a fact beyond the reach is invisible, and a fact inside the reach has to survive several hops.

In code: sliding_window_decode generates token by token with a buffer that never holds more than w keys, and returns the same outputs as masked attention with sliding_window_mask.

Everyday picture An open-plan office. You mostly talk to the people at the desks next to yours (local). Anyone can phone the front desk, and the front desk hears from everyone (a global token): any two people are at most two calls apart. Another design is the express train: you talk to your neighbours and also to every fourth desk down the row (strided), so a message can travel far in a few big jumps.

Tiny worked example Sixteen tokens, window 4.

• The window alone: 1 + 2 + 3 + 4 × 13 = 58 pairs.
• Make token 0 global: the 12 rows past the window (tokens 4 to 15) add a pair each: 70 pairs, against 136 for full causal attention.
• Strided with stride 4: the last 4 tokens, plus every 4th token before them. Token 13 reads 10, 11, 12, 13 and then 9, 5, 1. Over all rows that is 58 + 24 = 82 pairs.

flowchart LR
  G(("token 0<br/>global"))
  t3["token 3"] --- G
  t7["token 7"] --- G
  t11["token 11"] --- G
  t15["token 15"] --- G
  t14["token 14"] --- t15
  t13["token 13"] --- t14

s = 4
# row 13: which keys does token 13 read?
[j for j in range(16) if 13 - j >= 0 and (13 - j < s or (13 - j) % s == 0)]  # → [1, 5, 9, 10, 11, 12, 13]
# strided pairs over all 16 rows
sum(1 for i in range(16) for j in range(i + 1) if i - j < s or (i - j) % s == 0)  # → 82
# a window of 4 plus global token 0
sum(1 for i in range(16) for j in range(i + 1) if i - j < 4 or j == 0)  # → 70

Why it matters in practice. The Sparse Transformer introduced strided patterns; Longformer and BigBird combined windows with global tokens (BigBird added a few random links too) to read documents of thousands of tokens. StreamingLLM found that trained models park a lot of attention on the very first tokens (attention sinks); a plain window drops them and falls apart, while keeping the first few tokens plus a window lets a model stream indefinitely. One caution: a sparse pattern only saves time if the GPU kernel skips whole blocks of the score matrix, which is why real patterns are built from blocks.

In code: global_local_mask adds global rows and columns to a window, strided_mask builds the express-stop pattern, and pairs_computed counts what each one scores. Any of them can be passed as the mask to primer.ml.attention.scaled_dot_product_attention, and every skipped pair gets a weight of exactly 0.

Everyday picture A potluck soup. In softmax attention every new guest tastes every dish on the table, one by one, and then mixes a bowl: the more dishes, the longer it takes. In linear attention every guest pours their dish into one shared pot as they arrive, and a new guest takes a single ladle, seasoned to their own taste. The pot never grows, and the ladle costs the same for guest 3 as for guest 3 million.

The trick that makes the pot possible: softmax's score cannot be split into "a part that depends on q" times "a part that depends on k". Replace it with , a dot product of transformed vectors, and it can. Then the order of the matrix multiplies can be swapped: . The left side builds an n × n matrix; the right side builds a small d × d one (the pot) and never builds the big one. Matrix multiplication allows regrouping like this (associativity), which primer.notation covers under matrix multiply.

Tiny worked example Three tokens with 2-number keys (1, 0), (0, 1), (1, 1) and one-number values 2, 4, 6. The third token's query is (1, 0). Use the feature map φ(x) = elu(x) + 1, which for a positive number is simply x + 1 and for zero or a negative number is (always above zero).

1. φ of the keys: (2, 1), (1, 2), (2, 2).
2. Pour into the pot: S = (2, 1)·2 + (1, 2)·4 + (2, 2)·6 = (20, 22), and the running total of keys z = (5, 5).
3. Ladle with φ(q) = (2, 1): (2·20 + 1·22) / (2·5 + 1·5) = 62 / 15 = 4.13.

The slow way agrees: the weights φ(q)·φ(k) are 5, 4 and 6, so the output is (5·2 + 4·6 + 6·6) / 15 = 62/15. Softmax attention on the same numbers gives 4.00: linear attention is a different attention, not a faster copy of the same one.

flowchart LR
  subgraph T["each new token i"]
    K["φ(k_i)"]
    V["v_i"]
    Qi["φ(q_i)"]
  end
  K & V -->|"add φ(k_i) v_iᵀ"| S[("pot S<br/>d_k × d_v numbers")]
  K -->|"add φ(k_i)"| Z[("total z<br/>d_k numbers")]
  Qi --> R["output = φ(q_i)ᵀ S / φ(q_i)ᵀ z"]
  S --> R
  Z --> R

import math
def phi(v):
    # elu(x) + 1: x + 1 above zero, e^x at or below it
    return [x + 1 if x > 0 else math.exp(x) for x in v]
def dot(a, b):
    return sum(a_m * b_m for a_m, b_m in zip(a, b))
keys = [[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]]
values = [2.0, 4.0, 6.0]
q3 = [1.0, 0.0]
[phi(k) for k in keys]  # → [[2.0, 1.0], [1.0, 2.0], [2.0, 2.0]]
# S: the running sum of φ(k_j) v_j; z: the running sum of φ(k_j)
S = [sum(phi(k)[m] * v for k, v in zip(keys, values)) for m in range(2)]
S  # → [20.0, 22.0]
z = [sum(phi(k)[m] for k in keys) for m in range(2)]
z  # → [5.0, 5.0]
# o_3 = φ(q)ᵀS / φ(q)ᵀz
round(dot(phi(q3), S) / dot(phi(q3), z), 3)  # → 4.133
# the quadratic form: the weights φ(q)·φ(k_j), then their weighted average
weights = [dot(phi(q3), phi(k)) for k in keys]
weights  # → [5.0, 4.0, 6.0]
round(dot(weights, values) / sum(weights), 3)  # → 4.133
# softmax attention on the same numbers (unscaled), for contrast
e = [math.exp(dot(q3, k)) for k in keys]
round(dot(e, values) / sum(e), 3)  # → 4.0

Why it matters in practice. The work drops from n² × d to n × d², and generation needs only the pot, not a cache. The catch is recall: asked to copy back a specific token from far away, a fixed-size pot does worse than a full cache. Later work added forgetting (decay) to the pot, which turns out to be exactly the state-space models of the next section; the Mamba-2 paper shows the two views describe the same computation.

In code: feature_map is φ, linear_attention_quadratic builds all n × n weights (from linear_attention_weights) and blends the values, and linear_attention_recurrent gets the same outputs from the two running sums, returning the fixed-size S and z.

A recurrence that is also a convolution

Everyday picture A cup of tea cooling on a desk. Every minute it keeps some fraction of its heat and gains whatever hot water you pour in. Its temperature right now is a summary of everything you ever poured, with older pours counting less. That is a recurrent network's one-page summary (see primer.ml.cnn_rnn), with one difference: the rewrite rule is linear (multiply and add, nothing more), and linearity buys a second way to compute the same thing.

Tiny worked example The state keeps half of itself each step: A = 0.5, B = 1, C = 1. Inputs x = (1, 0, 0, 2).

Unroll the loop and each output is a weighted sum of all past inputs with weights 1, 0.5, 0.25, 0.125 for "now, 1 step ago, 2 steps ago, 3 steps ago": y₄ = 1·2 + 0.5·0 + 0.25·0 + 0.125·1 = 2.125. The same answer, with no loop.

flowchart LR
  X["inputs x_1 ... x_T"] --> R["recurrence<br/>h_t = A h_(t−1) + B x_t<br/>one step at a time"]
  X --> K["convolution<br/>slide the kernel (CB, CAB, CA²B, ...)<br/>all positions at once"]
  R --> Y["the same outputs y_1 ... y_T"]
  K --> Y

A, B, C = 0.5, 1.0, 1.0
x = [1.0, 0.0, 0.0, 2.0]
h, y = 0.0, []
for x_t in x:
    # keep half of the old state, add the new input
    h = A * h + B * x_t
    y.append(C * h)
y  # → [1.0, 0.5, 0.25, 2.125]
# the kernel C A^k B, for k = 0, 1, 2, 3
kernel = [C * A ** k * B for k in range(4)]
kernel  # → [1.0, 0.5, 0.25, 0.125]
# y_4 as a weighted sum of all four inputs, newest first
sum(kernel[k] * x[3 - k] for k in range(4))  # → 2.125

In code: ssm_recurrent runs the loop, ssm_kernel builds C A^k B, and ssm_convolution gets the same outputs from the kernel with a single NumPy convolution.

Selective: letting each token decide what to keep (Mamba)

Everyday picture A note-taker with a dial. For filler words ("um", "so", "anyway") they barely touch their notes. For a name or a number they wipe the relevant line and write the new fact. A fixed SSM uses the same dial setting for every word, so it must either write everything (and forget quickly) or write little (and never take in the important word properly). Mamba reads the dial setting off each token itself: that is what selective means.

The dial is a step size Δ. The parameters of the state update are derived from it at every step (this is called discretization: turning a continuous rate of change into one step's keep and write amounts):

• keep factor , with negative, so a big Δ makes it nearly 0 (forget) and a tiny Δ nearly 1 (keep);
• write factor , which goes the opposite way.

Tiny worked example One channel, a = −1, b = 1. A token that carries a marker gets Δ ≈ 5; an ordinary token gets Δ ≈ 0.0067.

The recall task: a sequence of small noise values with one marked 7 in third place, then nine more noise values. The selective SSM writes the 7 almost fully and then keeps it: at the end its state is 6.55. With one fixed Δ for every token, the best any setting manages is 0.26: a Δ big enough to write the 7 also lets the nine later tokens overwrite it.

flowchart LR
  U["token u_t"] --> D["Δ_t = softplus(w · u_t + β)<br/>how much this token matters"]
  D --> AB["Ā_t = e^(Δ_t a): keep<br/>B̄_t: write"]
  U --> XV["x_t: what to write"]
  AB --> H["h_t = Ā_t h_(t−1) + B̄_t x_t"]
  XV --> H
  HP["h_(t−1)"] --> H
  H --> Y["y_t = C h_t"]

import math
def softplus(x):
    return math.log(1 + math.exp(x))
a, b = -1.0, 1.0
# Δ for a marked token, then for an ordinary one
d_mark, d_plain = softplus(10 * 1 - 5), softplus(10 * 0 - 5)
round(d_mark, 3), round(d_plain, 4)  # → (5.007, 0.0067)
# Ā_t and B̄_t for the marked token: keep almost nothing, write almost everything
keep_mark, write_mark = math.exp(d_mark * a), (math.exp(d_mark * a) - 1) / a * b
round(keep_mark, 4), round(write_mark, 4)  # → (0.0067, 0.9933)
# and for an ordinary token: keep almost everything, write almost nothing
keep_plain = math.exp(d_plain * a)
round(keep_plain, 4)  # → 0.9933
# the 7 is written once, then kept through nine ordinary tokens
round(write_mark * 7 * keep_plain ** 9, 2)  # → 6.55

In code: discretize turns Δ into Ā and B̄, softplus and step_sizes compute each token's Δ, selective_ssm runs the per-token recurrence, and selective_recall and time_invariant_recall run the recall task.

Training in parallel when the kernel keeps changing: the scan

Everyday picture A relay race where each runner must know the total time so far. Done one after another it takes as long as the whole race. But "multiply by a, then add b" steps can be merged: two consecutive steps are themselves one step of the same shape. So pairs of runners merge their legs, then pairs of pairs, like a knockout tournament: log₂ n rounds instead of n (log₂ n is how many times n can be halved before reaching 1: 10 for 1,024).

Tiny worked example The four steps of the tea example are (a, b) = (0.5, 1), (0.5, 0), (0.5, 0), (0.5, 2), meaning "h becomes a·h + b". Round 1: every step merges with the one before it. Round 2: every step merges with the result two places before it. After 2 rounds (log₂ 4 = 2), the b parts are 1, 0.5, 0.25, 2.125: every state from the loop.

flowchart TB
  s1["(0.5, 1)"]
  s2["(0.5, 0)"]
  s3["(0.5, 0)"]
  s4["(0.5, 2)"]
  s1 --> r2["(0.25, 0.5)"]
  s2 --> r2
  s2 --> r3["(0.25, 0)"]
  s3 --> r3
  s3 --> r4["(0.25, 2)"]
  s4 --> r4
  s1 --> f3["(0.125, 0.25)"]
  r3 --> f3
  r2 --> f4["(0.0625, 2.125)"]
  r4 --> f4

def merge(first, then):
    a1, b1 = first
    a2, b2 = then
    return (a1 * a2, a2 * b1 + b2)
steps = [(0.5, 1.0), (0.5, 0.0), (0.5, 0.0), (0.5, 2.0)]
# round 1: each step absorbs the one just before it
r1 = [steps[0]] + [merge(steps[t - 1], steps[t]) for t in range(1, 4)]
r1  # → [(0.5, 1.0), (0.25, 0.5), (0.25, 0.0), (0.25, 2.0)]
# round 2: each absorbs the result two places before it
r2 = r1[:2] + [merge(r1[t - 2], r1[t]) for t in range(2, 4)]
[b for a, b in r2]  # → [1.0, 0.5, 0.25, 2.125]

Why it matters in practice. This is how Mamba trains on long sequences at GPU speed despite being recurrent: a parallel scan, written so the state stays in fast on-chip memory. At generation time it switches back to the plain loop, one token at a time, with a state that never grows.

In code: parallel_scan runs the rounds for any length and returns the states and the number of rounds: 10 for 1,024 steps.

Constant memory, and hybrids

Everyday picture An SSM travels with a backpack of fixed size; the KV cache is a suitcase that grows with every token. A hybrid model mostly uses backpacks but brings one suitcase for every few layers, so it can still look up an exact earlier token when it needs to.

Tiny worked example One SSM layer with 4,096 channels and a 16-number state per channel holds 4,096 × 16 × 2 bytes = 128 KB, at 10 tokens or at 10 million. One attention layer of the running 8B example holds 537 MB at 128k tokens. Build 32 layers as 4 attention layers and 28 SSM layers (one in eight, the ratio Jamba uses) and the cache at 128k drops from 17.2 GB to 2.15 GB.

flowchart TB
  I["tokens in"] --> M1["SSM layer<br/>fixed state"] --> M2["SSM layer"] --> M3["..."] --> A1["attention layer<br/>KV cache grows"]
  A1 --> M4["SSM layer"] --> M5["..."] --> A2["attention layer"] --> O["next-token scores"]

n, H_kv, d_h, b = 131_072, 8, 128, 2
D, N = 4096, 16
# one SSM layer's state, in KB, at any context length
D * N * b / 1024  # → 128.0
# one attention layer's KV cache at 128k tokens, in MB
n * 2 * H_kv * d_h * b / 1e6  # → 536.870912
# 32 layers: all attention, then 4 attention + 28 SSM, in GB
round(32 * n * 2 * H_kv * d_h * b / 1e9, 2)  # → 17.18
round((4 * n * 2 * H_kv * d_h * b + 28 * D * N * b) / 1e9, 2)  # → 2.15

Why it matters in practice. Pure SSMs are strong at language modelling but weaker than attention at copying and exact recall over long contexts, for the same reason as linear attention: a fixed-size state is a summary. Hybrids such as Jamba keep a few attention layers for those jobs and get most of the SSM's memory savings.

In code: ssm_state_bytes is the fixed backpack, and hybrid_cache_bytes adds up a mixed stack.

Fewer key/value heads: GQA and MQA, a recap

Everyday picture Colleagues sharing one reference binder instead of each keeping a personal copy: everyone still asks their own questions, but the shelf holds fewer binders.

Tiny worked example On the running example (32 layers, 128 numbers per head, 16-bit), the cache per token for different numbers of KV heads:

The diagram of shared heads and the full memory arithmetic live in primer.ml.attention and primer.ml.inference; everything below stacks on top of whichever of these a model uses.

Latent KV: cache the ingredients, cook on demand

Everyday picture A restaurant that stores ingredients, not finished dishes. Every head's keys and values can be cooked from a short list of ingredients (the latent vector) with a fixed recipe shared by every token. Store the ingredients, keep the recipe once, and cook when a query needs it. Better still, the recipe can be folded into the query itself, so nothing is ever cooked at all.

Tiny worked example A token x = (1, 0, 2, 1), two heads of width 2. Ordinary attention would cache 2 heads × 2 numbers × (key and value) = 8 numbers. Instead:

1. Squeeze: c = x W_down = (1 + 2, 0 + 1) = (3, 1). Only these 2 numbers are cached: 4 times smaller.
2. When needed, expand: k = c W_uk = (3, 1, 3, 1), both heads' keys.
3. Or fold the expansion into the query: for q = (1, 2, 0, 1), q · k = 6, and (q W_ukᵀ) · c = (1, 3) · (3, 1) = 6. Same score, and k was never built.

DeepSeek-V2 uses this at scale: its 128 heads of width 128 would cache 32,768 numbers per token per layer; its latent caches 512, plus 64 for a small key that carries position information: 576, about 57 times less.

flowchart LR
  X["token x<br/>(d_model numbers)"] -->|"W_down"| C[("cache: latent c<br/>d_c numbers")]
  C -->|"W_uk"| K["keys for every head"]
  C -->|"W_uv"| V["values for every head"]
  Qn["query q"] -->|"absorbed: q W_ukᵀ"| QL["query in latent space"]
  QL -->|"score directly against c"| C

x = [1, 0, 2, 1]
W_down = [[1, 0], [0, 1], [1, 0], [0, 1]]
W_uk = [[1, 0, 1, 0], [0, 1, 0, 1]]
def vecmat(v, M):
    # row vector times matrix: entry c is Σ_r v_r M[r][c]
    return [sum(v[r] * M[r][c] for r in range(len(v))) for c in range(len(M[0]))]
# c = x W_down: the only thing cached
c = vecmat(x, W_down)
c  # → [3, 1]
# k = c W_uk: both heads' keys, rebuilt on demand
k = vecmat(c, W_uk)
k  # → [3, 1, 3, 1]
q = [1, 2, 0, 1]
sum(q_m * k_m for q_m, k_m in zip(q, k))  # → 6
# absorbed: move q into latent space, then score against c
W_uk_T = [list(col) for col in zip(*W_uk)]
q_latent = vecmat(q, W_uk_T)
q_latent, sum(a * b for a, b in zip(q_latent, c))  # → ([1, 3], 6)
# DeepSeek-V2's shape: cached numbers per token per layer
2 * 128 * 128, 512 + 64, round(2 * 128 * 128 / (512 + 64), 1)  # → (32768, 576, 56.9)

Why can so few numbers stand in for so many? Because every head's keys and values are built from the same token, they are highly redundant. In the language of primer.notation, the key projection W^D W^UK has low rank: however many numbers it outputs, they all vary along only d_c independent directions. Storing those d_c coordinates loses nothing that projection can produce. One wrinkle: rotary position embeddings (see primer.ml.positional) rotate each key by its position, which breaks the absorption trick, so DeepSeek-V2 carries position in that separate small 64-number key.

In code: LatentKVAttention holds the four projections; LatentKVAttention.compress produces the cache, LatentKVAttention.attend expands it into keys and values, and LatentKVAttention.attend_absorbed gets the identical output without expanding. latent_kv_worked_example is the (3, 1) example and latent_kv_bytes_per_token the memory arithmetic.

Quantizing the cache

Everyday picture The same move as for weights in primer.ml.inference: write each number to the nearest tenth instead of the nearest thousandth, with one ruler per stored vector so one big number does not coarsen all the others.

Tiny worked example A cached key (0.7, −0.3, 0.2, 0.04) at 4 bits (codes −7 to 7). The step size is 0.7 / 7 = 0.1, the codes are (7, −3, 2, 0), and reading back gives (0.7, −0.3, 0.2, 0): the tiny 0.04 is lost. A 128-number head vector drops from 256 bytes to 64 bytes of codes plus a 2-byte scale: 66 bytes, 3.9 times smaller.

flowchart LR
  KV["new key or value<br/>16-bit numbers"] -->|"scale = max / 7<br/>code = round(x / scale)"| ST[("cache: 4-bit codes<br/>+ one scale per vector")]
  ST -->|"code × scale"| R["approximate key or value"]
  R --> A["attention as usual"]

k = [0.7, -0.3, 0.2, 0.04]
# one scale per vector: the largest magnitude lands on code 7
s = max(abs(k_j) for k_j in k) / 7
round(s, 3)  # → 0.1
codes = [round(k_j / s) for k_j in k]
codes  # → [7, -3, 2, 0]
[round(s * c, 2) for c in codes]  # → [0.7, -0.3, 0.2, 0.0]
L, H_kv, d_h = 32, 8, 128
# bytes per token: 16-bit, then 4-bit codes plus a 16-bit scale per vector
2 * L * H_kv * d_h * 16 / 8, 2 * L * H_kv * (d_h * 4 / 8 + 16 / 8)  # → (131072.0, 33792.0)

Why it matters in practice. On random test data, an 8-bit cache moves the attention output by under 1% and a 4-bit cache by about 12%; trained models tolerate this far better than random data suggests, and careful schemes go lower. Keys tend to have a few channels with consistently large values, so KIVI quantizes keys per channel and values per token and reaches 2 bits. Quantization stacks with everything above: a GQA model with a sliding window and a 4-bit cache enjoys all three savings.

In code: quantize_kv quantizes each cached vector with its own scale via primer.ml.inference.quantize, quantized_kv_bytes_per_token is the formula, and quantized_cache_error measures how far attention's output moves.

In code: cache_bytes_by_method computes every line in the figure for a given context length.

Answer each one out loud or on paper before you open it. If you can explain it, you know it.