当前位置：首页 > article >正文

ChunkKV：优化 KV 缓存压缩，让 LLM 长文本推理更高效

article 2025/2/10 14:24:37

每周跟踪AI热点新闻动向和震撼发展想要探索生成式人工智能的前沿进展吗？订阅我们的简报，深入解析最新的技术突破、实际应用案例和未来的趋势。与全球数同行一同，从行业内部的深度分析和实用指南中受益。不要错过这个机会，成为AI领域的领跑者。点击订阅，与未来同行！订阅：https://rengongzhineng.io/

在长文本推理中，大型语言模型（LLM）需要存储大量的 KV 缓存，这对 GPU 内存占用带来了极大挑战。传统的 KV 缓存压缩方法往往通过注意力分数筛选出“重要”Token，并裁剪掉相对不重要的部分。然而，这些方法通常单独评估每个 Token，忽视了Token 之间的依赖关系。结果就是，某些关键主题词被保留，但与之相关的重要上下文信息却可能被删减，导致模型语义理解出现偏差。因此，一种更结构化的 KV 缓存压缩方法迫在眉睫，它不仅要考虑Token 重要性，还要兼顾语义完整性。

近年来，一些动态 KV 缓存压缩策略被提出，以优化内存使用的同时维持模型性能。例如，H2O 和 SnapKV 通过注意力机制筛选关键 Token，而Chunking（文本分块） 方法则将文本划分为语义相关的段落，以确保上下文连贯性。这种分块技术在 NLP 预训练和检索任务中早已广泛应用。此外，还有像 LISA 和 DoLa 这样的分层策略，利用 Transformer 不同层的结构信息提升模型效率。然而，在 KV 缓存压缩中引入 Token 依赖意识，依然是提升长文本推理能力的关键。

ChunkKV：按块存储，提升 KV 缓存效率

香港大学的研究团队推出了ChunkKV，一种全新的 KV 缓存压缩方法，与传统方法不同的是，它不再单独评估 Token，而是将 Token 组织成“语义块”进行处理。通过按块存储，ChunkKV 减少了内存占用，同时确保语义信息的完整性。除此之外，该方法还引入分层索引复用（Layer-wise Index Reuse） 技术，在不同层间共享已压缩的索引，进一步提升计算效率。

在多个基准测试（如 LongBench、Needle-In-A-Haystack、GSM8K 和 JailbreakV）上的实验表明，ChunkKV 在高压缩率下仍能将模型准确率提高 10%，相比传统方法，它在保留上下文信息和提升内存效率方面表现更优，被认为是当前长文本推理的最优解决方案之一。

ChunkKV 如何优化 GPU 内存？

随着 LLM 支持的文本长度不断增加，KV 缓存的内存占用已经成为推理效率的核心瓶颈。ChunkKV 通过分块存储语义丰富的 Token 组，减少存储需求的同时，仍能保留关键信息。具体来说，它利用注意力分数，筛选出最具信息量的 Token 组，并使用分层索引复用 技术，提高跨层的计算效率。

实验数据显示，与 SnapKV 等现有方法相比，ChunkKV 在层间索引相似度上表现更优，确保了更稳定的语义保留。这一结构化 KV 缓存方法符合 LLM 的上下文学习机制，在优化内存使用的同时，保证了长文本推理的连贯性。

实验结果：ChunkKV 是否真的好用？

研究团队在两大关键任务上评估了 ChunkKV 的有效性：

上下文学习（ICL，In-Context Learning）：在 GSM8K、Many-Shot GSM8K 和 JailbreakV 基准测试中，ChunkKV 结合 LLaMA-3.1-8B-Instruct 和 DeepSeek-R1-Distill-Llama-8B 等模型，在不同压缩率下都保持了更高的准确率。
长文本任务（Long-Context Tasks）：在 LongBench 和 Needle-In-A-Haystack（NIAH） 数据集上的实验显示，ChunkKV 比现有方法更擅长保留关键信息。

此外，索引复用实验 进一步验证了 ChunkKV 的效率提升。在 A40 GPU 上，该方法减少了 20.7% 的推理延迟，并提升了 26.5% 的吞吐量，证明其在计算资源受限环境下的可行性。