通过3D堆叠架构实现寄存器级通信,显著降低大模型推理能耗与延迟。
From Buffers to Registers: Unlocking Fine-Grained FlashAttention with Hybrid-Bonded 3D NPU Co-Design
- 采用3D堆叠设计,用微米级垂直互连实现跨层寄存器直通通信。
- 在长序列任务中,能降低46%-93%能耗,速度提升1.4倍至7.6倍。
- 适合需要高能效、低延迟的Transformer大模型部署场景。
基于Transformer的模型主导现代AI工作负载,但其二次方注意力复杂度和不断增长的模型规模加剧了内存瓶颈。现有加速器如Groq和Cerebras通过大容量片上缓存缓解片外传输,而FlashAttention等算法创新通过融合算子避免生成大型注意力矩阵。然而,随着片外流量减少,我们测量发现片上SRAM访问在长序列任务中占总能耗超60%,成为新的瓶颈。本文提出3D-Flow,一种混合键合的3D堆叠空间加速器,支持跨垂直分区处理单元(PE)层级的寄存器到寄存器通信。相比受限于NoC路由传输的2D多阵列架构,3D-Flow利用亚10微米垂直TSV,实现周期级算子流水线,开销极小。在此架构基础上,设计3D-FlashAttention,一种细粒度调度方法,均衡各层级延迟,形成无气泡的垂直数据流,无需片上SRAM往返。在OPT和QWEN模型的Transformer工作负载评估中,该3D空间加速器相比先进2D与3D设计,能耗降低46%-93%,速度提升1.4倍至7.6倍。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。