在大模型推理中，Transformer的每一层自注意力都要生成query、key和value向量。如果不做缓存，每生成一个新token就要对全历史序列重新投影并计算注意力，计算量随序列长度平方增长。KV Cache把每一层已经算好的key和value保存下来，后续步只计算当前token的query，再与历史K、V做注意力。显存占用公式大致为2乘以层数乘以头数乘以头维度乘以序列长度乘以精度字节数。例如32层、32头、128维、4096长度、FP16下约2×32×32×128×4096×2字节，超过2GB。长上下文模型长度达到128K时单请求缓存就会到几十GB。
KV Cache为何成为显存瓶颈
KV Cache的显存增长与上下文长度成正比。以Llama这类模型为例，隐藏维度4096、32个注意力头、每头128维，KV通常合并缓存，每个token每层需要存储的键和值向量总数是4096个元素。FP16下每个元素占2字节，单层单token就是16KB，32层就是512KB。当长度达到8192时，单请求约4GB显存。如果并发8路请求，仅KV Cache就需要32GB，这还没有算模型权重和激活值。
更麻烦的是，KV Cache在推理的不同阶段行为不一样。预填充阶段可以并行计算整段输入的键值，显存分配较快；解码阶段每步只追加一个token，但显存已经被整段历史占用，无法释放。长文本、多轮对话、批处理场景下，KV Cache几乎成为吞吐量的硬限制。
要降低这个压力，常见思路有三类：一是减少缓存条目数量，比如窗口注意力、检索增强；二是降低每个条目的字节数，也就是量化；三是把不常用的缓存搬到更便宜的存储层级。量化直接作用于字节宽度，Attention Sink则给条目裁剪和层级卸载提供了依据，两者可以组合使用。
KV Cache量化：把显存压到INT8和INT4
KV Cache量化本质是把存储在显存中的键值张量从FP16或BF16压缩为低精度整数。量化不改变序列长度，只减少每个元素的位宽。INT8可以把缓存减半，INT4可以减到四分之一。具体做法通常有三步：先收集一批校准样本，统计各层键值的数值范围；再选择量化参数，例如对称量化的scale；推理时把浮点数除以scale后取整存入缓存，注意力计算前反量化为浮点。
量化的难点在于注意力对数值比较敏感。Key向量参与点积，Value向量参与加权求和，如果量化误差过大，softmax输出可能出现尖峰或偏移。实践中常用per-token或per-channel的粒度来分配scale，而不是对整个张量使用同一个尺度。per-channel对Key的每个头单独计算缩放系数，能保留不同头之间的数值差异；per-token则对每个token单独归一化，适合Value的离群值问题。更精细的分组量化按每32或64个元素一组计算scale，精度更高但需要的元数据也会增加。
import torch
def quantize_kv_int8(tensor, group_size=64):
# 按最后一维分组做对称量化，返回量化值和scale
orig_shape = tensor.shape
tensor = tensor.reshape(-1, group_size)
amax = tensor.abs().amax(dim=-1, keepdim=True).clamp_min(1e-6)
scale = amax / 127.0
q = torch.round(tensor / scale).clamp(-127, 127).to(torch.int8)
scale = scale.to(tensor.dtype)
return q.reshape(orig_shape), scale
def dequantize_kv_int8(q, scale):
return q.to(torch.float16) * scale
上面的代码按分组做对称量化。实际推理框架中，反量化通常和注意力矩阵乘法融合。例如Key量化后，在计算query和key的点积之前实时反量化成FP16，只产生临时浮点数据，显存常驻的仍是INT8或INT4。这样既节省缓存空间，又避免了全缓存反量化带来的额外显存峰值。
INT4比INT8压缩率更高，但对离群值非常敏感。少数绝对值很大的元素会拉高整组的scale，导致大部分正常值被量化到很小的整数范围，反量化后误差明显。解决方案包括保留离群值用高精度存储、改用非对称量化、或对Key和Value采用不同精度。比如Key对点积更敏感，可以保留INT8或FP16，Value用INT4。MLC、vLLM等推理系统已经在部分模型上提供KV Cache量化选项，常见的配置是K用INT8、V用INT4，或K/V都用分组INT4。
Attention Sink：开头的token为何权重特别高
检查长序列生成时的注意力分布，会发现一个稳定现象：前几个token，尤其是句子起始符或第一个真实token，会获得远超其语义贡献的注意力权重。即使这些token被替换成几乎无意义的占位符，后续token仍然会把很大一部分注意力分配给它们。这个现象被称为Attention Sink，是模型在训练过程中形成的偏置。
原因与softmax有关。注意力权重经过softmax后所有位置的概率之和必须为1。长序列里位置很多，但大多数token之间的相关性并不强。模型为了保持输出稳定，倾向把一部分多余的注意力概率固定放在某些绝对位置，开头token就成了这种默认槽位。这有点像内存中的保留区，不参与主要信息检索，却让注意力分布不至于被均匀稀释。
Attention Sink带来的好处是，可以显式保留少量Sink token，然后对中间缓存做激进裁剪。比如滚动窗口模型只保留最近的N个token和开头的4个Sink token，中间历史全部丢弃。StreamingLLM的实验表明，这种配置下数十万长度的文本生成仍能保持较低的困惑度，而单纯窗口注意力一旦长度超过窗口，模型很快崩溃。
import torch
def add_sink_mask(seq_len, sink_ids, window_size):
# 构建注意力掩码：保留sink token和最近window token
keep = set(sink_ids)
start = max(0, seq_len - window_size)
keep.update(range(start, seq_len))
mask = torch.zeros(seq_len, seq_len, dtype=torch.bool)
for i in keep:
mask[:, i] = True
return mask
上述掩码表示每个token只能关注保留位置。sink_ids通常取第0和第1个token，window_size根据显存预算调整。只要开头几个Sink token不被淘汰，模型在生成过程中就有了注意力概率的落点，即使历史信息不完整，输出质量也能维持在可接受范围。
缓存卸载与Sink结合的工程方案
如果量化后KV Cache仍然超出GPU显存，可以把部分缓存搬运到CPU内存。CPU容量远大于显存，但带宽低，直接搬运全部KV会在每一步解码时引入不小的延迟。Attention Sink此时提供了一个优先级策略：把Sink token和最靠近当前token的窗口缓存固定在GPU显存，把较早的中间缓存量化后卸载到CPU。生成新token时，只有当前层需要读取CPU缓存，其他层可以重叠传输和计算。
工程实现可以采用分层调度。每一层维护一个GPU缓存池和一个CPU缓存队列。当序列长度超过阈值，最旧的中间块被异步复制到CPU，GPU只保留Sink块和最近窗口块。注意力计算前判断当前query是否需要访问被卸载的块，如果命中，则通过预取流水线提前把对应块拉回GPU，或者直接在CPU上计算注意力分数，再把结果送回GPU。
class HierarchicalKVCache:
def __init__(self, sink_size, window_size, io_queue):
self.sink_size = sink_size
self.window_size = window_size
self.gpu_cache = {}
self.cpu_cache = {}
self.io_queue = io_queue
def evict(self, layer_id, seq_len):
# 只保留sink和window区间，其他块标记为可卸载
middle = seq_len - self.window_size
if middle > self.sink_size:
block = self.gpu_cache.pop((layer_id, self.sink_size, middle))
self.cpu_cache[(layer_id, self.sink_size, middle)] = block
self.io_queue.push(block)
这种分层缓存设计把量化、Sink裁剪和卸载串了起来。量化减少单块体积，Sink决定必须保留的锚点，卸载扩大整体容量。需要特别注意的是attention计算中的掩码要与缓存位置对应。如果中间块被移走，且模型没有显式Sink anchor，输出可能迅速退化。工程上通常把Sink块设置为不可淘汰，并在模型输入中显式添加专用Sink token。
效果评估与配置策略
评估KV Cache优化不能只看显存下降比例，还要观察长文本生成质量、困惑度和检索准确率。例如在PG19或LongBench这类测试上，INT4 KV Cache相比全精度缓存可能只有不到0.5的困惑度增加，但显存降到四分之一。如果引入Attention Sink并裁剪窗口，需要重点测试多跳推理和长距离依赖任务，因为中间token被裁掉后，模型可能失去跨段落的证据。
一个比较稳妥的策略是分阶段启用。先对Key做分组INT8、Value做分组INT4，显存紧张时再缩小窗口并引入Sink token卸载。每次调整后用同一批下游任务回归，记录困惑度、首token延迟、每token延迟和并发吞吐。对于客服问答、文档摘要等任务，缓存优化通常不会带来明显体验差异；但对于需要严格依据长文档细节的法律或医疗场景，建议保留更多窗口并降低量化激进程度。
此外，量化校准集必须与线上请求分布尽量一致。如果校准数据全是短文本，线上突然出现长代码或数学公式，Key和Value的数值分布可能超出校准范围，导致注意力分数偏移。可以在推理过程中加入动态校准或异常值保护，例如对超过scale范围的元素单独以FP16保存，牺牲少量显存换取稳定性。
综合来看，KV Cache量化和Attention Sink优化并不是彼此独立的技巧。量化解决每个缓存条目的字节宽度，Sink揭示模型注意力分布的静态偏好，使裁剪与卸载有了理论依据。两者组合能够在长上下文和高并发场景下同时降低显存占用与传输开销，是大模型推理工程中非常实用的一对优化手段。
KV Cache量化Attention Sink大模型推理修改时间：2026-09-22 22:44:23