Low-Cost FlashAttention with Fused Exponential and Multiplication Hardware Operators

📄 arXiv: 2505.14314v2 📥 PDF

作者: Kosmas Alexandridis, Vasileios Titopoulos, Giorgos Dimitrakopoulos

分类: cs.AR, cs.LG

发布日期: 2025-05-20 (更新: 2025-05-30)

备注: IEEE Computer Society Annual Symposium on VLSI (ISVLSI 2025)

💡 一句话要点

针对FlashAttention，提出融合指数运算和乘法运算的低成本硬件加速器

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: FlashAttention 硬件加速器 指数运算 向量乘法 低功耗设计 ASIC Transformer 大型语言模型

📋 核心要点

  1. Transformer和LLM中的注意力机制面临长序列计算的挑战，需要专门的硬件加速。
  2. 论文提出ExpMul硬件算子，融合指数运算和向量乘法，优化FlashAttention内核。
  3. 实验结果表明，该方法在28nm ASIC技术中，面积减少28.8%，功耗降低17.6%。

📝 摘要（中文）

注意力机制，尤其是在Transformer架构和大型语言模型（LLM）中，彻底改变了机器学习和人工智能应用中的序列建模。为了计算日益增长的长序列的注意力，已经提出了专门的加速器，以直接在硬件中执行关键的注意力步骤。在最近提出的各种架构中，基于FlashAttention算法变体的架构因其优化的计算、平铺能力和减少的内存流量而脱颖而出。本文重点在于优化基于浮点运算的FlashAttention内核，使用新的硬件算子融合指数运算和向量乘法运算，例如e^x, V。所提出的ExpMul硬件算子显著降低了基于FlashAttention的硬件加速器的面积和功耗成本。在28nm ASIC技术中实现时，与具有独立指数运算和向量乘法运算硬件算子的最先进硬件架构相比，平均面积改善了28.8%，功耗改善了17.6%。

🔬 方法详解

问题定义：现有FlashAttention硬件加速器设计中，指数运算和向量乘法通常由独立的硬件算子实现，导致面积和功耗开销较大。尤其是在处理长序列时，这些开销会显著影响整体性能和效率。因此，需要一种更高效的硬件架构来降低这些成本。

核心思路：论文的核心思路是通过融合指数运算和向量乘法运算，设计一种新的硬件算子ExpMul，从而减少硬件面积和功耗。这种融合的设计能够减少中间数据的存储和传输，并优化计算流程。

技术框架：该技术框架主要包括：1）分析FlashAttention算法中指数运算和向量乘法运算的计算模式；2）设计ExpMul硬件算子，将指数运算和向量乘法运算融合在一个硬件单元中；3）将ExpMul算子集成到FlashAttention硬件加速器中；4）在28nm ASIC技术中实现并评估性能。

关键创新：最关键的创新点在于ExpMul硬件算子的设计，它将原本独立的指数运算和向量乘法运算融合在一起，减少了中间数据的存储和传输，从而降低了面积和功耗。与现有方法相比，这种融合设计能够更有效地利用硬件资源。

关键设计：ExpMul算子的具体设计细节未知，但可以推测其关键在于如何高效地共享计算资源，例如乘法器和加法器，以及如何优化数据流，以减少中间数据的存储和传输。此外，可能还涉及到一些定制化的电路设计，以提高指数运算的效率。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，在28nm ASIC技术中，与具有独立指数运算和向量乘法运算硬件算子的最先进硬件架构相比，所提出的ExpMul硬件算子平均面积改善了28.8%，功耗改善了17.6%。这表明该方法在降低硬件成本和功耗方面具有显著优势。

🎯 应用场景

该研究成果可应用于各种需要加速注意力机制的场景，例如大型语言模型推理、机器翻译、语音识别等。通过降低硬件成本和功耗，可以更经济高效地部署这些模型，并推动边缘计算设备上AI应用的发展。此外，该技术还可以推广到其他涉及类似计算模式的领域。

📄 摘要（原文）

Attention mechanisms, particularly within Transformer architectures and large language models (LLMs), have revolutionized sequence modeling in machine learning and artificial intelligence applications. To compute attention for increasingly long sequences, specialized accelerators have been proposed to execute key attention steps directly in hardware. Among the various recently proposed architectures, those based on variants of the FlashAttention algorithm, originally designed for GPUs, stand out due to their optimized computation, tiling capabilities, and reduced memory traffic. In this work, we focus on optimizing the kernel of floating-point-based FlashAttention using new hardware operators that fuse the computation of exponentials and vector multiplications, e.g., e^x, V. The proposed ExpMul hardware operators significantly reduce the area and power costs of FlashAttention-based hardware accelerators. When implemented in a 28nm ASIC technology, they achieve improvements of 28.8% in area and 17.6% in power, on average, compared to state-of-the-art hardware architectures with separate exponentials and vector multiplications hardware operators.