मुख्य अंतर्दृष्टि (Dao et al., 2022): मानक attention पूरे N×N attention matrix को GPU HBM (high bandwidth memory) में materialize करता है, जो memory-intensive (sequence length में quadratic) और धीमा (HBM bandwidth bottleneck है) दोनों है। Flash Attention यह matrix कभी materialize नहीं करता। इसके बजाय, यह attention को tiles में गणना करता है, Q, K, V के छोटे blocks को तेज़ on-chip SRAM में लोड करता है, partial results की गणना करता है, और उन्हें accumulate करता है — एक तकनीक जिसे "tiling" या "kernel fusion" कहते हैं।
मानक attention N×N attention matrix स्टोर करता है, जिसके लिए O(N²) memory चाहिए। 128K context और 128 attention heads के साथ, यह सैकड़ों gigabytes है। Flash Attention softmax को incrementally compute करके और पूरी matrix कभी स्टोर न करके O(N) memory का उपयोग करता है। यही वह चीज़ है जिसने मौजूदा hardware पर 128K–1M context windows को संभव बनाया। FlashAttention-2 ने GPU thread blocks में बेहतर parallelization करके throughput को और बेहतर किया।
Flash Attention एक व्यापक सिद्धांत का उदाहरण है: आधुनिक hardware पर, bottleneck अक्सर compute नहीं बल्कि memory bandwidth होता है। GPUs प्रति सेकंड trillions operations कर सकते हैं लेकिन HBM से केवल सैकड़ों gigabytes प्रति सेकंड read/write कर सकते हैं। जो algorithms memory traffic को कम करते हैं (भले ही अतिरिक्त computation की कीमत पर) वे अक्सर जीतते हैं। यह "IO-aware" दृष्टिकोण प्रभावित कर रहा है कि पूरा क्षेत्र AI workloads के लिए algorithm design के बारे में कैसे सोचता है।