FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

Tri Dao; Daniel Y. Fu; Stefano Ermon; Atri Rudra; Christopher Ré

शोध पत्र · arXiv 2205.14135

FlashAttention: IO-अवेयरनेस के साथ तेज़ और मेमोरी-सक्षम सटीक अटेंशन (Fast and Memory-Efficient Exact Attention with IO-Awareness)

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré

प्रकाशित 2022-05-27cs.LG arXiv PDF

अनुवाद और AI के साथ पढ़ें

तीन पंक्तियों में सारांश

🎯 मानक सटीक अटेंशन धीमी GPU HBM में एक $O (N^{2})$ अटेंशन मैट्रिक्स को भौतिक रूप देता है (materializes), जिससे यह मेमोरी-बाउंड हो जाता है और लंबे संदर्भों तक स्केल करना कठिन हो जाता है।
🔧 FlashAttention तेज़ ऑन-चिप SRAM पर इनपुट की टाइलिंग करके, सॉफ़्टमैक्स को वृद्धिशील (incrementally) रूप से अपडेट करके, और बैकवर्ड पास के दौरान ऑन-चिप अटेंशन की पुनर्गणना करके एकल फ़्यूज्ड CUDA कर्नेल के भीतर सटीक अटेंशन की गणना करता है।
📈 FlashAttention GPT-2 पर 3x तक एंड-टू-एंड ट्रेनिंग स्पीडअप प्रदान करता है, MLPerf 1.1 BERT-large रिकॉर्ड से 15% तेज़ ट्रेनिंग प्राप्त करता है, और संदर्भ को 64K टोकन तक स्केल करता है।

सार

Transformers लंबी अनुक्रमों (sequences) पर धीमे और अत्यधिक मेमोरी का उपभोग करने वाले होते हैं, क्योंकि सेल्फ़-अटेंशन की समय और मेमोरी जटिलता अनुक्रम की लंबाई के साथ द्विघातीय (quadratic) होती है। सन्निकट अटेंशन (approximate attention) विधियों ने गणना जटिलता को कम करने के लिए मॉडल की गुणवत्ता से समझौता करके इस समस्या को हल करने का प्रयास किया है, लेकिन अक्सर वे वॉल-क्लॉक स्पीडअप प्राप्त नहीं कर पाते हैं। हमारा मानना है कि एक अनुपस्थित सिद्धांत अटेंशन एल्गोरिदम को IO-अवेयर बनाना है — यानी GPU मेमोरी के विभिन्न स्तरों के बीच रीड और राइट का ध्यान रखना। हम FlashAttention का प्रस्ताव करते हैं, जो एक IO-अवेयर सटीक अटेंशन एल्गोरिदम है जो GPU हाई बैंडविड्थ मेमोरी (HBM) और GPU ऑन-चिप SRAM के बीच मेमोरी रीड/राइट की संख्या को कम करने के लिए टाइलिंग का उपयोग करता है। हम FlashAttention की IO जटिलता का विश्लेषण करते हैं, जिससे पता चलता है कि इसे मानक अटेंशन की तुलना में कम HBM एक्सेस की आवश्यकता होती है, और यह SRAM आकारों की एक विस्तृत श्रृंखला के लिए इष्टतम है। हम FlashAttention को ब्लॉक-स्पार्स अटेंशन तक भी विस्तारित करते हैं, जिससे एक ऐसा सन्निकट अटेंशन एल्गोरिदम प्राप्त होता है जो किसी भी मौजूदा सन्निकट अटेंशन विधि से तेज़ है। FlashAttention मौजूदा बेसलाइनों की तुलना में Transformers को तेज़ी से प्रशिक्षित करता है: MLPerf 1.1 ट्रेनिंग स्पीड रिकॉर्ड की तुलना में BERT-large (अनुक्रम लंबाई 512) पर 15% एंड-टू-एंड वॉल-क्लॉक स्पीडअप, GPT-2 (अनुक्रम लंबाई 1K) पर 3 $\times$ स्पीडअप, और लॉन्ग-रेंज एरिना (अनुक्रम लंबाई 1K-4K) पर 2.4 $\times$ स्पीडअप। FlashAttention और ब्लॉक-स्पार्स FlashAttention Transformers में लंबे संदर्भ (context) को सक्षम बनाते हैं, जिससे उच्च गुणवत्ता वाले मॉडल (GPT-2 पर 0.7 बेहतर पर्पलेक्सिटी और लंबे-दस्तावेज़ वर्गीकरण पर 6.4 अंकों का सुधार) और पूरी तरह से नई क्षमताएं प्राप्त होती हैं: Path-X चुनौती (अनुक्रम लंबाई 16K, 61.4% सटीकता) और Path-256 (अनुक्रम लंबाई 64K, 63.1% सटीकता) पर संयोग-से-बेहतर प्रदर्शन प्राप्त करने वाले पहले Transformers।

मूल सार (अंग्रेज़ी)

Transformers are slow and memory-hungry on long sequences, since the time and memory complexity of self-attention are quadratic in sequence length. Approximate attention methods have attempted to address this problem by trading off model quality to reduce the compute complexity, but often do not achieve wall-clock speedup. We argue that a missing principle is making attention algorithms IO-aware -- accounting for reads and writes between levels of GPU memory. We propose FlashAttention, an IO-aware exact attention algorithm that uses tiling to reduce the number of memory reads/writes between GPU high bandwidth memory (HBM) and GPU on-chip SRAM. We analyze the IO complexity of FlashAttention, showing that it requires fewer HBM accesses than standard attention, and is optimal for a range of SRAM sizes. We also extend FlashAttention to block-sparse attention, yielding an approximate attention algorithm that is faster than any existing approximate attention method. FlashAttention trains Transformers faster than existing baselines: 15% end-to-end wall-clock speedup on BERT-large (seq. length 512) compared to the MLPerf 1.1 training speed record, 3 $\times$ speedup on GPT-2 (seq. length 1K), and 2.4 $\times$ speedup on long-range arena (seq. length 1K-4K). FlashAttention and block-sparse FlashAttention enable longer context in Transformers, yielding higher quality models (0.7 better perplexity on GPT-2 and 6.4 points of lift on long-document classification) and entirely new capabilities: the first Transformers to achieve better-than-chance performance on the Path-X challenge (seq. length 16K, 61.4% accuracy) and Path-256 (seq. length 64K, 63.1% accuracy).

पढ़ने की गाइड

समस्या और प्रेरणा

Transformers की अनुक्रम लंबाई $N$ में समय और मेमोरी जटिलता द्विघातीय (quadratic) होती है। यद्यपि सन्निकट अटेंशन विधियाँ सैद्धांतिक FLOP गणना को कम करती हैं, वे अक्सर वॉल-क्लॉक स्पीडअप प्राप्त करने में विफल रहती हैं क्योंकि आधुनिक GPU गणनाएँ मुख्य रूप से अंकगणितीय संचालन के बजाय मेमोरी ट्रांसफर (IO) द्वारा बाधित होती हैं।

मानक अटेंशन में, $S = Q K^{⊤}$ , $P = softmax (S)$ , और $O = PV$ की गणना करने पर $N \times N$ आकार के मध्यवर्ती मैट्रिक्स धीमी हाई बैंडविड्थ मेमोरी (HBM) में लिखे जाते हैं और उन्हें कई बार वापस पढ़ा जाता है। इससे $Θ (N d + N^{2})$ HBM एक्सेस और $O (N^{2})$ मेमोरी फ़ुटप्रिंट की आवश्यकता होती है, जो व्यावहारिक संदर्भ विंडो को सीमित करता है।

मुख्य विचार

IO-Awareness: GPU मेमोरी पदानुक्रम को स्पष्ट रूप से ध्यान में रखना: तेज़, छोटा ऑन-चिप SRAM (उदा., A100 पर ~19 TB/s बैंडविड्थ के साथ 192 KB/SM) बनाम बड़ा, धीमा HBM (1.5–2.0 TB/s बैंडविड्थ के साथ 40–80 GB)।
Tiling (ब्लॉक सॉफ़्टमैक्स गणना): इनपुट $Q, K, V$ को ऐसे ब्लॉकों में विभाजित करना जो पूरी तरह से SRAM के भीतर फ़िट हो जाएँ। रनिंग सॉफ़्टमैक्स रो-मैक्स (row-max) और सामान्यीकरण आँकड़ों को ट्रैक करके, HBM में पूर्ण $N \times N$ मैट्रिक्स को भौतिक रूप दिए बिना ब्लॉकों में सॉफ़्टमैक्स रिडक्शन की गणना वृद्धिशील रूप से की जाती है।
Recomputation in Backward Pass: बैकप्रॉपैगैशन के लिए मध्यवर्ती $N \times N$ अटेंशन मैट्रिक्स को संग्रहीत करने से बचना। इसके बजाय, केवल आउटपुट $O$ और सॉफ़्टमैक्स आँकड़े $(ℓ, m)$ को संग्रहीत किया जाता है, और बैकवर्ड पास के दौरान SRAM में ब्लॉक-दर-ब्लॉक ऑन-चिप $S$ और $P$ की पुनर्गणना की जाती है।
Kernel Fusion: अनावश्यक HBM राउंड-ट्रिप को समाप्त करने के लिए मैट्रिक्स गुणन, मास्किंग, सॉफ़्टमैक्स और ड्रॉपआउट को एकल CUDA कर्नेल में संयोजित (फ़्यूज़) करना।

विधि (मुख्य समीकरणों के साथ)

$Q, K, V \in R^{N \times d}$ दिए जाने पर, FlashAttention ब्लॉक आकार $B_{c} = ⌈ M / (4 d)⌉$ और $B_{r} = min (⌈ M / (4 d)⌉, d)$ निर्धारित करता है, जहाँ $M$ SRAM क्षमता है।

वृद्धिशील सॉफ़्टमैक्स ट्रैकिंग

उप-वेक्टरों $x^{(1)}$ और $x^{(2)}$ में विघटित एक वेक्टर के लिए, रनिंग आँकड़े निम्नलिखित के माध्यम से अपडेट होते हैं: $m (x) = max (m (x^{(1)}), m (x^{(2)})),$ $ℓ (x) = e^{m (x^{(1)}) - m (x)} ℓ (x^{(1)}) + e^{m (x^{(2)}) - m (x)} ℓ (x^{(2)}) .$

कुंजियों/मानों (keys/values) के ब्लॉक $j$ और प्रश्नों (queries) के ब्लॉक $i$ के लिए:

SRAM पर गैर-सामान्यीकृत स्कोर की गणना करें: $S_{ij} = Q_{i} K_{j}^{⊤} \in R^{B_{r} \times B_{c}}$ ।
ब्लॉक स्थानीय आँकड़ों की गणना करें: $\tilde{m}_{ij} = rowmax (S_{ij})$ , $\tilde{P}_{ij} = exp (S_{ij} - \tilde{m}_{ij})$ , और $\tilde{ℓ}_{ij} = rowsum (\tilde{P}_{ij})$ ।
रनिंग ग्लोबल सामान्यीकरणकर्ताओं को अपडेट करें: $m_{i}^{new} = max (m_{i}, \tilde{m}_{ij}), ℓ_{i}^{new} = e^{m_{i} - m_{i}^{new}} ℓ_{i} + e^{\tilde{m}_{ij} - m_{i}^{new}} \tilde{ℓ}_{ij} .$
रनिंग आउटपुट ब्लॉक $O_{i} \in R^{B_{r} \times d}$ को अपडेट करें: $O_{i} \leftarrow diag (ℓ_{i}^{new})^{- 1} (diag (ℓ_{i}) e^{m_{i} - m_{i}^{new}} O_{i} + e^{\tilde{m}_{ij} - m_{i}^{new}} \tilde{P}_{ij} V_{j}) .$

जटिलता विश्लेषण

IO जटिलता: मानक अटेंशन के लिए $Θ (N d + N^{2})$ HBM एक्सेस की आवश्यकता होती है। FlashAttention के लिए $Θ (N^{2} d^{2} M^{- 1})$ HBM एक्सेस की आवश्यकता होती है। लेखक साबित करते हैं कि कोई भी सटीक अटेंशन एल्गोरिदम सभी $M \in [d, N d]$ में स्पर्शोन्मुख रूप से (asymptotically) $o (N^{2} d^{2} M^{- 1})$ एक्सेस प्राप्त नहीं कर सकता है।
FLOPs और मेमोरी: FlashAttention $O (N^{2} d)$ FLOPs निष्पादित करता है (पुनर्गणना के कारण बैकवर्ड पास में थोड़े अधिक FLOPs के साथ) और इनपुट व आउटपुट के अलावा केवल $O (N)$ सहायक मेमोरी की आवश्यकता होती है।

विस्तार: Block-Sparse FlashAttention

पूर्वनिर्धारित स्पार्सिटी मास्क $M \in {0, 1}^{N / B_{r} \times N / B_{c}}$ के अनुसार गैर-शून्य ब्लॉकों तक मूल्यांकन को सीमित करने से $Θ (N d + N^{2} d^{2} M^{- 1} s)$ की IO जटिलता प्राप्त होती है, जहाँ $s$ गैर-शून्य ब्लॉकों का अंश है।

परिणाम (सबसे महत्वपूर्ण आंकड़ों के साथ)

ट्रेनिंग स्पीडअप:
- BERT-large (अनुक्रम लंबाई 512): MLPerf 1.1 ट्रेनिंग स्पीड रिकॉर्ड को 15% से पछाड़ता है ( $8 \times$ A100s पर 17.4 मिनट बनाम 20.0 मिनट)।
- GPT-2 (अनुक्रम लंबाई 1K): HuggingFace पर 3.5x तक और Megatron-LM पर 1.8x तक स्पीडअप प्राप्त करता है।
- Long-Range Arena (LRA): मानक अटेंशन की तुलना में 2.4x स्पीडअप प्रदान करता है; Block-sparse FlashAttention 2.8x स्पीडअप प्राप्त करता है।
लंबे संदर्भ की क्षमताएं:
- 4K संदर्भ के साथ GPT-2: 1K संदर्भ पर Megatron-LM की तुलना में तेज़ी से प्रशिक्षित होता है और वैलिडेशन पर्पलेक्सिटी में 0.7 अंकों का सुधार करता है (17.5 बनाम 18.2)।
- लंबा दस्तावेज़ वर्गीकरण: संदर्भ को 512 से 16K तक बढ़ाने पर MIMIC-III पर 4.3 micro- $F_{1}$ का लाभ (52.8 से 57.1) और ECtHR पर 8.5 तक का लाभ (8K पर 72.2 से 80.7) प्राप्त होता है।
- Path-X और Path-256: FlashAttention 61.4% सटीकता के साथ Path-X (अनुक्रम लंबाई 16K) को हल करने वाला पहला Transformer है। Block-sparse FlashAttention Path-256 (अनुक्रम लंबाई 64K) पर 63.1% सटीकता प्राप्त करता है।

सीमाएं और खुले प्रश्न

CUDA इंजीनियरिंग ओवरहेड: IO-अवेयर कर्नेल को लागू करने के लिए मानक PyTorch या कंपाइलर एब्स्ट्रैक्शन के बजाय कस्टम, निम्न-स्तरीय CUDA प्रोग्रामिंग की आवश्यकता होती है, जिससे विभिन्न हार्डवेयर आर्किटेक्चर में पोर्टेबिलिटी चुनौतीपूर्ण हो जाती है।
मल्टी-GPU अटेंशन: इष्टतम सिंगल-नोड IO शेड्यूल इंटरकनेक्ट्स (NVLink, InfiniBand) में वितरित मल्टी-GPU अटेंशन तक स्वचालित रूप से विस्तारित नहीं होता है।
सटीक अटेंशन निचली सीमा: सामान्य $M$ में इष्टतम होने के बावजूद, $M$ के संदर्भ में सूक्ष्म पैरामीटरयुक्त जटिलता निचली सीमाओं (lower bounds) को साबित करना एक खुला प्रश्न बना हुआ है।

मुख्य शब्द

IO-awareness: हार्डवेयर मेमोरी के विभिन्न स्तरों (उदा., तेज़ ऑन-चिप SRAM बनाम धीमा HBM) के बीच रीड और राइट का ध्यान रखने के लिए संचालन डिज़ाइन करने का एल्गोरिद्मिक सिद्धांत।
High Bandwidth Memory (HBM): बड़ी प्राथमिक GPU मेमोरी जो उच्च कुल क्षमता प्रदान करती है लेकिन ऑन-चिप कैश की तुलना में काफी कम बैंडविड्थ प्रदान करती है।
SRAM: प्रत्येक स्ट्रीमिंग मल्टीप्रोसेसर (SM) के भीतर स्थित आधुनिक GPUs पर तेज़, ऑन-चिप स्क्रैचपैड मेमोरी।
Tiling: इनपुट मैट्रिक्स को छोटे ब्लॉकों में विभाजित करके गणना का पुनर्गठन करना जो SRAM के भीतर फिट होते हैं और विभिन्न पासों में परिणामों को संचित करते हैं।
Selective gradient checkpointing: एक ऐसी तकनीक जो कैश्ड सॉफ़्टमैक्स आँकड़ों का उपयोग करके बैकवर्ड पास के दौरान ऑन-चिप पुनर्गणना द्वारा फ़ॉरवर्ड पास के दौरान बड़े मध्यवर्ती अटेंशन मैट्रिक्स को सहेजने से बचाती है।
Kernel fusion: HBM में बार-बार होने वाले रीड/राइट को समाप्त करने के लिए कई अलग-अलग ऑपरेशनों (मैट्रिक्स गुणा, मास्किंग, सॉफ़्टमैक्स, ड्रॉपआउट) को एकल GPU कर्नेल निष्पादन में संयोजित करना।
Block-sparse FlashAttention: FlashAttention का एक विस्तार जो ब्लॉक स्पार्सिटी मास्क द्वारा परिभाषित केवल गैर-शून्य टाइल ब्लॉकों की गणना करता है, जिससे स्पार्सिटी अंश के अनुपात में IO जटिलता कम हो जाती है।
Path-X / Path-256: लॉन्ग-रेंज एरिना से लंबी-अनुक्रम विज़ुअल रीज़निंग बेंचमार्क (16K और 64K अनुक्रम लंबाई) जिस पर पूर्ववर्ती Transformers संयोग-से-बेहतर सटीकता प्राप्त करने में विफल रहे थे।

प्रश्न

क्या FlashAttention अटेंशन मैट्रिक्स का सन्निकटन (approximation) करता है?

नहीं। FlashAttention संख्यात्मक फ़्लोटिंग-पॉइंट परिशुद्धता तक मानक अटेंशन के समान गणितीय आउटपुट के साथ सटीक सेल्फ़-अटेंशन की गणना करता है।

यदि FlashAttention बैकवर्ड पास में अधिक FLOPs की गणना करता है तो यह तेज़ क्यों है?

क्योंकि मानक अटेंशन मेमोरी-बाउंड है; HBM और SRAM के बीच $N \times N$ मैट्रिक्स को स्थानांतरित करने में तेज़ SRAM पर अतिरिक्त अंकगणितीय संचालन करने की तुलना में कहीं अधिक वॉल-क्लॉक समय लगता है।

FlashAttention बैकवर्ड पास के लिए कौन सी सहायक मेमोरी संग्रहीत करता है?

FlashAttention केवल फ़ॉरवर्ड आउटपुट $O$ , सॉफ़्टमैक्स सामान्यीकरण कारक $ℓ$ और $m$ ( $O (N)$ अतिरिक्त मेमोरी की आवश्यकता), और ऑन-चिप ड्रॉपआउट मास्क को पुन: उत्पन्न करने के लिए स्यूडो-रैंडम जनरेटर स्थिति संग्रहीत करता है।

मानक अटेंशन की तुलना में FlashAttention की IO जटिलता कैसी है?

मानक अटेंशन के लिए $Θ (N d + N^{2})$ HBM एक्सेस की आवश्यकता होती है, जबकि FlashAttention के लिए $Θ (N^{2} d^{2} M^{- 1})$ एक्सेस की आवश्यकता होती है, जहाँ $M$ SRAM का आकार है और $d$ हेड का आयाम है।

इस पृष्ठ की व्याख्या AI ने शोध पत्र से लिखी है और इसमें गलतियाँ हो सकती हैं। विवरण के लिए मूल शोध पत्र पढ़ें।

अनुवाद और AI के साथ पढ़ें