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

文章主要内容总结

1. 研究背景与问题：长上下文LLM应用中，自注意力模块在预填充阶段因二次时间复杂度成为瓶颈。现有稀疏注意力方法多采用粗粒度检查，导致模型精度显著损失。
2. SALE方法核心：
   • 三阶段处理流程：通过量化（4位查询-键乘积）、选择阶段（块稀疏注意力掩码构建）和计算阶段（仅计算重要块）实现高效稀疏注意力。
   • 相对注意力分数（Relative Attention Score）：基于“sink-local”区域（序列首尾）的注意力权重相对大小评估重要性，动态调整稀疏度。
   • 硬件优化：定制CUDA内核减少开销，量化操作利用低比特Tensor Core指令，降低全局内存访问。
3. 实验结果：在Llama-3.1-8B和Qwen-2.5-32B上，处理64K+序列时速度提升至少3.36×，精度损失可忽略，优于FlashAttention2、MInference等基线方法。

创新点

1. 细粒度稀疏注意力：通过4位量化的查询-键乘积实现细粒度注意力权重估计，构建高稀疏掩码的同时控制误差。
2. 相对注意力分数指标：基于sink-local区域的相对重要性评估，计算开销可忽略，自适应调整稀疏度。