用图示方法自动推导深度学习优化策略,提升GPU性能
FlashAttention on a Napkin: A Diagrammatic Approach to Deep Learning IO-Awareness
- 用神经电路图建模资源分布与任务调度
- 实现无需手动调优的高性能流水与分块策略
- 适合研究者构建可解释的硬件感知算法
当前深度学习算法优化依赖缓慢的手动推导,导致性能潜力未被充分挖掘。尽管FlashAttention通过避免不必要的数据传输,实现了比原生PyTorch高6倍的性能提升,但其开发历经三年三轮迭代。自动化编译方法始终落后。本文将神经电路图扩展至考虑资源使用和GPU层级的任务分配,展示如何通过简单重标记生成高层流水与分块优化策略及性能模型。该模型可轻松纳入量化与多级GPU架构的影响。我们提出用图示表示中间层伪代码的方法,实现硬件感知算法的逐步推导。最后,展示了该方法如何更好理解现有技术如FlashAttention。本工作建立理论框架,将对GPU行为的假设与性能结论关联,旨在推动GPU优化走向科学化,使实验能针对明确假设而非事后解释。
原文摘要 · Abstract (English)
Optimizing deep learning algorithms currently requires slow, manual derivation, potentially leaving much performance untapped. Methods like FlashAttention have achieved a x6 performance improvement over native PyTorch by avoiding unnecessary data transfers, but required three iterations over three years to be developed. Automated compiled methods have consistently lagged behind. This paper extends Neural Circuit Diagrams for deep learning models to consider resource usage and the distribution of tasks across a GPU hierarchy. We show how diagrams can use simple relabellings to derive high-level streaming and tiling optimization strategies along with performance models. We show how this high-level performance model allows the effects of quantization and multi-level GPU hierarchies to be readily considered. We develop a methodology for representing intermediate-level pseudocode with diagrams, allowing hardware-aware algorithms to be derived step-by-step. Finally, we show how our methodology can be used to better understand existing techniques like FlashAttention. This work uses a theoretical framework to link assumptions about GPU behaviour to claims about performance. We aim to lay the groundwork for a scientific approach to GPU optimization where experiments can address clear hypotheses rather than post-hoc rationalizations.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。