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 Systolic Array)를 제안합니다. 기존의 systolic array 기반 가속기는 FlashAttention의 연속되지 않은 작은 행렬 곱셈과 softmax 연산으로 인해 활용률이 낮았습니다. FSA는 SystolicAttention이라는 새로운 스케줄링 알고리즘을 통해 FlashAttention의 모든 연산을 단일 systolic array 내에서 수행하도록 설계되어 외부 벡터 유닛과의 데이터 교환을 없애고 활용률을 높였습니다. 합성 가능한 RTL로 구현된 FSA는 AWS NeuronCore-v2 및 Google TPUv5e 대비 각각 1.77배 및 4.83배 높은 attention FLOPs/s 활용률을 달성했으며, 면적 오버헤드는 약 10%에 불과했습니다.

시사점, 한계점

•
시사점:
◦
단일 systolic array 내에서 FlashAttention 알고리즘 전체를 실행 가능하게 함으로써, 기존 systolic array 기반 가속기의 낮은 활용률 문제를 해결했습니다.
◦
SystolicAttention 알고리즘을 통해 수치적 안정성을 유지하면서 systolic array 활용률을 크게 향상시켰습니다.
◦
기존 상용 가속기 대비 월등히 높은 성능을 달성했습니다.
◦
낮은 면적 오버헤드로 효율적인 설계임을 보여주었습니다.
•
한계점:
◦
현재 제시된 결과는 특정 상용 가속기와의 비교를 기반으로 하며, 다른 아키텍처나 알고리즘에 대한 일반화 가능성은 추가 연구가 필요합니다.
◦
FSA의 에너지 효율성에 대한 분석은 제시되지 않았습니다.
◦
다양한 크기의 Transformer 모델에 대한 성능 평가가 부족합니다.
◦
SystolicAttention 알고리즘의 일반적인 확장성 및 다른 attention 메커니즘으로의 적용 가능성에 대한 추가 연구가 필요합니다.
👍
Made with Slashpage