FlexAttention 是什么？几行代码写出各种高效注意力变体

FlexAttention 是 PyTorch 团队在 2024 年推出的一个灵活的注意力机制编程接口，于 2024 年 8 月前后随 PyTorch 2.5 的开发版本公开。它允许用几行典型的 PyTorch 代码实现许多不同的注意力变体，并通过 torch.compile 自动把它们编译成融合的高效内核，在不牺牲性能的情况下提供灵活性。此后研究团队也发表了相关论文对其进行介绍。

问题：注意力的变体太多了

注意力是 Transformer 的核心，为了提升效果或效率，研究者提出了大量的注意力变体：因果掩码，每个 token 只能看前面的内容；滑动窗口注意力，只看附近一段距离内的 token；ALiBi 和各种相对位置偏置，根据距离调整注意力分数；文档掩码，把多个文档打包在一起训练时，让不同文档之间互不可见；前缀语言模型掩码；软上限，限制注意力分数的大小；等等。

FlashAttention 等高度优化的注意力内核运行得非常快，但它们通常只支持少数几种固定的模式。研究者想试验一种新的注意力变体，往往面临两难：要么用普通的 PyTorch 代码直接实现，灵活但速度很慢、显存占用很大；要么自己动手写一个专门的 GPU 内核，速度快，但需要很深的专业知识和大量时间。这严重拖慢了研究的迭代速度。

FlexAttention 怎么做

FlexAttention 提供了两个简单的“钩子”，让用户用普通的 Python 函数描述想要的注意力行为。一个是分数修改函数（score_mod）：在注意力分数计算出来之后、softmax 之前，对每个分数进行修改，比如加上一个和距离相关的偏置，实现 ALiBi。另一个是掩码函数（mask_mod）：规定哪些查询和键之间可以计算注意力，比如实现因果掩码或滑动窗口。用户只需要写出这样的小函数，FlexAttention 就会借助 PyTorch 的编译器，把它们自动融合进一个基于 FlashAttention 思想的高效内核中。对于掩码，它还能自动识别出整块不需要计算的区域并跳过，利用稀疏性进一步加速。

效果

PyTorch 团队报告，FlexAttention 自动生成的内核，性能可以达到手工优化的 FlashAttention 的大部分水平，而在很多 FlashAttention 本身不支持的变体上，比普通实现快得多、省显存得多。它大大降低了探索新注意力机制的门槛，让研究者把精力集中在想法本身，而不是底层的 GPU 编程上。