【翻译】【PyTorch 奇技淫巧】FlexAttetion 基于Triton打造灵活度拉满的Attention

文章预览

blog链接：https://pytorch.org/blog/flexattention/ 代码示例：https://github.com/pytorch-labs/attention-gym/blob/main/examples/flex_attn.ipynb FlexAttention: PyTorch 的灵活性与 FlashAttention 的性能 by Team PyTorch: Horace He, Driss Guessous, Yanbo Liang, Joy Dong 理论上，Attention is All You Need。然而在实践中，我们还需要像FlashAttention这样的优化注意力实现。 尽管这些融合的注意力实现显著提高了性能并支持了长上下文，但这种效率是以灵活性为代价的。你不能再通过编写几个PyTorch操作符来尝试新的注意力变体——你通常需要编写一个新的自定义kernel！这对机器学习研究人员来说就像是一种“软件彩票”——如果你的注意力变体不适合现有的优化kernel之一，你注定会面临缓慢的运行时间和CUDA内存不足的问题。 一些注意力变体的例子包括因果关系、相对位置嵌入（https://paperswithcode.com/method/relative-posit ………………………………