提出FSA优化稀疏注意力计算,让小头数模型也能高效用上稀疏加速。
FSA: An Alternative Efficient Implementation of Native Sparse Attention Kernel
- 设计新内核实现,适配小头数分组查询注意力的稀疏计算
- 在主流大模型上平均提速1.6倍,生成阶段推理快1.11倍
- 开源代码,适合关注推理效率与落地部署的研究者
稀疏注意力机制在降低大语言模型长上下文训练与推理的计算开销方面展现出巨大潜力。现有先进方法原生稀疏注意力(NSA)虽能实现硬件对齐、高效且准确,但其内核实现强制特定循环顺序,仅在每组查询头数较多时高效,而当前主流大模型普遍采用较少的查询头数,导致该技术应用受限。本文提出闪速稀疏注意力(FSA),一种替代性内核实现,可在现代GPU上支持多种常见大模型中不同数量查询头的小头数组合下实现高效计算。实证表明,相比原始NSA内核,FSA在内核级延迟上最高降低3.5倍,平均降低1.6倍;在端到端训练中最高提速1.25倍,平均提升1.09倍;在生成式推理的预填充阶段最高提速1.36倍,平均提升1.11倍。源码已公开于https://github.com/Relaxed-System-Lab/Flash-Sparse-Attention。
原文摘要 · Abstract (English)
Recent advances in sparse attention mechanisms have demonstrated strong potential for reducing the computational cost of long-context training and inference in large language models (LLMs). Native Sparse Attention (NSA), one state-of-the-art approach, introduces natively trainable, hardware-aligned sparse attention that delivers substantial system-level performance boosts while maintaining accuracy comparable to full attention. However, the kernel implementation of NSA forces a loop order that is only efficient with a relatively large number of query heads in each Grouped Query Attention (GQA) group, whereas existing LLMs widely adopt a much smaller number of query heads in each GQA group -- such an inconsistency significantly limits the applicability of this sparse algorithmic advance. In this work, we propose Flash Sparse Attention (FSA), an alternative kernel implementation that enables efficient NSA computation across a wide range of popular LLMs with a varied, smaller number of heads in each GQA group on modern GPUs. Compared to vanilla NSA kernel implementation, our empirical evaluation demonstrates that FSA achieves (i) up to 3.5x and on average 1.6x kernel-level latency reduction, (ii) up to 1.25x and 1.09x on average end-to-end training speedup on state-of-the-art LLMs, and (iii) up to 1.36x and 1.11x on average for prefill-phase speedup in LLM generative inference. The source code is open-sourced and publicly available at https://github.com/Relaxed-System-Lab/Flash-Sparse-Attention.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。