用非线性哈希优化注意力缓存检索，提速超3倍
Spotlight Attention: Towards Efficient LLM Generation via Non-linear Hashing-based KV Cache Retrieval
- 通过非线性哈希重塑查询与键的分布，提升编码效率
- 哈希码长度缩短5倍以上，512K token检索<100μs
- 适配16GB显存训练，适用于长序列高效生成场景
降低大语言模型推理中键值（KV）缓存的负担可显著加速生成。动态选择解码过程中的关键KV缓存有助于保持性能。现有方法采用随机线性哈希识别重要标记，但因大语言模型中查询与键在两个窄锥内正交分布，导致效率低下。本文提出Spotlight Attention，利用非线性哈希函数优化查询与键的嵌入分布，提升编码效率与鲁棒性。同时设计轻量级、稳定的训练框架，基于Bradley-Terry排序损失，在16GB显存的GPU上8小时内完成非线性哈希模块优化。实验表明，Spotlight Attention 显著提升检索精度，哈希码长度至少缩短5倍。通过专用CUDA核实现位运算优势，单张A100 GPU下512K token的哈希检索耗时低于100μs，端到端吞吐量最高达原始解码的3倍。
原文摘要 · Abstract (English)
Reducing the key-value (KV) cache burden in Large Language Models (LLMs) significantly accelerates inference. Dynamically selecting critical KV caches during decoding helps maintain performance. Existing methods use random linear hashing to identify important tokens, but this approach is inefficient due to the orthogonal distribution of queries and keys within two narrow cones in LLMs. We introduce Spotlight Attention, a novel method that employs non-linear hashing functions to optimize the embedding distribution of queries and keys, enhancing coding efficiency and robustness. We also developed a lightweight, stable training framework using a Bradley-Terry ranking-based loss, enabling optimization of the non-linear hashing module on GPUs with 16GB memory in 8 hours. Experimental results show that Spotlight Attention drastically improves retrieval precision while shortening the length of the hash code at least 5$\times$ compared to traditional linear hashing. Finally, we exploit the computational advantages of bitwise operations by implementing specialized CUDA kernels, achieving hashing retrieval for 512K tokens in under 100$μ$s on a single A100 GPU, with end-to-end throughput up to 3$\times$ higher than vanilla decoding.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。