arXiv:2507.11331cs.ARcs.AI2025-07被引 8

将FlashAttention全量部署于流水线阵列,提升算力利用率

SystolicAttention: Fusing FlashAttention within a Single Systolic Array

  • 设计FSA架构,让注意力计算全程在流水线阵列内完成
  • 相较AWS Neuron-v2和Google TPUv5e,算力利用率分别提升1.77倍和4.83倍
  • 仅增加12%面积开销,适合高吞吐推理芯片设计

Transformer模型高度依赖缩放点积注意力(SDPA)操作,通常由FlashAttention实现。由于其频繁交错矩阵乘法与Softmax操作,FlashAttention无法充分使用面向连续大规模矩阵乘法的现代流水线阵列加速器的计算资源。为充分发挥流水线阵列在FlashAttention上的性能潜力,我们提出FSA——一种可在阵列内完整运行FlashAttention的增强型流水线阵列架构,无需外部向量单元。结合SystolicAttention优化核,FSA实现细粒度、逐元素重叠的FlashAttention运算,最大化阵列利用率并保持原始浮点运算顺序。我们在可综合RTL中实现FSA,评估其在主流流水线阵列加速器上的表现。结果表明,相较于AWS Neuron-v2和Google TPUv5e,FSA的注意力每秒浮点运算数(FLOPs/s)利用率分别提高1.77倍和4.83倍。在16 nm工艺下,工作频率1.5 GHz,FSA仅带来12%的面积开销,相比标准权值固定型流水线阵列。

原文摘要 · Abstract (English)

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.

注意力机制硬件加速流水线阵列

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。