在RISC-V向量处理器上加速注意力计算，不增指令集也能提升性能。
Vectorized FlashAttention with Low-cost Exponential Computation in RISC-V Vector Processors
- 用低开销近似法替代指数运算，减少软硬件负担。
- 向量化实现使注意力层处理速度显著提升，实测表现良好。
- 适合关注边缘AI推理、低功耗芯片优化的研究者和开发者。
注意力是众多机器学习与人工智能模型的核心操作。本文聚焦于在向量处理器（特别是基于RISC-V指令集架构的处理器）上加速FlashAttention算法。这是首次将FlashAttention向量化的工作，通过最小化标量代码并简化softmax中所需的指数计算复杂度来实现。我们采用浮点数运算中的低成本近似方法，降低指数函数的计算成本，无需扩展基础向量指令集即可完成。同时，探索了合适的分块策略以提升内存局部性。实验结果表明，该方法具有良好的可扩展性，在实际应用中处理注意力层时展现出显著性能提升。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。