Flash Attention 3 вышел!
[Статья][Блог]

Не прошло и года с выхода Flash Attention 2️⃣, как вышло продолжение, доведя серию до трилогии.

Метод

Основная идея первого Flash Attention в уменьшении передачи памяти 📝 между HBM и кэшом GPU, а второго - в отпимизации не matmul операций. Flash Attention 2 позволяет довольно эффективно использовать вычислительные возможности A100, но для более современной H100 все еще имеет место сильная недоутилизация - всего 35%.

В Flash Attention 3 отпимизируют Attention с учетом новых архитектурных особенностей Hopper+ - новых тензорных ядер с большим throughput, тензорным ускорителем памяти (Tensor Memory Accelerator) и поддержкой fp8.

Основных источника ускорения два:

1️⃣ Паралеллизация вычислений между GEMM (General Matrix Multiply) и иными операциями (softmax, RMSNorm).

Хоть по flops softmax и нормализации кажутся мизерными, они выполняются в 256 раз медленее на H100, потому могут занимать до половины времени матричного умножения. За их выполнение отвечают другие компоненты GPU, чем тензорные ядра для перемножения матриц, потому их можно обрабатывать параллельно.

2️⃣ Использование fp8.

fp8 позволяет почти удвоить скорость. Однако есть нюанс - выбросы, приводящие к большим ошибкам квантования. Чтобы избавиться от них используют incoherence processing из QuIP. Благодаря этому ошибка квантизации уменьшается до 2.6 раз.

Результаты

С fp16 удается достичь ускорения до 1.6-1.8 раз по сравнению с Flash Attention 2.

И под 1.2 PFLOPs c fp8 на больших последовательностях и с большими головами трансформера.

Обещают в будущем накатить в торч.

Вывод

Сильно! Ждем через год Flash Attention 4️⃣!