SystolicAttention: Fusing FlashAttention within a Single Systolic Array

TL;DR

This paper introduces FSA, a novel systolic array architecture that fully accelerates FlashAttention operations within a single array, significantly improving hardware utilization and performance for transformer models.

Contribution

The paper proposes FSA, an enhanced systolic array design that executes entire FlashAttention operations internally, optimizing resource use and maintaining operation order, with implementation and performance validation.

Findings

FSA achieves 1.77x higher FLOPs/s utilization than AWS Neuron-v2.

FSA achieves 4.83x higher FLOPs/s utilization than Google TPUv5e.

FSA has only 12% area overhead in 16 nm technology.

Abstract

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.