Какво е FlashAttention от гледна точка на програмирането на ядра?

FlashAttention е оптимизиран CUDA ядро за изчисляване на механизма на внимание (attention) в трансформърите, което значително ускорява изчисленията и намалява потреблението на памет.

От гледна точка на програмирането на ядра, FlashAttention реализира изчисляването на attention чрез ефективно управление на паметта и паралелизъм на GPU. Основните характеристики:

• Използване на блокове споделена памет за съхранение на междинни данни, което намалява достъпите до бавната глобална памет.
• Изчисляване на attention в нишки с минимално дублиране на данни.
• Оптимизация на изчисляването на softmax с оглед на числовата стабилност.
• Обработка на големи последователности без необходимостта да се зареждат всички данни в паметта едновременно.

Това позволява по-бързо изпълнение на attention и по-малко използване на памет в сравнение с класическите реализации, което е критично за обучението на големи модели.

Например, ядро зарежда блокове ключове и заявки, изчислява скаларни произведения, прилага softmax и умножава по стойностите, използвайки споделена памет за ускоряване и минимизиране на глобалните достъпи.