🤖 AI Summary
Hand-tuned deep learning optimizations (e.g., FlashAttention) suffer from low efficiency and poor interpretability, while existing automatic compilation approaches lack generalizability and hardware awareness. Method: We propose an I/O-aware Neural Circuit Graph (NCG) optimization framework that explicitly models GPU microarchitecture, resource constraints, and memory hierarchy as a relabelable graph structure—enabling joint derivation of streaming computation, tiling schedules, and higher-order effects (e.g., quantization, multi-level storage). It establishes a formal bridge from hardware behavioral assumptions to performance assertions and supports graphical modeling of IR-level pseudocode alongside GPU-driven performance modeling. Contribution/Results: NCG is the first framework to systematically reverse-engineer the intrinsic mechanisms of state-of-the-art techniques like FlashAttention. It establishes a GPU algorithm design paradigm that is interpretable, evolvable, and formally verifiable—outperforming both manual tuning and black-box auto-compilation by significant margins.
📝 Abstract
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.