Applied AI · advanced · Session 18
Exercises with solutions — KV cache, recurrent memory, MLA, and low rank
← Back to courseFrançaisMarkdown source

Exercises with solutions — KV cache, recurrent memory, MLA, and low rank

General instruction: every answer must show data, transformation, result, one check, and one limitation. A bare number or copied definition is insufficient.

Starting data: 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.

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

Exercise 1 — Calculated trace — Standard KV cache

Reproduce and annotate the chain bytes≈tokens×layers×2×heads×d_head×bytes/value. Double context length from 4,096 to 8,192 tokens while keeping 32 layers, 8 KV heads, d_head=128, and BF16. Recompute the simplified raw cache.

Deliverable: a data → operation → result → interpretation table, plus two sentences about the changed value.

Worked solution

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.

Solved variant: Cache grows from 536,870,912 bytes (512 MiB) to 1,073,741,824 bytes (1 GiB). Growth is linear in token count; MLA or recurrent state can reduce different dimensions but also changes fidelity or retrieval behavior.

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. 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. The minimum check covers dimensions, sign, and order of magnitude. If the observed change contradicts the prediction, locate the first operation whose direction changes instead of fixing only the final line.

Rubric Exercise 1 — /10

Criterion Points
Explicit data and shapes 2
Traceable calculation 3
Prediction before variation 2
Interpretation and check 2
Named limitation 1

Exercise 2 — Diagnose a seductive explanation — MLA

A colleague claims: « MLA proves the system will be accurate, fast, and stable in every context. »

  1. Separate mechanism, assumption, observation, and conclusion.
  2. Name two correct lesson elements and two unsupported extrapolations.
  3. Propose a bounded experiment with controlled variable, metric, and stop threshold.
  4. Rewrite the claim as one defensible sentence.
Reasoned solution

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. 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. Simplified formulas omit alignment, quantization, buffers, and sharing details. They compare trends, not promise real footprint.

The claim mixes a local relation with a global guarantee. A defensible version states only the observed mechanism, test conditions, and measured metric. Stop the test if shapes become invalid, the metric crosses the declared degradation threshold, or another variable changes.

Rubric Exercise 2 — /10

2 points per element: separation, lesson grounding, extrapolations, protocol, and rewrite.

Exercise 3 — Architecture decision and transfer — MLA is not fixed memory

You must reproduce the 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.” under two conditions. Option A uses the full chain through “MLA is not fixed memory.” Option B is a transparent baseline that retains “Standard KV cache,” calculates the expected output directly, and does not use the compression or adjustment mechanism studied. Build a decision record containing:

Elements of a strong solution

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. 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. Mixed: established mechanisms + source-reported Kimi K3-style choices.

A strong answer does not make the newer mechanism the default winner. It retains a measurable baseline, sets thresholds before testing, and separates component cost from whole-system behavior. The verdict names what remains uncertain and the next evidence that could change it.

Rubric Exercise 3 — /15

Criterion Points
Framing and baseline 3
Compared causal chains 4
Protocol and metrics 4
Rollback threshold 2
Bounded verdict 2

Extension

Repeat Exercise 3 after reversing the dominant constraint. If you optimized memory, impose a strict quality floor; if you optimized fidelity, halve the memory budget. Identify the first part of the verdict that changes and the evidence required.

Review before submission

Review the packet as if another group had to reproduce it without speaking to you. Are all starting values or assumptions present? Are shapes or roles stated before operations? Does the prediction truly precede the observation? Is the result translated into behavior rather than left as an isolated number? Did you test a boundary value and identify a stop condition? Does the procedure or architecture choice retain a measurable baseline and a rollback threshold set before the test? Finally, highlight one sentence that states what is established, one that remains a hypothesis, and one measurement that could change your verdict. If any element is missing, the work is not reproducible.

Reference appendix for correction

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

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

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

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.

Sources and evidence boundary

Scope: Mixed: established mechanisms + source-reported Kimi K3-style choices. These references support the session frame; they do not turn a reported product choice into an independently verified result.