26× Faster Inference with Layer-Condensed KV Cache for Large Language Models
arxiv.org
arxiv.org
Initial result -- those KV caches in lower layers matter, and output suffered.
Updated plan -- cull half the KV layers! This works 'nearly' as well as keeping all of them, with memory and compute savings.
Downside - triple the training, worse out of band / long context performance.
This feels to me like a technique you'd use on a particular architecture deployed at the edge where compute matters and you have a little extra room on performance. Phi-3 on raspberry pi, basically.
Interesting! As always, I wish models showed prompt output in their papers, not just perplexity numbers. But, here we are.
From what I understand, not quite. It looks like the cost of training might be similar, but less parallelisable within a specific token sequence. This is because they have to compute the KV of token T before they can use it in T+1 whereas in a regular training process you can compute the KV at each layer for every subsequence. You're right that it took 2.7x longer to train the smallest model but I wouldn't be surprised if the GPU utilisation was proportionally lower too.
If we split a model layer-wise for inference across X GPUs, does the cache split as well?
1. The correct title is Layer-Condensed KV Cache for Efficient Inference of Large Language Models.
2. The paper does make a 26x claim later in the introduction, but it’s an outlier.
26x is for only one benchmark and that benchmark is CPU based, not GPU based like 99% of transformer loads actually run on.
If you look at GPU only workloads, the improvements range from 1.4x to 4.7x.
Edit: I apologize. The table was cut off on mobile and I didn't see that they sneaked in GPU+CPU offloading for the 25x result.
Best way to address this is to email hn@ycombinator.com.
It is not possible to read all posts to HN.
(Also, the convention is to use 'dang rather than @dang, borrowing from a lisp convention.)
"Layer-Condensed KV Cache for Efficient Inference of Large Language Models"
Not sure what this means? Would this work for a Mistral model as well?
It’s not hard to find ways to speed up transformers if you’re willing to give up quality.
You could argue some tradeoffs are worth it and that’s true sometimes but I don’t see that they’ve made the case for it here.
You want to store your KV cache in the same processor that does the rest of your matmuls.
This is unavoidable by the way. Any method to avoid quadratic attention will necessarily have to degrade accuracy, because less than quadratic means you will have to look at "less than all the tokens" in a given pass. You're bound to miss at least some of them. When you consider how simple the classical attention mechanism is, you realize that there is not much you can do. Any preattention pass is necessarily going to be more complicated than your quadratic attention mechanism.
What they do is just get rid of it entirely after w layers. This saves memory consumption, which allows you to fit more context in the same amount of VRAM. In the paper they decided to run more batches in parallel, but they could have also advertised a massive increase in context length. 128k context needs 16GB on my machine for llama3 7b. Their approach would allow at least 784k context in the same 16GB.