🤖 machine learning

Models Take Notes at Prefill: KV Cache Can Be Editable and Composable

This paper reveals that during prefill, models primarily store field-conditioned conclusions in downstream KV cache notes rather than relying on the field's own vectors, enabling a novel approach where caches can be efficiently edited via chain-of-thought and composed across contexts to achieve near-perfect accuracy with significantly reduced latency compared to full recomputation.

Original authors: Bojie Li

Published 2026-06-17

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Bojie Li

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Idea: The Model's "Notebook"

Imagine a large language model (like a very smart robot assistant) is reading a long document to answer a question. Usually, when the robot reads a sentence, it "thinks" about it and then moves on. If the sentence changes later, the robot has to re-read the whole document from the beginning to understand the new context. This is slow and expensive.

This paper discovers something surprising: The robot doesn't just read the document; it takes notes on a separate "notebook" (the KV Cache) as it reads.

Crucially, the robot writes its conclusions into this notebook, not just the raw words. Once it figures out, "Oh, the order status is 'shipped,' so I must deny the refund," it writes that conclusion down in the notebook. Later, when it needs to make a decision, it doesn't look back at the original sentence; it just reads its own notes.

The Problem: The "Stale Note" Trap

The researchers tried a simple fix when a piece of information changed (e.g., the order status changed from "shipped" to "pending"). They thought: "If I just update the specific word in the text, the robot will see the change and update its answer."

They were wrong.

Because the robot had already written its conclusion ("Deny the refund") into the notebook based on the old information, simply changing the original word didn't work. The robot ignored the change and kept reading its old note. It was like trying to change a math problem from $2+2=4$ to $2+2=5$ on a piece of paper, but the student had already written "The answer is 4" in their notebook and was just copying that down. The student wouldn't notice the change in the problem unless you forced them to re-read the whole page.

The Solution 1: The "Erratum" (The Sticky Note)

Since the robot relies on its notes, the researchers found a way to fix the mistake without re-reading the whole book.

Instead of trying to surgically erase and rewrite the old note (which fails), they simply append a new, loud note at the end of the document.

The Analogy: Imagine the robot is reading a contract. You can't easily cross out a clause in the middle of the contract without messing up the page numbers. Instead, you stick a bright yellow "ERRATUM" note at the very end that says: "Ignore the previous note. The status is now 'Pending'. The answer is 'Approve'."
The Result: The robot sees this new, loud note, updates its decision, and ignores the old note. This is incredibly fast and works almost perfectly.

The Solution 2: "Composing" Skills (The Lego Brick)

The second discovery is about reusing work.

Imagine you have a long, complex instruction manual for a specific task (like "How to book a flight"). Usually, every time you ask the robot to book a flight, it has to read that whole manual from scratch.

The researchers found that because the robot writes its conclusions into the notebook, you can pre-write the notes for that manual once, save them, and then "paste" them into a new conversation.

The Analogy: Instead of reading a 50-page instruction manual every time you need to build a chair, you have a pre-assembled "Chair Kit" (the notes). You just grab the kit and drop it into your workspace.
The Magic: You can move this kit to a different spot in the conversation (reposition it) and it still works perfectly. The robot doesn't need to re-read the manual; it just picks up the pre-written notes and continues. This makes the process 14 times faster for long documents.

Why This Matters

It's Editable: If a user changes a detail (like their name or an order status), you don't need to restart the whole conversation. You just add a "correction note" at the end, and the robot instantly updates its behavior.
It's Composable: You can build a library of "skills" (like a recipe for booking flights or a guide for handling returns). You can pre-compute these skills and paste them into any conversation instantly, saving massive amounts of time and computer power.
It Works Everywhere: The researchers tested this on many different types of AI models (from small ones to huge ones, and even models that look at images), and it worked for all of them.

Summary in One Sentence

The paper reveals that AI models write their conclusions into a hidden notebook as they read; by realizing this, we can fix mistakes by adding a "correction note" at the end and speed up tasks by pasting in pre-written "skill notes," rather than forcing the AI to re-read everything from scratch.

Technical Summary: Models Take Notes at Prefill: KV Cache Can Be Editable and Composable

1. Problem Statement

Modern Large Language Model (LLM) agents frequently re-read long, mostly-static instructions (system policies, tool specifications, retrieved documents) across multiple turns. While Key/Value (KV) caching allows for the reuse of prefill computations, current production systems (e.g., vLLM, SGLang) rely on exact prefix matching.

The core problem arises when a mutable field within a reused prefix changes (e.g., a user ID, order status, or timestamp). Because every subsequent token in the context attends to the changed token, the keys and values of the entire downstream region become invalid. The de-facto workaround—hoisting all mutable content to the end of the prompt—forces application-layer constraints, preventing the clean reuse of nested sub-agent prompts or dynamically assembled contexts.

A naive attempt to surgically refresh only the changed field's KV vectors while reusing the rest of the cache fails completely: the model reverts to the old decision, ignoring the edit. This paper investigates why this failure occurs and proposes a mechanism to make the KV cache both editable and composable.

2. Core Discovery: Memoized Inference

The paper establishes a causal mechanism explaining why naive KV editing fails. Through four causal probes (locality patching, suffix concentration, linear probing, and circuit knockout) across four model families (Qwen3, Llama-3.1, Gemma-2, Mistral), the authors demonstrate that:

Prefill-Time Memoization: During the prefill phase, the model does not merely cache raw token representations. Instead, it computes the field-conditioned conclusion (e.g., "status is shipped $\to$ action is deny") and writes this conclusion onto the KV vectors of specific downstream aggregator tokens (such as punctuation, newlines, or section breaks).
Decision Reading: At decode time, the decision token does not re-derive the answer from the field itself. Instead, it attends back to these downstream "note" positions.
Causal Contribution: The field's own KV vectors drive less than 1% of the final decision. The decision is driven almost entirely by the "notes" written downstream.

Consequently, refreshing only the field's KV (the "field-only" edit) is ignored because the decision relies on the stale notes downstream. The KV cache is reframed not as a frozen byproduct, but as a notebook of memoized conclusions.

3. Methodology and Capabilities

Based on the "notebook" mechanism, the paper proposes two capabilities: Editable and Composable KV caches.

A. Editable KV Cache

To update a decision when a field changes, one must amend the "notes" rather than just the field. The paper evaluates several strategies:

Naive Field-Only Edit: Fails (recovery $\approx$ 0) unless the model uses Chain-of-Thought (CoT). With CoT, the model re-reads the field, making the field-only edit sufficient (recovery $\approx$ 1.0). Without CoT, it is ignored.
Selective Recompute (Field + Top-K): Recomputing the field plus the top- $K$ downstream tokens with the highest causal effect on the decision. This is reliable but requires model-specific tuning of $K$ and is less robust.
The Erratum (Append-Only): The most robust and cheap solution. Appending a salient correction line (e.g., [STATE UPDATE] field $\to$ new; overrides any earlier value) acts as a new, authoritative note. This matches the performance of a full recompute (Oracle) without rewriting the prompt. It composes with existing prefix caching systems because it does not invalidate the static prefix.

B. Composable KV Cache

If notes are localized and position-portable, precompiled "skills" (long policy blocks or tool specs) can be reused across contexts without re-prefilling.

Mechanism: Precompute the KV cache of a skill in isolation. When inserting it into a new context, RoPE-reposition the keys (rotating them to the new absolute positions) and splice the values (which are position-free) into the target context.
Seam Repair: A small number of boundary tokens at the start of the spliced chunk may need recomputation to account for the missing prefix attention, but this is a constant cost.
Complexity: This reduces Time-to-First-Token (TTFT) from $O(L^2)$ (full prefill) to $O(L)$ (re-rotation + splicing).

4. Key Results

Editing Performance

Accuracy: The "Erratum" method achieves 1.00 decision recovery (matching full recompute) across models, whereas naive field-only edits achieve 0.00 without CoT.
Efficiency: The Erratum approach requires appending only one line of text, incurring negligible compute compared to full recompute.
Comparison to Weight Editing: Unlike weight editing methods (ROME, LoRA) which are global, slow (seconds), and cause cross-request contamination, KV editing is local, fast (~100ms), and request-specific.

Composition Performance

Fidelity: Transplanted skills match full recompute with logit cosine similarity of 0.90–0.999 across 12 models (ranging from 1.7B to 70B parameters, including MoE and quantized variants).
Speedup: TTFT speedups scale with skill length, reaching 13.9× at 32k tokens on an 8B model.
Generality: The method works for rules, RAG passages, and agentic tool-calling. It holds for multimodal inputs (images), where image KV can be spliced with re-rotation of M-RoPE.

Systems Benchmarks (vLLM)

Prefix Cache Alignment: In an online vLLM benchmark with Poisson load, the append-only erratum keeps the prefix cache-aligned, achieving a 98.5% prefix cache hit rate (vs. 1% for in-prefix edits).
Latency & Throughput: The p90 TTFT is reduced by 53–398×. Throughput gains grow with load, reaching 14.5× at saturation compared to a baseline that re-prefills on every change.

User Memory Application

The paper applies these concepts to user memory (large, mutable profiles).

Placement: Placing memory at the end of the prompt (after the trajectory) allows for $O(L)$ editing and transplanting.
Trade-off: While late placement incurs a small accuracy cost due to the lack of "pre-digestion" (prefill-time integration), this is offset by the massive efficiency gains.
Validation: On the LoCoMo benchmark (real long-conversation QA), transplanting memory is statistically equivalent to full recompute in QA accuracy.

5. Significance and Claims

The paper claims to provide the first causal account of why in-place KV field editing fails and offers a unified framework for programmable KV caches.

Mechanism: It identifies that transformers write conclusions to downstream tokens during prefill, making the cache a structured record of intermediate reasoning rather than just raw data.
Programmability: It demonstrates that systems can read, write, and rearrange this "notebook" without retraining the model.
Practical Impact: The approach enables editable and composable caching that is compatible with existing production infrastructure (vLLM, SGLang) and attention mechanisms (FlashAttention, GQA, MLA, M-RoPE).
Future Vision: The authors posit that while current models exploit this mechanism "for free," future models could be trained to expose these notes explicitly, turning the KV cache from a linear log into a programmable memory for agents.

The work is validated across scale (0.6B to 70B), quantization (FP8, 4-bit), architecture families (MoE, standard), and modalities (text, vision), establishing a robust foundation for efficient, stateful LLM agents.