BasisKV: Value Latents for Communication and Key Routing
作者: Lixun Zhang, Zishan Shao, Kangning Cui, Yu Liang, Fan Yang, Yiran Chen
分类: cs.AI
发布日期: 2026-11-02
🔗 代码/项目: GITHUB
💡 一句话要点
提出BasisKV以解决长上下文LLM推理中的KV缓存问题
🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 长上下文推理 KV缓存压缩 低秩投影 张量并行 大型语言模型 分布式推理 效率提升
📋 核心要点
- 现有方法在分布式推理中未充分考虑张量并行通信的成本,导致KV缓存压缩效果不佳。
- 提出BasisKV,通过联合优化每个GPU的低秩因子,减少Value缓存宽度和张量并行通信。
- 在多个LLM家族中,BasisKV实现了比现有低秩基线更好的质量-效率平衡,提升了推理速度。
📝 摘要(中文)
长上下文的大型语言模型(LLM)推理受到维护不断增长的KV缓存成本的限制。低秩K/V投影为以最小质量损失压缩KV缓存提供了有前景的方法,但之前的工作主要在孤立的情况下优化KV缓存压缩,导致分布式推理,尤其是张量并行计算的研究相对不足。本文提出BasisKV,一种针对通信密集型长上下文LLM推理的低秩KV缓存压缩方法,能够在多个GPU上优化低秩因子,显著提高效率和质量。实验结果表明,BasisKV在多个LLM家族和系统设置下,达到了比现有低秩基线更好的质量-效率平衡,在8-GPU张量并行推理中实现了最高1.42倍的端到端加速。
🔬 方法详解
问题定义:本文旨在解决长上下文LLM推理中KV缓存的高成本问题,现有方法未能有效整合分布式推理与KV缓存压缩的优化。
核心思路:BasisKV通过联合优化每个GPU的低秩因子,减少Value缓存的宽度,同时降低张量并行通信的成本,从而提高推理效率。
技术框架:整体架构包括低秩KV缓存压缩模块和Key路由预测模块,前者负责压缩Value缓存,后者则通过预测部分Key信息来优化Key的访问。
关键创新:最重要的创新在于将Value压缩与输出投影关联,使得同一低秩Value表示可以同时减少Value缓存宽度和张量并行通信,从而实现更高效的推理。
关键设计:关键参数包括低秩因子的选择,损失函数设计为联合优化每个GPU的注意力分区,网络结构则采用了适应性残差机制来存储无法从Value恢复的Key信息。
📊 实验亮点
实验结果显示,BasisKV在多个LLM家族和系统设置下,达到了比现有低秩基线更好的质量-效率平衡,在8-GPU张量并行推理中实现了最高1.42倍的端到端加速,显著提升了推理性能。
🎯 应用场景
该研究的潜在应用领域包括大规模自然语言处理任务,如机器翻译、文本生成和对话系统等。通过提高长上下文LLM的推理效率,BasisKV能够在资源受限的环境中实现更快速的响应,具有重要的实际价值和未来影响。
📄 摘要(原文)
Long-context large language model (LLM) inference is increasingly constrained by the cost of maintaining an ever-growing KV-cache. Low-rank K/V projections provide a promising way to compress the KV-cache with minimal quality loss, but prior work has largely optimized KV-cache compression in isolation, leaving distributed inference, especially with tensor parallelism where attention computation is partitioned across multiple GPUs, largely underexplored. We make two observations: (i) tensor-parallel communication is an important cost when deploying low-rank KV-cache compression in distributed inference; and (ii) Value compression can be associated with the output projection, allowing the same low-rank Value representation to reduce both Value-cache width and tensor-parallel communication. Based on these observations, we propose BasisKV, a low-rank KV-cache compression method for communication-intensive, long-context LLM inference. BasisKV jointly optimizes the low-rank factors across per-GPU attention partitions against the complete attention output, rather than fitting each partition in isolation. The compressed Value cache is further reused to help identify which Key pages should be accessed: it predicts part of the Key information needed for routing, while a Key-specific residual stores what cannot be recovered from Value. The original Keys are then retrieved only for the selected pages to compute exact QK scores. Across multiple LLM families and system settings, BasisKV achieves a better quality--efficiency tradeoff than existing low-rank baselines, reaching up to 1.42x end-to-end speedup under 8-GPU tensor-parallel inference. Our code is available at https://github.com/alexz949/BasisKV