从缓冲区到寄存器：基于混合键合 3D NPU 的细粒度 FlashAttention
硬件体系结构
2026-02-12 v1 人工智能
摘要
Transformer 模型主导现代 AI 工作负载，但其二次注意力复杂度和不断增大的模型尺寸加剧了内存瓶颈。现有加速器如 Groq 和 Cerebras 通过大型片上缓存来缓解芯片外流量，算法创新如 FlashAttention 通过融合算子来避免实例化大型注意力矩阵。然而，随着芯片外流量降低，我们的测量显示，在长序列工作负载中，片上 SRAM 访问占能耗的 60%以上，使缓存访问成为新的瓶颈。我们提出 3D-Flow，一种混合键合、3D 堆叠的空间加速器，通过垂直分区的 PE 层实现寄存器到寄存器的通信。不同于受 NoC 基于路由器到路由器传输限制的 2D 多数组架构，3D-Flow 利用亚 10 微米垂直 TSV 以最小开销维持周期级算子流水线。在此架构之上，我们设计 3D-FlashAttention，一种细粒度调度方法，在层之间平衡延迟，形成无气泡的垂直数据流，无需片上 SRAM 往返。针对 Transformer 工作负载（OPT 和 QWEN 模型），我们的 3D 空间加速器相较于最先进的 2D 和 3D 设计可减少 46-93% 能耗，实现 1.4 倍至 7.6 倍的加速。
引用
@article{arxiv.2602.11016,
title = {From Buffers to Registers: Unlocking Fine-Grained FlashAttention with Hybrid-Bonded 3D NPU Co-Design},
author = {Jinxin Yu and Yudong Pan and Mengdi Wang and Huawei Li and Yinhe Han and Xiaowei Li and Ying Wang},
journal= {arXiv preprint arXiv:2602.11016},
year = {2026}
}
备注
Accepted to DATE 2026