[论文解读] Faster Causal Attention Over Large Sequences Through Sparse Flash Attention

Matteo Pagliardini, Daniele Paliotta+2 more|arXiv (Cornell University)|Jun 1, 2023

Topic Modeling被引用 5

一句话总结

本文提出稀疏因果闪存注意力（SCFA），作为FlashAttention的GPU内核扩展，支持因果自注意力中的动态、非规则稀疏模式，实现长序列上的高效计算。SCFA通过优化稀疏内核执行，高效处理基于哈希和查询/键丢弃的稀疏模式，在不牺牲模型困惑度的前提下，相较于FlashAttention实现数倍加速——16k token序列的训练速度最高提升3.3倍。

ABSTRACT

Transformer-based language models have found many diverse applications requiring them to process sequences of increasing length. For these applications, the causal self-attention -- which is the only component scaling quadratically w.r.t. the sequence length -- becomes a central concern. While many works have proposed schemes to sparsify the attention patterns and reduce the computational overhead of self-attention, those are often limited by implementations concerns and end up imposing a simple and static structure over the attention matrix. Conversely, implementing more dynamic sparse attentions often results in runtimes significantly slower than computing the full attention using the Flash implementation from Dao et al. (2022). We extend FlashAttention to accommodate a large class of attention sparsity patterns that, in particular, encompass key/query dropping and hashing-based attention. This leads to implementations with no computational complexity overhead and a multi-fold runtime speedup on top of FlashAttention. Even with relatively low degrees of sparsity, our method improves visibly upon FlashAttention as the sequence length increases. Without sacrificing perplexity, we increase the training speed of a transformer language model by $2.0 imes$ and $3.3 imes$ for sequences of respectively $8k$ and $16k$ tokens.

研究动机与目标

• 解决长序列Transformer模型中因果自注意力的计算瓶颈。
• 实现对动态和非规则注意力稀疏模式（如哈希或查询/键丢弃产生的模式）的高效、高性能推理与训练。
• 将FlashAttention的效率扩展至非三角形因果掩码，这类掩码在现代稀疏注意力机制中很常见。
• 在不增加计算复杂度或牺牲模型质量的前提下，实现相对于FlashAttention的实际加速。
• 提供一个灵活、开源的GPU内核，支持广泛类别的稀疏模式，且实现开销极低。

提出的方法

• 扩展FlashAttention以支持任何可表示为每个查询对应键范围的稀疏模式，从而支持非规则因果掩码。
• 引入一个GPU内核，以零计算复杂度开销实现稀疏因果注意力计算，利用Triton进行底层优化。
• 通过几何哈希（如Reformer风格的LSH）实现动态稀疏性，但采用精确计算，无近似。
• 实现细粒度的每头查询和键丢弃，实现计算量的按比例减少，而无需完整剪枝整头。
• 采用分块内存访问模式，即使在稀疏访问模式下也能保持高内存带宽利用率。
• 采用自定义内核融合策略，最小化内核启动开销，并最大化现代GPU上的占用率。

Figure 3: Comparing several hash-based sparse attention implementations with FlashAttention. Similarly to QK-dropping-based sparsity in Fig. 7 , due to the non-triangular causal mask resulting from re-ordering the tensors based on the hash buckets (see Fig. 1 ), a naive implementation would force th

实验结果

研究问题

• RQ1 我们能否将FlashAttention扩展以支持非三角形因果掩码，从而实现对非规则稀疏模式的高效计算？
• RQ2 我们能否在不增加额外计算复杂度的前提下，实现相对于FlashAttention在稀疏注意力机制中的实际加速？
• RQ3 动态、细粒度的稀疏性（如查询/键丢弃）是否能带来可测量的训练加速，同时保持模型性能？
• RQ4 SCFA通过哈希实现的注意力在速度和覆盖范围方面与原始Reformer LSH相比如何？
• RQ5 随着序列长度增加（尤其是超过8k token），SCFA能否保持或提升训练效率？

主要发现

• SCFA在8k和16k token序列上分别实现最高2.0倍和3.3倍的训练加速，且未牺牲困惑度。
• 通过SCFA实现的哈希稀疏注意力，在8192 token序列下，使用8和16个桶时，相比FlashAttention基线分别将训练时间减少1.4倍和1.8倍。
• 该方法在训练过程中保持一致的加速效果，H-LM模型在训练初期即实现加速，并持续保持高吞吐量。
• SCFA的基于哈希的注意力实现精确计算，运行时间快于Reformer的LSH，避免了其覆盖率低和近似误差的问题。
• 通过SCFA实现的查询和键丢弃可按丢弃对的比例减少计算量，实现细粒度的效率控制。
• 运行时性能提升随序列长度和稀疏度增加而增强，但当桶数量超过某一临界点后，收益开始递减。

(a) Forward attn. runtimes