arXiv:2505.24179cs.LGcs.AI2025-05被引 4

用4比特量化实现高效低精度注意力,显著加速长文本推理

SALE : Low-bit Estimation for Efficient Sparse Attention in Long-context LLM Prefilling

  • 通过4比特量化查询键乘积,精细估计注意力权重
  • 在64K以上序列上达3.36倍加速,模型质量几乎不变
  • 无需训练,可无缝接入现有系统,适合长文本应用

大量先进大语言模型应用需处理长上下文,但自注意力模块在推理预填充阶段因序列长度的二次复杂度成为瓶颈。现有稀疏注意力方法通过跳过注意力图中不重要的区域加速计算,但通常采用粗粒度检查,导致模型精度明显下降。本文提出SALE,一种细粒度稀疏注意力方法,在几乎不影响模型准确率的前提下加速长上下文预填充。SALE通过4比特量化查询键乘积实现快速精准的注意力权重估计,并结合块稀疏注意力加速计算。为评估查询-键对的重要性,我们引入相对注意力得分(Relative Attention Score)指标,显著提升框架效率。我们实现了一个定制的CUDA内核以优化硬件性能,额外开销仅约全注意力延迟的11%。SALE无需参数训练,可通过极小代码修改无缝集成到现有系统。在长上下文基准测试中,该方法在准确性与效率权衡上优于现有方案,在序列长度超过64K时,Llama-3.1-8B模型的推理速度至少提升3.36倍。

原文摘要 · Abstract (English)

Many advanced Large Language Model (LLM) applications require long-context processing, but the self-attention module becomes a bottleneck during the prefilling stage of inference due to its quadratic time complexity with respect to sequence length. Existing sparse attention methods accelerate attention computation by skipping less significant regions of the attention map. However, these approaches typically perform coarse-grained inspection of the attention map, rendering considerable loss in model accuracy. In this paper, we propose SALE, a fine-grained sparse attention method that accelerates the long-context prefilling stage of LLM with negligible loss in model accuracy. SALE achieves fast and accurate fine-grained attention weight estimation through 4-bit quantized query-key products, followed by block-sparse attention to accelerate prefilling computations. For importance evaluation for query-key pairs, we adopt our Relative Attention Score metric, which offers significantly higher efficiency within our framework. We implement a custom CUDA kernel optimized for our approach for hardware efficiency, reducing the additional overhead to approximately 11% of the full attention latency. Notably, SALE requires no parameter training and can be seamlessly integrated into existing systems with trivial code modifications. Experiments on long-context benchmarks demonstrate that our method outperforms existing approaches in accuracy-efficiency trade-offs, achieving at least 3.36x speedups on Llama-3.1-8B for sequences longer than 64K while maintaining model quality.

稀疏注意力长文本生成4比特量化推理加速

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