Low-Cost FlashAttention with Fused Exponential and Multiplication Hardware Operators

2025-05-20

AI Summary

To address the high area and power overhead of long-sequence attention computation in Transformers and large language models, this work proposes a low-power, area-efficient hardware accelerator tailored for FlashAttention. The core innovation is the ExpMul fused operator, which jointly executes exponentiation (e^x) and vector multiplication (V) in hardware—breaking the energy-efficiency bottleneck of conventional separated implementations. This is realized through a custom floating-point ExpMul unit and a tile-aware, memory-centric architecture, optimized for FlashAttention kernels in 28 nm ASIC technology. Experimental results demonstrate that, compared to state-of-the-art hardware accelerators, the proposed design reduces chip area by 28.8% and power consumption by 17.6% on average, significantly improving the energy efficiency of long-sequence attention computation.

Technology Category

• Machine Learning: Hardware-aware ML
• Search and Optimization: Learning to Search
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• Graph Algorithms and Modeling for the Web: Algorithms and analysis for heterogeneous, signed, attributed, multi-relational, temporal, higher-order, and annotated Web-related graphs
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications
• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines

Abstract

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.

Problem

• Optimize FlashAttention kernel with fused hardware operators
• Reduce area and power costs in attention accelerators
• Improve efficiency of exponential and multiplication computations

Innovation

• Fused exponential and multiplication hardware operators
• Optimized FlashAttention kernel for efficiency
• Reduced area and power costs significantly