Sep 19, 2024 · 53m · latent-space
[Paper Club] Writing in the Margins: Chunked Prefill KV Caching for Long Context Retrieval
0:00 / 0:00
▶ Watch on YouTube →
gold bands on the timeline = statements, start to end. Hover to read, click to jump. CC turns on captions
Machine learning researcher Umar Jamil presents 'Writing in the Margins,' an inference-only technique that leverages chunked KV cache prefilling in transformers to generate intermediate annotations and significantly enhance long-context retrieval without model fine-tuning.
How this conversation actually went
Every chapter scored 0–10 on four independent dynamics. Hover any point for the reasoning behind the score. How this is scored →
The hosts as informed peer 1.3 Guest teaching 2.7 Guest disagreement 0.3 The hosts pushing back 0.4
speaking balance: gold is the hosts, purple is the guest (3 minute bins)
Sharpest disagreement ▶ 32:52 Refuting chunk-isolated attention masking
Umar directly rejects the premise of masking preceding chunks, explaining that KV cache representations are inherently contextualized and isolating them forces the model out of distribution.
Hardest push from the hosts ▶ 45:10 Host pushes on short context utilityVibhu challenges the scope of the technique by presenting a small-scale paragraph and sentence highlighting scenario to test where the cost-benefit trade-off collapses.
Biggest teaching moment ▶ 32:52 Contextualized KV tokens masterclassUmar clearly breaks down why autoregressive KV cache states depend on prior tokens, educating the group on why simple causal sub-masking cannot be applied without degrading model performance.
The host holds their own ▶ 46:45 Host outlines hierarchical multi-step chunking architectureVibhu demonstrates deep system architecture familiarity by extending the paper's core margin concept into a multi-tier recursive map-reduce pipeline for constrained context windows.
the scores for every segment, with the reasoning behind each
| Chapter | Topic | The hosts as informed peer | Guest teaching | Guest disagreement | The hosts pushing back | Why |
|---|---|---|---|---|---|---|
| Transformers, KV Caching, and Chunked Prefilling | 0 | 2 | 0 | 0 | Umar gives a monologue introduction explaining autoregressive transformers, KV caching, prefill quadratic complexity, and the book-reading analogy for chunked prefilling. The hosts act solely as session facilitators. | |
| Mechanics of the Writing in the Margins Technique | 0 | 2 | 0 | 0 | Umar walks through the core mechanism of Writing in the Margins, including extractive intermediate summary generation, cache eviction from the tail, and overcoming Lost in the Middle. | |
| Performance Benchmarks and Token Efficiency | 0 | 2 | 0 | 0 | Umar presents token cost comparisons, illustrating how naive re-prompting incurs double the prefill cost while Writing in the Margins stays near single-pass cost. | |
| Video Demo, Memory Implementation, and Code Repository | 0 | 2 | 0 | 0 | Umar shows a demo video and the GitHub repository, highlighting O(1) PagedAttention tensor resizing, overlapping margin classification, and human-in-the-loop progress updates. | |
| Technical Q&A: Attention Masking and Computation Overlapping | 3 | 5 | 1 | 1 | Eugene and audience members ask technical questions about attention masking and pipeline overlap. Umar directly clarifies why masking preceding chunks breaks causal autoregressive modeling due to contextualized tokens. | |
| Technical Q&A: Prefill Quadratic Dynamics and Trade-Offs | 4 | 4 | 1 | 2 | Audience members and host Vibhu probe chunk sizing trade-offs, break-even context thresholds, and multi-tier chunking. Umar explains compute-bound prefilling versus memory-bound token generation dynamics. | |
| Future Research Directions and Paper Club Announcements | 2 | 2 | 0 | 0 | Umar outlines future research directions focusing on attention softmax distribution issues like Sink Attention and Sigmoid Attention, followed by host announcements for upcoming Paper Club sessions. |
Statements from this episode (11)
Jamil: Chunked prefill is experimental in vLLM, likely used by majors
“This is called the chunked pre-fill, and it's an experimental feature that has been recently introduced in VLLN, but it's probably used in more sophisticated inference engines at major companies.”
Jamil: 'Writing in the Margins' works on any transformer without fine-tuning
“So it can be used with any transformer model without fine-tuning, just by doing it, just by doing this inference differently.”
Jamil: Smaller models gain more improvement from Writing in the Margins
“As you can see, for example smaller models have a better more more improvement.”
Jamil: Writing in the Margins avoids re-prefilling tokens, halving compute cost
“Again, to the language model to generate the answer, and it would cost you another million, because the model has to reprocess this prefilling again of one million tokens, so it would cost you two million tokens, but with writing in the margins, it would cost …”
Jamil: Writing in the Margins solves the lost-in-the-middle problem
“It improves the ability of any language model to extract relevant information, so solving the lost in the middle problem”
Writing in the Margins requires inference engine modifications but zero fine-tuning
“We are not doing any change to the model architecture, so you don't have to fine-tune anything, you don't have to change anything, like the can you use this stuff with, like, a Lank pane? No, because it requires a modification on how the inference engine is us…”
Masking out prior KV cache tokens pushes autoregressive transformers out of distribution
“The token number two in the KVCache is a contextualized version of the token zero, one, and two. So if you tell the model to only look at the last tokens you are creating an autoregressive model that is generating the logits of a P of let's say X, but only loo…”
Jamil: KV Cache Prefilling Scales Quadratically While Token Generation Is Linear
“Token generation, which means generating one token using whatever is in the kvcash, is linear with respect to whatever is in the side of the kvcash. Prefilling the kvcash is quadratic, and mostly because it's quadratic, it's very expensive, so we are talking a…”
Jamil: OpenAI and Cohere Overlap Prefill and Generation to Maximize GPU Utilization
“Token generation is memory bound means that the limitation is only given by how much your KVCache can hold. So the memory can hold in terms of KVCache. While prefilling is compute bound, so to maximize the GPU utilization, whenever you work with OpenAI or Cohe…”
Jamil: Ablations Show Supplying Both Context and Margins Outperforms Either Alone
“And we also prove in the ablation studies that actually it's always convenient to send the context plus the margins, never just the margins. So you can see here, this ablation contest compression. So if you only send the margins or only the context, it's alway…”
Jamil: LLM Attention Allocates Most Weight to Initial Tokens
“So we have seen with the paper called sync attentions that actually the language model allocates a lot of a lot of, because when you do the attention mechanism, you are doing a weighted sum over the tokens, and each token is given a weight, and we see that mos…”