advanced · Session 18

KV cache, recurrent memory, MLA, and low rank

Compare three memory budgets and distinguish per-token compression, fixed state, and low-rank factorization.

120 min6 mechanismscausal lab

What you will be able to do

Open the lab

Study method

Study this session as a causal investigation. Before every formula or interaction, write what you expect to change and what must remain fixed. During the calculation, retain units, shapes, and intermediate values so an error can be located without restarting at random. After the result, translate the number or state into one sentence about system behavior. Always finish with a counterexample or boundary value. This discipline separates understanding a mechanism from merely recognizing its vocabulary and makes the lab reproducible by another learner. Record the evidence that changed your initial prediction.

Build the mechanism step by step

1. Standard KV cache

The problem: Your service targets 131,072 tokens of context. Before any optimization, you need the raw number: at 32 layers, 8 KV heads, d_head=128 and BF16, each token costs 128 KiB of cache — 16 GiB per request at full length. And that bill returns with every request.

The idea: The cache keeps every past token’s keys and values, per layer: bytes ≈ tokens × layers × 2 × heads × d_head × bytes. Reads are faithful — exact attention over the whole past — and the arithmetic can be redone factor by factor, no calculator.

bytes≈tokens×layers×2×heads×d_head×bytes/value

Why / at what price: Fidelity is total: exact recall of any token. The price is structural: strictly linear growth — ×32 tokens = ×32 memory — with re-read bandwidth to match. No setting changes the slope, only the coefficient.

Understanding check

Name the input, transformed state, output, and one required assumption. Then compare your chain with the explanation above.

2. Fixed recurrent state

The problem: 16 GiB per request rules out most deployments. Sessions 13 to 16 built the radical alternative: what if the past fit in a fixed-size matrix, whatever the token count?

The idea: The state S — 32 layers × 128 × 128 × BF16 = 1 MiB — summarizes the whole past: 4,096 or 524,288 tokens, still 1 MiB. Growth disappears; this is where the previous sessions’ memories were heading.

Why / at what price: A constant, negligible budget — 16,000× less than the cache at 131,072 tokens. The price, familiar from sessions 13-16: compression and interference — recall is no longer exact and degrades with length, even though the memory itself never moves.

Understanding check

Name the input, transformed state, output, and one required assumption. Then compare your chain with the explanation above.

3. MLA

The problem: Between an exact 16 GiB and an approximate 1 MiB, the gap is brutal. Is there a middle ground: keep one entry PER token — structured fidelity — but pay less per entry?

The idea: Multi-head Latent Attention compresses each token into a latent vector c_t (512 values) — the only thing cached — then reconstructs K and V through W_UK and W_UV at read time. Width ÷ 4 ⇒ 32 KiB/token, 4 GiB at 131,072 tokens.

Why / at what price: The cache keeps its per-token structure and the bill is divided by 4. The price: a reconstruction on every read — memory traded for compute, mostly at decode — and growth stays O(n): the gain is a coefficient, not a change of asymptote.

Understanding check

Name the input, transformed state, output, and one required assumption. Then compare your chain with the explanation above.

4. Low-rank factorization

The problem: MLA’s compression rests on a pure algebra question: when does replacing a large matrix W by a product of two small ones actually save anything? Ill-chosen, the bottleneck r saves nothing at all.

The idea: W (d×m) ≈ A(d×r)·B(r×m) costs r(d+m) parameters instead of d·m. For 4096×4096: r=512 divides by 4; r=2048 gives 16,777,216 — exactly the original cost. The break-even threshold is r = d·m/(d+m) = 2048 here.

W≈AB, A∈R^{d×r}, B∈R^{r×m}

Why / at what price: Below the threshold the saving is real and the compute faster. The price: the projection’s capacity is capped at rank r — any transformation requiring a higher rank is structurally out of reach, no matter the training.

Understanding check

Name the input, transformed state, output, and one required assumption. Then compare your chain with the explanation above.

5. MLA is not fixed memory

The problem: Two announcements sound alike: “4× compressed cache” and “constant memory”. A team budgets 4 GiB “forever” with MLA — and discovers 16 GiB at 524,288 tokens. Where did the promise go?

The idea: There never was one: MLA compresses each entry, it does not merge tokens. One entry per token ⇒ linear growth with a smaller coefficient. Only a recurrent state merges the past into a fixed-size object.

Why / at what price: Separating the two avoids the production capacity mistake: MLA bounds the coefficient, the state bounds the growth. The price of confusing them is written in the trace: MLA’s ×4 gain is eaten by ×4 tokens — length still rules.

Understanding check

Name the input, transformed state, output, and one required assumption. Then compare your chain with the explanation above.

6. Low rank is not LoRA

The problem: The same formula W ≈ AB in two contexts: MLA’s architectural factorization and LoRA adaptation. A hurried reader concludes “MLA is built-in LoRA” — and proposes “removing the adapter” from a model that has none.

The idea: Opposite roles: in MLA, A and B ARE the normal path, trained from scratch, immovable. LoRA adds a low-rank delta BESIDE frozen weights, to adapt after the fact — mergeable or removable at will.

Why / at what price: The distinction prevents absurd decisions — freezing MLA’s “adapter”, or believing LoRA is required at inference. The price: permanent vigilance; the same algebra serves architectures and adaptation procedures, and only context decides the meaning.

Understanding check

Name the input, transformed state, output, and one required assumption. Then compare your chain with the explanation above.

Development from the course source

Chapter 13 — Full records versus compressed memory

13.1 Two memory strategies

Goal: compare exact token records with summaries. Standard key-value caching stores per-token records. Recurrent linear memory compresses many tokens into a fixed-size state.

Intuition: A video archive preserves every frame; meeting notes preserve a compact summary. Each is useful for different questions.

Step by step

  • Per-token cache: grows with context; preserves more direct access to individual tokens.

  • Fixed state: does not grow with token count; risks interference and information loss.

  • Hybrid designs can use both mechanisms in different layers or roles.

Worked example: For 1,000 tokens, a per-token method keeps 1,000 records per relevant layer. A fixed-state method keeps one state of predetermined shape. This does not prove it uses less total memory in every implementation, but it explains the scaling difference.

Why it matters: Architecture is a trade-off among fidelity, speed, memory traffic and trainability.

Quick check: Which strategy is more likely to retrieve one exact old token? Answer: the per-token record, although retrieval quality still depends on learned attention.

13.2 Latent compression and MLA

Goal: reduce cache size without collapsing all history into one recurrent matrix. A latent vector is a smaller learned representation. Multi-head Latent Attention, abbreviated MLA, stores compressed latent information and reconstructs head-specific keys or values when needed.

Intuition: Store a zipped folder instead of several expanded copies, then unpack the view needed by each worker.

Step by step

  • Compress hidden representation x into c = xW_down.

  • Cache the smaller latent c.

  • Use learned up-projections to create the key/value forms needed by attention heads.

  • A head is one parallel attention subspace; multiple heads can learn different relations.

Worked example: Worked shape example: x has 8 coordinates. Compressing to 2 gives c with 2 coordinates. Expanding c back to an 8-coordinate key does not magically preserve every possible 8-D vector; it restricts the key to patterns learnable through the 2-D bottleneck.

Why it matters: MLA is compressed per-token memory, not the same as one fixed recurrent state for the whole past.

Quick check: Does MLA’s cache normally remain constant as token count grows? Answer: No. It can be smaller per token, but still grows with the number of cached tokens.

13.3 Low-rank factorization and LoRA

Goal: understand narrow intermediate spaces. A full 8×8 matrix has 64 entries. Replacing it with an 8×2 matrix followed by a 2×8 matrix uses 16+16=32 entries in this simplified count.

Intuition: A narrow hallway limits how many independent flows can pass at once.

Step by step

  • Down-project: h = xA.

  • Up-project: y = hB.

  • The combined map AB has rank at most the narrow width.

  • Low-Rank Adaptation, LoRA, usually adds a trainable low-rank update to a frozen base weight; it is related mathematics but not automatically the same use as architectural compression.

Worked example: This distinction prevents a common confusion: every low-rank factorization is not necessarily a LoRA fine-tuning adapter.

Why it matters: Low rank trades flexibility for fewer parameters, less storage or cheaper computation, depending on where it is used.

Quick check: How many entries are in 10×3 plus 3×10 matrices? Answer: 30+30=60.

Complete worked case

With 32 layers, 8 KV heads, d_head=128, BF16, and 4,096 tokens, simplified raw cache is 4,096×32×2×8×128×2 bytes = 536,870,912 bytes = 512 MiB (0.5 GiB). Quartering latent width reduces the per-token term, not linear growth.

Reading method: write the data, state every object shape, perform one transformation, and interpret the result before continuing.

Validity boundary

Simplified formulas omit alignment, quantization, buffers, and sharing details. They compare trends, not promise real footprint.

Evidence status: Mixed: established mechanisms + source-reported Kimi K3-style choices.

Quick checks

1. What does Standard KV cache?

The cache keeps every past token’s keys and values, per layer: bytes ≈ tokens × layers × 2 × heads × d_head × bytes. Reads are faithful — exact attention over the whole past — and the arithmetic can be redone factor by factor, no calculator.

2. What does Fixed recurrent state?

The state S — 32 layers × 128 × 128 × BF16 = 1 MiB — summarizes the whole past: 4,096 or 524,288 tokens, still 1 MiB. Growth disappears; this is where the previous sessions’ memories were heading.

3. What does MLA?

Multi-head Latent Attention compresses each token into a latent vector c_t (512 values) — the only thing cached — then reconstructs K and V through W_UK and W_UV at read time. Width ÷ 4 ⇒ 32 KiB/token, 4 GiB at 131,072 tokens.

4. What does Low-rank factorization?

W (d×m) ≈ A(d×r)·B(r×m) costs r(d+m) parameters instead of d·m. For 4096×4096: r=512 divides by 4; r=2048 gives 16,777,216 — exactly the original cost. The break-even threshold is r = d·m/(d+m) = 2048 here.

Sources and evidence boundary

Scope: Mixed: established mechanisms + source-reported Kimi K3-style choices.