FlashAttention (EL)

FlashAttention — είναι ένας επαναστατικός αλγόριθμος υπολογισμού του μηχανισμού attention, που αναπτύχθηκε για σημαντική επιτάχυνση της εκπαίδευσης και του inference μεγάλων γλωσσικών μοντέλων (LLM) με διατήρηση πλήρους ακρίβειας υπολογισμών. Ο αλγόριθμος παρουσιάστηκε για πρώτη φορά το 2022 από ομάδα ερευνητών του Πανεπιστημίου Stanford υπό την καθοδήγηση του Tri Dao.

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

Προβληματική του τυπικού attention

Ο τυπικός μηχανισμός self-attention στα transformers υπολογίζεται με τον τύπο: $\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{Q{K}^T}{\sqrt{d_k}}\right)V$ όπου Q, K, V είναι οι πίνακες ερωτημάτων, κλειδιών και τιμών.

Το κύριο πρόβλημα αυτής της προσέγγισης είναι η τετραγωνική πολυπλοκότητα ως προς τον χρόνο και τη μνήμη (O(N²)) σε σχέση με το μήκος ακολουθίας N. Σε μια απλή υλοποίηση είναι απαραίτητο να υπολογίζεται και να αποθηκεύεται στη μνήμη GPU ο πλήρης πίνακας attention S μεγέθους N×N, γεγονός που οδηγεί σε δύο κρίσιμα προβλήματα:

1. Υψηλή κατανάλωση μνήμης: Η αποθήκευση πίνακα N×N καθίσταται αδύνατη κατά την εργασία με μεγάλα πλαίσια.
2. Λειτουργίες εισόδου/εξόδου (IO): Το κύριο σημείο συμφόρησης δεν είναι ο αριθμός των αριθμητικών πράξεων, αλλά οι συνεχείς προσβάσεις στη βραδεία μνήμη του GPU.

Ιεραρχία μνήμης GPU

Για την κατανόηση του προβλήματος είναι σημαντικό να διακρίνουμε δύο τύπους μνήμης στο GPU (με παράδειγμα το NVIDIA A100):

• SRAM (στατική μνήμη): Γρήγορη ενδοεπεξεργαστική μνήμη μικρής χωρητικότητας (~20 MB) με τεράστιο εύρος ζώνης (έως 19 TB/s).
• HBM (μνήμη υψηλού εύρους ζώνης): Βραδεία μνήμη μεγάλης χωρητικότητας (40–80 GB) με πολύ μικρότερο εύρος ζώνης (περίπου 1.5 TB/s).

Αυτή η ασυμμετρία καθιστά τον τυπικό αλγόριθμο attention περιορισμένο από το εύρος ζώνης μνήμης (memory-bound), καθώς διαβάζει και γράφει συνεχώς μεγάλους πίνακες από τη βραδεία HBM, κάτι που αποτελεί την κύρια πηγή καθυστερήσεων.

Βασικές καινοτομίες του FlashAttention

Το FlashAttention είναι ένας IO-aware αλγόριθμος που επιλύει το πρόβλημα μέσω ελαχιστοποίησης των προσβάσεων στην HBM. Αυτό επιτυγχάνεται με τρεις βασικές τεχνικές.

Tiling και επεξεργασία σε μπλοκ

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

Online υπολογισμός Softmax

Βασική τεχνική ανακάλυψη υπήρξε ο «online» υπολογισμός Softmax. Το τυπικό Softmax απαιτεί γνώση όλων των στοιχείων του διανύσματος εισόδου για κανονικοποίηση. Το FlashAttention χρησιμοποιεί έναν τροποποιημένο αλγόριθμο που επιτρέπει τον υπολογισμό Softmax τμηματικά. Διατηρεί δύο ενδιάμεσες τιμές (το τρέχον μέγιστο και το άθροισμα εκθετικών), οι οποίες ενημερώνονται καθώς επεξεργάζονται νέα μπλοκ, επιτρέποντας την επίτευξη ακριβούς αποτελέσματος χωρίς πρόσβαση σε ολόκληρο τον πίνακα.

Συγχώνευση πράξεων σε ένα CUDA kernel

Όλες οι πράξεις attention (πολλαπλασιασμός πινάκων QKᵀ, masking, Softmax, πολλαπλασιασμός με V) ενοποιούνται σε ένα ενοποιημένο fused CUDA kernel. Αυτό μειώνει δραστικά τον αριθμό πράξεων ανάγνωσης/εγγραφής στην HBM: αντί για επανειλημμένες διαβάσεις ολόκληρου του πίνακα, ο αλγόριθμος φορτώνει ένα μπλοκ στη SRAM μία φορά, εκτελεί όλους τους υπολογισμούς και γράφει μόνο το τελικό αποτέλεσμα.

Θεωρητική και πρακτική αποδοτικότητα

Πολυπλοκότητα και βελτιστότητα

Το FlashAttention μειώνει την κατανάλωση μνήμης από O(N²) σε O(N), εξασφαλίζοντας γραμμική κλιμάκωση. Έχει αποδειχθεί ότι η IO-πολυπλοκότητα του αλγόριθμου είναι θεωρητικά βέλτιστη για τον υπολογισμό attention σε διαβαθμισμένη ιεραρχία μνήμης δύο επιπέδων, δηλαδή δεν είναι δυνατό να εκτελεστεί ακριβής attention ταχύτερα χωρίς αλλαγή του υλικού.

Εμπειρικά αποτελέσματα

Η πρώτη έκδοση του FlashAttention επέδειξε σημαντικές βελτιώσεις:

• Επιτάχυνση:
  • BERT-large (μήκος 512): 15% επιτάχυνση εκπαίδευσης.
  • GPT-2 (μήκος 1K): 3-πλάσια επιτάχυνση.
  • Εργασίες Long-Range Arena (1K-4K): 2.4-πλάσια επιτάχυνση.
• Εξοικονόμηση μνήμης: Έως 20-πλάσια εξοικονόμηση μνήμης σε σύγκριση με τις ακριβείς βασικές υλοποιήσεις.
• Βελτίωση ποιότητας μοντέλων: Χάρη στη δυνατότητα εργασίας με μεγαλύτερα πλαίσια, το FlashAttention όχι μόνο δεν χάνει αλλά και βελτιώνει την ποιότητα των μοντέλων. Για παράδειγμα, η perplexity του GPT-2 βελτιώθηκε κατά 0.7 μονάδες, ενώ η ακρίβεια σε εργασίες ταξινόμησης μεγάλων εγγράφων αυξήθηκε κατά 6.4 μονάδες.

Εξέλιξη και περαιτέρω ανάπτυξη

Η επιτυχία του FlashAttention έδωσε το έναυσμα για μια ολόκληρη σειρά αλγόριθμων προσανατολισμένων στο υλικό.

FlashAttention-2 (2023)

Η δεύτερη έκδοση στόχευε σε πληρέστερη αξιοποίηση των πόρων GPU. Στο αρχικό FlashAttention η απόδοση στο NVIDIA A100 ήταν μόλις 25–40% του μεγίστου. Το FlashAttention-2 εισήγαγε βελτιώσεις στην παραλληλοποίηση των υπολογισμών, που επέτρεψαν:

• Επίτευξη διπλάσιας επιτάχυνσης σε σύγκριση με την πρώτη έκδοση.
• Αύξηση της χρήσης GPU έως 50–73% του θεωρητικού μεγίστου.
• Επέκταση υποστήριξης για κεφαλές attention μεγέθους 256, καθώς και για αρχιτεκτονικές Multi-Query Attention (MQA).

FlashAttention-3 (2024)

Η τρίτη έκδοση βελτιστοποιήθηκε ειδικά για την αρχιτεκτονική GPU NVIDIA Hopper (H100). Αξιοποιεί νέες δυνατότητες υλικού, όπως η ασυγχρονία Tensor Cores και η υποστήριξη FP8, επιτρέποντας:

• Επίτευξη επιπλέον 1.5–2-πλάσιας επιτάχυνσης σε σύγκριση με το FlashAttention-2.
• Επίτευξη απόδοσης έως 740 TFLOPS σε FP16 και κοντά στα 1.2 PFLOPS σε FP8.

Εξειδικευμένες λύσεις

Οι ιδέες του FlashAttention αναπτύχθηκαν περαιτέρω σε άλλα έργα:

• FlashInfer (2025): Παραμετροποιήσιμη μηχανή attention, βελτιστοποιημένη ειδικά για εργασίες inference LLM. Εστιάζει στην αποδοτική εργασία με το KV-cache σε λειτουργία ροϊκής παραγωγής.
• FlashMLA (2024): Υλοποίηση attention με συμπίεση του πλαισιακού cache (latent attention), που επιτρέπει εξοικονόμηση μνήμης σε πολύ μεγάλες ακολουθίες με ελάχιστη απώλεια πληροφορίας.

Επίδραση στη βιομηχανία και το οικοσύστημα

Το FlashAttention αποτέλεσε θεμελιώδη ανακάλυψη και μετατράπηκε γρήγορα σε βιομηχανικό πρότυπο για αποδοτική εκπαίδευση και inference LLM. Ενσωματώθηκε σε βασικές βιβλιοθήκες όπως PyTorch και Hugging Face, και χρησιμοποιείται στα περισσότερα μεγάλα γλωσσικά μοντέλα (LLaMA, MPT, Falcon, Claude κ.ά.).

Ακριβώς το FlashAttention και οι επόμενες εκδόσεις του διαδραμάτισαν καθοριστικό ρόλο στην αύξηση των πλαισίων περιεχομένου των γλωσσικών μοντέλων: από 2–4 χιλ. tokens (GPT-3) έως 128 χιλ. tokens (GPT-4) και ακόμη έως εκατομμύρια tokens σε πειραματικά μοντέλα. Ο αλγόριθμος εξάλειψε ένα από τα κύρια εμπόδια στην κλιμάκωση των transformers, ανοίγοντας νέες δυνατότητες για εφαρμογές ΤΝ, από την ανάλυση μεγάλων εγγράφων έως την πολυτροπική κατανόηση.

Παραπομπές

• Επίσημο αποθετήριο FlashAttention στο GitHub

Βιβλιογραφία

• Dao, T. et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. arXiv:2205.14135.
• Dao, T. (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. arXiv:2307.08691.
• Shah, J. et al. (2024). FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-Precision. arXiv:2407.08608.
• Kwon, W. et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. arXiv:2309.06180.
• Hong, K. et al. (2023). FlashDecoding++: Faster Large Language Model Inference on GPUs. arXiv:2311.01282.
• Ye, Z. et al. (2025). FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving. arXiv:2501.01005.
• Dege, P. et al. (2025). FlashMLA-ETAP: Efficient Transpose Attention Pipeline for Accelerating MLA Inference on NVIDIA H20 GPUs. arXiv:2506.01969.
• Wang, G. et al. (2025). FlashMask: Efficient and Rich Mask Extension of FlashAttention. OpenReview wUtXB43Chi.
• Dao, T. et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (OpenReview version). OpenReview H4DqfPSibmx.
• Gholami, A. et al. (2024). FlashAttention on a Napkin: A Diagrammatic Approach to Deep Learning IO-Awareness. OpenReview pF2ukh7HxA.