ZYFA, kernel de atención para AMD MI355X, mejora rendimiento y reduce error

Zyphra presentó ZYFA, un kernel de atención hasta 1,76 veces más rápido en AMD MI355X, según datos proporcionados. El kernel acelera hasta 1,76 veces la atención FP8 frente a AITER en una AMD MI355X y reduce el error en comparación con AITER. ZYFA es de código abierto y está diseñado específicamente para AMD Instinct MI355X, con licencia Apache 2.0.

Según Zyphra, ZYFA sustituye el cálculo exponencial del softmax por una aproximación y, según la empresa, logra menos error FP8 que AITER. El kernel está pensado para la etapa de prefill en modelos de inteligencia artificial e implementa atención causal hacia adelante en FP8 e4m3 y BF16.

En activaciones reales, Zyphra reporta mejoras FP8 de 1,16 a 1,76 veces frente a AITER. En pruebas gaussianas, ZYFA alcanza 2.672 TFLOP/s FP8, equivalente a 53% del pico teórico FP8 de la tarjeta. En pruebas de Mistral-7B, ZYFA reporta una ventaja FP8 de 1,16 veces frente a AITER y una eficiencia de 1,94 TFLOP por joule.

El kernel no incluye kernels para decode ni soporte para caché KV paginada, según Zyphra. ZYFA no se presenta como un kernel genérico para cualquier GPU de AMD, requiere ROCm 7.2.3 y PyTorch 2.12, y no se comparó con GPU de NVIDIA en la tabla comparativa.

Zyphra destacó que ZYFA no es una solución completa para todas las fases de inferencia y recomienda evaluar el kernel con activaciones reales y frecuencias sostenidas. ZYFA está disponible bajo la licencia Apache 2.0 y ofrece una interfaz compatible con flash_attn_func.