FlatAttention: Accelerated Attention

• FlatAttention is an optimized co-design approach that minimizes off-chip HBM traffic by leveraging hardware-accelerated collective primitives and two-level tile blocking in multi-head attention.
• It employs group flattening and collective scheduling to maximize on-chip data reuse, reducing memory bandwidth demands by up to 16× compared to traditional methods.
• The technique generalizes to various attention variants, delivering significant throughput and latency improvements in transformer inference on large-scale, tile-based accelerator systems.

FlatAttention denotes a class of optimized dataflow and algorithm–hardware co-design techniques for accelerating attention mechanisms—specifically multi-head attention (MHA) and its modern variants—on scalable tile-based accelerator architectures. Its core contribution is the systematic minimization of off-chip high-bandwidth memory (HBM) traffic via the exploitation of hardware-accelerated collective primitives (e.g., row/col multicast and all-reduce) in network-on-chip (NoC) fabrics, combined with a two-level tile-group flattening/blocking of the attention computation. This approach achieves near-peak utilization of on-chip compute resources, dramatic reductions in HBM bandwidth demands (up to 16× vs. established baselines), and generalizes seamlessly to autoregressive, multi-query, and latent attention variants. End-to-end evaluation on large wafer-scale systems demonstrates substantial throughput and latency improvements over both FlashAttention-derived GPU kernels and conventional transformer execution engines (Zhang et al., 2 Apr 2026, Zhang et al., 24 May 2025).

1. Tile-Based Accelerator Model and Attention Formulation

FlatAttention is designed for domain-specific accelerators organized as large 2D meshes of tiles (e.g., 32×32 arrays). Each tile provides the following architectural features:

• High-throughput matrix engine (for GEMM), vector engine (for elementwise and Softmax ops), at least one scalar RISC-V control core
• L1 scratchpad and programmable DMA for local–HBM transfers
• 2D-mesh NoC providing O(1)-latency hardware collectives: row-wise multicast, column-wise multicast, and all-reduce

Off-chip data movement is serviced by multiple HBM stacks located on mesh boundaries (e.g., 2–4 TB/s per die). The standard MHA computation per head is given by

$$\operatorname{Attn}_i = \operatorname{Softmax}\bigg( \frac{Q_i K_i^\top}{\sqrt{d}} \bigg) V_i$$

$$\operatorname{MHA}(Q,K,V) = [\operatorname{Attn}_1; \ldots; \operatorname{Attn}_H] W^O$$

where $Q, K, V \in \mathbb{R}^{S \times d}$ and $H$ is the head count. The quadratic cost in sequence length $S$ and exploding HBM traffic necessitate specialized dataflows for tractable large-scale inference (Zhang et al., 2 Apr 2026, Zhang et al., 24 May 2025).

2. FlatAttention Dataflow: Group Flattening and Collective Scheduling

FlatAttention introduces a novel group-based dataflow that partitions the attention computation across $G_x \times G_y$ tile-groups, with each group jointly handling a $(B_r \times B_c)$ SuperBlock of the $QK^\top$ score matrix.

• Tiling and Mapping: The model sequence is blocked so that each group can hold an entire $(B_r \times B_c)$ block of intermediates in their collective L1s, with each tile owning a slice of $(B_r/G_y) \times (B_c/G_x)$.
• Collectives Protocol: Processing of a SuperBlock proceeds as follows:
  1. Diagonal tiles load Q-slices and row-multicast them.
  2. For each K/V-slice, diagonal tiles load from HBM and column-multicast. Each tile computes its local QK segment.
  3. Softmax normalization is performed using row-wise all-reduce (max for row normalization, sum for denominators), all using hardware collectives.
  4. Each tile accumulates its O, ending with row-wise all-reduce for final output storage.

• Asynchronous Double Buffering: Matrix-engine operations for block $$\operatorname{MHA}(Q,K,V) = [\operatorname{Attn}_1; \ldots; \operatorname{Attn}_H] W^O$$0 are overlapped with vector-engine/DMA/Softmax operations on block $$\operatorname{MHA}(Q,K,V) = [\operatorname{Attn}_1; \ldots; \operatorname{Attn}_H] W^O$$1, maximizing compute concurrency.

The key advantage lies in maximizing on-chip data reuse and minimizing global memory traffic; all non-input/output communications are handled by collective operations in the NoC, eliminating redundant HBM fetches and reducing effective I/O per attention block by orders of magnitude (Zhang et al., 2 Apr 2026, Zhang et al., 24 May 2025).

3. Hardware Collectives and Traffic Reduction

Central to FlatAttention is its reliance on hardware-accelerated collectives implemented directly within NoC routers. The essential primitives include:

• Row-wise multicast: Distributes data along a row in O(1) router steps.
• Column-wise multicast: Analogous operation along the column.
• Row-wise all-reduce: Partial aggregates propagate and are reduced tree-wise, then results are broadcast.

Table: Key NoC Collective Operations in FlatAttention

Collective	Data Direction	Complexity
Row-wise multicast	Row	O(1) flit-level
Column-wise multicast	Column	O(1) flit-level
Row-wise all-reduce	Row	O(1) tree-based

Each primitive obviates N-stage software messaging, reducing overall inter-tile traffic by up to 30× and collective latency by 5–10× relative to baseline P2P (Zhang et al., 2 Apr 2026). As a result, HBM usage is reserved primarily for bulk initial/final Q, K, V, and O transfers. The effective off-chip I/O for FlatAttention with $$\operatorname{MHA}(Q,K,V) = [\operatorname{Attn}_1; \ldots; \operatorname{Attn}_H] W^O$$2 tiles, per layer, scales as

$$\operatorname{MHA}(Q,K,V) = [\operatorname{Attn}_1; \ldots; \operatorname{Attn}_H] W^O$$3

in contrast to the baseline

$$\operatorname{MHA}(Q,K,V) = [\operatorname{Attn}_1; \ldots; \operatorname{Attn}_H] W^O$$4

which yields an HBM traffic reduction factor of $$\operatorname{MHA}(Q,K,V) = [\operatorname{Attn}_1; \ldots; \operatorname{Attn}_H] W^O$$5 when $$\operatorname{MHA}(Q,K,V) = [\operatorname{Attn}_1; \ldots; \operatorname{Attn}_H] W^O$$6 (Zhang et al., 24 May 2025, Zhang et al., 2 Apr 2026).

4. Performance Metrics and Empirical Results

Comprehensive benchmarking across simulated and physical tile-meshes demonstrates substantial gains:

• Matrix Engine Utilization: Up to 92.3% on a 32×32 tile array (GH200-scale) vs. ~20% for FlashAttention-3; 4.6× increase.
• HBM Bandwidth Utilization: For memory-bound tasks, FlatAttention achieves ~78% sustained bandwidth (comparable to GH200).
• Speedup: 4.1× over FlashAttention-3 on matched tile-based hardware, and overall 1.9× end-to-end speedup over FlashAttention/FlashMLA on NVIDIA GH200.
• End-to-End DeepSeek-v3 FP8 Decoding: On an 8×8 wafer mesh (64 chips), FlatAttention achieved 1.9× system throughput improvement and 1.4× reduction in per-token user latency compared to the baseline, matching or exceeding solutions with 1.5× more aggregate peak FLOP/s (Zhang et al., 2 Apr 2026).

Table: Representative Performance Metrics

Metric	FlashAttn-3 (32×32)	FlatAttention w/ HW	FlatAttn Speedup	GH200 Avg.	Flat vs. GH200
Matrix-util (prefill)	~20%	92.3%	4.6×	26–64%	1.9×
HBM BW util (decode)	~80%	~78%	–	~80%	Comparable
Throughput (DS671B)	–	–	1.9×	–	1.9×
Latency (T_POT)	–	–	1.4× lower	~50 ms	1.4× lower

(Zhang et al., 2 Apr 2026, Zhang et al., 24 May 2025)

5. Generalization Across Attention Variants

FlatAttention generalizes efficiently to a broad range of modern attention mechanisms used in LLMs and inference engines:

• Autoregressive Decode: Collapses group-row $$\operatorname{MHA}(Q,K,V) = [\operatorname{Attn}_1; \ldots; \operatorname{Attn}_H] W^O$$7 for sequence length $$\operatorname{MHA}(Q,K,V) = [\operatorname{Attn}_1; \ldots; \operatorname{Attn}_H] W^O$$8; KV block reuse is maximized.
• Multi-Query/Grouped-Query Attention (MQA/GQA): Tiles the query block against a shared KV as a long (1×G) block, maintaining the dataflow.
• Multi-Head Latent Attention (MLA): Absorbed projection weights reduce the computation to a compressed MQA form, handled identically in FlatAttention.
• Speculative Decoding: Flexible $$\operatorname{MHA}(Q,K,V) = [\operatorname{Attn}_1; \ldots; \operatorname{Attn}_H] W^O$$9 group shapes can be used to map blocks of speculative decode versus KV length.

All these variants benefit from the hardware collectives and group flattening, requiring no microarchitectural changes (Zhang et al., 2 Apr 2026).

6. Algorithm–Architecture Co-Exploration and Design Implications

Extensive exploration of hardware–software trade-offs in simulators such as SoftHier has yielded several key design insights:

• Two-level Blocking: The $Q, K, V \in \mathbb{R}^{S \times d}$0 parameter exposes a key balance between reducing global I/O and maintaining per-tile compute workload; excessive flattening (large group sizes with small per-tile slices) can underutilize cores on small $Q, K, V \in \mathbb{R}^{S \times d}$1.
• HBM Channel Placement: A moderate fan-out of HBM controllers (e.g., $Q, K, V \in \mathbb{R}^{S \times d}$2 on mesh edges vs. $Q, K, V \in \mathbb{R}^{S \times d}$3) optimally utilizes aggregate bandwidth and area.
• NoC Architecture: Hardware-embedded collectives in the router datapath are essential; host-software-based collectives incur prohibitive overheads at large scale.
• Die Size and Efficiency: On 32×32 tile mesh, FlatAttention reduced aggregate HBM BW needs by 40% and total die area by 1.8× relative to NVIDIA H100 at comparable peak TFLOPS (Zhang et al., 24 May 2025).
• Extensibility: The approach naturally extends to next-generation LLM accelerators and multi-die/wafer systems.

A plausible implication is that future AI hardware will further emphasize programmable, collective-rich NoC fabrics and distributed scratchpad control, with FlatAttention-like dataflows setting the foundation for large-scale transformer inference (Zhang et al., 2 Apr 2026, Zhang et al., 24 May 2025).