FLASH-D: FlashAttention with Hidden Softmax Division

TL;DR

FLASH-D introduces a simplified, hardware-efficient formulation of FlashAttention that reduces computational cost, area, and power consumption while maintaining core properties and numerical stability for faster transformer attention computations.

Contribution

It presents a mathematically equivalent reformulation of FlashAttention that simplifies implementation and enhances hardware efficiency without sacrificing accuracy or performance.

Findings

Achieves 22.8% reduction in hardware area

Achieves 20.3% reduction in power consumption

Maintains numerical stability and core properties of FlashAttention

Abstract

The transformer's attention mechanism has revolutionized AI and machine learning, with its efficient computation being crucial to its performance. However, calculating attention involves matrix operations interspersed with softmax rescaling, which inherently slows down computation and requires processing the entire input sequence. Building on online softmax computation, FlashAttention integrates softmax calculation with matrix arithmetic, enabling tiled computation independent of sequence length. While optimized for GPUs, FlashAttention's simplicity makes it amenable to direct hardware acceleration. This work re-evaluates the core FlashAttention kernel, presenting FLASH-D a mathematically equivalent, yet simplified, formulation that achieves: (a) hiding softmax division within other non-linear function evaluations; (b) inherently numerically stable computation of exponentials, eliminating the need for maximum value subtraction; and (c) a reduction in computational cost without introducing numerical approximations to the FlashAttention kernel. Importantly, the essential FlashAttention properties that facilitate efficient tiled implementation are fully preserved. Hardware implementation results at 28nm demonstrate that this proposed formulation achieves a 22.8% reduction in area and a 20.3% reduction in power, on average, compared to state-of-the-art parallel hardware architectures without any performance penalty.