SageAttention2: Технический отчет

В современном мире обработки данных, где внимание играет ключевую роль в моделях машинного обучения, особенно в LLM (Large Language Models), эффективность вычислений становится критически важной. С увеличением длины последовательностей в реальных приложениях, потребность в оптимизации внимания возрастает. В этой статье мы рассмотрим SageAttention2, новый подход к ускорению вычислений внимания с использованием 4-битной квантизации, который значительно улучшает производительность и точность по сравнению с предшественниками.

Основные концепции SageAttention2

1. Квантизация матриц Q и K в INT4

В SageAttention2 матрицы запросов (Q) и ключей (K) квантизируются в 4-битный целочисленный формат (INT4) на уровне warp (группа потоков в GPU). Это позволяет использовать более быстрые операции матричного умножения (Matmul), доступные на современных GPU, что в свою очередь приводит к значительному ускорению вычислений.

2. Сглаживание матриц Q и V

Для улучшения точности внимания с использованием INT4 для QK и FP8 для PV, SageAttention2 применяет методы сглаживания. Это включает в себя вычитание среднего значения по каналам из матриц Q и V перед квантизацией, что помогает уменьшить ошибки, вызванные выбросами данных.

3. Адаптивная квантизация

Анализ точности квантизации на разных временных шагах и слоях модели показал, что в некоторых случаях требуется более высокая точность. SageAttention2 вводит адаптивную квантизацию, где для критических слоев и временных шагов используется 8-битная квантизация (INT8+FP8), а для остальных — 4-битная (INT4+FP8).

Ускорение и точность

SageAttention2 превосходит FlashAttention2 и xformers по скорости примерно в 3.1 раза и 5.4 раза соответственно на GPU RTX4090, достигая пиковой производительности в 485 TOPS (тераопераций в секунду). Эти улучшения достигаются за счет использования более эффективных операций квантизации и матричного умножения.

Точность

• Сглаживание Q и V: Применение сглаживания значительно улучшает точность, особенно при использовании низкобитной квантизации. Это позволяет SageAttention2 сохранять высокую точность даже при использовании INT4 для QK.

• Адаптивная квантизация: Использование смешанной точности для различных слоев и временных шагов обеспечивает минимальные потери в метриках конечного результата, что делает SageAttention2 универсальным решением для различных моделей и задач.

Экспериментальные результаты

Эксперименты проводились на разнообразных моделях, включая те, что используются для обработки текста, генерации изображений и видео. Результаты показывают:

• Минимальные потери в метриках: Независимо от задачи, SageAttention2 демонстрирует несущественные потери в точности по сравнению с полноценным вниманием.

• Ускорение: На L20 GPU SageAttention2 достигает 288 TOPS, что делает его значительно быстрее, чем предыдущие решения.

Заключение и будущие работы

SageAttention2 представляет собой значительный шаг вперед в оптимизации внимания для LLM и других моделей, предлагая ускорение вычислений без значительных потерь в точности. Благодаря своим методам квантизации и сглаживания, SageAttention2 становится важным инструментом для ускорения вывода моделей в реальных приложениях. В будущем планируется реализация FP8 MatMul с использованием FP16 аккумулятора для еще большего ускорения на архитектуре Hopper.

SageAttention2 не только улучшает производительность, но и делает возможным использование более сложных моделей на устройствах с ограниченными ресурсами, открывая новые возможности для применения AI и LLM в различных областях.