Nunchux AI представила VC-Attention: низкоразрядное ядро внимания без обучения, ускоряющее видео Diffusion Transformers
Nunchux AI выпустила VC-Attention — низкоразрядное ядро внимания без обучения, созданное для видеотрансформеров диффузии (DiT). Оно одновременно решает 2 проблемы: ошибку квантования значений и медленный этап softmax.
Почему внимание является узким местом для видео
Видеомодели DiT разворачивают клип в 1 последовательность пространственно-временных токенов и выполняют полное self-attention на каждом слое. 5-секундный клип Wan2.2-14B в разрешении 720p содержит около 70 тыс. токенов. На RTX 5090 на внимание приходится более 64% времени генерации. Исследовательская команда сообщает, что на одной B200 внимание занимает около двух третей каждого шага денойзинга MiniMax-H3.
Тензорные ядра низкой разрядности ускоряют 2 матричных произведения — QK и PV. Остаются 2 препятствия. Во-первых, предыдущие методы, такие как SageAttention2, сглаживают запросы и ключи. После сглаживания и вращения QK на компонент значений приходится 82% ошибки результата в Wan2.2. Во-вторых, softmax между произведениями по-прежнему выполняется в FP32. На B200 и H200 экспонента и её преобразование в FP8 становятся самым долгим этапом конвейера.
V-Smooth: устранение выбросов в значениях
Выбросы значений сосредоточены в небольшом числе токенов, а их каналы меняются между головами, слоями и шагами. Вращение Адамара сохраняет нормы токенов, поэтому не устраняет выбросы. Вращение V изменяет ошибку значений всего на 0,2%.
V-Smooth использует другой подход:
- Группировка: Онлайн-алгоритм k-means кластеризует токены значений для каждого пакета и каждой головы. Ключи и значения перемещаются совместно, поэтому результат некаузального внимания не меняется.
- Удаление среднего: Каждый аппаратный блок из 128 токенов вычитает своё среднее значение. Квантуется только остаток — с использованием E4M3 по каналам при 8 битах или NVFP4 при 4 битах.
- Восстановление: Среднее значение добавляется обратно с использованием суммы строки, которую уже поддерживает онлайн-softmax. Второй проход или дополнительный буфер не требуются.
В среднем по 100 головам Wan2.2 среднее блока устраняет 8% энергии блока при порядке токенов в последовательности. При использовании статического куба DeltaQuant этот показатель составляет 12%, а после сортировки — 36%. Каждое среднее добавляет 0,125 бита на элемент значения.
Группировка выполняется только на первых 25% шагов денойзинга. Перестановка повторно используется для 4 соседних шагов. В среднем по всему расписанию группировка занимает от 3 до 4% времени работы внимания.
ExpCast-FP8: устранение узкого места softmax
Байт E4M3 уже близок к логарифму хранимого в нём значения. Если прочитать его как целое число, он примерно равен 8 log2(v) + 56. Поэтому ExpCast-FP8 записывает байт непосредственно из оценки в логарифмической области с помощью 1 объединённой операции умножения и сложения. Константа β = -0,35 центрирует остаточную ошибку, и для каждой модели отдельная константа не подбирается.
Прямой путь записывает тот же байт, что и путь с вычислением экспоненты в FP32 и последующим преобразованием, в 79,6% случаев каждого удвоения. В остальных случаях результат отличается на 1 код. В статье доказывается ограничение полной вариации для каждой строки менее 3,64% плюс возможный хвост недополнения. Для 204,8 тыс. строк внимания Wan2.2 измеренное среднее составляет 1,6%. ExpCast-FP8 применяется только к 8-битному ядру, поскольку у NVFP4 нет единого аффинного отображения из логарифма в код.
Написанное вручную объединение цепочки предварительной обработки в CuTe/CUDA сокращает время 1 вызова V-Smooth с 42,2 мс до 4,8 мс на B200.
Объяснение: как работает VC-Attention
Результаты тестирования
Тесты охватывают 4 видеомодели DiT с открытыми весами: Wan2.2-T2V-A14B, LongCat-Video, HunyuanVideo-1.5 и MiniMax-H3. Качество оценивается относительно результатов BF16 FlashAttention-4 по 100 запросам.
| GPU (Wan2.2) | Точность | Ускорение внимания | Сквозное ускорение |
|---|---|---|---|
| B200 | 8 бит | 1,59× | 1,19× |
| H200 | 8 бит | 1,46× | 1,13× |
| RTX PRO 6000 | 4 бита | 2,27× | 1,36× |
| RTX 5090 | 4 бита | 3,58× | 1,70× |
На B200 VC-Attention работает в 6,02 раза быстрее SageAttention2, в составе которого нет ядра для Blackwell. На H200 разрыв составляет 1,16 раза. На рабочих станциях 4-битный V-Smooth не уступает SageAttention3 на RTX PRO 6000. На RTX 5090 разница не превышает 5%, поэтому модели разделяет качество.
Результаты по качеству:
- При 8 битах V-Smooth повышает PSNR на 2,3 дБ по сравнению с SageAttention2 на Wan2.2 и на 2,8 дБ на HunyuanVideo-1.5.
- Добавление ExpCast-FP8 возвращает от 0,7 до 2,1 дБ, но результат всё равно превосходит SageAttention2 на всех 4 моделях.
- При 4 битах V-Smooth превосходит SageAttention3 на 2,9 дБ на Wan2.2 и на 3,6 дБ на LongCat-Video.
- При работе без обучения Attn-QAT уступает SageAttention2 от 3,4 до 6,7 дБ.
На MiniMax-H3 в разрешении 1344×768 внимание работает в 1,60 раза быстрее, чем BF16 FlashAttention-4 на B200. PSNR составляет 20,2 дБ против 19,9 дБ у SageAttention2. На B300 в статье сообщается об ускорении 1,47 раза против 1,31 раза у обычного ядра FP8. На графике в блоге указано ускорение 1,51 раза для B300.
Nunchux Attention, собственное расширение компании, обеспечивает ускорение 1,91 раза на B200 и 1,83 раза на B300 для внимания MiniMax-H3.
Метод изменяет только стоимость отдельных взаимодействий. Поэтому его можно комбинировать с разреженным вниманием, например с Sparse VideoGen и Radial Attention, а также с дистилляцией и выполнением на нескольких GPU. Nunchux сообщает, что бесплатный доступ к MiniMax-H3 появится через список ожидания Modelverse.
Главные выводы
- VC-Attention — низкоразрядное внимание без обучения для видеомоделей DiT от Nunchux AI.
- V-Smooth группирует токены значений, а затем квантует только остатки после вычитания среднего блока.
- ExpCast-FP8 заменяет экспоненту FP32 и преобразование в число одной операцией умножения и сложения.
- Внимание Wan2.2 работает в 1,59 раза быстрее на B200 и в 3,58 раза быстрее на RTX 5090.
- Публичного релиза ядра пока нет; Nunchux использует в своём стеке собственное расширение.
Ознакомьтесь с исследовательской статьёй и техническими подробностями. Вся заслуга принадлежит исследователям этого проекта. Также подписывайтесь на нас в Twitter и не забудьте присоединиться к нашему сабреддиту о машинном обучении с аудиторией более 150 тыс. участников и подписаться на нашу рассылку. Постойте! Вы есть в Telegram? теперь вы также можете присоединиться к нам в Telegram.
Хотите сотрудничать с нами для продвижения своего репозитория GitHub, страницы на Hugging Face, релиза продукта, вебинара и т. д.? Свяжитесь с нами
Переведено автоматически с английского. Оригинал статьи — по ссылке ниже.