🌟 SageAttention: метод квантования механизма внимания в архитектурах трансформеров. Внимание - ключевой компонент трансформеров, но его квадратичная сложность вычислений становится проблемой при обработке длинных последовательностей. Квантование успешно применяется для ускорения линейных слоев, но оно мало изучено применительно к механизму внимания. SageAttention - экспериментальный метод, который использует 8-битное квантование механизма внимания для ускорения вычислений и сохранения точности модели. Метод не требует специального обучения и конвертации моделей в какой-либо формат, он применяется к существующим трансформеным моделям в режиме "plug-and-play". Ключевые особенности метода: 🟢Для уменьшения ошибки квантования используется сглаживания матртицы К (среднее значение K вычитается по всем токенам); 🟢Квантование Q и K в INT8; INT8 в четыре раза быстрее, чем в FP16, и в два раза быстрее, чем в FP8. 🟢Matmul PV выполняется с FP16-накопителем; Умножение матриц в высокой разрядности позволяет ускорить вычисления без потери точности. 🟢Адаптивное квантование; Для каждого слоя внимания выбирается наиболее быстрый вариант квантования. SageAttention реализован с использованием Triton и оптимизирован для GPU RTX4090 и 3090. Метод превосходит FlashAttention2 и xformers по скорости примерно в 2,1 и 2,7 раза соответственно. Тестирование на Llama2, CogvideoX, Unidiffuser и TIMM подтвердило сохранение метрик точности при использовании SageAttention. ⚠️ Использование SageAttention рекомендуется с версиями: 🟠python>=3.11; 🟠torch>=2.4.0; 🟠triton-nightly. ⚠️ SageAttention оптимизирован для RTX4090 и RTX3090. На других архитектурах GPU прирост производительности может быть незначительным. ▶️Пример использования:

# Install sageattention
pip install sageattention

# How to use
from sageattention import sageattn
attn_output = sageattn(q, k, v, is_causal=False, smooth_k=True)

# Plug-and-play example with Cogvideo
# add the following codes and run
from sageattention import sageattn
import torch.nn.functional as F

F.scaled_dot_product_attention = sageattn

# Specifically
cd example
python sageattn_cogvideo.py

📌Лицензирование: BSD-3-Clause license. 🟡Arxiv 🖥GitHub @ai_machinelearning_big_data #AI #ML #SageAttention #Transformers