KV Cache

在LLM推理时，在attention层中计算attention score时，需要计算query和key的点积，然后进行softmax归一化，最后与value进行加权求和。由于自回归的特性，每次生成一个token都需要和之前的所有token进行attention计算，这个时候就需要将之前的key和value进行缓存，以便下一次计算时直接使用。是一种用来加速推理的技术。

但是KV Cache占用显存空间很大，尤其对于Long Context的inference，一般可以到 30% 以上。因此需要对KV Cache进行压缩，以减少显存占用。

KV Cache压缩和计算优化

压缩和计算优化是两个不同的概念，压缩是为了减少显存占用，计算优化是为了减少计算量。压缩和计算优化可以同时进行，也可以单独进行。

KV Cache压缩的方法有很多，主要分为两类，一类是量化(有机会再单独写一篇文章，现在我也整不明白)，一类是剪枝。量化是将浮点数转换为低精度的整数，剪枝是将不重要的token/channel/layer进行eviction 或者分配较少的budget。

首先我们看一下KV Cache的结构。

其中，

• B是batch size
• N是sequence length
• H是head number
• L是attention layers number
• d是 embedding dimension

一般来说Batch size无法进行优化，因此我们主要关注这四个维度。

因此，我们可以从这四个维度进行压缩。

1. N维度压缩，即对sequence length进行压缩，可以使用剪枝的方法，将不重要的token进行eviction。一般有以下几种
   1. 随机剪枝，随机选择一些token进行eviction。
   2. Sliding window，只保留最近的N个token。
2. H维度压缩，即对head number进行压缩，可以使用剪枝的方法，将不重要的head进行eviction。某些head对于推理的结果影响不大，我们可以将这些head进行eviction或分配较少的Cache Budget。
   1. RazorAttention 通过仅为少数检索头保留全量缓存、对其他头丢弃远程 tokens 并用补偿 token 恢复信息，实现了 >70% 的 KV‑Cache 缩减、无性能损失且无需重训、兼容 FlashAttention。
3. L维度压缩，即对attention layers number进行压缩，可以使用剪枝的方法，将不重要的layer进行eviction。比如说某些layer对于推理的结果影响不大，我们可以将这些layer进行eviction或分配较少的Cache Budget。
   1. SqueezeAttention
   2. CAKE CAKE 将 KV 缓存驱逐抽象为“蛋糕切分”问题，通过结合各层的时空注意力动态按层级级联分配缓存资源并引入新的时序驱逐指标，在仅用 3.2% 缓存的前提下保持性能并显著加速长序列推理。
4. d维度压缩，即对embedding dimension进行压缩，一般是对某个一些不重要的维度进行删除，比如说 256 维度的embedding，我们可以删除一些不重要的维度，比如说删除128维度，这样就可以将embedding dimension从256降到128。
5. Low Rank Projection: 将KV Cache进行低秩分解，然后只保留低秩分解的结果，这样就可以减少显存的使用。
   1. Palu Palu 通过对生成 Key/Value 的投影层做低秩分解，仅缓存低秩中间结果并在注意力计算时动态重构全维 KV，从而大幅降低内存占用并提升推理效率。

还有可以通过CPU offloading来减少显存的使用，将一些不重要的token的KV Cache放到CPU上，这样就可以减少显存的使用。

KV Cache的计算优化

TopK Selection：在计算attention score的时候，我们可以只选择topK个score，这样就可以减少计算量。

这里计算量的减少主要是在加权求和值的计算上，因为 的维度变小了，从而减少了计算量。之前的计算量是 ，现在的计算量是 ，其中 是序列长度， 是embedding dimension， 是topK的个数。(这里计算的是一个token的计算量)

对于Selection还有一些其他的做法，比如说通过Sampling的方式来选择，可以保留部分整体的信息比如MagicPIG 通过 LSH 采样加上 CPU 异构执行，以带理论保证的方式近似注意力计算，为超长上下文 LLM 推理提供高效、低延迟的解决方案。