Lightning-AI
Support for KV caching and batched inference
Adds abstraction for key-value caches, implements batched inference.
I am also adding two baseline KV caches, the default one from before (all KV are stored) and a last-recent one.
The abstraction contains methods not used by these baselines, but they are required to implement more advanced KV caches such as Heavy Hitter Oracle (H2O).
I have implemented some of these, but I may not be allowed to contribute them here (working for a company). I'll see what I can do.
Hey, great work @mseeger .
Can we decouple things a lot, though?
Some initial thoughts:
- I would prefer if we kept the KVCache initialization as in the current version (i.e. that you initialize the model, then potentially adjust the max seq len and then initialize the KVCache) in this PR. Adding this to the init parameters seems orthogonal to the other changes.
- We do have batched generation today. Can we please split changes to batched generation from the KVCache improvements. We probably don't want to do batching via lists of tensors. I'm currently looking at passing in "packed" input/input_pos sequences but these changes. Changing the existing tests should be a bit of a red flag, as it will screw existing users to change the API (we can do this if we need to, but TBH I am not convinced this is the case).
- In general, can we be very conservative with adding arguments? For optional arguments, we should look into making them keyword-only unless there is a good reason not to.
- We do keep control flow simple.
self._default_kv_cache = Falseis not a good idea. - I'm not sure I understand the
both_in_parallel. Maybe the right time to add it and the associated refactors is when they are used? - I'm generally a bit weary of the amount of data structures and cases that are being passed around here, those add a lot of complexity. To my mind, this likely means that the right abstraction has not yet been found. Maybe integrating KVCache and SDPA more could be a thing, but I am not sure.
- In general, we do not want to do the cache setup during the forward. Please keep the initialization separate. I think we are rather seeing movement towards less of a distinction between pre-fill and next token, so this seems a bit in the wrong direction.
Again, super good stuff in the PR! I think there are a few things to split out and consider individually and then maybe we can have a video call about the core KVCache things, wdyt?
Thanks for the initiative for better KVCacheing!
Hello, sure we can have a call, I am in the central Europe (Germany) time zone.
My impression was that batched generation is not really there. But if it is, I don't ask to change it.
One thing is important through. KV caches really work by filling positions sequentially. So, you filled positions 0:(T-1), you need to continue with T, or with T:(T+k). The current API of just passing some position indexes is really not going to work.
Also, the implementation right now allows you to send in KV cache objects from the start. If you do not do that, it will create them by default. This is done by set_kv_cache. If you also do not do that, it is done in the first forward with for_prefill=True.
Note that prefill here means that I can do a single pass, and the cache can take it all, without having to evict anything. It does not mean that this will encode even the shortest prompt in the batch. If prompts are longer than the max prefill length, you need to do it sequentially in chunks.
Maybe there is an easier way, we can discuss.
It is annoying I cannot show you the KV cache code I have. But in a talk, I could explain why a few things are the way they are. Of course, I am not on top of other constraints you guys have.
You may ask why KVCache.prefill? The main reason is that you want to use SDPA whenever you can, but SDPA cannot return the attention weights, which some KV cache algorithms (H2O) need in order to decide what to evict next.
We can do things so the very first call to the model, with input_pos=0, is doing this. So, instead of
model(x, for_prefill=True)
you'd call
model(x, input_pos=0)
This I could do. That would indeed be a little simpler.
@t-vi Let me know what the next steps here should be. If I understand correctly, I could:
- Get rid of
for_prefillparameter, and useinput_pos=0instead - Don't create default KV cache in
forwardand rather fail the call ifinput_posis used, s.t. user needs to callset_kv_cache - You don't seem to approve of passing the KV caches at construction (if user does not want to use default ones). Would you rather use
set_kv_cachefor that?
Hi, so I think we should try to break things down.
We could either start with the core caching itself and try to see how to integrate it with minimal changes or see what is the deal with batching and prefill first. I sent to your gmail address to find a good time to discuss.
Hello @t-vi , let me try to break things down. Changes are these:
KVCacheand its implementations. This replaces the default cache, which just stores everything. No behavior changes.- Caches for each layer can be passed when model is created. Before, there is
set_kvcache, which creates the default caches. If nothing is done at all, default caches are created when first needed. This is a change. Before, it would create an exception. - Refactoring of generation code: This works for batch generation now, and single sequence generation is a special case. Inside, this properly supports large prompts by splitting generation into prefill (as large as caches allow), and then aequential blocks of desired length.
If I understand you correctly, you complain about 2., especially the automatic creation of default cache when nothing is done, and the change of __init__ of GPT. This, I can work on. I could to the following:
- Allow passing KV caches per layer in
set_kvcache(or have another method?) - Create default KV caches by calling
set_kvcache. If this is not done, callingforwardfor inference fails, so no cache is created automatically
Would that be what you prefer?
As for 1. and 3., in the end, they go together, but I can try split it into two. I'd first do 1., keeping the generation code in place, which would however not work for batches and not support the sequential processing of prompts properly.
First doing 3. is not really sensible, because it requires things from 1.
What do you think?
Note that with DeepSeek (I am involved trying to bring this to Hugging Face), there is a lot of movement now not to ignore KV caching in the future. They even released a paper now how they can train with large contexts.
@t-vi Maybe I can change your mind about first keeping the current generation code in place, and only contribute the KV cache support?
This is quite a bit of extra work for me, and new code of mine has a number of improvements. in particular, the current code does not really do batch generation, it is marked with several TODO and is not used.
If we could have a chat, I'd appreciate that.
Out of curiosity: Why do you object to batch prompts being a list of tensors? In the end, they can have wildly different lengths, and there is not much you can do against that (sure, if you get lots of requests, you can maybe cluster them, but doing this too much delays requests, so increases latency).
Also, you really don't want to push PAD tokens into models just because a prompt in a batch happens to be shorter than others. The model, not being trained on this, would certainly get confused. And since you need to start token-by-token forward for generation, you really gain nothing by padding prompts.
I always thought if this as some kind of TensorFlow artefact when all tensors had to be allocated up front, etc. But I thought we have overcome this with PyTorch.
Hey, sorry, I am totally swamped, still want to have a video call to chat.
Out of curiosity: Why do you object to batch prompts being a list of tensors? In the end, they can have wildly different lengths, and there is not much you can do against that (sure, if you get lots of requests, you can maybe cluster them, but doing this too much delays requests, so increases latency).
Because lists are a lot less nice to work with in various setups passing to kernels, cudagraphs etc.
For somewhat homogeneous seq lengths, padding works fine. We are using it in production, so I'm doubting claims that it does not work. It does have limitations with the inhomogenous sequence lengths, which we want to support.
But the proper way to support this is packed sequences, i.e. pass in flat input_tokens, input_pos (i.e. 1d shape, no batch index) and then batch_seq_lens of shape batch_size.
batch_seq_lens gives lengths for each batch item (and might even be 0).
This is hugely more flexible. It needs FlexAttention or somesuch https://pytorch.org/docs/stable/nn.attention.flex_attention.html to make it work efficiently in stock PyTorch.
After our call, I think I understand more what you mean. Something like an abstraction in multi-head attention, where the input are keys, values, query for the current input chunk, all the same size, but then this is bundled:
- Take in keys, values and replace with KV-cached ones, so now keys, values are larger
- Do the SDPA computation
- Feed attention weights back to KV cache if needed
- Return MHA outputs before final linear mapping
This makes a lot of sense, and is quite elegant.
Your CI system seems to be broken still.
@mseeger shall be fixed now, thank you for your patience :)
As discussed with @t-vi , I'll refactor this as stated in the comment above. Makes total sense
@t-vi , is this what you had in mind? KVCache is now very simple, and so the code in CausalSelfAttention is not polluted by details of KV caching. I added DefaultKVCache which most KV caches will use.
I'd be OK to take out my batched inference part here, but this means there won't be any. Do you have plans to add the batched inference code you talked about any time soon?
OK, I've taken out the batched inference code. Still working on fixing the tests (and need to refactor speculative decoding), but this is essentially it.
@t-vi , it would be great to get some feedback on this one, before I spent time on fixing tests for code which I need to change afterwards anyway.
BTW: Even vLLM does not have a consistent support of different KV caching techniques and strategies. They just offer some blunt stuff like tensor parallelism and quantization.
This could be a real differentiating feature of LitGPT.
If you know another open source library that indexes on KV caching, and which you'd like to integrate with instead, please let me know.
@t-vi , @Borda : Any change there will be some progress here?
I recognize this is a big PR. On the other hand, a decent support for selective (sparse) KV caching could be a real differentiator for LitGPT, in that none of the other OS libraries spend effort on it, not even vLLM.
I made quite some progress, also on fine-tuning with long contexts. I am trying to get approval to open source this. In my team, we start to use LitGPT quite seriously (HF being just too messy).
One thing still missing is good support for batch inference without excessive padding.
@t-vi , @Borda: Any sign of life on this PR? A few things:
- This is not meant to be a single PR, but I am happy to chop it into smaller pieces. At the moment, I am simply using this PR to get my work done on the extra repo I am writing
- I made substantial progress: I now have a method to compute gradients (!) with KV caching inside. To my best knowledgem, this does not exist so far in any OS repo. This could be a real strong point for
LitGPT - I am getting closer to being allowed to open source my code (which is a library on top of
LitGPT), but this needs some integration into your code, which is what I am proposing here - Finally, I'd argue my code is also cleaning a few things up