Zyphra выпустила FP8-ядра внимания для AMD MI355X

Zyphra представила ZYFA — библиотеку с открытым исходным кодом, содержащую ядра внимания prefill в форматах FP8 и BF16 для ускорителей AMD Instinct MI355X. Проект распространяется по лицензии Apache 2.0 и рассчитан на каузальное forward-внимание при работе с длинным контекстом. Поддерживаются FP8 e4m3, BF16, размерность головы 128 и формы MLA 192/128. Библиотека работает на MI355X с ROCm 7.2.3 и PyTorch 2.12 и предоставляет совместимый с flash_attn_func интерфейс.

По данным Zyphra, на реальных активациях моделей FP8-ядра работают в 1,16-1,76 раза быстрее AITER — стандартной библиотеки ядер AMD. Производительность BF16 сопоставима с AITER и составляет 1,01-1,02 от его показателей. В тесте с 32K токенов ядро для размерности головы 128 достигло 2672 TFLOP/s, или 53% пикового значения FP8 по спецификации, тогда как AITER показал 1944 TFLOP/s.

При размерности головы 128 ошибка FP8 у ZYFA была на 62-78% ниже, чем у AITER. Реализация заменяет экспоненту softmax приближением на основе одного умножения и сложения, применяет двухволновое планирование и масштабирование блоков. Результат зависит от нагрузки: на активациях Mistral-7B ускорение составило 1,16 раза. ZYFA поддерживает только prefill и работает лишь на MI355X.