Τι είναι το FlashAttention από την άποψη του προγραμματισμού πυρήνων;

FlashAttention είναι ένας βελτιστοποιημένος πυρήνας CUDA για τον υπολογισμό του μηχανισμού προσοχής (attention) στους μετασχηματιστές, που επιταχύνει σημαντικά τους υπολογισμούς και μειώνει την κατανάλωση μνήμης.

Από την άποψη του προγραμματισμού πυρήνων, το FlashAttention υλοποιεί τον υπολογισμό της προσοχής χρησιμοποιώντας αποτελεσματική διαχείριση μνήμης και παράλληλο επεξεργαστή στην GPU. Τα κύρια χαρακτηριστικά:

• Χρήση μπλοκ κοινόχρηστης μνήμης για την αποθήκευση ενδιάμεσων δεδομένων, μειώνοντας τις κλήσεις σε αργή παγκόσμια μνήμη.
• Υπολογισμός της προσοχής σε νήματα με ελάχιστη διπλή κατανομή δεδομένων.
• Βελτιστοποίηση του υπολογισμού softmax λαμβάνοντας υπόψη τη αριθμητική σταθερότητα.
• Επεξεργασία μεγάλων ακολουθιών χωρίς την ανάγκη φόρτωσης όλων των δεδομένων στη μνήμη ταυτόχρονα.

Αυτό επιτρέπει την ταχύτερη εκτέλεση της προσοχής και τη μείωση της κατανάλωσης μνήμης σε σύγκριση με τις κλασικές υλοποιήσεις, κάτι που είναι κρίσιμο για την εκπαίδευση μεγάλων μοντέλων.

Για παράδειγμα, ο πυρήνας φορτώνει μπλοκ κλειδιών και ερωτημάτων, υπολογίζει εσωτερικά προϊόντα, εφαρμόζει softmax και πολλαπλασιάζει με τις τιμές, χρησιμοποιώντας κοινόχρηστη μνήμη για επιτάχυνση και ελαχιστοποίηση των κλήσεων σε παγκόσμια μνήμη.