Question: Details of KV cache compression ratios and settings for all models

[bg51717]

Hello, thank you for sharing this great work!

I have some detailed questions about the KV cache compression experiments:

1. How exactly are the KV cache compression ratios calculated?

2. For all reported models and all compression ratios in the paper/experiments, could you share the specific parameter settings?

   • qk_rope_head_dim
   • kv_lora_rank

Thanks a lot for your time and support!

[Stanleytowne]

Hi,

Thank you for your interest in our work and for your detailed questions! I'm happy to provide more specifics on the KV cache compression experiments.

1. How the KV Cache Compression Ratio is Calculated

The compression ratio can be calculated directly by comparing the KV cache size per token per layer before and after applying our method.

For a standard model with Grouped-Query Attention (GQA), the original KV cache size per token per layer is:
2 * num_kv_heads * head_dim
(The 2 accounts for both the Key and Value caches).

With our TransMLA method, the KV cache is reparameterized. The new size per token per layer becomes the sum of the LoRA rank for the non-positional KV (nope_kv) and the dimension of the single RoPE-encoded Key head (rope_k):
kv_lora_rank + qk_rope_head_dim

Therefore, the compression ratio is calculated as follows:
$$\text{Compression Ratio} = 1 - \frac{\text{Compressed Size}}{\text{Original Size}} = 1 - \frac{\text{kv lora rank} + \text{qk rope head dim}}{2 \times \text{num kv heads} \times \text{head dim}}$$

2. Specific Parameter Settings

The experiments in our paper follow this calculation. For the Llama-2-7B model, the baseline KV cache size is 2 * 32 * 128 = 8192 per token per layer.

The parameter settings for the compression ratios reported in the paper are as follows:

• -68.75% compression (New size: 2560): kv_lora_rank = 2432, qk_rope_head_dim = 128
• -81.25% compression (New size: 1536): kv_lora_rank = 1408, qk_rope_head_dim = 128
• -87.5% compression (New size: 1024): kv_lora_rank = 896, qk_rope_head_dim = 128
• -92.97% compression (New size: 576): kv_lora_rank = 448, qk_rope_head_dim = 128

3. Note on vLLM Acceleration Experiments

For the vLLM acceleration benchmarks, the fused kernel only supported the configuration of kv_lora_rank=512 and qk_rope_head_dim=64.

This setting results in a total compressed size of 512 + 64 = 576, which is the same total cache footprint as the -92.97% setting (448 + 128) mentioned above. This is the configuration we used for the vLLM-related results in the paper.

Hope this clarifies things!

[bg51717]

I also want to know the setting for SmolLM-1.7B. Thanks!