mlx-lm · Lesson 6

Chat & the KV Cache

Multi-turn conversation, and the one mechanism that makes it fast: the KV cache. You'll meet mlx_lm.chat for an interactive REPL, then see — with real token counts from your Mac — why a model doesn't re-read the whole conversation on every turn, and how to save that work to disk with cache_prompt.

🎯 The last big runtime concept — why decode stays fast as context grows
The one idea

The KV cache stores what the model already computed for past tokens, so each new token is cheap. Without it, generating token 500 would re-process tokens 1–499 from scratch. With it, the model keeps each token's key/value tensors and only does work for the new token. It's why a long chat doesn't grind to a halt — and why you can cache a big prompt once and reuse it.

01 The REPL: mlx_lm.chat

Lesson 1's generate is one-shot — it forgets everything after answering. mlx_lm.chat gives you a conversation that remembers within the session:

~/projects/learn/public/courses/mlx-lm/practice/.venv/bin/mlx_lm.chat \
  --model mlx-community/Llama-3.2-1B-Instruct-4bit \
  --system-prompt "You are a terse assistant."

>> What's the capital of France?
Paris.
>> And its population?            # "its" works — the REPL kept the context
About 2.1 million.
💡 What "keeps context" means mechanically The REPL holds a growing list of messages and the KV cache for everything said so far. Each new turn only feeds your new message through the model — the prior turns are already cached. Useful flags: --adapter-path (chat with your Lesson-4 pirate!), --temp, --system-prompt, and --max-kv-size (cap the cache so very long chats don't grow unbounded — it evicts the oldest tokens).

02 Why it's fast — the KV cache, drawn

Inside the model, every token produces a key and value tensor at each layer; generating the next token attends to all previous keys/values. The cache keeps them so they're computed once, not re-computed every step.

Without KV cache With KV cache step 3: t1 t2 t3 step 4: t1 t2 t3 t4 recomputes t1..t3 every step — O(n²) work step 4: t1 t2 t3 t4 t1..t3 read from cache (free) only t4 is computed — O(n) work KV cache (keys + values)
Solid = computed this step; faded = reused from cache. The cache turns per-token work from "re-read the whole context" into "process just the new token." The trade is memory: the cache grows with context length — that's what --max-kv-size and peak-memory watch.

03 Cache a prompt to disk — measured on your Mac

If you query the same big context repeatedly (a long system prompt, a document, a few-shot preamble), don't re-process it each time. mlx_lm.cache_prompt runs the prefill once and saves the KV cache to a file; generate loads it and skips straight to your new question.

# 1) cache a ~1,200-token reference context once:
mlx_lm.cache_prompt --model …4bit --prompt "<long reference text>" \
  --prompt-cache-file ~/projects/learn/public/courses/mlx-lm/practice/ctx.safetensors      # → a 40 MB cache file

# 2) reuse it for many questions — the context isn't reprocessed:
mlx_lm.generate --model …4bit \
  --prompt-cache-file ~/projects/learn/public/courses/mlx-lm/practice/ctx.safetensors \
  --prompt "Summarize the reference in 5 words."

The payoff, real numbers from your machine — same context, with and without the cache:

Prompt tokens prefilledWhat ran
No cache1257 tokensthe whole context, re-read every call
With cache15 tokensjust your new question — context loaded from the 40 MB file
💡 When this matters Any repeated-context workload: a fixed system prompt you reuse across requests, a document you ask many questions about, a few-shot template. Cache the constant part once; pay only for the variable part. (The server from Lesson 3 does the in-memory version of this automatically while it stays warm.)

04 Check yourself

Think first, then expand. (Curious how this interacts with your pirate adapter or the server? Ask me.)

Q: Why does generation stay roughly constant-speed per token even as a chat gets long?

A: The KV cache means each new token only computes itself and reads prior keys/values from the cache, instead of recomputing the whole history. Per-token work scales with one token, not with the full context length.

Q: With the cache, the prompt showed only 15 tokens instead of 1257. Where did the other ~1242 go?

A: They were prefilled once by cache_prompt and saved to the cache file. generate loaded those keys/values from disk, so only the 15 new tokens of your question needed processing.

Q: What does the KV cache cost, and what flag bounds it?

A: Memory — it grows with context length (part of your peak-memory number). --max-kv-size caps it by evicting the oldest tokens, trading unbounded memory for a sliding window of context.

Q: You have one long system prompt reused across 100 requests. Cheapest approach?

A: Cache the system prompt once (cache_prompt → file, or a warm server that keeps it in memory) and reuse it, so you prefill those tokens once instead of 100 times.

05 Going deeper

Keep handy: Glossary & quick reference (see KV cache, prefill/decode) · ← Lesson 5. Next up → Lesson 7: the capstone — wire everything you've learned into a real local-LLM tool on your machine.