From Buffers to Registers: Unlocking Fine-Grained FlashAttention with Hybrid-Bonded 3D NPU Co-Design

Paper: 2602.11016 Authors: Jinxin Yu, Yudong Pan, Mengdi Wang, Huawei Li, Yinhe Han, Xiaowei Li, Ying Wang Categories: cs.AR, cs.AI

Abstract

Transformer models have become the backbone of modern AI, but their quadratic attention complexity creates severe memory bottlenecks as models scale. While existing solutions like large on-chip caches (Groq, Cerebras) and algorithmic optimizations (FlashAttention) have reduced off-chip memory traffic, this paper identifies a new critical bottleneck: on-chip SRAM accesses now consume over 60% of energy in long-sequence workloads. The authors propose 3D-Flow, a hybrid-bonded 3D-stacked spatial accelerator that fundamentally rethinks data movement by enabling direct register-to-register communication across vertically partitioned processing element (PE) tiers using sub-10μm through-silicon vias (TSVs). Combined with 3D-FlashAttention, a fine-grained scheduling method that eliminates pipeline bubbles, the system achieves 46-93% energy reduction and 1.4x-7.6x speedup over state-of-the-art designs on OPT and QWEN models.

Key Contributions

• Identification of on-chip SRAM as the new bottleneck: Empirical measurements showing that SRAM accesses dominate energy consumption (>60%) in attention workloads once off-chip traffic is minimized
• 3D-Flow architecture: A hybrid-bonded 3D spatial accelerator with vertical PE tier partitioning that enables cycle-level register-to-register data transfer via dense TSV arrays
• 3D-FlashAttention scheduling: A fine-grained operator fusion and scheduling method that balances latency across tiers to create bubble-free vertical dataflow without SRAM roundtrips
• Comprehensive evaluation: Demonstrated 46-93% energy savings and 1.4x-7.6x performance improvements on real Transformer workloads compared to 2D and existing 3D designs

The Memory Hierarchy Problem in Attention

The attention mechanism’s $O(N^2)$ complexity with sequence length $N$ creates a fundamental memory challenge. Traditional approaches have focused on two strategies:

Hardware-centric solutions like Groq’s LPU and Cerebras’ WSE deploy massive on-chip SRAM (up to 40GB) to cache intermediate attention matrices, reducing DRAM bandwidth pressure. However, this approach hits diminishing returns as SRAM access energy becomes dominant.

Algorithm-centric solutions like FlashAttention restructure computation to avoid materializing the full $N \times N$ attention matrix by fusing operations and using tiling. While this reduces memory footprint, implementations still rely heavily on SRAM for intermediate results.

The key insight of this work is that both approaches converge on the same bottleneck: once off-chip traffic is minimized, the energy cost shifts to on-chip SRAM accesses. For a 4096-token sequence, the authors measure that SRAM reads/writes account for 62% of total energy, while computation is only 23%. This motivates moving data through registers instead of buffers.

3D-Flow Architecture Design

3D-Flow employs hybrid bonding technology to create ultra-dense vertical interconnects between stacked PE tiers. Key architectural features include:

Vertical PE partitioning: The attention computation is decomposed across multiple physical tiers (typically 4-8 layers), with each tier handling a portion of the query, key, or value matrices. Unlike traditional 3D stacking that simply stacks memory on compute, 3D-Flow distributes the compute itself vertically.

Register-to-register communication: TSVs with <10μm pitch enable direct data transfer between PE registers across tiers without routing through SRAM. This achieves single-cycle latency for vertical data movement, comparable to register file access within a tier.

Minimal NoC overhead: By handling most data movement vertically through TSVs, the design reduces pressure on the 2D network-on-chip, which becomes a bottleneck in multi-array 2D architectures. Horizontal NoC traffic is reserved primarily for control and boundary data.

The architecture supports three data movement patterns:

• Vertical streaming: Sequential data flows up/down tiers (e.g., Q matrix to attention computation)
• Vertical broadcast: Same data distributed to multiple tiers (e.g., K matrix reuse)
• Vertical reduction: Partial results accumulated across tiers (e.g., softmax normalization)

3D-FlashAttention: Fine-Grained Scheduling

Building on the 3D-Flow hardware, the authors develop 3D-FlashAttention, which extends FlashAttention’s operator fusion principles to the vertical dimension:

Tier-aware tiling: The algorithm partitions Q, K, V matrices not just spatially (as in FlashAttention) but also across tiers. Each tier processes a subset of attention heads or sequence chunks, with careful load balancing to prevent pipeline stalls.

Bubble-free pipelining: By analyzing the latency of each operation (matrix multiply, softmax, etc.) and the vertical communication cost, the scheduler arranges operations so that each tier has work ready when the previous operation completes. This eliminates idle cycles (“bubbles”) in the pipeline.

SRAM bypass: Critical intermediate values flow directly between tiers via registers. For example, the attention scores $QK^T$ computed on one tier feed directly into the softmax on the next tier without writing to SRAM. Only final outputs and small control data touch SRAM.

The scheduling algorithm solves a constrained optimization problem: $\min \sum_{t=1}^{T} E_{\text{tier}_t} \quad \text{s.t.} \quad \max_{t} L_{\text{tier}_t} \leq L_{\text{target}}$

where $E_{\text{tier}_t}$ is energy per tier, $L_{\text{tier}_t}$ is latency per tier, and the constraint ensures balanced execution.

Experimental Results and Analysis

The evaluation compares 3D-Flow against multiple baselines:

• 2D spatial accelerators: Multi-array designs with NoC interconnect
• 3D memory-on-logic: Traditional 3D stacking with HBM on compute die
• 3D compute-on-compute: Prior work on 3D stacked processing arrays

Energy efficiency: On OPT-6.7B with 4096-token sequences, 3D-Flow achieves:

• 93% energy reduction vs. 2D baseline
• 78% reduction vs. 3D memory-on-logic
• 61% reduction vs. 3D compute-on-compute

The energy breakdown shows that register-to-register transfers consume only 8% of total energy, compared to 62% for SRAM in the baseline.

Performance: Speedups range from 1.4x (short sequences, small models) to 7.6x (long sequences, large models). The benefit scales with sequence length because longer sequences amplify the SRAM access bottleneck in baseline designs.

Scalability: The architecture scales efficiently to 8 tiers before vertical communication latency begins to impact performance. Beyond 8 tiers, the benefit diminishes due to increased TSV routing complexity.

Implications for Future AI Accelerators

This work demonstrates that memory hierarchy optimization must extend beyond DRAM-SRAM to SRAM-register. As AI models continue scaling, the energy cost of data movement at every level becomes critical.

3D integration as enabler: Hybrid bonding technology, with its sub-10μm pitch, makes register-level communication across dies practical. This is qualitatively different from earlier 3D approaches using micro-bumps (40μm+ pitch) that could only efficiently move data between large memory blocks.

Co-design necessity: The gains come from tight coupling between architecture (vertical PE partitioning), technology (hybrid bonding), and algorithms (3D-FlashAttention scheduling). Neither hardware nor software innovation alone would achieve these results.

Broader applicability: While demonstrated on Transformer attention, the principles apply to any compute pattern with producer-consumer relationships that can be vertically partitioned—potentially including convolutions, graph neural networks, and sparse operations.

Takeaways

1. On-chip SRAM access has become the dominant energy cost in attention computation, accounting for over 60% of energy once off-chip traffic is minimized through caching or algorithmic optimizations.

2. Hybrid-bonded 3D stacking enables a new data movement paradigm: register-to-register communication across vertical PE tiers with single-cycle latency, bypassing SRAM entirely for intermediate results.

3. 3D-FlashAttention extends operator fusion to the vertical dimension, creating bubble-free pipelines through tier-aware scheduling that balances computation and communication latency.

4. The 3D-Flow architecture achieves 46-93% energy reduction and 1.4x-7.6x speedup on Transformer workloads compared to state-of-the-art 2D and 3D accelerators, with benefits scaling with sequence length.

5. Future AI accelerator design must consider the full memory hierarchy from DRAM to registers, with 3D integration providing a path to eliminate intermediate storage bottlenecks through direct die-to-die communication.