Zyphra ZYFA: FP8 FlashAttention для AMD MI355X — ускорение до 1.76x

Zyphra представила ZYFA (Zyphra FlashAttention on AMD) — набор высокопроизводительных ядер внимания (attention kernels) для архитектуры CDNA4 (AMD Instinct MI355X). Проект открыт под лицензией Apache 2.0 и интегрируется в стек Zyphra Optimization Engine (ZOE). Решение нацелено на оптимизацию фазы префилла (prefill) в моделях с длинным контекстом, используя точность FP8 e4m3 и BF16.

01Почему FP8 на MI355X отстает от BF16?

На бумаге FP8 матричные инструкции работают на пиковой частоте, в 2 раза превышающей производительность BF16. Однако в механизме внимания (Attention) вычисление softmax происходит на векторных процессорах (vector units), а не в матричных блоках (matrix pipes). Это создает узкое горлышко:

• Матричная работа (QKᵀ и PV): 1024 такта для FP8.
• Softmax: 896 тактов векторного вычисления (одинаково для FP8 и BF16).

Библиотека AITER (стандарт от AMD) простаивает в матричном конвейере 31.7% времени, так как MI355X не отслеживает завершение инструкций аппаратно. ZYFA решает эту проблему перепланированием потоков.

02Технические инновации ZYFA

Zyphra применяет три метода для устранения узких мест и повышения утилизации GPU:

1. Lean P (Lean Probability): Замена экспоненциальной функции softmax на аппроксимацию. Поскольку битовая паттерн e4m3 растет на 8 при удвоении значения, код $2^x$ можно выразить как $8x + C$. Это сокращает векторные такты на блок с 896 до 384 за счет одного умножения-сложения и преобразования типов.
2. Двухволновая пинг-понг-схема: Два потока (waves) делят один SIMD. Пока один поток выдает матричные инструкции, другой вычисляет вероятности. При использовании Lean P загрузка матричного конвейера возрастает с 60% до 86%.
3. Масштабирование по блокам: Каждый блок из 64 ключей масштабируется степенью двойки перед округлением. Это позволяет использовать инструкцию масштабированного умножения CDNA4 без дополнительных затрат, а сумма округленных вероятностей вычисляется одной дополнительной матричной инструкцией.

03Производительность и бенчмарки

Тестирование проводилось на реальных активациях моделей и синтетических данных (Gaussian) при длине контекста 32K токенов. ZYFA демонстрирует значительный прирост пропускной способности по сравнению с AITER.

На реальных моделях прицел (sustained clock) зависит от количества переключающихся битов. AITER генерирует больше нулей и субнормальных чисел, что позволяет поддерживать более высокую тактовую частоту (2108 МГц против 1628 МГц у ZYFA для Mistral-7B). Тем не менее, ZYFA завершает вычисления быстрее за счет меньшего количества тактов.

04Точность и совместимость

Точность ZYFA близка к эталону fp64. Ошибка FP8 в ZYFA на 62-78% ниже, чем в AITER при head dimension 128. В полной модели Mistral-7B расхождение KL-divergence от fp64 составляет $3.4 \times 10^{-3}$ (ZYFA) против $5.9 \times 10^{-3}$ (AITER).

Для дальнейшего снижения ошибки квантования Q, K и V рекомендуется использовать вращение Хадамара (Hadamard rotation), что добавляет лишь ~0.1% к времени вычислений.

Установка и запуск

Решение требует ROCm 7.2.3 и PyTorch 2.12. Интерфейс совместим с библиотекой flash_attn.

pip install zyfa

Использование в коде аналогично стандартному FlashAttention:

from zyfa import flash_attn_func

# Пример вызова для префилла
out = flash_attn_func(q, k, v, causal=True, softmax_scale=1.0)

05Сравнение с другими решениями

06Кому подойдёт / что запустится

• Обязательное железо: Только AMD Instinct MI355X (CDNA4, gfx950). На MI300X/310X не работает.
• Стек: ROCm 7.2.3, PyTorch 2.12.
• Сценарий: Оптимизация фазы префилла (prefill) для длинных контекстов (до 32K+ токенов) в LLM. Не подходит для фазы декодирования (decode) или задач с paged KV-cache.
• Модели: Mistral-7B, Qwen3-8B, DeepSeek-V2-Lite и архитектуры с MLA.

Источник: MarkTechPost ↗