基于 RISC-V 向量处理器的低成本指数计算的向量化 FlashAttention
机器学习
2025-10-09 v1 分布式、并行与集群计算
性能
摘要
注意力是众多机器学习和人工智能模型中的核心操作。本工作聚焦于在向量处理器（特别是基于 RISC-V 指令集架构（ISA）的处理器）上使用 FlashAttention 算法加速注意力内核。本工作代表了首次对 FlashAttention 进行向量化，最小化标量代码，简化评估 softmax 所需指数函数计算的计算复杂度。通过采用浮点算术中低成本近似值来计算指数函数，无需扩展基础向量 ISA 即可降低指数函数计算成本。同时，探索了合适的铺幕策略，以提高内存局部性。实验结果突出了该方法的可扩展性，表明在处理实际应用中的注意力层时，向量化实现在性能上具有显著优势。
引用
@article{arxiv.2510.06834,
title = {Vectorized FlashAttention with Low-cost Exponential Computation in RISC-V Vector Processors},
author = {Vasileios Titopoulos and Kosmas Alexandridis and Giorgos Dimitrakopoulos},
journal= {arXiv preprint arXiv:2510.06834},
year = {2025}
}