Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Chi Zhang, Luca Colagrande, Renzo Andri, Thomas Benz, Gamze Islamoglu, Alessandro Nadalini, Francesco Conti, Yawei Li, Luca Benini

TL;DR

FlatAttention tackles the memory bottleneck of multi-head attention on tile-based many-PE accelerators by co-designing a dataflow with on-chip collective primitives. By grouping tiles and sharing on-chip memory, it reduces off-chip HBM traffic and increases utilization, achieving up to 4.1× speedup over the FlashAttention-3 dataflow and up to 89.3% utilization on large tile meshes. An optimization framework (SoftHier) coupled with GVSoC modeling identifies a BestArch configuration (32×32 fabric, 16×2 HBM) that matches Nvidia’s H100 peak throughput with 40% less HBM bandwidth and a ~1.8× die-area reduction. The work demonstrates that collective-enabled dataflows and architecture co-design enable scalable, energy-efficient MHA acceleration, with implications for end-to-end LLM inference on multi-chiplet systems and future memory hierarchies.

Abstract

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.1x performance speedup over FlashAttention-3 dataflow on tile-based accelerators whilst reducing HBM traffic by 16x. Through algorithm-architecture co-exploration, we identify an optimal configuration for a large scaled-out tile-based accelerator featuring a 32x32 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.3x higher utilization over FlashAttention-3 on the H100 GPU. Meanwhile, this tile-based accelerator configuration requires 40% less HBM bandwidth compared to the H100, enabling a 1.8x reduction in die size, estimated on the same technology node.

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

TL;DR

FlatAttention tackles the memory bottleneck of multi-head attention on tile-based many-PE accelerators by co-designing a dataflow with on-chip collective primitives. By grouping tiles and sharing on-chip memory, it reduces off-chip HBM traffic and increases utilization, achieving up to 4.1× speedup over the FlashAttention-3 dataflow and up to 89.3% utilization on large tile meshes. An optimization framework (SoftHier) coupled with GVSoC modeling identifies a BestArch configuration (32×32 fabric, 16×2 HBM) that matches Nvidia’s H100 peak throughput with 40% less HBM bandwidth and a ~1.8× die-area reduction. The work demonstrates that collective-enabled dataflows and architecture co-design enable scalable, energy-efficient MHA acceleration, with implications for end-to-end LLM inference on multi-chiplet systems and future memory hierarchies.

Abstract

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.1x performance speedup over FlashAttention-3 dataflow on tile-based accelerators whilst reducing HBM traffic by 16x. Through algorithm-architecture co-exploration, we identify an optimal configuration for a large scaled-out tile-based accelerator featuring a 32x32 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.3x higher utilization over FlashAttention-3 on the H100 GPU. Meanwhile, this tile-based accelerator configuration requires 40% less HBM bandwidth compared to the H100, enabling a 1.8x reduction in die size, estimated on the same technology node.

Paper Structure

This paper contains 12 sections, 2 equations, 7 figures, 2 tables.

Table of Contents

  1. Introduction
  2. Reference Tile-Based Many-PE Architecture
  3. FlatAttention Dataflow
  4. Motivation
  5. Detailed FlatAttention Dataflow
  6. Asynchronous FlatAttention
  7. Modeling and Analysis Methodology
  8. Experimental Results
  9. FlashAttention vs. FlatAttention
  10. Tile Group Scale Trade-offs for FlatAttention
  11. Co-exploration of Architecture and Algorithm Parameters
  12. Conclusion

Figures (7)

  • Figure 1: Tile-Based Many-PE Architecture Template
  • Figure 2: (a) Parametric definition of FlatAttention. (b) Detailed FlatAttention dataflow, with each step corresponding to the line numbers in Algorithm 2. (c) FlatAttention dataflow optimization.
  • Figure 3: Runtime breakdown (bars) and average HBM BW utilization (star markers) for different implementations and layer sizes. $^+$Runtime not overlapped with RedMulE. $^{++}$Runtime not overlapped with either Spatz or RedMulE. *Implementations without double buffering.
  • Figure 4: Runtime breakdown for different (square) flattening scales and layer sizes. Percentage labels above the bars indicate the average utilization of the RedMulE units when active. $^+$Runtime not overlapped with RedMulE. $^{++}$Runtime not overlapped with either Spatz or RedMulE.
  • Figure 5: (a) Heatmap of utilization with best group size. (b) Comparison with FlashAttention-3 on H100 solutions, with absolute performance in TFLOPS labeled above the bar. (c) Comparison of GEMMs, including FFN layer in Meta's LLaMA 70B gemmbench, between BestArch and H100.
  • ...and 2 more figures