FlashAttention on a Napkin: A Diagrammatic Approach to Deep Learning IO-Awareness

本文是LLM系列文章，针对《FlashAttention on a Napkin: A Diagrammatic Approach to Deep Learning IO-Awareness》的翻译。

摘要

优化深度学习算法目前需要缓慢的手动推导，可能会留下许多未开发的性能。FlashAttention等方法通过避免不必要的数据传输，实现了比原生PyTorch高6倍的性能提升，但需要在三年内进行三次迭代。自动编译方法一直落后。GPU受到处理器传输和可用计算的限制，传输带宽的改善速度要慢得多。传输带宽已经占GPU能源成本的46%。这表明，能源和资本效率算法的未来依赖于对传输成本（IO感知）的改进考虑和推导优化算法的系统方法。在本文中，我们提出了一种深度学习模型的图解方法，该方法通过简单的重新标记，得出考虑低级记忆的最佳实现和性能模型。图表在GPU层次结构中进行了概括，为比较硬件和量化选择提供了一个通用的性能模型。图表生成伪代码，揭示了硬件特定功能的应用，如合并内存访问、张量核心操作和重叠计算。我们提出了Ampere和Hopper的注意力算法，Ampere每SM适合13次扭曲（FlashAttention适合8次），Hopper改进了重叠，可能达到1.32 PFLOP。

1 引言

2 深度学习算法示意图</