Zyphra выложила в открытый доступ ZYFA (Zyphra FlashAttention on AMD) — набор ядер attention для префилл-фазы (causal forward) в форматах FP8 e4m3 и BF16. Код распространяется под лицензией Apache 2.0, а работает всё это только на AMD Instinct MI355X (архитектура CDNA4, gfx950) при ROCm 7.2.3 и PyTorch 2.12. На реальных активациях моделей FP8-ядра ZYFA обгоняют AITER, штатную библиотеку ядер AMD, в 1,16–1,76 раза.
Это не самостоятельный продукт, а движок attention, на котором работает Zyphra Optimization Engine (ZOE), обслуживающий Zyphra Cloud. BF16-ядра написаны на HIP, FP8-ядра — вручную распланированный ассемблер, который генерирует Python-скрипт. Интерфейс совместим с flash_attn_func, то есть подмена происходит без переписывания кода вокруг.
Почему FP8 на MI355X упирается в softmax, а не в матрицы
Матричные инструкции FP8 на MI355X работают вдвое быстрее BF16, но ускоряют только два матричных произведения — QKᵀ и PV. Softmax между ними выполняется на векторном блоке с одинаковой скоростью в обеих точностях. Для 32 запросов против 128 ключей матричная работа в FP8 занимает 1 024 такта, а точный FP8-softmax — 896 векторных тактов. Чтобы матричные блоки не простаивали, нужно почти идеальное перекрытие вычислений.
По трассировкам потоков Zyphra видно, что у AITER матричный конвейер простаивает 31,7% тактов. Аппаратно MI355X не отслеживает завершение матричных инструкций, поэтому перекрытие приходится создавать порядком инструкций вручную.
Три приёма, которые дали ускорение
- Lean P. В формате e4m3 битовый шаблон, прочитанный как целое число, растёт на 8 при каждом удвоении значения, поэтому код 2^x примерно равен 8x плюс константа. ZYFA записывает каждую FP8-вероятность одним умножением-сложением и одним целочисленным преобразованием. Векторные такты softmax на блок падают с 896 до 384.
- Двухволновой ping-pong. Две волны делят один SIMD: одна выдаёт матричные инструкции, пока другая считает вероятности, затем они меняются местами. С Lean P матричный конвейер занят 86% тактов в установившемся режиме против 60% при точном softmax.
- Поблочное масштабирование и округлённая сумма строк. Каждый блок из 64 ключей масштабируется степенью двойки перед округлением, а масштабированная матричная инструкция CDNA4 снимает множитель бесплатно. Деление выполняется на сумму уже округлённых вероятностей — на это уходит одна дополнительная матричная инструкция на блок.
Цифры: от 1,16x до 1,76x на живых моделях
На реальных активациях при 16K токенов и устойчивых частотах ускорение FP8 составило 1,16x на Mistral-7B, 1,17x на Qwen3-8B и 1,76x на DeepSeek-V2-Lite. На гауссовых бенчмарках при 32K ядро с head dim 128 выдаёт 2 672 TFLOP/s против 1 944 у AITER — 53% от паспортного пика FP8 против 39%. На форме MLA результат 2 464 против 1 237, то есть 1,99x.
На захваченных слоях Mistral-7B и Qwen3-8B ZYFA даёт 1,94 и 1,91 TFLOP/Дж, AITER — 1,60 и 1,56. По BF16 результаты практически равны: 1,01–1,02x.
Почему на моделях выигрыш меньше, чем на бенчмарке
Частота MI355X зависит от того, сколько бит операндов переключается. AITER позволяет большему числу малых вероятностей стать нулями или субнормалями — биты переключаются реже, частота выше. На слое Mistral-7B ZYFA работает на 1 628 МГц, а AITER — на 2 108 МГц. ZYFA всё равно тратит меньше тактов в сумме. Команда Zyphra рекомендует мерить на реальных активациях при устойчивых частотах, а не на синтетике.
Точность: ошибка FP8 ниже, чем у AITER
При head dim 128 ошибка FP8 у ZYFA на 62–78% ниже, чем у AITER. Относительно attention в fp64 ошибка составляет 0,22x от ошибки AITER на Mistral-7B и 0,38x на Qwen3-8B; до Composable Kernel и Triton — в пределах нескольких процентов. У полной Mistral-7B расхождение KL с fp64 равно 3,4×10⁻³ у ZYFA и 5,9×10⁻³ у AITER. Выходы ZYFA оказались побитово идентичны на 7 200 запусках.
Основная часть ошибки полной модели в FP8 приходит от квантования Q, K и V. Поворот Адамара для Q и K примерно вдвое снижает эту входную ошибку, добавляя к времени attention около 0,1% на Mistral-7B.
С чем сравнивают и чего ZYFA не умеет
ZYFA заявлена как ядро только для префилла: декодирование и paged KV не поддерживаются, а целевая платформа — единственная модель MI355X. AITER покрывает диапазон от MI300X до MI355X и умеет префилл, декод, MLA и paged KV. FlashAttention-3 от Dao-AILab и соавторов нацелена на NVIDIA H100/H800, FlashAttention-4 — на B200 и Hopper. Кроссвендорные цифры получены на разных входах и частотах, и GPU NVIDIA Zyphra не измеряла.
Для российского админа практический вывод скромный: MI355X — редкая в локальных стойках платформа, а ROCm 7.2.3 и PyTorch 2.12 задают узкий коридор совместимости. Если инференс на AMD Instinct уже стоит, ZYFA стоит проверить как замену AITER на префилле длинного контекста — но только с замером на своих активациях. Смежная тема — как инференс ускоряют на уровне компиляции: Together Compiler перекомпилирует граф под нагрузкой, а контроль расходов агентов разбирает плагин TokenEconomizer от Microsoft.
Частые вопросы
На каком железе работает ZYFA?
Только на AMD Instinct MI355X (CDNA4, gfx950). Нужны ROCm 7.2.3 и PyTorch 2.12. На MI300X и на GPU NVIDIA ядра не запустятся.
Что такое AITER и почему с ним сравнивают?
AITER — библиотека ядер AMD, используемая по умолчанию. Zyphra мерила свои FP8-ядра против неё на одних и тех же реальных активациях и получила 1,16–1,76x.
Это готовый сервер инференса?
Нет. Это библиотека ядер с интерфейсом flash_attn_func, реализующая только causal forward attention для префилла. Декодирование и paged KV в неё не входят.
Под какой лицензией код?
Apache 2.0. Для сравнения: AITER распространяется под MIT, FlashAttention-3 и FlashAttention-4 — под BSD-3-Clause.
Почему на реальных моделях ускорение ниже бенчмарка?
Из-за частоты: у AITER больше малых вероятностей превращаются в нули и субнормали, меньше бит переключается, частота выше. На слое Mistral-7B это 2 108 МГц против 1 628 МГц у ZYFA.
Насколько ZYFA точнее AITER в FP8?
При head dim 128 ошибка ниже на 62–78%. На полной Mistral-7B KL-дивергенция с fp64 — 3,4×10⁻³ против 5,9×10⁻³ у AITER.
Что даёт поворот Адамара?
Применённый к Q и K, он примерно вдвое снижает ошибку от квантования входов и добавляет около 0,1% ко времени attention на Mistral-7B.
Можно ли ставить ZYFA в продакшен вместо AITER?
Только после замера на своих активациях при устойчивых частотах: на синтетике выигрыш выглядит больше, чем на живых слоях моделей. И помните об ограничении — префилл и только MI355X.
Источник: marktechpost.com