If you want 8 different answers to the same prompt, the naive way costs you 8x the prefill.
Every one of those 8 forward passes re-reads the prompt, re-computes attention over the whole context, re-fills a KV cache from scratch. On a 4k-token system prompt with a 30B-parameter model, that is not a small tax. It is most of the wall-clock time before the first token comes out.
Andrej Karpathy's nanochat has one of the cleanest workarounds for this I have read. The core move is a 12-line method called prefill on the KVCache object, and it turns the 8x prefill bill into a 1x prefill plus an 8-way memcpy. This post walks through why that works, why it depends on Flash Attention 3's flash_attn_with_kvcache API, and where the design bites you.
What the KV cache looks like on disk
nanochat/engine.py allocates the cache up front, one tensor per attention side:
# k_cache and v_cache shape: (n_layers, batch_size, seq_len, n_kv_head, head_dim)
The tricky index is n_kv_head, not n_head. Under GQA (grouped-query attention), multiple query heads share a single KV head. nanochat only allocates enough cache slots for the reduced set. That is not an optimization the code goes out of its way to advertise, but it is why the cache stays inside VRAM budgets that would otherwise not fit.
cache_seqlens (shape (batch_size,)) tracks how far each row has written into the buffer. Flash Attention 3's flash_attn_with_kvcache API writes the new K/V slices in place at the right offsets and does attention against the full cache, all in a single kernel launch. Python does not touch the cache tensor except to increment cache_seqlens at the end of the step.
That single-kernel guarantee is what makes the whole design tractable. If Python had to re-assemble K/V from N shards between steps, the memory bandwidth alone would eat the win.
The prefill trick, 12 lines
Here is the method that does the interesting work:
def prefill(self, other):
"""
Copy cached KV from another cache into this one.
Used when we do batch=1 prefill and then want to generate multiple samples in parallel.
"""
assert self.get_pos() == 0, "Cannot prefill a non-empty KV cache"
assert self.n_layers == other.n_layers and self.n_heads == other.n_heads and self.head_dim == other.head_dim
assert self.max_seq_len >= other.max_seq_len
other_pos = other.get_pos()
self.k_cache[:, :, :other_pos, :, :] = other.k_cache[:, :, :other_pos, :, :]
self.v_cache[:, :, :other_pos, :, :] = other.v_cache[:, :, :other_pos, :, :]
self.cache_seqlens.fill_(other_pos)
The flow is: run the prompt through the model with batch size 1, get a KVCache back that has other_pos tokens filled. Allocate a fresh cache sized for N sampling rows. prefill copies the batch-1 K/V state into every row of the batch-N cache, then updates cache_seqlens so all N rows agree they are already other_pos tokens into their history.
From that point, each of the N rows samples its own next token independently, writes its own K/V into its own row of the cache, and diverges from the others. The prompt itself was only ever processed once, and by an attention kernel that got the full luxury of contiguous batch-1 memory access.
The cost model changes shape:
| Sampling mode | Prefill FLOPs | Decode FLOPs |
|---|---|---|
| Naive (N independent runs) | N x prompt_len^2 | N x output_len |
| Prefill + clone | 1 x prompt_len^2 | N x output_len |
For a 2000-token prompt and N=8, that is 32M prefill units against 4M. Almost an order of magnitude, and you get it back before the first decode token.
I know this because I spent an evening tuning the wrong kernel before finding it. Prefill was eating 40% of my TTFT budget and I was busy micro-optimizing the decode loop. An evening I would like back.
Where the trick would not work
Two conditions have to hold for this to be safe.
First, the KV cache tensors have to be the same shape across rows. That is the assert on n_layers == other.n_layers etc. If your inference engine hot-reloads models mid-flight (say, for a routing layer), you cannot reuse the prefill.
Second, no state that is per-row can leak into the prompt encoding. If your prompt embedding depended on, for example, a rotary embedding phase that was per-sample, the copy would put every row into the same phase and downstream tokens would drift. nanochat sidesteps this because rotary embeddings are position-based, and positions are the same for all N rows at prefill time.
The engine.py module actually copies prev_embedding alongside the K/V, exactly to preserve one piece of per-step state that would otherwise diverge. It is a small correctness detail that is easy to miss on a first read. I missed it. My second read was because my clone worked but the 8 outputs looked suspiciously like 8 copies of the same one.
The generator loop and one queue
Engine.generate is a Python generator that yields (tokens, mask) per step. Internally each row has a RowState object holding current_tokens, a completed flag, and a queue called forced_tokens.
That queue is where the tool-call machinery lives.
When the model emits <|python_start|>, the row switches into "collecting an expression" mode and stops emitting anything to the caller. When <|python_end|> comes out, the collected tokens get decoded into a string, passed to use_calculator, and if that returns a value, the encoded result gets pushed into forced_tokens wrapped in <|output_start|> / <|output_end|> markers.
if next_token == python_end and state.in_python_block:
state.in_python_block = False
if state.python_expr_tokens:
expr = self.tokenizer.decode(state.python_expr_tokens)
result = use_calculator(expr)
if result is not None:
result_tokens = self.tokenizer.encode(str(result))
state.forced_tokens.append(output_start)
state.forced_tokens.extend(result_tokens)
state.forced_tokens.append(output_end)
On the next step, the loop pops from forced_tokens before checking what the model sampled. The mask value 0 means "we overrode the model's choice." The mask value 1 means "the model actually picked this." The forward pass still runs on every step (the model gets to keep contributing to its own internal state), but its output gets discarded while the queue drains.
Instead of a second event loop, a state machine, or an interrupt handler, tool calls are one queue check per token. That is roughly 40 lines of Python doing the work that Ollama or vLLM would split across three modules.
The calculator is a very small calculator
use_calculator is not a Python interpreter. It accepts numeric expressions and .count() string calls. Anything containing __, import, or eval gets rejected before it reaches eval. There is an eval_with_timeout wrapper that fires SIGALRM at 3 seconds to catch runaway expressions.
The full-Python execution path (execute_code in nanochat/execution.py) exists, but only gets called from tasks/humaneval.py for benchmark scoring. Even that one has a docstring that says explicitly it is not a security sandbox, just a crash-prevention wall. It spawns a subprocess, scrubs the environment, and limits memory to 256 MB. There is no network isolation and no kernel-level jail.
If you were expecting the chat CLI to run arbitrary Python for you, it does not. It runs a calculator. The gap between "calculator" and "code execution" is exactly the gap between "I trust this to run in-process" and "I do not."
What breaks between turns
One design choice worth flagging: the KV cache does not survive across turns of the same conversation.
scripts/chat_cli.py maintains a conversation_tokens list that grows by appending <|user_start|> / <|user_end|> / <|assistant_start|> blocks per turn. When you send turn 6, engine.generate gets handed the entire 5-turn history and re-prefills all of it. The KV state from turn 5 was garbage-collected the moment the generator returned.
For a chat CLI this is fine. Turn latency stays acceptable up to a few thousand tokens of history, and the code stays boring. For a serving system this would be a scaling wall, and you would want to hold cache blocks per conversation and reattach them. nanochat is explicit that it is not that system, and the choice is a good example of picking the boring option that ships.
The right way to read this file
If you are debugging your own inference engine and hit the "N samples of the same prompt" problem, engine.py is worth reading end-to-end. Roughly 350 lines. The interesting patterns compose:
-
KVCache.prefillfor the batch-1-to-batch-N clone. -
RowState.forced_tokensfor tool call injection without a second loop. -
sample_next_tokenfor the top-k + temperature that is deliberately not top-p (the argument being that top-p is not worth the extra code for this use case). - The GQA-aware
n_kv_headallocation, which is one of those details that only matters until it saves you 4 GB of VRAM.
Karpathy's design principle here reads as: put the complexity in the shape of the data, not in the control flow. The queue is one line. The clone is 4 assignments. The tool call system is a state variable and a check at the top of the loop. Everything that could have been an event system or a plugin architecture stayed as a plain function.
Notes
- Repository: karpathy/nanochat. The file to read is
nanochat/engine.py. - Flash Attention 3 requires SM90 (Hopper) GPUs. The
flash_attn_with_kvcacheAPI is documented in the Dao-AILab/flash-attention repo. - The chapter this article is adapted from also covers the tool-call state machine,
use_calculatorvsexecute_code, andscripts/infer_bench.py(TTFT/TPOT/MBU/MFU measurement): 8000行でわかる大規模言語モデル, chapter 7 (Japanese).
Top comments (0)