Flex Attention：生成优化注意力内核的编程模型
机器学习
2024-12-10 v1 性能
编程语言
摘要
过去 7 年间，注意力机制已成为深度学习中最重要的基本运算之一。优化注意力的主要方法是 FlashAttention，通过融合运算显著提升运行时性能和内存消耗。然而，FlashAttention 的重要性及其单体化特性给追求新型注意力变体的研究者带来了问题——即“软件抽奖”问题。编写高效融合注意力内核的困难也使得传统编译器方法难以解决。我们引入 FlexAttention，这是一种新型编译器驱动的编程模型，允许通过几行恰当的 PyTorch 代码实现大多数注意力变体。我们展示，许多现有注意力变体（如 Alibi、Document Masking、PagedAttention 等）可通过 FlexAttention 实现，并实现了与这些手写内核相当的性能。最后，我们表明 FlexAttention 使得注意力变体的易于组合成为可能，解决了注意力变体的组合爆炸问题。
引用
@article{arxiv.2412.05496,
title = {Flex Attention: A Programming Model for Generating Optimized Attention Kernels},
author = {Juechu Dong and Boyuan Feng and Driss Guessous and Yanbo Liang and Horace He},
journal= {arXiv preprint arXiv:2412.05496},
year = {2024}
}