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.