Vectorized FlashAttention with Low-cost Exponential Computation in RISC-V Vector Processors

2025-10-08

AI Summary

This work addresses the low computational efficiency of attention mechanisms on RISC-V vector processors. We present the first end-to-end fully vectorized implementation of FlashAttention—entirely within the standard RISC-V Vector Extension (RVV) without ISA extensions. Methodologically, we replace costly exact exponential computations in Softmax with a low-overhead floating-point exponent approximation, and integrate fine-grained tiling with memory-access optimizations to significantly improve vector unit utilization and data locality. Our key contributions are: (1) the first FlashAttention vectorization scheme relying solely on standard RVV instructions; (2) a 67% reduction in exponent-operation latency within Softmax, achieved without custom instructions; and (3) 2.1×–3.4× throughput improvement for attention layers on a real RISC-V vector core, demonstrating strong scalability across vector lengths and batch sizes.

Technology Category

• Machine Learning: Hardware-aware ML
• Computer Vision: Large Vision Models
• Search and Optimization: Learning to Search

Application Category

• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines
• Security and Privacy: Large-scale security measurements
• User Modeling, Personalization and Recommendation: Attacks and countermeasures in recommendation systems

Abstract

Attention is a core operation in numerous machine learning and artificial intelligence models. This work focuses on the acceleration of attention kernel using FlashAttention algorithm, in vector processors, particularly those based on the RISC-V instruction set architecture (ISA). This work represents the first effort to vectorize FlashAttention, minimizing scalar code and simplifying the computational complexity of evaluating exponentials needed by softmax used in attention. By utilizing a low-cost approximation for exponentials in floating-point arithmetic, we reduce the cost of computing the exponential function without the need to extend baseline vector ISA with new custom instructions. Also, appropriate tiling strategies are explored with the goal to improve memory locality. Experimental results highlight the scalability of our approach, demonstrating significant performance gains with the vectorized implementations when processing attention layers in practical applications.

Problem

• Vectorizing FlashAttention algorithm for RISC-V processors
• Simplifying exponential computation in softmax with low-cost approximation
• Improving memory locality through optimized tiling strategies

Innovation

• Vectorized FlashAttention algorithm on RISC-V processors
• Low-cost exponential approximation for softmax computation
• Memory locality optimization through tiling strategies