Mixed precision flashattention boosts efficiency on dedicated chips
LampAttention: Look-Ahead Mixed-Precision FlashAttention for Dedicated Accelerators
Machine Learning
Summary
Attention mechanisms in AI usually do calculations using numbers with the same precision. The authors found that many of these calculations can be done in lower precision without making the results worse. They created a method that mostly uses 8-bit numbers but switches to 16-bit numbers only for the tricky parts. They also designed a plan for hardware that can run this method efficiently. Tests show that this selective approach keeps the AI's accuracy while using less computing power.
What this means in practice
- •For hardware accelerator designers: Design chips that run attention models using mixed precision to improve speed and energy efficiency without losing accuracy.
- •For machine learning engineers: Implement mixed-precision FlashAttention to optimize transformer-based models for faster inference and training on custom hardware.
Authors
Stanislav Budzinskiy, Marian Gloser, Tolunay Yilmaz, Ying Hong Tham, Yuanyi Lin, Wenyi Fang, Fan Wu, Philipp Petersen
Abstract
While most attention logits can be computed in low precision without degrading numerical stability, current attention kernels fail to exploit this phenomenon. We introduce a novel hardware-algorithm co-design in the form of mixed-precision FlashAttention. Our method accumulates key-query products and evaluates their exponentials in 8-bit formats, then adaptively identifies sensitive sub-blocks and recomputes them in 16-bit formats. We propose the specifications for a dedicated accelerator capable of executing this pipeline efficiently. Simulated experiments with Qwen3 and Gemma 3 show that rerouting a selective minority of sub-blocks to high precision is sufficient to recover the baseline model performance.