[论文解读] From Buffers to Registers: Unlocking Fine-Grained FlashAttention with Hybrid-Bonded 3D NPU Co-Design

Jinxin Yu, Yudong Pan+5 more|arXiv (Cornell University)|Feb 11, 2026

Interconnection Networks and Systems被引用 0

一句话总结

该论文提出了 3D-Flow，一种混合绑定的三维堆栈 NPU 协同设计，使 FlashAttention 实现寄存器到寄存器的数据流，形成无气泡的垂直流水线，相较于 2D/3D 基线在能耗和速度上有显著提升。

ABSTRACT

Transformer-based models dominate modern AI workloads but exacerbate memory bottlenecks due to their quadratic attention complexity and ever-growing model sizes. Existing accelerators, such as Groq and Cerebras, mitigate off-chip traffic with large on-chip caches, while algorithmic innovations such as FlashAttention fuse operators to avoid materializing large attention matrices. However, as off-chip traffic decreases, our measurements show that on-chip SRAM accesses account for over 60% of energy in long-sequence workloads, making cache access the new bottleneck. We propose 3D-Flow, a hybrid-bonded, 3D-stacked spatial accelerator that enables register-to-register communication across vertically partitioned PE tiers. Unlike 2D multi-array architectures limited by NoC-based router-to-router transfers, 3D-Flow leverages sub-10 um vertical TSVs to sustain cycle-level operator pipelining with minimal overhead. On top of this architecture, we design 3D-FlashAttention, a fine-grained scheduling method that balances latency across tiers, forming a bubble-free vertical dataflow without on-chip SRAM roundtrips. Evaluations on Transformer workloads (OPT and QWEN models) show that our 3D spatial accelerator reduces 46-93% energy consumption and achieves 1.4x-7.6x speedups compared to state-of-the-art 2D and 3D designs.

研究动机与目标

• 在长序列 Transformer 工作负载中，证明片上 SRAM 作为离芯片通信的能量瓶颈将成为新的瓶颈。
• 提出 3D-Flow：一个混合绑定、3D 堆栈的逐鹿阵列，支持跨层寄存器到寄存器通信。
• 开发 3D-FlashAttention：一种细粒度调度策略，平衡垂直层之间的延迟，形成无气泡的数据流。
• 展示垂直堆叠的处理单元（PE）和在 PE 内的 softmax/流水线，能够降低片上通信并提升 LLM 推理的能效。

提出的方法

• 提出 3D-Flow 架构：四层垂直堆叠的 PE，通过混合绑定的 TSV（<10 μm）实现寄存器到寄存器的数据流。
• 为每层设计面向 FlashAttention 子算子的 PE 单元（QK^T、max/subtract、exp/RowSum、PV/scaling）。
• 开发 3D-FlashAttention 调度，将连续的 FlashAttention 运算映射到跨层的延迟平衡数据流。
• 实现无气泡的垂直管线，使中间数据直接通过 TSV 传递，而非 SRAM 循环传输。
• 在循环精确的仿真与 RTL 验证的 4 层 3D-Stack 上进行评估，基于 16nm 制程假设。
• 与 OPT 与 Qwen 模型在长序列长度下的 2D-Unfused、2D-Fused（FuseMax、FLAT、TileFlow）、Dual-SA 和 3D-Base 基线进行对比。

Figure 1 : Energy breakdown of operator fusion and unfusion with different sequence lengths for OPT.

实验结果

研究问题

• RQ1 3D 集成搭配混合绑定是否能够在没有 SRAM 交换的情况下实现 FlashAttention 的逐级算子流水线？
• RQ2 通过跨垂直堆叠的 PE 将 FlashAttention 子算子映射为寄存器到寄存器通信，可以获得多少能耗和吞吐量的提升？
• RQ3 在为 Transformer 注意力负载部署 4 层 3D-Flow 堆栈时，能耗、面积和热性能的权衡如何？

主要发现

• 相较于 2D 基线，能耗降低 46% 到 93%，相较于基线在不同序列长度下平均降低 32.7% 到 64.2%（序列长度从 1K 到 64K）。
• 推理速度在平均层面上较 2D-Unfused 提升 7.62×，较 2D-Fused 提升 1.46×，较 Dual-SA 提升 2.36×，较 3D-Base 提升 1.43×。
• 在测试的序列长度下，PE 利用率平均为 87%，由最小化的内存流量和良好平衡的垂直流水线驱动。
• 中间结果通过垂直堆叠的 PE 之间通过 TSV 直接流动，消除了 SRAM 循环传输，实现无气泡执行。
• 热分析表明，在合理封装条件下四层堆栈可安全工作，内部温升在 128×128 PE 阵列中约为 2.8°C。

Figure 2 : Overview of 3D-stacked PE array architecture and the operator mapping of each layer.