SALE：面向长上下文LLM预填充的高效稀疏注意力低比特估计
机器学习
2025-06-02 v1 人工智能
摘要
许多先进的大型语言模型（LLM）应用需要长上下文处理，但由于自注意力模块相对于序列长度的二次时间复杂度，它在推理的预填充阶段成为瓶颈。现有的稀疏注意力方法通过跳过注意力图中不重要的区域来加速注意力计算。然而，这些方法通常对注意力图进行粗粒度检查，导致模型精度出现显著损失。在本文中，我们提出了SALE，一种细粒度稀疏注意力方法，能够以可忽略的模型精度损失加速LLM的长上下文预填充阶段。SALE通过4比特量化的查询-键乘积实现快速且准确的细粒度注意力权重估计，随后进行块稀疏注意力以加速预填充计算。对于查询-键对的重要性评估，我们采用了我们的相对注意力分数指标，该指标在我们的框架内提供了显著更高的效率。我们为我们的方法实现了优化的自定义CUDA内核以提升硬件效率，将额外开销降低至仅约为完整注意力延迟的11%。值得注意的是，SALE无需参数训练，并且可以通过极少的代码修改无缝集成到现有系统中。在长上下文基准上的实验表明，我们的方法在精度与效率的权衡上优于现有方法，在Llama-3.1-8B上对长于64K的序列实现了至少3.36倍的加速，同时保持了模型质量。
引用
@article{arxiv.2505.24179,
title = {SALE : Low-bit Estimation for Efficient Sparse Attention in Long-context LLM Prefilling},
author = {Xiaodong Ji and Hailin Zhang and Fangcheng Fu and Bin Cui},
journal= {arXiv preprint arXiv:2505.24179},
year = {2025}
}