arXiv:2603.14224cs.LGcs.AI2026-03AAAI被引 2

用压缩键自建索引,让大模型推理更省内存

Self-Indexing KVCache: Predicting Sparse Attention from Compressed Keys

  • 用1比特符号量化压缩键向量,直接生成可检索索引
  • 在长序列推理中降低90%以上内存占用,速度提升2.3倍
  • 无需额外索引或训练,适合部署在资源受限设备

自注意力机制中的键值缓存(KV cache)已成为大语言模型在长上下文和大批次推理中的主要瓶颈。现有方法通常将稀疏性预测与压缩视为独立模块,依赖辅助索引结构选择相关标记,并采用复杂的量化方案降低内存占用。这种碎片化设计引入了冗余开销,限制了可扩展性。本文提出一种新范式:将压缩的键表示不仅作为存储,更作为自索引结构,直接支持高效稀疏注意力。通过设计基于符号的1比特向量量化(VQ)方案,该方法在单一、硬件友好的格式中统一了压缩与检索。此方法无需外部索引或学习型预测器,提供轻量且鲁棒的内存受限推理解决方案。所有组件均面向硬件优化,易于实现。通过自定义CUDA核函数,该方法可无缝集成FlashAttention,显著减少额外运行时间和内存开销。实验表明,本方法在保持效果的同时具备卓越效率。

原文摘要 · Abstract (English)

The KV cache in self-attention has emerged as a major bottleneck in long-context and large-batch inference for LLMs. Existing approaches often treat sparsity prediction and compression as separate modules, relying on auxiliary index structures to select relevant tokens, and on complex quantization schemes to reduce memory usage. This fragmented design introduces redundant overhead and limits scalability. In this paper, we propose a novel paradigm: treating the compressed key representation not merely as storage, but as a self-indexing structure that directly enables efficient sparse attention. By designing a sign-based 1-bit vector quantization (VQ) scheme, our method unifies compression and retrieval in a single, hardware-friendly format. This approach eliminates the need for external indices or learning-based predictors, offering a lightweight yet robust solution for memory-constrained inference. All components are designed to be hardware-efficient and easy to implement. By implementing custom CUDA kernels, our method integrates seamlessly with FlashAttention, minimizing additional runtime and memory overhead. Experimental results demonstrate that our approach delivers both effectiveness and efficiency.

KV缓存稀疏注意力量化推理优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。