FlashAttention-3: असिंक्रोनस और लो-प्रिसिजन से तेज़ और अधिक सटीक Attention

 (together.ai)
1 पॉइंट द्वारा GN⁺ 2024-07-12 | 1 टिप्पणियां | WhatsApp पर शेयर करें
  • Transformer में Attention bottleneck ने लंबे context वाले LLM की training और inference speed को सीमित किया है, और FlashAttention-3 Hopper GPU features का उपयोग करके इसे और अधिक घटाने की कोशिश करने वाला नया version है
  • मुख्य बात Tensor Core और TMA की asynchronicity का उपयोग करके computation और data movement को overlap करना, और matmul व softmax को interleave करके GPU idle time घटाना है
  • H100 पर FlashAttention-2 theoretical maximum FLOPS के केवल 35% utilization तक सीमित था, लेकिन FlashAttention-3 FP16 में अधिकतम 740 TFLOPS, यानी H100 theoretical peak के 75% तक पहुंचता है
  • FP8 low precision में LLM activation के outliers error बढ़ाते हैं, इसलिए Hadamard transform आधारित incoherent processing से baseline FP8 attention की तुलना में error 2.6 गुना घटाया गया है
  • FlashAttention-3 GitHub पर open source किया गया है, और जैसे-जैसे Attention cost घटती है, लंबे context processing तथा LLM training/inference efficiency सुधारने की गुंजाइश बढ़ती है

FlashAttention-3 का लक्ष्य और performance

  • Attention Transformer की core layer है, लेकिन बड़े language models और लंबे context applications में यह मुख्य bottleneck बन जाता है
  • FlashAttention और FlashAttention-2 ने GPU memory reads/writes घटाकर Attention को तेज़ बनाने की approach शुरू की, और आज अधिकतर libraries इसे Transformer training और inference acceleration के लिए उपयोग करती हैं
  • इस approach ने पिछले 2 वर्षों में LLM context length को GPT-3·OPT के 2~4K से GPT-4 के 128K और Llama 3 के 1M तक बढ़ाने में योगदान दिया है
  • FlashAttention-2 H100 GPU पर theoretical maximum FLOPS के केवल 35% utilization तक पहुंचा था, लेकिन FlashAttention-3 Hopper GPU के नए features का उपयोग करके इसे बढ़ाता है
  • प्रकाशित FlashAttention-3 performance इस प्रकार है
    • FP16 में FlashAttention-2 से 1.5~2.0 गुना तेज़
    • FP16 में अधिकतम 740 TFLOPS
    • H100 theoretical maximum FLOPS का 75% utilization
    • FP8 में लगभग 1.2 PFLOPS
    • baseline FP8 attention की तुलना में 2.6 गुना कम error

FlashAttention approach का recap

  • FlashAttention Attention computation order को rearrange करता है और tiling व recomputation का उपयोग करके speed बढ़ाता है, साथ ही sequence length के संदर्भ में memory usage को quadratic से linear कर देता है
  • input blocks को HBM से SRAM में load किया जाता है, उन blocks पर Attention किया जाता है, और फिर output को HBM में update किया जाता है
  • बड़े intermediate Attention matrix को HBM में नहीं लिखा जाता, इसलिए memory reads/writes घटते हैं और real runtime के आधार पर 2~4 गुना speedup संभव होता है
  • tiling और softmax rescaling को साथ इस्तेमाल करने पर block-wise processing करते हुए भी approximation के बिना सही output मिल सकता है

Hopper GPU features: WGMMA, TMA, FP8

  • FlashAttention-2 Ampere A100 GPU पर theoretical maximum FLOPS के 70% तक पहुंच सकता है, लेकिन Hopper GPU के नए features का पर्याप्त उपयोग नहीं कर पाता
  • FlashAttention-3 Hopper के तीन features का उपयोग करता है
    • WGMMA: Hopper के नए Tensor Core का उपयोग करने वाला warpgroup matrix multiply-accumulate feature, जिसकी throughput Ampere के mma.sync से अधिक है
    • TMA: global memory और shared memory के बीच data transfer को accelerate करने वाली dedicated hardware unit, जो index computation और out-of-bound predication संभालकर register usage घटाती है
    • FP8: FP16 की तुलना में Tensor Core throughput को दोगुना कर सकता है, लेकिन कम bits में floating-point values represent करता है, इसलिए accuracy के साथ tradeoff होता है
  • FlashAttention-3 Hopper features का उपयोग करने के लिए NVIDIA CUTLASS के abstractions का उपयोग करता है
  • FlashAttention को नए features इस्तेमाल करने के लिए rewrite करने भर से FP16 forward pass performance FlashAttention-2 के लगभग 350 TFLOPS से बढ़कर 540~570 TFLOPS स्तर तक पहुंचती है

Asynchronicity से GEMM और softmax overlap करना

  • Attention के मुख्य operations Q-K, P-V के बीच GEMM और softmax हैं
  • आधुनिक accelerators में non-matmul operations matmul से काफी धीमे होते हैं, और softmax का exponential जैसे special functions floating point multiply-add या matrix multiply-add से अलग units पर process होते हैं
  • H100 SXM5 FP16 matrix multiply में 989 TFLOPS देता है, लेकिन special function throughput 3.9 TFLOPS है, यानी 256 गुना कम
  • head dimension 128 में matmul FLOPS exponential से 512 गुना अधिक होने पर भी, exponential matmul की तुलना में time का 50% ले सकता है
  • FP8 में matmul FLOPS दोगुने तेज़ हो जाते हैं, लेकिन exponential speed वही रहती है, इसलिए matmul और softmax को parallel में चलाना और महत्वपूर्ण हो जाता है

  • warpgroups के बीच pingpong scheduling

    • GPU warp scheduler, जब कुछ warps GEMM result का इंतज़ार करते हैं, तब दूसरे warps चलाकर कुछ overlap अपने-आप करता है
    • FlashAttention-3 synchronization barrier का उपयोग करके दो warpgroups के GEMM और softmax को manually बेहतर तरीके से overlap करता है
    • warpgroup 1 पहले एक iteration का GEMM1 और अगले iteration का GEMM0 करता है
    • इसके बाद जब warpgroup 2 GEMM करता है, तब warpgroup 1 softmax process करता है
    • यह pingpong schedule softmax को दूसरे warpgroup के GEMM execution time के पीछे छिपाने का तरीका है
    • वास्तविक scheduling diagram की तरह पूरी तरह साफ-सुथरी नहीं होती, लेकिन FP16 attention forward pass में head dimension 128, sequence length 8K के आधार पर यह लगभग 570 TFLOPS को 620 TFLOPS तक बढ़ाती है

  • warpgroup के अंदर overlap

    • एक ही warpgroup के अंदर भी उस warpgroup के GEMM चलते समय softmax का कुछ हिस्सा execute किया जा सकता है
    • यह pipelining FP16 attention forward throughput को लगभग 620 TFLOPS से 640~660 TFLOPS तक बढ़ाती है
    • इसके बदले GEMM accumulator और softmax input/output को साथ रखना पड़ता है, जिससे register pressure बढ़ता है
    • कुल मिलाकर यह technique लाभदायक tradeoff देती है

FP8 low precision और incoherent processing

  • LLM activation में कुछ outlier हो सकते हैं जिनका magnitude बाकी features से बहुत अधिक होता है
  • outlier quantization को कठिन बनाते हैं और quantization error को काफी बढ़ा देते हैं
  • FlashAttention-3 QuIP जैसी quantization literature में इस्तेमाल की गई incoherent processing का उपयोग करता है
  • query और key को random orthogonal matrix से multiply करके outliers को फैलाया जाता है और quantization error घटाया जाता है
  • implementation में random sign वाला Hadamard transform इस्तेमाल होता है
    • head dimension को d मानने पर इसे O(d²) नहीं बल्कि O(d log d) time में हर attention head पर किया जा सकता है
    • Hadamard transform memory-bandwidth bound है, इसलिए इसे rotary embedding जैसे पहले के memory-bandwidth bound operation के साथ fuse करने पर extra cost के बिना process किया जा सकता है
  • Q, K, V को standard normal distribution से generate करके और entries के 0.1% में बड़ा magnitude डालकर outlier simulate करने वाले experiment में, incoherent processing ने quantization error को 2.6 गुना घटाया

Benchmark और release status

  • FlashAttention-3 की तुलना FlashAttention-2 के साथ-साथ Hopper GPU के नए hardware features पहले से उपयोग करने वाले Triton और cuDNN implementations से भी की गई
  • FP16 में यह FlashAttention-2 की तुलना में लगभग 1.6~1.8 गुना speedup दिखाता है
  • FP8 में यह लगभग 1.2 PFLOPS तक पहुंचता है
  • FlashAttention-3 GitHub repository उपलब्ध है
  • paper भी उसी flash-attention repository में देखा जा सकता है

बाकी optimizations और आगे integration

  • paper में blog में बताए गए हिस्सों के अलावा variable length sequence, persistent kernel, FP8 के लिए in-kernel transpose जैसी optimizations भी शामिल हैं
  • execution hardware के हिसाब से algorithms design करने पर बड़ी efficiency improvement और लंबे context जैसी नई model capabilities खुल सकती हैं
  • आगे के काम में LLM inference optimization और techniques को अन्य hardware architectures पर generalize करना शामिल है
  • FlashAttention-3 के future PyTorch releases में integrate होने की उम्मीद है

संबंधित पढ़ाई

1 टिप्पणियां

 
GN⁺ 2024-07-12
Hacker News की राय

  • कोड comments देखकर लगता है कि Tri Dao Hopper/H100 की घोषणा के तुरंत बाद, अप्रैल 2022 से FA3 पर काम कर रहे थे
    यह थोड़ा दिलचस्प है कि कोड आज सार्वजनिक होने में 2 साल से ज़्यादा लग गए; शायद इसलिए कि बेहतर समाधान तैयार हो रहे हों
    Tri के हालिया papers का रुझान SSM और Mamba-family architectures की ओर है। FlashAttention में sequence length के लिए quadratic time complexity होती है, लेकिन नए algorithms sub-quadratic हैं, यानी वे वही computation सिर्फ़ ज़्यादा efficient तरीके से नहीं करते, बल्कि computation की मात्रा ही काफ़ी कम कर देते हैं
    Dao और Gu ने इस साल एक लंबे paper में दिखाया कि Mamba/SSM को भी ऐसे formalize किया जा सकता है कि उन्हें Transformer को फायदा देने वाली hardware primitive operations से accelerate करना आसान हो

    • जब तक Strong Exponential Time Hypothesis (SETH) सिद्ध या खंडित नहीं हो जाती, तब तक quadratic cost ज़रूरी होगी या फिर कुछ छोड़ना पड़ेगा। आखिरकार यह exhaustive search की cost है
      SETH को सिद्ध या खंडित करने पर P बनाम NP समस्या भी हल हो जाएगी, इसलिए इसके जल्द होने की उम्मीद करना मुश्किल है
      मुख्य बात यह है कि कोई खास use case इस cost को झेल सकता है या नहीं

  • मैं सोच रहा हूं कि FlashAttention algorithm hardware से कितना बंधा हुआ है
    उदाहरण के लिए, इस घोषणा में कहा गया है कि यह H100 GPU की asynchronous features का इस्तेमाल करता है, तो लगता है कि H series के अलावा अन्य cards पर वह speedup नहीं मिलेगा
    साथ ही वास्तविक FlashAttention library को CUDA चाहिए, लेकिन algorithm Metal पर port किया गया लगता है[^0]। अगर algorithm लगभग pure function जैसा है, तो उसे किसी भी GPU/machine learning framework में implement किया जा सकना चाहिए, ऐसा लगता है
    [0]: https://github.com/philipturner/metal-flash-attention

    • अच्छे जवाब बहुत हैं, लेकिन संक्षेप में कहें तो “practical तौर पर, काफी हद तक” यह hardware से बंधा है। नीचे एक अच्छा example है

      https://github.com/karpathy/nanoGPT/blob/master/model.py#L45
      Karpathy का nanoGPT torch.nn.functional.scaled_dot_product_attention मौजूद है या नहीं, यह check करके FlashAttention को call करता है
      https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html
      docs देखने पर असल में ज़्यादातर मामलों में आप FA2 को call करना चाहेंगे, और FA2 device के kernels को optimize करके triangular matrix के Softmax operation को split करता है, और non-essential floating-point batches को GPU और CPU के बीच round-trip कराने की जरूरत कम करता है
      https://arxiv.org/pdf/2307.08691
      FA2 paper लगभग पूरी तरह उस hardware के नजरिए से लिखा गया है जिस पर वह चलता है

    • FlashAttention में algorithmic improvement मुख्य रूप से attention के Softmax हिस्से को split और merge करने में है, और यह अपने-आप में बिल्कुल नया idea नहीं है। इसका जबरदस्त योगदान इस method और बारीक details को Nvidia hardware पर efficiently implement करने में है
    • मूल FlashAttention में hardware dependency लगभग नहीं थी
      latest version abstraction के level पर निर्भर करता है। ThunderKittens[0] article में बताई गई चीज़ के समान FA2 के मुकाबले 1.3~2x speedup देता है, और फिर भी GPU ecosystem में अपेक्षाकृत general रूप से लागू होता है
      हर नए hardware में hardware-specific features हो सकते हैं जिनसे extra performance निकाली जा सकती है। आम तौर पर vendors उन features को अपनाते हैं जो उन्हें आगे रखते हैं, लेकिन CUDA में भी जैसा पहले से है, APIs और libraries fragmented हो जाती हैं
      [0]: https://hazyresearch.stanford.edu/blog/2024-05-12-tk
    • conceptually थोड़ा, लेकिन practical implementation के नजरिए से काफी बंधा है। standard Python implementation भी internally खास hardware के लिए kernels compile करता है
    • practical angle से जोड़ें तो, AMD hardware पर अभी flash-attention-2 का सही implementation पर्याप्त नहीं है। ROCm धीरे-धीरे उपयोग लायक हो रहा है, लेकिन CUDA से तुलना करने लायक अभी नहीं पहुंचा है

  • कंपाइलर वाले लोगों से पूछना चाहूंगा। क्या कोई संभावना है कि कंपाइलर FlashAttention जैसी optimizations को खुद खोज पाए? TVM और tinygrad उस दिशा में जाते लगते हैं, लेकिन यह मानना मुश्किल है कि यह संभव है

    • सैद्धांतिक रूप से संभव है। गणितीय algebraic गुणों की वजह से बड़े स्तर पर reordering संभव है, और इसमें अपेक्षाकृत structured polyhedral loop tiling जोड़नी होगी
      हालांकि cost बड़ी है, इसलिए उस search result को cache करना होगा
      e-graph optimization इस क्षेत्र के लिए अच्छी fit लगती है। लेकिन optimization pass को process करने के तरीके में बड़ा paradigm shift चाहिए, इसलिए कुछ niche tools के अलावा यह लगभग deploy नहीं हुआ है। उदाहरण के लिए, यह पारंपरिक call graph के साथ ठीक से fit नहीं बैठता, इसलिए basic blocks और for loops के बाहर/बीच से गुजरते हुए e-graph deploy करने के लिए control flow को काफी बदलना पड़ेगा, और break व return भी support नहीं हैं
    • यह बेहद कठिन, लेकिन असंभव नहीं समस्या लगती है
      compiler optimization की state-of-the-art data layout और processor utilization को अधिकतम करने के मामले में कहां तक पहुंची है, यह मुझे ठीक से नहीं पता
      पहले optimization से जुड़ा एक video देखा था जिसमें छोटी optimization ने speed बढ़ाई थी, लेकिन उस optimization या यहां तक कि random changes से बने memory layout के अंतर के कारण आने वाले speed variations की तुलना में उसका असर बहुत छोटा था
      वह talk noise में signal अलग करने पर ज्यादा focused थी, लेकिन वही noise अपने-आप में इस बात का संकेत है कि compiler यहां कही गई समस्या से कहीं सरल रूपों को भी अच्छी तरह handle नहीं कर पाता
      cache और access patterns speed को प्रभावित करते हैं—सिर्फ ऐसी CPU-memory architecture भी complex है; इसमें GPU architecture जोड़ दें तो यह काफी unexplored area लगता है
      किसी दिन संभव हो सकता है। यह AI field है, इसलिए सवाल उठता है कि क्या sufficiently smart AI यह कर सकती है, लेकिन यह “sufficiently” के मानदंड पर निर्भर है
      AI models के बहुत high-level test के रूप में, micrograd जैसा कुछ देकर, वही interface बनाए रखते हुए torch से तेज चीज बनाने को कहना सोचा जा सकता है। अभी हम उसके आसपास भी नहीं हैं, लेकिन अगर यह संभव हुआ तो दिलचस्प होगा
    • मुझे नहीं लगता। इसे अलग algorithm की तरह सोचना चाहिए। केवल गणित पर विचार करने के बजाय hardware की shape को ध्यान में रखकर algorithm design करना है
      TVM समझ में आता है। सख्ती से कहें तो वह अलग काम करता है, लेकिन काफी करीब का क्षेत्र है
      लेकिन tinygrad के बारे में आपको ऐसा क्यों लगा, यह समझ नहीं आया
    • https://github.com/uwplse/tensat
    • Python जैसी wrapped language से high-level operators call करना काफी tricky है

  • अगर कोई इसे ROCm / AMD MI300x पर port करना चाहता है, तो hello@hotaisle.xyz पर संपर्क करें। spam बिल्कुल नहीं भेजूंगा
    इस काम के लिए compute time donate कर सकता हूं

    • AMD accelerator server company है! बढ़िया काम है, उम्मीद है कोई इसे ले लेगा :)
    • बदतमीजी का इरादा नहीं है, लेकिन इस offer का उद्देश्य जानना चाहता हूं। सिर्फ hardware access पाकर कौन यह porting मुफ्त में करेगा? उस व्यक्ति को क्या फायदा होगा?

  • FlashAttention-3 is optimized for Hopper GPUs (e.g. H100).
    FA3 3090 और 4090 जैसे consumer GPUs पर कितना performance देता है?

  • TMA (Tensor Memory Accelerator). This is a special hardware unit that accelerates the transfer of data between global memory and shared memory, taking care of all index calculation and out-of-bound predication. This frees up registers, which is a valuable resource to increase tile size and efficiency.
    मेरी समझ में TMA registers घटाने में मदद करता है, यह भी सही है, लेकिन ज्यादा महत्वपूर्ण बात यह है कि address generation hardware से करवाता है। आसपास के operations तेज होते जाएं तो address generation bottleneck बन सकता है

  • यह पूरे AI में सबसे महत्वपूर्ण improvements में से एक है। वही hardware इस्तेमाल करके ज्यादा चीजें और तेज चला पाने देता है, और अधिकांश AI users को लगभग बिना trade-off के फायदा देता है

    • H100 रखने वाले users के लिए तो होगा

  • FlashAttention में variable masking इस्तेमाल करने पर, न इस्तेमाल करने की तुलना में यह लगभग 5 गुना धीमा क्यों हो जाता है, यह जानना चाहता हूं। अच्छी masking support की कमी हो तो optimization का असर लगभग खत्म हो जाता है

    • वह benchmark आप कहां देख रहे हैं?

  • अच्छा होगा अगर कोई विशेषज्ञ कुछ सवालों के जवाब दे दे :)
    क्या FlashAttention, LLM के attention ऑपरेशन का बस एक drop-in replacement है? जहाँ भी “attention” ऑपरेशन इस्तेमाल होता है, क्या वहाँ इसे इस्तेमाल किया जा सकता है, या FA इस्तेमाल करने के लिए LLM को अलग से train करना पड़ता है?
    FA का GQA, sliding window attention जैसी strategies से क्या संबंध है? क्या ये एक-दूसरे से orthogonal concepts हैं, या हर strategy के लिए अलग FA implementation चाहिए?
    हाल ही में llama.cpp ने FlashAttention support जोड़ा है; क्या इसका मतलब है कि उसने FlashAttention द्वारा दिए जाने वाले CUDA kernels जैसी चीज़ों का इस्तेमाल शुरू किया है?
    आखिर में, इस लेख में FlashAttention और Triton की तुलना की गई है। क्या Triton कोई abstraction layer जैसा नहीं है? क्या FA को Triton में implement नहीं किया जा सकता? “FlashAttention बनाम Triton” वाली अभिव्यक्ति ठीक से समझ नहीं आ रही

      1. लगभग सही है। mathematically equivalent है। software की तरफ़ से दिक्कतें बस dependency version management या memory में data format जैसी चीज़ें हैं, और FlashAttention 2 पहले से HuggingFace और कई popular libraries में शामिल है। FlashAttention 3 भी जल्द शामिल होने की संभावना है, लेकिन चलाने के लिए H100 GPU चाहिए
      2. FlashAttention 2 ने पहले के version update में GQA support जोड़ा था:
        https://github.com/Dao-AILab/flash-attention
      3. यहाँ pure CUDA C++ में लिखे इस FlashAttention implementation की तुलना Triton में लिखे समान algorithm के Triton implementation से की जा रही है: https://triton-lang.org/main/getting-started/tutorials/06-fused-attention.html
    • FlashAttention, LLM के attention ऑपरेशन को drop-in रूप में replace कर सकता है
      FlashAttention attention के Softmax(QK^T)V हिस्से को calculate करने का तरीका है, और GQA Q, K, V matrices calculate करने का तरीका है। sliding window attention के बारे में पक्का नहीं, लेकिन यह attention mask बदलने का तरीका है जो control करता है कि कौन-सी query कौन-सी key देख सकती है
      मैंने llama.cpp इस्तेमाल नहीं किया है, लेकिन यह explanation कि उसने CUDA kernels इस्तेमाल करने शुरू किए हैं, कुल मिलाकर सही लगता है
      आख़िरी सवाल पहले Triton में लिखे FlashAttention implementation की बात कर रहा है

  • इस लेख में कहा गया है कि sigmoid जैसे operations बहुत slow होते हैं, इसलिए जिज्ञासा हुई
    modern LLMs SiLU, Swish, SOLU जैसे sigmoid या Softmax वाले activation functions बहुत इस्तेमाल करते हैं
    क्या ReLU में performance loss कम होता है? अगर ऐसा है, तो क्या पुराने अच्छे ReLU पर लौटना बेहतर हो सकता है?

    • ReLU सचमुच एक linear function है जो किसी point पर 0 पर कट जाता है, इसलिए जिन चीज़ों में exponential function होता है उनकी तुलना में इसका computation काफ़ी कम है। हालांकि इतने simple activation function से competitive results पाना मुश्किल लगता है