Jamil: KV Cache Prefilling Scales Quadratically While Token Generation Is Linear
Umar Jamil · [Paper Club] Writing in the Margins: Chunked Prefill KV Caching for Long Context Retrieval · Sep 19, 2024 · at 42:03
AI researcher Umar Jamil explains the computational trade-offs between autoregressive generation and prefilling attention matrices for long contexts in LLMs.
0:00 / 0:14exact quote · 14.2s
▶ Watch the full episode on YouTube → 720p mp4 · rendered on demand · StarZero watermark
“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 about something that is linear with, Quadratic.”
quote is from the automated transcript, cleaned for reading: filler sounds and stutters are removed, nothing is rephrased. names can be misheard (the analysis reads context, assessments check outside sources). how →
More from Umar Jamil
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”
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: 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: 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 …”
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…”