FlatAttention: Dataflow and Fabric Collectives Co-Optimization for Efficient Multi-Head Attention on Tile-Based Many-PE Accelerators

Multi-Head Attention (MHA) is a critical computational kernel in transformer-based AI models. Emerging scalable tile-based accelerator architectures integrate increasing numbers of tightly-packed processing elements (PEs) with tensor units. MHA dataflow mapping is crucial for achieving high utilization of the available units. We propose FlatAttention, a new dataflow for MHA on tile-based many-PE accelerators, minimizing costly main memory (HBM) accesses by leveraging collective primitives integrated into the on-chip network fabric. FlatAttention achieves up to $89.3 \%$ utilization, and $4.1 \times$ performance speedup over FlashAttention-3 dataflow on tile-based accelerators whilst reducing HBM traffic by $16 \times$. Through algorithm-architecture co-exploration, we identify an optimal configuration for a large scaled-out tile-based accelerator featuring a $32 \times 32$ tile mesh with 1024 TFLOPS @ FP16 peak performance, comparable to the state-of-the-art Nvidia H100 GPU. FlatAttention in this configuration achieves up to $1.3 \times$ higher utilization over FlashAttention3 on the H100 GPU. Meanwhile, this tile-based accelerator configuration requires $40 \%$ less HBM bandwidth compared to the H100 GPU, enabling a $1.8 \times$ reduction in die size, estimated on the same technology node.