PQ-HSA：复用乘积量化分数实现混合稀疏-近似注意力
PQ-HSA: Reusing Product-Quantized Scores for Hybrid Sparse-Approximate Attention
浏览论文内容
中文总结 AI 辅助
PQ-HSA通过复用IVF-PQ索引的近似分数，将未选词元作为背景纳入softmax，在长上下文下提升稀疏注意力精度，并实现比FlashAttention-3更快的解码速度。
中文摘要 AI 辅助
在每一步解码中，语言模型都会关注之前每个词元的键值（KV）缓存，因此在长上下文下，注意力调用受限于内存带宽。稀疏注意力仅读取由廉价分数估计选择的键子集，大多数方法给未读取的词元赋予零权重。这样输出仅依赖于KV缓存的一小部分，在小预算下精度下降，尤其是在跨上下文聚合信息的任务上。对缓存键建立倒排文件乘积量化（IVF-PQ）索引，可为每个索引词元计算近似分数以进行排序；排序后，这些分数近似于被遗漏词元的注意力对数。PQ-HSA（混合稀疏-近似注意力）使用所选词元的原始键和值进行注意力计算，而未选词元（即背景）通过相同分数进入同一softmax，按倒排列表求和并乘以该列表的平均值。在128K上下文和1-2%检索预算下，PQ-HSA在Llama-3.1-8B和Qwen3-30B-A3B上比Quest和SnapKV更准确，且宏观准确率接近全注意力；使用相同选择器时，背景项将8B模型的宏观准确率从0.71提升至0.83。在相同128K设置下，在单个NVIDIA H20上的vLLM中，解码注意力调用比FlashAttention-3内核快1.6倍；加速比随上下文长度增长，基于8B到30B模型拟合的成本模型给出了加速开始的上下文长度。一个vLLM插件可在两个引擎版本上运行PQ-HSA，无需修改引擎源码；代码见https URL。
英文摘要
At each decoding step a language model attends over the key-value (KV) cache of every earlier token, so at long context the attention call is bounded by memory bandwidth. Sparse attention reads only a subset of keys chosen by a cheap score estimate, and most methods give the unread tokens zero weight. The output then draws on only a small fraction of the KV cache, and accuracy drops at small budgets, most on tasks that aggregate information across the context. An inverted-file product-quantization (IVF-PQ) index over the cached keys computes an approximate score for every indexed token in order to rank them; after ranking, those scores approximate the attention logits of the tokens left out. PQ-HSA (hybrid sparse-approximate attention) attends the selected tokens with their original keys and values, and the unselected tokens, the background, enter the same softmax through those scores, summed per inverted list and multiplied by the list's mean value. At 128K and a 1-2% retrieval budget, PQ-HSA is more accurate than Quest and SnapKV on Llama-3.1-8B and Qwen3-30B-A3B and stays close to full attention in macro accuracy; with the same selector, the background term raises macro accuracy on the 8B model from 0.71 to 0.83. In the same 128K setting, inside vLLM on one NVIDIA H20, the decode attention call runs 1.6x faster than the FlashAttention-3 kernel; the speedup grows with context length, and a cost model fitted on 8B to 30B models gives the context length at which it begins. A vLLM plugin runs PQ-HSA on two engine versions without changes to the engine source; code is available at https://github.com/KunmingSHAO/pqhsa_release.
发表机构
- The Hong Kong University of Science and Technology（香港科技大学）
- Huawei Technologies Ltd.（华为技术有限公司）
- Sun Yat-sen University（中山大学）
- Nanyang Technological University（南洋理工大学）
机构由 AI 辅助整理，请以论文原文为准。
