让注意力机制高效处理长序列，支持复杂遮蔽且速度提升3倍以上。
FlashMask: Efficient and Rich Mask Extension of FlashAttention
- 用列稀疏表示法重构注意力遮蔽，实现线性内存开销
- 在128K上下文下吞吐提升1.65至3.22倍，峰值算力达理论值62.3%
- 适用于大模型微调，已开源并集成于PaddleNLP
标准注意力计算和内存开销随序列长度N呈平方增长，制约了长序列建模。FlashAttention通过输入输出感知的内存优化，消除了O(N²)内存依赖，降低了延迟。但其对特定注意力遮蔽类型支持有限，难以满足复杂遮蔽需求。此前方法多采用稠密遮蔽，导致内存复杂度仍为O(N²)。本文提出FlashMask，作为FlashAttention的扩展，引入列稀疏形式的注意力遮蔽表示，高效支持多种遮蔽类型，并促进优化核实现。该方法将内存复杂度降至O(N)，适用于长上下文建模。同时利用遮蔽稀疏性消除冗余计算，不损失精度，显著提升效率。在SFT、LoRA、DPO和RM等大模型微调与对齐训练中验证，相比现有FlashAttention稠密方法，端到端速度提升1.65x至3.22x。内核级对比显示，较最新FlexAttention提升12.1%至60.7%的内核TFLOPs/s，A100上达到理论峰值的37.8%至62.3%。代码已开源于PaddlePaddle，集成于PaddleNLP，支持超1000亿参数模型，上下文长达128K tokens。
原文摘要 · Abstract (English)
The computational and memory demands of vanilla attention scale quadratically with the sequence length $N$, posing significant challenges for processing long sequences in Transformer models. FlashAttention alleviates these challenges by eliminating the $O(N^2)$ memory dependency and reducing attention latency through IO-aware memory optimizations. However, its native support for certain attention mask types is limited, and it does not inherently accommodate more complex masking requirements. Previous approaches resort to using dense masks with $O(N^2)$ memory complexity, leading to inefficiencies. In this paper, we propose FlashMask, an extension of FlashAttention that introduces a column-wise sparse representation of attention masks. This approach efficiently represents a wide range of mask types and facilitates the development of optimized kernel implementations. By adopting this novel representation, FlashMask achieves linear memory complexity $O(N)$, suitable for modeling long-context sequences. Moreover, this representation enables kernel optimizations that eliminate unnecessary computations by leveraging sparsity in the attention mask, without sacrificing computational accuracy, resulting in higher computational efficiency. We evaluate FlashMask's performance in fine-tuning and alignment training of LLMs such as SFT, LoRA, DPO, and RM. FlashMask achieves significant throughput improvements, with end-to-end speedups ranging from 1.65x to 3.22x compared to existing FlashAttention dense method. Additionally, our kernel-level comparisons demonstrate that FlashMask surpasses the latest counterpart, FlexAttention, by 12.1% to 60.7% in terms of kernel TFLOPs/s, achieving 37.8% to 62.3% of the theoretical maximum FLOPs/s on the A100 GPU. The code is open-sourced on PaddlePaddle and integrated into PaddleNLP, supporting models with over 100 billion parameters for contexts up to 128K tokens.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。