GPU implementation

[zhou9402]

Recently, KV cache compression has emerged as a critical optimization technique for large language models (LLMs). The KV cache exhibits strong temporal and spatial locality—similar to time-series data—where adjacent tokens and attention heads often share redundant patterns. Given these characteristics, does this algorithm (or approach) effectively adapt to KV cache compression while maintaining efficient GPU execution?

[azimafroozeh]

Hi,
Thanks for bringing this up — I wasn’t aware of KV cache compression.
So, to answer your question, it’s hard to give a definitive answer without any benchmarks. Based on your description of the KV cache characteristics and a bit of research I’ve done, it does seem like ALP could be a good fit.

As for GPU execution, we’re planning to release a GPU-optimized version of ALP soon, which may be especially interesting to try. You can follow the latest updates and join the discussion on our Discord: https://discord.gg/JsyHHWGY.

Overall, it’s very exciting to explore using ALP for compressing the KV cache!

[zhou9402]

Thanks! And one more thing is in LLM most of the data is bf16 or even fp8(E4M3 and E5M2), can ALP also extend to support these?

[azimafroozeh]

We’ve extended ALP’s double-precision implementation to single-precision (float), and it works well. Based on the formats you mentioned, I believe ALP could also work for bf16 and fp8. However, it’s hard to say for sure without benchmarks. If you can provide a way to generate a dataset similar to what you’d expect in the KV cache, we can give you a more definitive answer very quickly.

[zhou9402]

kv_cache_original.txt

diff_data.txt

Chang txt -> bin

First row is original and second is diff between different rows, and this is the analysis I get:

Below is the distribution of diff data

[azimafroozeh]

Thanks for sharing the dataset! I'll look into it and get back to you.

[azimafroozeh]

I'm looking at kv_cache_original.txt — is the floating point format BF16?

[zhou9402]

I'm looking at kv_cache_original.txt — is the floating point format BF16?

No, it is float32 but transformed from bf16. original is original value, and diff is difference between consecutive values, both of them are float32 transformed from bf16.

[azimafroozeh]

I checked both data files you shared. FastLanes compresses kv_cache_original to 421 KB, giving a compression ratio of approximately 2.4× (1 MB / 421 KB), and compresses data_diff to 612 KB, compression ratio of about 1.6×. Check cwida/FastLanes#17 — I've added your dataset to the FastLanes repo for testing purposes. Also check https://github.com/azimafroozeh/reproduce_LLM_KV_cache_compression_result to reproduce the results using PyFastLanes!