Суть инновации: ZYFA против AITER
Компания Zyphra представила ZYFA (Zyphra FlashAttention on AMD) — open-source библиотеку (лицензия Apache 2.0) для оптимизации внимания в префилле (prefill) на GPU AMD Instinct MI355X (архитектура CDNA4, gfx950). Решение работает в форматах FP8 (e4m3) и BF16 при размерности головы (head dimension) 128 и поддерживает MLA-формат (192/128).
Ключевая проблема, которую решает ZYFA: на MI355X операция softmax является узким местом (bottleneck), а не матричные умножения. Нативная библиотека AMD AITER простаивает в матричном конвейере до 31.7% времени из-за отсутствия аппаратного отслеживания завершения инструкций. ZYFA использует ручную раскладку (hand-scheduled assembly) для идеального перекрытия вычислений.
Технические детали: Lean P и Ping-Pong
Zyphra применяет три основные техники для максимизации эффективности:
- Lean P: Замена экспоненты в softmax на одну операцию умножения-сложения (multiply-add) и преобразование целого числа. Это сокращает векторные циклы на блок с 896 до 384.
- Two-wave ping-pong schedule: Два потока (waves) делят каждый SIMD-блок. Пока один выполняет матричные инструкции, другой вычисляет вероятности, повышая загрузку матричного конвейера до 86%.
- Per-block scaling: Масштабирование блоков вероятностей перед округлением, что позволяет использовать масштабированные матричные инструкции CDNA4 без дополнительных затрат.
Производительность и метрики
На реальных активациях моделей ZYFA демонстрирует значительный прирост пропускной способности FP8 по сравнению с AITER. Ниже приведено сравнение ключевых показателей:
| Метрика / Модель | ZYFA FP8 | AITER FP8 (Baseline) | Ускорение |
|---|---|---|---|
| DeepSeek-V2-Lite (реальные активации) | — | — | 1.76x |
| Qwen3-8B (реальные активации) | — | — | 1.17x |
| Mistral-7B (реальные активации) | — | — | 1.16x |
| Пиковая производительность (32K токенов, head dim 128) | 2,672 TFLOP/s | 1,944 TFLOP/s | — |
| Доля от пиковой мощности (FP8) | 53% | 39% | — |
| Энергоэффективность (Mistral-7B) | 1.94 TFLOP/J | 1.60 TFLOP/J | — |
Точность и стабильность
Несмотря на оптимизации, точность FP8 в ZYFA выше, чем в AITER. Ошибка FP8-ядер ZYFA на 62-78% ниже, чем у AITER при head dimension 128. Расхождение Кульбака-Лейблера (KL divergence) от fp64 для Mistral-7B составляет 3.4×10⁻³ против 5.9×10⁻³ у AITER. Результаты ZYFA являются битово идентичными (bitwise identical) при 7,200 запусках.
Ограничения и совместимость
Решение работает только на AMD Instinct MI355X (gfx950) с ROCm 7.2.3 и PyTorch 2.12. ZYFA поддерживает только causal forward prefill (без decode, без paged KV cache). Для достижения максимальной производительности Zyphra рекомендует использовать sustained clocks, так как на реальных моделях частота GPU может падать (например, до 1,628 МГц у ZYFA против 2,108 МГц у AITER в Mistral-7B из-за переключения битов), но общее время выполнения всё равно выигрывает за счет меньшего количества циклов.
Источник: MarkTechPost ↗