Reproduce “FlatAttention: Dataflow and Fabric Collectives Co-Optimization for Large Attention-Based Model Inference on Tile-Based Accelerators” on OpenResearch
Turn your favorite coding agents into research agents that review literature, develop hypotheses, and run experiments.
Universal app for Apple silicon and Intel. Requires macOS 11 or later.
About this paper
Chi Zhang, Luca Colagrande, Renzo Andri, Luca Benini
FlatAttention is a dataflow designed for efficient large attention-based model inference on tile-based accelerators, co-optimized with hardware-supported collective communication. This approach reduces HBM traffic by up to 16x and achieves up to a 4.1x speedup over FlashAttention-3 for MHA, leading to a 2.1x end-to-end decoding speedup for DeepSeek-v3-671B on wafer-scale multi-die systems.
Key insights
- Hardware-supported collective communication primitives on a 32x32 tile mesh offer substantial speedups over software-based implementations (e.g., 30.7x for row-wise multicast, 67.3x for row-wise sum reduction), crucial for efficient inter-tile data movement.
- FlatAttention's reduced I/O complexity and effective use of on-chip collectives enable significant reductions in HBM access time and higher compute utilization compared to memory-bound FlashAttention variants on tile-based accelerators.
- A general tiling strategy that optimizes per-tile efficiency first, then scales the group size based on the attention score matrix shape, is essential for maintaining high utilization across various sequence lengths and preventing performance degradation from 'over-flattening'.
Summary by alphaXiv.
Tip: change arxiv.org to autoarxiv.org in any arXiv abstract link to open the paper here.