[논문 리뷰] From Buffers to Registers: Unlocking Fine-Grained FlashAttention with Hybrid-Bonded 3D NPU Co-Design

Jinxin Yu, Yudong Pan외 5명|arXiv (Cornell University)|2026. 02. 11.

Interconnection Networks and Systems인용 수 0

한 줄 요약

본 논문은 등록-대-등록 데이터 흐름을 가능하게 하는 하이브리드 본딩 방식의 3D-스택 NPU 공동 설계인 3D-Flow를 제시하여 FlashAttention에 대해 버블-프리(vertical) 파이프라인을 달성하고 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 개발: 수직 계층 간 지연을 균형 있게 맞춰 버블-프리 데이터플로를 형성하는 미세한 스케줄링 전략.
• 수직으로 적층된 PEs 및 PE 내 softmax/파이프라인이 온칩 트래픽 감소를 가능하게 하고 LLM 추론의 에너지 효율을 향상시킴을 보여준다.

제안 방법

• 3D-Flow 아키텍처 도입: 네 층으로 구성된 수직 적층 PEs를 하이브리드 본딩으로 연결하고 sub-10 μm TSV를 통해 등록-대-등록 데이터플로를 구현한다.
• 각 층의 PE 유닛을 FlashAttention 서브 연산자(QK^T, max/subtract, exp/RowSum, PV/scaling)에 맞춰 설계한다.
• 계층 간 연속 FlashAttention 연산자를 매핑하여 지연균형 데이터플로를 구현하는 3D-FlashAttention 스케줄링을 개발한다.
• 중간 데이터가 SRAM 왕복 대신 TSV를 통해 직접 전달되어 버블-프리 수직 파이프라인을 구현한다.
• 사이클 정확도 시뮬레이션과 RTL-검증된 4-layer 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-layer 3D-Flow 스택을 배치할 때 에너지, 면적, 열 거동의 트레이드오프는 무엇인가?

주요 결과

• 에너지 소비는 2D 기준선과 비교하여 46%에서 93%까지 감소하고, 시퀀스 길이가 1K에서 64K에 이르는 동안 평균적으로 기준선 대비 32.7%에서 64.2% 감소한다.
• 추정 결과는 7.62× vs 2D-Unfused, 1.46× vs 2D-Fused, 2.36× vs Dual-SA, 1.43× vs 3D-Base의 평균 속도향상을 보인다.
• PE 활용률은 시험된 시퀀스 길이에 걸쳐 평균 87%로, 메모리 트래픽 최소화와 잘 균형 잡힌 수직 파이프라인에 의해 주도된다.
• 중간 결과는 TSV를 통해 수직으로 적층된 PEs 사이에서 직접 흐르며 SRAM 왕복을 제거하고 버블-프리 실행을 가능하게 한다.
• 발열/열 해석은 합리적인 패키징을 가진 4층 스택에서 안전한 동작 온도를 나타내며, 128×128 PE 배열에서 내부 온도 상승은 약 2.8°C이다.

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