FlexAttention：兼具PyTorch的灵活与FlashAttention的性能

文章预览

点击下方 卡片 ，关注“ AI生成未来 ” Pytorch官方Blog：FlexAttention: The Flexibility of PyTorch with the Performance of FlashAttention 简介 理论上，Attention is All You Need。然而在实践中，我们还需要像FlashAttention这样的优化的注意力机制实现。 尽管这些融合的注意力实现在性能上有了显著提升，并使得长序列上下文成为可能，但这种效率的提升是以牺牲灵活性为代价的。你不能再通过简单地编写几个PyTorch运算符来尝试新的 注意力 变体，而通常需要重新写一个新的自定义的Kernel，即使是使用triton等工具，也并不简单！这为机器学习研究人员创造了一种“Software Lottery”（这个词来源于谷歌的论文The Hardware Lottery，在机器学习领域中用来描述某些研究想法因适合现有的软硬件环境而成功，而非因为这些想法在本质上优于其他研究方向）——如果你的 注意力 变体不适用于 ………………………………