FlatAttention: Dataflow and Fabric Collectives Co-Optimization for Large Attention-Based Model Inference on Tile-Based Accelerators

Abstract: 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.

• The paper introduces a novel dataflow that aggregates tile groups for large attention models, reducing HBM I/O by up to 6.6×.
• It leverages hardware collectives to achieve over 92% compute utilization and up to a 4.1× speedup compared to GPU-based methods.
• The approach unifies multiple attention variants with scalable tiling, guiding future architecture-algorithm co-design for wafer-scale systems.

FlatAttention: Dataflow and Fabric Collectives Co-Optimization for Large Attention-Based Model Inference on Tile-Based Accelerators

Introduction and Motivation

FlatAttention introduces a fundamentally re-architected dataflow mechanism for efficient inference of large attention-based models on tile-based accelerator platforms. As mixture-of-experts (MoE) LLMs move towards wafer-scale deployments and employ advanced attention mechanisms such as MHA, GQA, and MLA, attention kernels increasingly bottleneck system throughput due to their high bandwidth requirements, particularly for off-chip high-bandwidth memory (HBM) accesses. Traditional GPU-oriented dataflows such as FlashAttention-3 and FlashMLA deliver suboptimal utilization—only 26–64% of peak on NVIDIA GH200—even with advances like asynchronous execution (Figure 1).

Figure 1: (a) FLOP breakdown for LLMs shows attention dominating compute in recent architectures; (b) Roofline plot revealing FlashAttention-3 and FlashMLA on GH200 operate far from hardware limits.

Simultaneously, emerging accelerator architectures, particularly those based on dense tile meshes interconnected via high-radix on-chip networks (NoCs), enable direct, hardware-accelerated collective communication primitives (e.g., multicast, reduction), which are underutilized by standard GPU-style dataflows. Exploiting such collectives, FlatAttention redefines how attention blocks are mapped and parallelized, targeting extreme data reuse and minimization of global HBM traffic while maximizing matrix engine utilization.

Tile-Based Architectures and Modern Attention Variants

Tile-based many-PE architectures, as exemplified in recent wafer-scale systems, comprise a 2D mesh of compute tiles each integrating matrix/vector/scalar engines and local scratchpad memories. Tiles communicate through a NoC fabric with hardware support for collective operations. Multiple dies can be interconnected ranging up to wafer-scale, enabling scaling in both memory and compute (Figure 2).

Figure 2: (a) Tile-based many-PE template; (b) Row-wise multicast with hardware collectives; (c) Wafer-scale multi-die mesh system.

Modern LLMs demand efficient support for numerous attention variants:

• MHA: Classic, with per-head key/value projections.
• GQA/MQA: Key/value projections are shared across groups or all heads, reducing cache pressure.
• MLA: Employs compression/decompression of key/value tensors, further optimizing KV cache usage.

These variants share a core: large-scale matrix- or vector-matrix operations with memory footprints that rapidly saturate on-chip resources if not carefully mapped.

FlatAttention Dataflow: Structure and Implementation

FlatAttention’s central innovation is forming logical groups of tiles (“tile groups”) to collectively process larger attention blocks that do not fit within individual tile memories. Each group aggregates its collective L1 to host larger data slices, dramatically increasing on-chip data reuse and reducing HBM I/O complexity.

This comes at the cost of required intra-group collective communication, efficiently supported in hardware (Figure 3).

Figure 3: (a) Parametric group definition for FlatAttention; (b) Detailed dataflow with block assignments; (c) Naive schedule; (d) Optimized asynchronous schedule.

Key principles include:

• Block aggregation: Tile groups process $\left(N \cdot M, N \cdot M\right)$ blocks (for group size $N \times N$, block size $M$), reducing off-chip I/O compared to per-tile FlashAttention dataflow.
• Row/column-wise collectives: Diagonal tiles load data from HBM, which is then distributed via efficient row/column multicasts and reductions.
• Asynchronous pipeline: Softmax and data movement are overlapped with compute-intensive matrix multiplications through dual-head scheduling, closely matching the utilization envelope of matrix engines.

Compared to naive GPU FlashAttention mapping, FlatAttention’s I/O complexity decreases by up to $6.6\times$ for typical LLM configurations, with the bottleneck shifting from HBM to on-chip network efficiency.

Optimization and Tiling Strategies

Performance is a function of group scale and tiling configuration, leading to a key tradeoff: larger groups improve reuse but can “over-flatten” slices, underutilizing individual matrix engines especially for short sequences.

Experimental results demonstrate:

• Utilization: Optimal per-tile slice size is $128 \times 128$, maximizing RedMulE utilization ($>$95%) without exceeding the L1 budget (Figure 4).
• Collective primitive acceleration: Hardware collectives deliver $5-67\times$ speedup over software collectives for intra-group communication (Figure 2b).
• Scalability across attention variants: FlatAttention unifies MHA, GQA, and MLA handling through tunable tiling/grouping strategies.

  Figure 4: General FlatAttention tiling/group-scaling strategy addressing computation/memory tradeoff for MHA, GQA, MLA.

Quantitative Results

• On a $32\times32$ tile architecture (GH200-class peak), FlatAttention achieves up to 92.3% compute utilization and 4.1$\times$ speedup over FlashAttention-3, with up to 16$\times$ lower HBM traffic.
• Generalized across attention variants, average utilization is 86% (compute-bound) and 78% HBM bandwidth (memory-bound).
• Against NVIDIA GH200’s best attention kernels, FlatAttention delivers an average 1.9$N \times N$0 speedup across prefill and decode.
• In end-to-end DeepSeek-v3-671B decoder inference on a wafer-scale multi-die system, FlatAttention improves overall system throughput by 1.9$N \times N$1 and per-user latency by 1.4$N \times N$2, despite 1.5$N \times N$3 lower system peak performance compared to 96×H800 GPU deployments (Figure 5).

  Figure 5: FlatAttention on tile-based accelerator vs. GH200 GPU attention baselines across variants and regimes.

Comparative Analysis and Implications

Relative to prior art, FlatAttention is the first to unify efficient fused-attention dataflow and fabric-level collective primitive utilization. Unlike recent efforts confined to single-tile/single-GPC fusion (e.g., FlashFuser, ClusterFusion), FlatAttention scales cleanly to mesh- or wafer-scale systems. While works such as COMET and Zen-Attention exploit collectives, they do not approach the full co-optimization of tiling, scheduling, and hardware-supported communication exposed in FlatAttention.

Critically, FlatAttention enables architectural codesign: the same performance analysis guides both attention dataflow and tile-based accelerator design (matrix engine width, memory organization, collective support), informing future wafer-scale LLM inference engines.

Practical and Theoretical Implications

Practically, FlatAttention enables:

• Sustained high utilization and throughput for attention-bound inference even at hundreds of billions of parameters, eliminating the attention bottleneck present in MoE architectures.
• Significant reduction in HBM energy and cost per token, crucial at data center scale.
• Efficient support for future speculative decoding and advanced attention mechanisms.

Theoretically, it demonstrates:

• The necessity of architecture-algorithm co-design for domain-specific AI workload acceleration.
• The return of collective communication to the core of on-chip ML dataflow, an inversion of recent GPU-exclusive designs.
• Scalability to system limits with the potential to generalize to broader classes of tensor algebra where on-chip data reuse dominates.

Future Directions

Future system design will benefit from refined codesign of NoC topologies for collective bandwidth, adaptive group scaling at runtime, and integration with more diverse inferencing workloads. Potential architectural extensions include tighter coupling with chiplet-based HBM and custom collective-aware instruction sets.

Conclusion

FlatAttention, via aggressive dataflow and hardware co-optimization, reclaims attention inference efficiency at wafer scale. By exploiting hierarchical tile-grouping and hardware collectives, it sets a new benchmark in utilization and throughput per Watt, and closes the performance gap between practical LLM deployment and hardware-limited peak.

References to Figures

• (Figure 1): Motivation: FLOP breakdown and Roofline gap
• (Figure 2): Architecture overview, fabric collective demonstration, wafer-scale system
• (Figure 3): FlatAttention group definition/dataflow/schedules
• (Figure 4): FlatAttention general tiling strategy
• (Figure 5): GH200 vs. FlatAttention benchmarking across attention variants