SystolicAttention: 在单个流体数组中融合FlashAttention
SystolicAttention: Fusing FlashAttention within a Single Systolic Array
AI总结：
SystolicAttention通过在单个流体数组中融合FlashAttention，提升了计算效率并优化了资源利用率。
摘要：
Transformer模型严重依赖于缩放点积注意力（SDPA）运算，通常以FlashAttention实现。其特点是频繁地交错矩阵乘法和softmax运算，FlashAttention未能充分利用现代基于流体数组的加速器的计算资源，这些加速器专为连续和大规模矩阵乘法设计。为了充分释放流体数组在FlashAttention中的性能潜力，我们提出了FSA，一种增强的流体数组架构，可以在数组上运行整个FlashAttention，而无需外部向量单元。结合SystolicAttention，一种为FSA优化的内核，实现了FlashAttention操作的细粒度和逐元素重叠，FSA在最大化数组利用率的同时保持了FlashAttention原始的浮点运算顺序。我们将在可综合的RTL中实现FSA，并将其性能与最先进的基于流体数组的加速器进行评估。我们的结果表明，FSA在AWS Neuron-v2和Google TPUv5e上分别实现了1.77倍和4.83倍更高的注意力FLOPs/s利用率。我们将在16nm工艺下以1.5GHz的速度综合FSA，结果表明与标准的权重固定的流体数组相比，仅增加了12%的面积开销。
英文摘要：
Transformer models rely heavily on the scaled dot-product attention (SDPA) operation, typically implemented as FlashAttention. Characterized by its frequent interleaving of matrix multiplications and softmax operations, FlashAttention fails to fully utilize the compute resources of modern systolic-array-based accelerators designed for consecutive and large matrix multiplications. To fully unleash the performance potential of systolic arrays for FlashAttention, we propose FSA, an enhanced systolic array architecture that runs the entire FlashAttention on the array without external vector units. Combined with SystolicAttention, an optimized kernel for FSA that achieves fine-grained and element-wise overlapping of FlashAttention operations, FSA maximizes array utilization while preserving the original floating-point operation order of FlashAttention. We implement FSA in synthesizable RTL and evaluate its performance against state-of-the-art systolic-array-based accelerators. Our results show that FSA achieves 1.77x and 4.83x higher attention FLOPs/s utilization compared to AWS Neuron-v2 and Google TPUv5e, respectively. We synthesize FSA in a 16 nm technology at 1.5 GHz, and results indicate only a 12% area overhead compared to a standard weight-stationary systolic array.
机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。
