提出更简洁的FlashAttention变体,加速计算并降低硬件开销。
FLASH-D: FlashAttention with Hidden Softmax Division
- 将softmax除法融入非线性计算,简化核心运算流程。
- 无需减去最大值即可稳定计算指数,避免数值误差。
- 硬件实现节省22.8%面积、20.3%功耗,性能无损。
Transformer的注意力机制推动了人工智能的发展,其计算效率至关重要。但注意力计算涉及矩阵运算与softmax归一化交替进行,天然导致速度瓶颈且需处理完整序列。基于在线softmax计算,FlashAttention将softmax与矩阵运算融合,支持与序列长度无关的分块计算。虽已针对GPU优化,其简洁性也利于硬件加速。本文重新审视核心FlashAttention内核,提出FLASH-D——一种数学等价但更简化的公式:(a) 将softmax除法嵌入其他非线性函数中;(b) 原生数值稳定地计算指数,无需最大值减法;(c) 在不引入近似的情况下降低计算成本,同时完全保留支持高效分块实现的关键特性。28nm硬件实现结果表明,相比现有先进并行架构,本方案平均减少22.8%面积和20.3%功耗,且无性能损失。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。