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\times$ and $3.3\times$ for sequences of respectively $8k$ and $16k$ tokens.

翻译：基于Transformer的语言模型已广泛应用于需要处理日益增长序列长度的多种场景。在这些应用中，因果自注意——唯一计算复杂度随序列长度呈二次增长的核心组件——成为关键瓶颈。尽管已有诸多研究提出通过稀疏化注意力模式来降低自注意的计算开销，但这些方法常受限于实现层面的约束，最终仅能在注意力矩阵上施加简单且固定的结构。反之，实现更动态的稀疏注意力机制往往导致运行时效率显著低于Dao等人（2022）提出的Flash全注意计算方案。我们将FlashAttention扩展至支持包含键/查询丢弃及哈希注意力在内的一大类稀疏注意力模式。由此实现的方案在计算复杂度上无额外负担，并在FlashAttention基础上实现多倍的运行加速。即使稀疏度相对较低，随着序列长度增加，我们的方法也能显著优于FlashAttention。在不牺牲困惑度指标的前提下，对于包含8k和16k令牌的序列，我们的方法分别将Transformer语言模型的训练速度提升至原方法的2.0倍和3.3倍。