FLASH-D：隐藏软最大除法的 FlashAttention
机器学习
2025-05-21 v1 人工智能
硬件体系结构
摘要
Transformer 的注意力机制已彻底改变 AI 与机器学习领域，其高效计算对性能至关重要。然而，注意力计算涉及交织在矩阵运算中的 softmax 重缩放，这本质上会减慢计算速度并要求处理整个输入序列。基于在线 softmax 计算，FlashAttention 将 softmax 计算与矩阵运算集成，实现与序列长度无关的块状计算。虽为 GPU 优化，FlashAttention 的简洁性使其易于直接硬件加速。本 work re-evaluates 核心 FlashAttention kernel，提出 FLASH-D 一种在数学上等价却更简化的表述，实现：(a) 将 softmax 除法隐藏在其他非线性函数评估中；(b) 内在数值稳定的指数计算，无需最大值为减法；(c) 在不引入数值近似的前提下降低计算成本。重要的是，促进高效块状实现的关键 FlashAttention 属性完全得以保留。28nm 硬件实现结果表明，在无性能惩罚的情况下，本提出方法相较于最先进的并行硬件架构平均实现面积降低 22.8%，功耗降低 20.3%。
引用
@article{arxiv.2505.14201,
title = {FLASH-D: FlashAttention with Hidden Softmax Division},
author = {Kosmas Alexandridis and Vasileios Titopoulos and Giorgos Dimitrakopoulos},
journal= {arXiv preprint arXiv:2505.14201},
year = {2025}
}
备注
IEEE/ACM International Symposium on Low Power Electronics and Design (ISLPED) 2025