FlashAttention 4: Γρηγορότερη και Αποδοτική σε Μνήμη Προσοχή για LLMs

Εισαγωγή

Η προσοχή με κλιμακωτό εσωτερικό γινόμενο (SDPA) κυριαρχεί στον χρόνο εκτέλεσης και τον ενεργειακό προϋπολογισμό των μεγάλων γλωσσικών μοντέλων (LLMs): σχεδόν όλες οι εκτελέσεις τελεστών πραγματοποιούνται μέσα σε ένα μόνο πρωτόγονο. Η προσοχή παίρνει τα ερωτήματα (Q) και τα πολλαπλασιάζει με τα κλειδιά (K). Στη συνέχεια, κανονικοποιεί το αποτέλεσμα μέσω μιας λειτουργίας softmax πριν το πολλαπλασιάσει με τις τιμές (V) για να παραχθεί η έξοδος. Είναι περιορισμένη από τη μνήμη: διαβάζει Q/K/V πολλαπλές φορές από μνήμη υψηλής ταχύτητας και μετά γράφει επανειλημμένα ενδιάμεσα τμήματα.

Οι οικογένειες FlashAttention μειώνουν αυτό το σημείο συμφόρησης διατηρώντας περισσότερα δεδομένα στο τσιπ για μεγαλύτερο χρονικό διάστημα αξιοποιώντας τους πολλαπλούς επεξεργαστές ροής (SMs) της GPU. Κυκλοφόρησε ανοιχτός κώδικας για πρώτη φορά το 2022, η FlashAttention ανασχεδίασε τη ρουτίνα προσοχής ώστε να υπολογίζει τη softmax σε πραγματικό χρόνο. Μόνο λίγες γραμμές του Q × K αποθηκεύονται στο τσιπ ανά πάσα στιγμή.

Αυτό το άρθρο λειτουργεί ως μια πρακτικά προσανατολισμένη σε βάθος ανάλυση του FlashAttention 4 (FA4). Θα συνοψίσουμε τις αλλαγές σε σχέση με τις προηγούμενες εκδόσεις και θα παρέχουμε κάποιες οδηγίες σχετικά με την υιοθέτηση/αξιολόγηση. Στόχος μας είναι να εξοπλίσουμε τους μηχανικούς υποδομών LLM, τους kernel engineers και τις ομάδες πλατφόρμας ML με αρκετές πληροφορίες για να γνωρίζουν πότε – και αν – πρέπει να χρησιμοποιηθεί το FA4 στο σύστημά τους.

Κύρια Συμπεράσματα

• Το A4 έχει σχεδιαστεί ως πυρήνας προσοχής πρώτης για Blackwell, με στόχο την ανάπτυξη σε GPU SM100. Διαθέτει έναν εξειδικευμένο για warp 5-σταδίων αγωγό, που επιτρέπει αυξημένη επικάλυψη και επαναχρησιμοποίηση στο chip σε σύγκριση με προηγούμενες γενιές του FlashAttention.
• Η αποδοτικότητα βελτιώνεται στο softmax μέσω δύο κύριων τεχνικών: την εκτέλεση της συνάρτησης exp2() στο λογισμικό απευθείας στους πυρήνες CUDA για μείωση της συμφόρησης στα SFU, και την προσαρμοστική διαδικτυακή επανακλίμακα για την αποφυγή περιττών υπολογισμών επανακλίμακας ενώ διατηρείται η σταθερότητα.
• Υπάρχουν ακόμα ορισμένες ελλείπουσες λειτουργίες: Το FA4 είναι προσανατολισμένο προς τα εμπρός και προς το παρόν δεν διαθέτει υλοποιήσεις για την αντίστροφη διέλευση, το varlen και το GQA/MQA, κάτι που περιορίζει τη χρήση του για εκπαίδευση και ορισμένες αρχιτεκτονικές μοντέλων.
• Υιοθετήστε το FA4 με πρακτικό τρόπο: χρησιμοποιήστε το πρώτα για συμπερασματολογία Blackwell, τοποθετήστε το πίσω από δοκιμές benchmarking + ελέγχους ορθότητας + εναλλακτικές λύσεις, και κρατήστε τα FA3 (Hopper)/FA2 (Ampere/Ada) ως τα σταθερά προεπιλεγμένα για ευρείες ανάγκες παραγωγής.

Γιατί οι πυρήνες προσοχής εξακολουθούν να κυριαρχούν στο κόστος των LLM

Παρά την υψηλή απόδοση FLOPS που παρέχεται από τις GPU, τα μοντέλα μετασχηματιστών συχνά περιορίζονται από το εύρος ζώνης μνήμης, επειδή ο μηχανισμός προσοχής χρειάζεται να φορτώσει τα Q, K και V από την κύρια μνήμη πολλές φορές. Για να υπολογιστεί η έξοδος της προσοχής για κάθε διακριτικό ερωτήματος, το καθένα πρέπει να εκτελέσει ένα προϊόν κουκκίδων με έως και χιλιάδες διακριτικά κλειδιών, να κανονικοποιήσει και να επανακλιμακώσει αυτό το αποτέλεσμα και, τέλος, να συγκεντρώσει τις τιμές. Απλές υλοποιήσεις υπολογίζουν ολόκληρο τον πίνακα Q×K και τον αποθηκεύουν στη μνήμη. Αυτό υπερβαίνει γρήγορα το μέγεθος της μνήμης στο τσιπ, οδηγώντας σε περιορισμένη απόδοση από αναγνώσεις εκτός τσιπ.

Το FlashAttention μετριάζει αυτό το πρόβλημα τοποθετώντας τη μήτρα προσοχής. Αντί να δημιουργήσει ολόκληρο το Q×K, μεταφέρει μικρά πλακίδια Q και K σε κοινόχρηστη μνήμη, υπολογίζει το softmax στο πλακίδιο και διατηρεί μόνο το μέγιστο και το άθροισμα για κάθε σειρά ερωτήματος. Αυτό μειώνει σημαντικά την κίνηση της μνήμης, καθιστώντας τον υπολογισμό του πυρήνα δεσμευμένο.

Η Εξέλιξη του FlashAttention με μια Ματιά

Είναι χρήσιμο να κατανοήσουμε πώς έχει εξελιχθεί το FlashAttention με την πάροδο του χρόνου. Κάθε έκδοση αξιοποίησε νέες λειτουργίες υλικού για να βελτιώσει την ταχύτητα ή την απόδοση. Το FA4 είναι ειδικά βελτιστοποιημένο για GPUs γενιάς Blackwell. Ακολουθεί ένας πίνακας με κάθε έκδοση του FlashAttention, τις κύριες βελτιώσεις που έφερε και την επίδρασή τους στην απόδοση:

Τι νέο υπάρχει στο FlashAttention 4

Το FlashAttention 4 συνεχίζει το παράδειγμα υπολογισμού με πλακάκια που καθιερώθηκε από προηγούμενες εκδόσεις του FlashAttention: οι τεσονές Q/K/V επεξεργάζονται πλακάκι προς πλακάκι (μπλοκ προς μπλοκ) για μέγιστη τοπικότητα δεδομένων στον επεξεργαστή και ελάχιστη κίνηση προς τη γενική μνήμη. Οι μεμονωμένες περιπτώσεις πυρήνων (που ονομάζονται συνεργατικός πίνακας νημάτων ή CTA) συνεργάζονται για να υπολογίσουν ένα ή περισσότερα πλακάκια εξόδου, μεταφέροντας αντίστοιχα πλακάκια των Q, K και V από τη γενική μνήμη και στη συνέχεια προβάλλοντάς τα στις τελικές εξόδους προσοχής.

Τα εσωτερικά στοιχεία του FA4 αποτελούνται από ένα πολύ σωληνοειδές πρόγραμμα που επικαλύπτει διάφορα στάδια του υπολογισμού της προσοχής για να βελτιώσει την απόκρυψη καθυστέρησης και να μεγιστοποιήσει τη χρήση πόρων της GPU. Σε αντίθεση με το FlashAttention 3, το οποίο έχει ουσιαστικά μία επικαλυπτόμενη 2-σταδίου διαδικασία φόρτωσης/υπολογισμού, το FA4 διασπά τον υπολογισμό σε αρκετά σύνολα ταυτόχρονων σταδίων, με διαφορετικά μέρη του σωλήνα να διαχειρίζονται ξεχωριστές ομάδες warp. Αυτές οι ομάδες warp ειδικεύονται η καθεμία σε μία μόνο ευθύνη κατά τη διάρκεια του "κύκλου ζωής" ενός block μέσω της λειτουργίας προσοχής.

Αυτό το μοντέλο εξειδίκευσης warp αναθέτει warp 32 νημάτων σε συγκεκριμένες εργασίες, όπως:

• Μεταφορά δεδομένων (φορτώσεις): ανάκτηση πλακιδίων ερωτήματος, κλειδιού και τιμής από τη μνήμη παγκοσμίως σε γρηγορότερη μνήμη στο τσιπ.
• Στάδια υπολογισμού: Εκτέλεση μερικών βαθμολογιών προσοχής χρησιμοποιώντας υλικό πολλαπλασιασμού-συσσώρευσης μητρών με πυρήνες tensor.
• Softmax και λογική κανονικοποίησης: Εκθετικές + λειτουργίες κανονικοποίησης για τα βάρη softmax της προσοχής.
• Λογική επανακαθορισμού κλίμακας και μείωσης: εφαρμογή αριθμητικών προσαρμογών ή συνδυασμός μερικών αποτελεσμάτων πριν από τη συσσώρευση.
• Επίλογος/λειτουργίες αποθήκευσης: η εγγραφή των ολοκληρωμένων πλακιδίων του τελικού τανυστή πίσω στη γενική μνήμη.

Η διαίσθηση πίσω από το βασικό μοντέλο εκτέλεσης είναι η διατήρηση ασύγχρονων σωληνώσεων για κάθε ομάδα warp, έτσι ώστε ο προγραμματιστής warp της GPU να μπορεί να εναλλάσσεται γρήγορα μεταξύ των ομάδων warp καθώς οι διαθέσιμοι τελεστές πηγής γίνονται διαθέσιμοι.

Αποκορύφωση Λογισμικού μέσω Πυρήνων CUDA

Ένα ιδιαίτερα εντυπωσιακό νέο κόλπο στο FA4 είναι ο τρόπος που υπολογίζει τη συνάρτηση εκθετικού που χρησιμοποιείται στο softmax. Κανονικά, οι πυρήνες εκμεταλλεύονται τα αφιερωμένα SFU στην GPU για να εκτελούν δαπανηρές μαθηματικές πράξεις όπως το exp(). Ωστόσο, επειδή υπάρχουν σχετικά λίγα SFU ανά SM, μπορούν εύκολα να γίνουν σημείο συμφόρησης στην ουρά.

Το FA4 παρακάμπτει αυτόν τον περιορισμό χρησιμοποιώντας κανονικούς πυρήνες CUDA για την εκτέλεση προσεγγιστικών εκθετικών συναρτήσεων λογισμικά. Συγκεκριμένα, το FA4 περιέχει μια προσαρμοσμένη υλοποίηση μιας πολυωνυμικής προσέγγισης exp2(x) με ακρίβεια υλισμικού. Το FA4 προσομοιώνει τον υπολογισμό της εκθετικής συνάρτησης χρησιμοποιώντας ένα κυβικό πολυώνυμο, το οποίο στη συνέχεια μπορεί να εκτελεστεί παράλληλα σε πολλούς πυρήνες CUDA αντί να συγκεντρώνεται σε λίγες SFU.

Το αποτέλεσμα: δεν χρειάζεται πια να περιμένουμε τις SFU για να εκτελέσουν softmax. Με την επεξεργασία της exp() στους κεντρικούς πυρήνες γενικής χρήσης, το FA4 επιτρέπει στην υποομάδα που υπολογίζει το softmax να ακολουθεί καλύτερα τον ρυθμό του pipeline.

Έξυπνος επανακανονισμός Softmax στο Διαδίκτυο (Νέο στο FA4)

Μια ακόμα καινοτομία στο FA4 σχετίζεται με τον τρόπο που η συσσώρευση softmax διατηρεί τη μαθηματική σταθερότητα. Όπως κάνουν όλες οι εκδόσεις του FlashAttention, η «online» προσέγγιση softmax – συσσωρεύει μερικές τιμές max() και αθροίσματα καθώς περνούν μέσα από τη σειρά (για να αποφευχθεί η υπερχείλιση). Οι προηγούμενες εκδόσεις ενημέρωναν συνεχώς τον τρέχοντα συντελεστή κλίμακας (συνήθως το μέγιστο logit μέχρι στιγμής). Επανεπεναρμοζόντουσαν συνεχώς – κάθε φορά που έβλεπαν μια νέα μέγιστη τιμή, επανακλιμάκωναν για να διατηρήσουν τη σταθερότητα του softmax. Το FA4 εισάγει μια προσαρμοστική πολιτική επανακλιμάκωσης. Ανακλιμάκωνεται μόνο αν η νέα μέγιστη τιμή που βλέπετε είναι σημαντικά μεγαλύτερη από την προηγούμενη. Αυτό, σύμφωνα με αναφορές, μειώνει τη συχνότητα των λειτουργιών επανακλιμάκωσης κατά παράγοντα 10. Ας εξετάσουμε το ακόλουθο διάγραμμα:

Το διάγραμμα συγκρίνει δύο εκδοχές επανακλιμάκωσης που χρησιμοποιούνται στο FlashAttention, επιτρέποντας σταθερό και αποδοτικό σε μνήμη υπολογισμό της προσοχής. Η αρχική στρατηγική επανακλιμακώνει την έξοδο κάθε μπλοκ εάν εμφανιστεί η μέγιστη τιμή, εξασφαλίζοντας ότι τα μικρά αθροίσματα παραμένουν ακριβή κανονικοποιώντας τις συνεισφορές βήμα προς βήμα. Με την προσαρμοστική επανακλιμάκωση, παρακολουθούμε τα τοπικά και παγκόσμια μέγιστα, επανακλιμακώνοντας μόνο όταν συμβαίνουν απότομες αλλαγές ή τα μπλοκ μεγαλώνουν. Αυτό μειώνει το κόστος επεξεργασίας ενώ εξακολουθεί να αποτρέπει το υπερπλήρωμα και επιτρέπει τον αριθμητικά σταθερό αθροισμό όρων με μικρές τιμές.

Συμβατότητα και τρέχουσα κατάσταση

Αυτός ο πίνακας παρέχει μια επισκόπηση της αναμενόμενης συμβατότητας του FlashAttention 4 (FA4) και του «τι λειτουργεί σήμερα» σε όλα τα επίπεδα:

Η προώθηση του FA4 είναι ήδη ευρέως διαθέσιμη (δεσμευμένη) στο δημόσιο δέντρο πηγαίου κώδικα. Υπάρχουν ανοιχτές αναρτήσεις σε έργα όπως το αποθετήριο FlashAttention του Lab όπου οι χρήστες ζητούν πλήρη υποστήριξη για την αντίστροφη πορεία/varlen/group/GQA. Αυτό υποδηλώνει ότι εξακολουθεί να είναι ένα ανοιχτό δημόσιο ζήτημα.

Η υποστήριξη ακολουθιών μεταβλητού μήκους (αποτελεσματικός batching ακολουθιών διαφορετικών μηκών χωρίς padding) και το GQA (Grouped-Query Attention) δεν υποστηρίζονται ακόμη στη διαδικασία πίσω διέλευσης. Το GQA είναι ένα κόλπο προσοχής όπου πολλαπλές κεφαλές προσοχής χρησιμοποιούν τις ίδιες προβολές κλειδιού/τιμής (μερικές φορές χρησιμοποιείται σε LLMs για μείωση της χρήσης μνήμης). Οι προηγούμενες εκδόσεις του FlashAttention υποστήριζαν διαμορφώσεις τύπου MQA/GQA. Ωστόσο, ο κύριος πυρήνας προόδου του FA4 πιθανότατα θα υποθέσει ότι χρησιμοποιείται ο ίδιος αριθμός κεφαλών για κάθε ερώτημα. Έτσι, πιθανότατα δεν υποστηρίζει σωστά ακόμη τις ομαδοποιημένες κεφαλές. Τα μοντέλα με grouped-query attention (όπως ορισμένες παραλλαγές του Llama2) δεν θα υποστηρίζονται από το FA4 μέχρι να υλοποιηθεί αυτό.

Ενσωμάτωση στο οικοσύστημα: Το FA4 δεν έχει ακόμη υιοθετηθεί από πολλά υψηλού επιπέδου πλαίσια/στρώματα επειδή είναι καινούριο. Ένα παράδειγμα είναι η μηχανή εξυπηρέτησης vLLM (υψηλά βελτιστοποιημένη για εκτέλεση LLM), η οποία αυτή τη στιγμή δεν υποστηρίζει FA4 – υποστηρίζονται μόνο τα FlashAttention 2 και 3. Η επιλογή του FA4 στο vLLM θα έχει ως αποτέλεσμα σφάλμα «μη υποστηριζόμενη έκδοση», επειδή ο επιλέκτης εκδόσεων τους σας επιτρέπει επί του παρόντος να επιλέξετε μόνο μεταξύ των εκδόσεων 2 ή 3.

Πώς να Υιοθετήσετε το FlashAttention – Οδηγός Αποφάσεων

Χρησιμοποιήστε αυτόν τον πίνακα για να καθορίσετε ποια έκδοση του FlashAttention είναι βέλτιστη για τις ανάγκες σας.

Συχνές Ερωτήσεις

Τι πρόβλημα προσπαθεί να λύσει το FlashAttention 4 (FA4);

Μειώνει τη μεγάλη κίνηση προσοχής μνήμης HBM υπολογίζοντας την προσοχή τμήμα προς κομμάτι σε τσιπ, αποφεύγοντας την υλοποίηση ολόκληρου του πίνακα QK⊤.

Ποια είναι η κύρια αρχιτεκτονική αλλαγή στο FA4 σε σύγκριση με το FA3;

Το FA4 χρησιμοποιεί έναν αγωγό εξειδικευμένο για warp, περίπου 5 σταδίων (φόρτωση, υπολογισμός, softmax/κανονικοποίηση, επανακλιμάκωση/μείωση, αποθήκευση) για να αυξήσει την επικαλυπτικότητα και την επαναχρησιμοποίηση στον chip στο Blackwell (SM10.x).

Γιατί η FA4 υπολογίζει εκθετικές τιμές σε πυρήνες CUDA αντί σε SFUs;

Δεδομένου ότι οι SFU είναι σπάνιες και μπορούν εύκολα να γίνουν σημείο συμφόρησης κατά τον υπολογισμό του softmax, το FA4 καταφεύγει στη συνέχεια στον υπολογισμό μιας λογισμικής προσέγγισης του exp2() στους πυρήνες CUDA.

Τι αλλάζει το «προσαρμοστικό διαδικτυακό επανακλιμάκωση» στο softmax;

Αντί να επανακλιμακώνεται κάθε φορά που εμφανίζεται ένα νέο μέγιστο, η FA4 επανακλιμακώνεται μόνο όταν το μέγιστο αυτό αυξηθεί σημαντικά. Αυτό μειώνει το κόστος επανακλιμάκωσης ενώ διατηρεί τη αριθμητική σταθερότητα.

Πότε θα πρέπει ένας μηχανικός να υιοθετήσει το FA4 στην παραγωγή σήμερα;

Ξεκινήστε με Blackwell μόνο για συμπερασματική λειτουργία, με σταθερό μήκος, τυπική προσοχή πίσω από ένα feature flag· διατηρήστε cuDNN/SDPA (και FA3/FA2 σε παλαιότερες GPUs) ως εναλλακτικές λύσεις μέχρι να είναι ώριμη η υποστήριξη για backward/varlen/GQA.

Συμπέρασμα

Το FlashAttention 4 βασίζεται στον στόχο του FlashAttention να μειώσει τις απαιτήσεις της μνήμης-εύρους ζώνης της προσοχής κάνοντας περισσότερη δουλειά στα Q/K/V εντός του chip. Ωστόσο, σε αντίθεση με τις προηγούμενες εκδόσεις, το FA4 το κάνει αυτό με έναν ειδικό για το Blackwell 5-σταδιακό αγωγό warp και δύο σημαντικές βελτιώσεις στο softmax (χρησιμοποιώντας εκθετικούς λογισμικού στους πυρήνες CUDA για να μειωθούν τα εμπόδια της SFU· και χρησιμοποιώντας προσαρμοστική επανακανονικοποίηση σε πραγματικό χρόνο για την εξάλειψη περιττών υπολογισμών σταθερότητας). Το FA4 έχει δείξει αξιοπρεπείς βελτιώσεις στη ροή μετάδοσης κατά το forward-pass σε μεγάλες ακολουθίες σε σύγκριση με το cuDNN, αλλά παραμένει προσανατολισμένο στο forward και μόνο για το Blackwell. Επιπλέον, η FA4 σήμερα έχει σημαντικά κενά χαρακτηριστικών (λειτουργίες όπως backward, varlen, GQA κ.λπ. λείπουν) και η υποστήριξη του οικοσυστήματος δεν είναι πλήρως ώριμη. Το συμπέρασμα προς το παρόν είναι να εξεταστούν οι δρόμοι της FA4 για την εκτέλεση Blackwell ως μια ενδεχομένως πολλά υποσχόμενη προσέγγιση που απαιτεί υιοθέτηση μετά από αυστηρή αξιολόγηση, σωστότητα και μηχανισμούς επαναφοράς. Ταυτόχρονα, οι εργασίες εκπαίδευσης και το ευρύτερο υλικό θα πρέπει να συνεχίσουν να χρησιμοποιούν τους ώριμους δρόμους κώδικα FA3/FA2 για την παραγωγή.

Αναφορές

• FlashAttention-3: Γρήγορη και Ακριβής Προσοχή με Ασύγχρονη Λειτουργία και Χαμηλή Ακρίβεια
• FlashAttention 4: Σπάζοντας το Φράγμα των Πεταφλοπ στην Προσοχή GPU
• Πώς να εφαρμόσετε το FA4 στο B200;
• Υποστήριξη για την αντίστροφη διαδικασία FA4 στο Blackwell

Ευχαριστούμε που μάθατε με την Κοινότητα Linux-Console.net.