英伟达开源KVzap:一种近乎无损的KV缓存剪枝方法

4 天前·来源:MarkTechPost
英伟达KV缓存剪枝方法TransformerAI优化

英伟达开源了KVzap,这是一种先进的KV缓存剪枝方法,能实现2到4倍的压缩。KVzap使用小型代理模型替代昂贵的评分机制,直接基于隐藏状态预测重要性分数。该方法在保持模型准确性的同时,显著减少了内存占用,适用于长上下文部署。

随着上下文长度达到数万或数十万令牌,Transformer解码器中的键值缓存成为主要部署瓶颈。缓存存储每个层和头的键值,形状为(2, L, H, T, D)。对于Llama1-65B这样的标准Transformer,在128k令牌和bfloat16精度下,缓存约达335 GB,直接限制了批处理大小并增加了首令牌生成时间。

生产模型已沿多个轴压缩缓存。分组查询注意力在多个查询间共享键值,在Llama3中实现4倍压缩,GLM 4.5中12倍,Qwen3-235B-A22B中高达16倍,均沿头轴进行。DeepSeek V2通过多头潜在注意力压缩键值维度。混合模型将注意力与滑动窗口注意力或状态空间层结合,以减少维护完整缓存的层数。

这些变化未沿序列轴压缩。稀疏和检索式注意力在每个解码步骤仅检索缓存子集,但所有令牌仍占用内存。因此,实用的长上下文服务需要删除对后续令牌影响可忽略的缓存条目的技术。

英伟达的KVpress项目在一个代码库中收集了二十多种此类剪枝方法,并通过Hugging Face上的公共排行榜公开。H2O、Expected Attention、DuoAttention、Compactor和KVzip等方法均以一致方式评估。

KVzip是当前KVpress排行榜上最强的缓存剪枝基线。它使用复制粘贴预文本任务为每个缓存条目定义重要性分数。模型在扩展提示上运行,要求精确重复原始上下文。对于原始提示中的每个令牌位置,分数是重复段中任何位置分配回该令牌的最大注意力权重,在使用分组查询注意力时跨同一组中的头。低分条目被逐出,直到满足全局预算。

KVzip+改进了此分数。它将注意力权重乘以值贡献到残差流的范数,并通过接收隐藏状态的范数归一化。这更好地匹配了令牌在残差流中引起的实际变化,并提高了与下游准确性的相关性。

这些oracle分数有效但昂贵。KVzip需要在扩展提示上预填充,这使上下文长度加倍,对生产环境来说太慢。它也无法在解码期间运行,因为评分过程假设固定提示。

KVzap用小型代理模型替代oracle评分,直接操作隐藏状态。对于每个Transformer层和每个序列位置t,模块接收隐藏向量hₜ并输出每个键值头的预测对数分数。考虑两种架构:单线性层(KVzap Linear)和两层MLP,带GELU且隐藏宽度等于模型隐藏大小的八分之一(KVzap MLP)。

训练使用Nemotron预训练数据集样本中的提示。研究团队筛选27k个提示,长度在750到1,250令牌之间,每个子集采样最多500个提示,然后每个提示采样500个令牌位置。对于每个键值头,他们获得约120万训练对和23k验证对。代理模型学习从隐藏状态回归到对数KVzip+分数。

背景阅读

在大型语言模型中,Transformer架构的键值缓存用于存储注意力机制中的键和值,以加速自回归生成过程。随着模型上下文长度的增加,如从数千扩展到数十万令牌,缓存的内存占用急剧增长,成为部署中的主要瓶颈。这限制了批处理大小,增加了延迟,并提高了硬件成本。为了解决这一问题,业界开发了多种压缩技术,包括分组查询注意力、稀疏注意力和状态空间模型,但这些方法通常沿头或维度轴压缩,而未有效处理序列轴。KV缓存剪枝技术旨在删除缓存中对未来预测影响较小的条目,从而减少内存使用而不显著降低模型性能。英伟达的KVpress项目整合了多种剪枝方法,提供了一个标准化评估平台,推动了该领域的研究和实际应用。KVzap作为最新进展,通过轻量级代理模型实现高效评分,代表了在长上下文服务中优化内存效率的重要一步。

评论 (0)

登录后参与评论

加载评论中...