버퍼에서 레지스터까지: 하이브리드 본딩 3D NPU로 구현하는 미세 입자 FlashAttention

초록

본 논문은 장기 시퀀스 트랜스포머 추론에서 온칩 SRAM 접근이 에너지 병목이 되는 현상을 분석하고, 서브 10 µm TSV를 활용한 하이브리드 본딩 3D 스택형 PE 배열 ‘3D‑Flow’를 제안한다. 3D‑Flow 위에 레지스터‑투‑레지스터 파이프라인을 구현한 ‘3D‑FlashAttention’ 스케줄링을 적용해 SRAM 라운드트립을 제거하고, OPT·QWEN 모델에서 46‑93 % 에너지 절감 및 1.4‑7.6× 속도 향상을 달성한다.

상세 분석

본 연구는 기존 2D systolic array 기반 가속기와 FlashAttention 알고리즘이 결합될 때 발생하는 두 가지 근본적인 한계를 정확히 짚어낸다. 첫 번째는 연산 단계 간 중간 결과를 온칩 SRAM에 저장·로드해야 하는 구조적 제약이다. SRAM 접근은 FP‑MAC 연산 대비 10‑20배 높은 에너지를 소모하므로, 시퀀스 길이가 2 k 이상으로 늘어날 경우 전체 에너지의 60 % 이상을 차지한다는 실험 결과는 매우 설득력 있다. 두 번째는 연산 부하 불균형이다. Q·Kᵀ는 대규모 매트릭스 곱으로 PE 활용도가 거의 100 %인 반면, row‑max, softmax, row‑sum, P·V 등은 스칼라·벡터 유닛에 의존해 병목이 된다. 기존 2D NoC 기반 라우터‑투‑라우터 전송은 다 사이클 지연을 유발해 파이프라인에 버블을 만들고, 결국 연산 유닛이 유휴 상태에 빠진다.

이를 해결하기 위해 저자는 하이브리드 본딩 기술을 이용한 초고밀도 TSV(피치 <10 µm, 1‑cycle 전송)로 수직 PE 층을 직접 연결하는 3D‑Flow 아키텍처를 설계한다. 각 층은 서로 다른 연산(예: Q·Kᵀ, row‑max/comparator, exp/row‑sum, P·V) 전용 마이크로아키텍처를 갖추고, 동일 위치 PE 간 레지스터가 TSV로 직접 연결돼 데이터가 SRAM을 거치지 않고 바로 위층으로 전달된다. 이는 “레지스터‑투‑레지스터” 흐름을 가능하게 하여, 연산 단계마다 발생하던 SRAM 라운드트립을 완전히 제거한다.

3D‑FlashAttention 스케줄링은 연산 부하를 층별로 균등하게 배분하도록 설계되었다. 각 연산의 사이클 수를 정밀히 모델링하고, 파이프라인이 완전 채워졌을 때 한 반복당 2 d 사이클(여기서 d는 PE 배열의 행/열 크기)만 소요된다. 이는 기존 2D 배열이 Q·Kᵀ 연산에만 3 d 사이클이 필요하던 것에 비해 33 % 이상의 속도 향상을 의미한다. 또한, 수직 데이터 흐름은 전력 소모가 낮은 TSV를 통해 이루어지므로, 전체 시스템 전력 밀도는 41 W/cm² 수준에 머물면서 열 상승도 2.8 °C에 불과해 열 관리가 용이함을 입증한다.

실험 결과는 OPT‑1.3B, QWEN‑7B 등 실제 LLM 워크로드에 적용했을 때, 3D‑Flow 기반 가속기가 기존 2D Groq, Cerebras 및 최신 3D HBM‑stack 가속기 대비 에너지 효율을 46‑93 % 개선하고, 처리량을 1.4‑7.6배 가속한다는 점에서 설계의 실효성을 강력히 뒷받침한다. 특히, 시퀀스 길이가 16 k에 달하는 경우 SRAM 접근 비중이 급증하는 상황에서도 3D‑FlashAttention은 버블‑프리 파이프라인을 유지해 성능 저하 없이 에너지 절감을 달성한다.

요약하면, 본 논문은 (1) 온칩 메모리 접근이 새로운 병목임을 정량적으로 규명하고, (2) 하이브리드 본딩 기반 초고밀도 TSV를 활용한 3D 수직 PE 스택으로 레지스터‑레지스터 데이터 흐름을 구현했으며, (3) 부하‑균형 스케줄링을 통해 연산 파이프라인을 완전 채우는 방법을 제시한다. 이러한 하드웨어‑알고리즘 공동 설계는 차세대 LLM 추론 가속기의 에너지·성능 한계를 크게 확장시킬 잠재력을 가진다.