Материал подготовлен автоматически по первоисточникам: ссылки на них — в конце статьи.
Robert Hu переработал вычисление механизма внимания для 4-битных чисел с плавающей точкой FP4 в препринте, который не проходил рецензирования. По замерам автора, новый путь достигает до 2,13 пропускной способности варианта на BF16 при прямом проходе на NVIDIA GB200. Работа важна разработчикам собственных ядер инференса и обучения для Blackwell: простая замена формата не даёт такого ускорения.
Что сделали
В механизме внимания два матричных умножения разделены нормализацией оценок в вероятности. Тензорные ядра Blackwell быстрее выполняют умножения в FP4, но не ускоряют вычисление экспонент, сокращение строк, упаковку вероятностей и передачу масштабов. После ускорения матричных операций именно этот промежуточный этап оказывается на критическом пути.
Метод Direct-P преобразует оценки сразу в коды MXFP4, которые потребляет второе матричное умножение, не создавая промежуточные вероятности высокой точности. Нормализатор рассчитывается по тем же округлённым значениям. Внешнее расписание FlashAttention-4 сохраняется, а упакованные вероятности занимают освобождённую память оценок.
Для причинного обучения, где позиция видит только предыдущие токены, прямой проход передаёт обратному квантованные запросы и ключи, их масштабы и нормализатор. По замерам автора, полный шаг обновления модели на 8 млрд параметров ускорился до 1,14 раза. При этом в распределённом обучении вероятности и значения пришлось оставить в FP8: все проверенные варианты с MXFP4 для них расходились.
Что это значит
Работа показывает, что производительность низкоточного внимания определяется не только скоростью тензорных ядер. Существенны зависимости между операциями и владение тензорной памятью: пока вероятности занимают прежний участок, следующее умножение не может записать туда оценки. Поэтому перенос Direct-P в другой стек потребует воспроизвести расписание и размещение данных, а не только формат квантования.
Ограничения. Direct-P для непричинного режима проверяли на фиксированных входах, а не в предобучении или дообучении. Сопоставленное распределённое обучение охватывает 100 млрд токенов, но для каждого маршрута выполнена одна траектория, поэтому разброс между запусками не измерен. Аппаратные выводы относятся к GB200, B300 и преимущественно размерности головы 128. Сопоставимого контроля качества с BF16 автор не провёл: использованные рецепты обучения различались.
Источники
Иллюстрация: рисунок из статьи «Hardware-Aware FP4 FlashAttention-4», Robert Hu, CC BY 4.0
Похоже на вашу задачу?
Расскажите, что собираете. За полчаса разложим на этапы и назовём сроки — это бесплатно и ни к чему не обязывает.
