Проблема: узкое место MUFU в FP8
На архитектуре NVIDIA Blackwell (B200) использование формата FP8 удваивает пропускную способность тензорных ядер. Однако специализированный блок MUFU, вычисляющий экспоненты, не получил такого ускорения. В стандартных ядрах внимания одна операция экспоненты на каждый скор-токен ограничивает производительность уровнем около 2.53 PFLOP/s — лишь половина от пиковой производительности FP8.
Решение: ExpCast и отказ от экспонент
Open-VC базируется на идее VC-Attention (Li et al., 2026), но переосмысливает её реализацию. Ключевая инновация — ExpCast: код вероятности softmax в формате FP8 (E4M3) является почти аффинной функцией от логарифма значения. Это позволяет вычислять экспоненту через одну операцию FMA (fused multiply-add), полностью исключая вызов MUFU и конвертацию в FP32.
В отличие от оригинальной статьи, MachGen отказалась от метода V-Smooth (кластеризация значений), так как на FP8 накладные расходы на перестановку данных превышают выгоду от снижения ошибки квантования. Вместо этого внедрена V-residual repair — точечное исправление ошибок для «плохо округленных» токенов.
Архитектурные оптимизации
- Среднеоконный обход (Mid-window traversal): Порядок посещения блоков ключей изменен для раннего обнаружения максимума, что критично для предотвращения клиппинга скор при ExpCast.
- Warp-specialized pipeline: Разделение задач между варпами: одни загружают данные, другие выполняют умножение матриц, третьи — обработку softmax. Это позволяет скрыть задержки памяти.
- Слияние подготовки ввода: Квантование Q, K и V объединено в один CUDA-граф, добавляя всего 0.86 мс (1.5%) к вызову внимания при длине последовательности 73,397 токенов.
Результаты производительности
Open-VC демонстрирует значительное ускорение на B200. Ниже приведено сравнение времени выполнения полного вызова внимания (включая подготовку и опциональное исправление остатков) для активаций с 56 головами:
| Конфигурация | Относительная ошибка L2 | Общее время (мс) | Снижение ошибки | Доп. время |
|---|---|---|---|---|
| Open-VC (без repair) | 2.980% | 58.14 | — | — |
| Open-VC (0.5% repair) | 2.924% | 59.05 | 1.90% | 1.55% |
| Open-VC (1% repair) | 2.919% | 59.57 | 2.07% | 2.46% |
| Open-VC (2% repair) | 2.912% | 60.81 | 2.28% | 4.58% |
| BF16 (базовый) | — | 114.36 | — | — |
Использование 0.5% токенов для исправления остатков (repair) дает оптимальный баланс: снижение ошибки почти на 2% при увеличении времени всего на 1.55%. Это позволяет достичь скорости, в 2 раза превышающей BF16 FlashAttention-4, сохраняя при этом высокое качество генерации видео.
Почему это важно
Открытый код доступен на GitHub. MachGen подчеркивает, что открытые ядра (kernels) и рецепты их использования могут стать фундаментом для всей индустрии, превращая узкое преимущество одной компании в общее достояние. Это особенно критично для генерации видео, где внимание является самой ресурсоемкой частью процесса.
Источник: Machgen ↗