4-bit KV Cache

[Dampfinchen]

Turboderp, developer of Exllama V2 has made a breakthrough: A 4 bit KV Cache that seemingly performs on par with FP16. Here are his words:

"I'm working on some benchmarks at the moment, but they're taking a while to run. Preliminary results show the Q4 cache mode is more precise overall than FP8, and comparable to full precision. HumanEval tests are still running."

https://www.reddit.com/r/LocalLLaMA/comments/1b9571u/80k_context_possible_with_cache_4bit/

This is huge. Sadly, Llama.cpp doesn't even have full 8 bit cache right now (only K cache). So in that aspect, there's a lot of potential for improvement.

[puppetm4st3r]

thats is huge! hope you can implement in llama.cpp :)

[BarfingLemurs]

Here is some benchmarks and information - https://github.com/turboderp/exllamav2/blob/master/doc/qcache_eval.md

[Dampfinchen]

Here is some benchmarks and information - https://github.com/turboderp/exllamav2/blob/master/doc/qcache_eval.md

Very interesting. What do @ggerganov @JohannesGaessler and @slaren think about these results? The current consensus is that 4 bit KV cache isn't worth it as the uptick in perplexity would be too severe. However, that isn't the case with Turboderp's implementation. I wonder what llama.cpp can learn from that.

[JohannesGaessler]

The MMQ kernels all use the same template with quantization type specific functions for allocation, data loading, and the actual dot products. I meant the latter when I was talking about the MMQ code. I'm thinking the way to go is to write a new template for FA and re-use the functions.

[henk717]

I have had overwhelming feedback from users that they are eager for 4-bit cache because it helps them run larger models they currently need to resort to EXL2 for. They prefer Llamacpp based solutions but when it comes down to being able to run a model and not being able to run a model on their limited VRAM they prefer to be able to run a model.

So from a user demand perspective at least on our community demand has been high for this and its something I expect the majority of our users to turn on if implemented (and given the nature of the request I expect this to be the same in the entire ecosystem) . I fully understand the "if it makes sense to add" is from a coding perspective and the end result of implementations can sway it one way or the other, but I did want to ensure that the interest from the users is known to help guide any decisions / evaluations if it's down to user interest.

[JohannesGaessler]

I mean, in September of last year I had already made a PR for quantizing the KV cache so I definitely agree that it would be a very good feature to have. But that PR was never merged due to introducing too many changes so for a new implementation I'll have to be careful to lay sufficient groundwork in advance.

[Dampfinchen]

On that note, what happened to your enhanced mmq code with tensor core support, Johannes? I remember being pretty impressed by the speed increase and low VRAM usage back then.

#4801 (comment)

Perhaps this would be useful for large batch sizes when you are working on the 4 bit kv cache?

[JohannesGaessler]

I at some point noticed that there are numerical issues with my implementation that are not easily fixable (I just never posted about this). After talking to an NVIDIA engineer I'm also thinking that I'll need to overhaul the implementation in general. But I have definitely not forgotten about this. My next three major goals are FlashAttention, n-gram lookup, and int8 tensor core support. After that I'll look into training.

[kir-gadjello]

Integral part of good performance of turboderp's KV-cache quantization is Hadamard transform for smoothing the kv distribution:

See https://arxiv.org/pdf/2404.00456 and this part of turboderp's implementation turboderp/exllamav2@324404e

[henk717]

Now that #5021 is merged i'd love to see this, we have a lot of users who are eager for 4-bit KV cache to help save vram.

[strawberrymelonpanda]

Count me as one of them, but the new Flash Attention flag really helps already. Combined with this, sounds amazing.

[WiseFarAI]

This might be useful:

1. https://github.com/ClubieDong/QAQ-KVCacheQuantization
2. https://arxiv.org/abs/2404.00456