KV Cache Quantization about aqlm HOT 3 CLOSED

Hi! I am not an author, but a contributor, and I have some familiarity with the issue.

As you correctly describe, the AQLM does not do cache quantization, relying on standard transformers code.
If you plug in any data-free cache quantization, e.g. 8-bit KV compression, it is likely that the impact will be the same as when using KV quantization in any other model.
As for 4-bit variant - I did not try that specific one, but according to your description and code, it should be easy to implement with AQLM.

One way you can do this is by extending transformers Cache class roughly as follows:

1. at __init__, create a storage of quantized KVs, similarly to StaticCache
2. during update, de-quantize cache items, return de-quantized cache to the user
3. as a side-effect of update, quantize user's key_states and value_states and write them to your KV storage

This should give you the expected memory saving since only one attention layer is de-quantized at a time.
As for speed-ups, this is unlikely to work any faster than 16bit attention. If you want speedups, you will need to use custom kernels for attention that accept KV inputs in 4 or 8 bit precision.

Since the neurips deadline is soon, it is unlikely that paper authors will be able to write this anytime soon. However, if you try this and share your observations, we'd be glad to take a look. In turn, if you have any issues with AQLM while doing so, please tell us.

from aqlm.

This issue is stale because it has been open for 30 days with no activity.

from aqlm.

For anyone curious, huggingface has done it: https://huggingface.co/blog/kv-cache-quantization

from aqlm.