VC-Attention от Nunchux AI: ускорение видео-диффузии без дообучения — внимание перестало быть узким местом

Видео-диффузионные трансформеры (DiTs) сталкиваются с проблемой: полное self-attention на каждом слое обрабатывает десятки тысяч токенов. Например, для 5-секундного клипа 720p на Wan2.2-14B это около 70K токенов, и на RTX 5090 attention занимает более 64% времени генерации. На B200 внимание составляет примерно две трети каждого шага денойзинга MiniMax-H3.

VC-Attention решает две ключевые проблемы. Первая — ошибка квантизации значений (V). Предыдущие методы сглаживали запросы и ключи, но после этого на долю значений приходилось 82% ошибки на Wan2.2. Nunchux предложили V-Smooth: онлайн k-means группирует токены значений по батчам и головам, затем для каждого блока из 128 токенов вычитается среднее, квантизуется только остаток (8-бит E4M3 или 4-бит NVFP4), а среднее восстанавливается через уже имеющийся online softmax. Группировка выполняется только на первых 25% шагов денойзинга и занимает 3–4% времени attention.

Вторая проблема — softmax, который до сих пор выполняется в FP32 и на B200/H200 становится самым долгим этапом. ExpCast-FP8 использует свойство E4M3: его байт при чтении как целое число примерно равен 8 log2(v) + 56. Это позволяет записывать результат напрямую из логарифмической области одной fused multiply-add операцией. На 79,6% значений результат совпадает с FP32-путем, средняя ошибка — 1,6%. Метод применим только для 8-битного kernel, так как NVFP4 не имеет однозначного логарифмического отображения.

Благодаря ручной фьюжн-оптимизации на CuTe/CUDA один вызов V-Smooth на B200 сократился с 42,2 мс до 4,8 мс. VC-Attention не требует дообучения и может быть интегрирован в существующие видео-диффузионные пайплайны для значительного ускорения инференса.