💡
原文中文，约2400字，阅读约需6分钟。
📝
内容提要
PyTorch团队引入了FlexAttention，一个灵活的API，允许用户使用几行PyTorch代码实现多个注意力变体。通过torch.compile将其降低到一个融合的FlashAttention内核中，生成了一个不会占用额外内存且性能可与手写内核相媲美的FlashAttention内核。FlexAttention具有令人惊讶的表达能力，可以满足大多数用户对注意力变体的需求。
🎯
关键要点
- FlexAttention 是一个灵活的 API，允许用户用几行 PyTorch 代码实现多个注意力变体。
- 通过 torch.compile，FlexAttention 被降低到一个融合的 FlashAttention 内核，性能可与手写内核相媲美且不占用额外内存。
- 现有的注意力机制在性能提升的同时失去了灵活性，导致用户面临运行缓慢和 CUDA 内存不足的问题。
- FlexAttention 允许用户定义 score_mod 函数，以满足对注意力变体的需求。
- FlexAttention 动态计算偏差值，显著提高内存和性能，支持相对位置编码等变体。
- FlexAttention 的性能接近手写的 Triton 内核，前向传播实现了 FlashAttention2 性能的 90%，反向传播实现了 85%。
- 研究者计划改进 FlexAttention 的反向算法，以缩小与 FlashAttention2 的性能差距。
➡️
