FlexAttention让研究人员用几行代码快速实现高效注意力变体。
Flex Attention: A Programming Model for Generating Optimized Attention Kernels
- 基于编译器的编程模型，支持用原生PyTorch代码实现注意力变体。
- 可高效实现Alibi、PagedAttention等主流注意力结构，性能接近手工优化版本。
- 支持注意力变体自由组合，解决新结构开发中的复杂性问题。
过去七年中，注意力机制已成为深度学习中最关键的计算原语之一。主流优化方法FlashAttention通过操作融合显著提升了运行效率和内存消耗。然而，FlashAttention的封闭性给研究者尝试新型注意力结构带来了困难，形成‘软件彩票’困境。此外，高效融合注意力内核的编写难度高，传统编译优化手段难以奏效。本文提出FlexAttention，一种由编译器驱动的编程模型，使大多数注意力变体仅需少量原生PyTorch代码即可实现。我们验证了包括Alibi、文档掩码、PagedAttention在内的多种现有结构均可通过FlexAttention实现，且性能与手工优化内核相当。最后，展示了其在注意力变体组合上的强大能力，有效缓解了注意力结构的组合爆炸问题。
原文摘要 · Abstract (English)
Over the past 7 years, attention has become one of the most important primitives in deep learning. The primary approach to optimize attention is FlashAttention, which fuses the operation together, drastically improving both the runtime and the memory consumption. However, the importance of FlashAttention combined with its monolithic nature poses a problem for researchers aiming to try new attention variants -- a "software lottery". This problem is exacerbated by the difficulty of writing efficient fused attention kernels, resisting traditional compiler-based approaches. We introduce FlexAttention, a novel compiler-driven programming model that allows implementing the majority of attention variants in a few lines of idiomatic PyTorch code. We demonstrate that many existing attention variants (e.g. Alibi, Document Masking, PagedAttention, etc.) can be implemented via FlexAttention, and that we achieve competitive performance compared to these handwritten kernels. Finally, we demonstrate how FlexAttention allows for easy composition of attention variants, solving the combinatorial explosion of attention variants.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。