FlashAttention-4 решает проблемы производительности тензорных ядер 4-битной плавающей точки (FP4) архитектуры Blackwell, внедряя новые пути для некаузального вывода и каузального обучения. Метод, названный Direct-P, напрямую отображает оценки в вероятности FP4, чтобы обойти накладные расходы на преобразование softmax, которые обычно доминируют при уменьшении размеров матричных произведений.

  • Direct-P достигает пропускной способности прямого распространения до 2.13x по сравнению с bfloat16 (BF16) на NVIDIA GB200 для некаузального вывода.
  • Каузальный путь передает прямое квантование непосредственно в обратное распространение, восстанавливая вероятности из сохраненных квантованных запросов и ключей с использованием операндов градиентов 8-битной плавающей точки (FP8).
  • Этот каузальный подход ускоряет полное обновление параметров объемом 8 миллиардов на одном GPU до 1.14x.
  • При распределенном обучении сохраняются вероятности и значения FP8, поскольку все протестированные траектории обучения с MXFP4 для вероятностей/значений расходятся.

Авторы считают это важным, поскольку это обеспечивает эффективное аппаратно-ориентированное внимание FP4 на архитектурах Blackwell, преодолевая зависимости, препятствующие автоматическому ускорению от тензорных ядер FP4.