NVIDIA 研发新技术 DMS，实现 8 倍 KV 缓存压缩，大幅提升 LLM 推理效率

近日，英伟达（NVIDIA）和爱丁堡大学的研究人员提出了一种名为动态内存稀疏化（Dynamic Memory Sparsification, DMS）的新方法，该方法能够显著压缩Transformer大语言模型中的键值缓存（KV cache），在不影响模型准确性的情况下实现推理时的超级扩展。 KV 缓存问题 在处理推理密集型任务时，大型语言模型（LLM）需要生成更长的序列或并行的推理链。然而，推理性能受到KV缓存内存占用的严重限制。KV缓存用于存储过去的令牌表示，以便自回归生成时重复使用。它的内存消耗随序列长度和宽度（并行线程数）线性增长，导致GPU内存大量消耗和频繁内存访问，从而减慢了推理速度。 现有技术的局限 目前优化KV缓存的方法主要有两种：一种是基于启发式的无需重新训练的技术，如根据注意力权重进行令牌逐出；另一种是需要复杂后训练改造的技术，如动态内存压缩（Dynamic Memory Compression, DMC）。前者虽然简单，但往往会损害模型准确性，后者则计算开销巨大。 DMS 的创新之处 DMS通过混合方法解决了这些局限。它像传统剪枝方法一样稀疏化KV缓存，但在训练过程中使用了一种基于Gumbel-sigmoid采样的机制，使得逐出决策可微分。被标记为未来逐出的令牌在一段滑动窗口时间内仍然可用，从而保留重要上下文信息，避免精度突降。 DMS的另一个优势在于其高效的改装能力。与DMC相比，DMS不需要为每个注意力头添加额外参数，而是重用注意力机制中的一小部分（单个神经元）来预测逐出。因此，即使只有大约1000步的训练，也能实现8倍KV缓存压缩，并保持甚至提升模型在多种推理任务上的性能。 实验结果 研究人员测试了DMS在多种模型大小（Qwen-R1 1.5B、7B和32B）上的效果。结果表明，在AIME、GPQA和LiveCodeBench等推理密集型基准测试中，DMS大幅提高了精确匹配性能，分别提升了9.1、7.6和9.6个百分点，且没有增加运行时或内存开销。 在与顶级基线模型如Quest和TOVA的对比中，DMS无论是从KV缓存读效率还是峰值内存使用，都表现出更好的帕累托前沿。 多样化应用场景 此外，DMS在非推理任务中的表现也相当出色。在MMLU、GSM8K和HellaSwag等短上下文基准测试中，DMS在4倍压缩比下仅出现了约3.5个百分点的轻微性能下降。在长上下文任务如Needle-in-a-Haystack和Variable Tracking中，DMS的性能甚至超过了未压缩的原生模型，显示出其能够有效缓解信息过压问题的潜力。 结论 总的来说，DMS 提供了一种实用且可扩展的解决方案，可以显著提高Transformer基础语言模型的推理效率。通过智能压缩KV缓存和最少的再训练，DMS使模型能够在不增加运行时间和内存需求的情况下处理更长的序列或并行推理。随着大语言模型在资源受限环境中的日益普及，DMS为实际应用中的压缩、准确性和易集成提供了有力支持，展示出了广泛的应用前景。 业内专家认为，DMS不仅解决了当前KV缓存优化技术的瓶颈，还在保持模型性能的同时大幅度降低了资源消耗。这对于推动大规模语言模型的实际应用具有重要意义。英伟达作为全球领先的高性能计算和人工智能技术公司，此次研究进一步巩固了其在这一领域的领先地位。