Attention accounts for an increasingly dominant fraction of total computation during inference for mixture-of-experts (MoE) models, making efficient acceleration critical. Emerging domain-specific accelerators for large model inference are shifting toward chip-scale and wafer-scale tile-based architectures. Tiles contain large matrix and vector engines and are connected through on-chip interconnects, which support tile-to-tile traffic to reduce the tile-to-main-memory traffic bottleneck. Hence, dataflow management is crucial to achieve high utilization. We propose FlatAttention, a dataflow for modern attention variants on tile-based accelerators. FlatAttention minimizes expensive high-bandwidth memory (HBM) accesses by exploiting collective primitives integrated into the on-chip network fabric, achieving up to 92.3% utilization, 4.1x speedup over FlashAttention-3, and 16x lower HBM traffic. On a 32x32 tile configuration with peak performance comparable to NVIDIA GH200, FlatAttention generalizes across multiple attention variants, achieving an average of 86% utilization for compute-bound attentions and 78% HBM bandwidth utilization for memory-bound ones, resulting in an average 1.9x speedup over attention implementations on GH200. Finally, we evaluate end-to-end DeepSeek-v3 FP8 decoding with FlatAttention on a wafer-scale multi-die system, achieving a 1.9x improvement in system throughput and a 1.4x reduction in per-user token output latency, despite operating with 1.5x lower peak system performance compared to the state-of-the-art solution.

翻译：注意力计算在混合专家（MoE）模型推理中占总计算量的比例持续上升，高效加速变得至关重要。面向大规模模型推理的新兴领域专用加速器正转向芯片级与晶圆级瓦片式架构。瓦片内集成大型矩阵与向量引擎，通过片内互连网络连接，支持瓦片间数据传输以减少主存访问瓶颈。因此，数据流管理对实现高利用率至关重要。本文提出FlatAttention——一种面向瓦片加速器现代注意力变体的数据流方案。FlatAttention通过利用集成于片内网络架构的集合通信原语，最小化高带宽内存（HBM）的昂贵访问，实现了高达92.3%的利用率、相比FlashAttention-3的4.1倍加速比以及16倍HBM流量降低。在峰值性能与NVIDIA GH200相当的32×32瓦片配置下，FlatAttention可泛化至多种注意力变体，计算密集型注意力平均利用率达86%，内存密集型注意力HBM带宽利用率达78%，相比GH200注意力实现方案平均实现1.9倍加速比。最后，我们在晶圆级多芯片系统上使用FlatAttention对DeepSeek-v3 FP8解码进行端到端评估，尽管峰值系统性能比现有最优方案低1.5倍，仍实现了1.9倍系统吞吐量提升和1.4倍单用户令牌输出延迟降低。