[공지사항]을 빙자한 안부와 근황 
Show more

Daily Arxiv

전 세계에서 발간되는 인공지능 관련 논문을 정리하는 페이지 입니다.
본 페이지는 Google Gemini를 활용해 요약 정리하며, 비영리로 운영 됩니다.
논문에 대한 저작권은 저자 및 해당 기관에 있으며, 공유 시 출처만 명기하면 됩니다.

SystolicAttention: Fusing FlashAttention within a Single Systolic Array

Created by
  • Haebom

저자

Jiawei Lin, Guokai Chen, Yuanlong Li, Thomas Bourgeat

개요

본 논문은 FlashAttention 알고리즘을 사용하는 Transformer 모델의 효율적인 가속화를 위해 새로운 systolic array 기반 아키텍처인 FSA(Flash Attention Systolic Array)를 제안합니다. 기존 systolic array 기반 가속기는 FlashAttention의 연속적이지 않은 행렬 곱셈과 softmax 연산으로 인해 활용률이 낮았으나, FSA는 SystolicAttention이라는 새로운 스케줄링 알고리즘을 통해 FlashAttention 알고리즘 전체를 단일 systolic array 내에서 실행하여 이 문제를 해결합니다. SystolicAttention은 세밀한 요소 단위의 겹침을 통해 array 활용률을 높이고, 원래의 부동 소수점 연산 순서를 유지하여 수치적 안정성을 보장합니다. 합성 가능한 RTL로 FSA를 구현하여 AWS NeuronCore-v2 및 Google TPUv5e와 비교 평가한 결과, 약 10%의 면적 오버헤드로 1.77배(NeuronCore-v2 대비) 및 4.83배(TPUv5e 대비) 높은 attention FLOPs/s 활용률을 달성했습니다.

시사점, 한계점

시사점:
단일 systolic array 내에서 FlashAttention 알고리즘 전체를 실행 가능하게 함으로써 기존 systolic array 기반 가속기의 성능 한계를 극복.
기존 상용 가속기 대비 월등히 높은 attention FLOPs/s 활용률 달성 (AWS NeuronCore-v2 대비 1.77x, Google TPUv5e 대비 4.83x).
낮은 면적 오버헤드 (약 10%)로 높은 성능 향상을 달성.
SystolicAttention 알고리즘은 다른 유사한 알고리즘에도 적용 가능성이 높음.
한계점:
현재는 FlashAttention 알고리즘에 특화된 아키텍처로, 다른 attention 메커니즘에는 적용이 어려울 수 있음.
실제 응용 프로그램에서의 성능 평가가 부족함. (논문에서는 합성 가능한 RTL 기반 평가만 제시)
에너지 효율에 대한 분석이 부족함.
SystolicAttention 알고리즘의 일반화 및 확장성에 대한 추가 연구가 필요함.
👍