KI-Hardware und Software-Stack
Was ist die Grundidee schneller Attention-Implementierungen wie FlashAttention?
- ADie Berechnung wird in Kacheln (Tiles) unterteilt und im On-Chip-Speicher ausgeführt, ohne die große n×n-Zwischentabelle in den HBM zu schreiben, wodurch Speicherzugriffe reduziert werden
- BDie Attention-Berechnung wird ausgelassen und ein Näherungsergebnis ausgegeben
- CDie Attention wird auf der CPU berechnet
Antwort und Erläuterung anzeigen
Antwort: A. Die Berechnung wird in Kacheln (Tiles) unterteilt und im On-Chip-Speicher ausgeführt, ohne die große n×n-Zwischentabelle in den HBM zu schreiben, wodurch Speicherzugriffe reduziert werden
Bei der naiven Implementierung wird die n×n-Tabelle in den HBM geschrieben, für den Softmax erneut gelesen und für das Produkt mit Value noch einmal gelesen; die Speicherzugriffe bestimmen damit die Leistung. FlashAttention teilt Query sowie Key und Value in Blöcke, lädt sie in den Shared Memory und führt die Berechnung in einem Durchgang aus, wobei der Softmax schrittweise aktualisiert wird; die Tabelle wird nicht in den HBM geschrieben. Bei gleichem Rechenaufwand sinken die Speicherzugriffe stark, und die Verarbeitung wird schneller.