FlashAttention for Scalable Vector Architectures
이 논문은 인터헤드 병렬성(inter-head parallelism)과 효율적인 메모리 액세스를 활용하여 CPU 상의 트랜스포머 추론을 크게 가속화함으로써 프리필(prefill)에서 최대 42배, 디코드(decode)에서 11배의 속도 향상을 달着하는 동시에 롱 벡터 실행을 위한 현재 양자화 포맷의 구조적 병목 현상을 강조하는, 확장 가능한 벡터 아키텍처에 최적화된 블록형 FlashAttention 구현체인 FlashAttention-V를 소개한다.
원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기
이메일을 작성하는 챗봇부터 소프트웨어의 버그를 수정하는 코딩 어시스턴트에 이르기까지, 현대의 인공지능 시스템은 트랜스포머(transformer)라고 불리는 특정한 유형의 컴퓨터 프로그램에 의존합니다. 이 프로그램들은 문장 내에서 단어들이 서로 어떻게 연관되는지를 살펴봄으로써 언어를 이해하도록 구축되었습니다. 이를 위해 그들은 '어텐션(attention)'이라 불리는 메커니즘을 사용하는데, 이는 마치 스포트라이트처럼 작동하여 시스템이 나머지 부분을 무시하고 텍스트의 가장 관련 있는 부분에 집중할 수 있게 해줍니다. 이러한 시스템은 흔히 강력한 그래픽 카드를 갖춘 거대한 데이터 센터에서 실행되지만, 노트북, 태블릿, 심지어 사물인터넷(IoT)에서 발견되는 특수 마이크로칩과 같은 더 작고 일상적인 기기에서 실행해야 할 필요성이 커지고 있습니다. 이러한 더 작은 기기들은 종មាន 데이터를 긴 줄 형태로 처리함으로써 한 번에 많은 계산을 수행하도록 설계된 벡터 아키텍처(vector architectures)라는 다른 종류의 프로세서를 사용하는 경우가 많습니다. 그러나 주요한 장애물이 등장했습니다. 바로 어텐션 메커니즘은 메모리를 매우 많이 소모하며, 프로세서가 방대한 양의 데이터를 끊임없이 가져오고 저장해야 하므로 모든 동작을 느리게 만든다는 점입니다.
찰머스 공과대학교(Chalmers University of Technology)와 글래스고 대학교(University of Glasgow)의 연구진은 이러한 벡터 프로세서에서 어텐션 메커니즘이 작동하는 방식을 재설계함으로써 이 문제를 해결했습니다. 그들은 FlashAttention-V라고 불리는 새로운 방법을 만들어냈는데, 이는 컴퓨터가 이러한 프로세서의 독특한 형태에 맞게 작업 방식을 조직하는 방식을 변경합니다. 프로세서가 한 번에 하나의 작은 데이터 조각을 처리하도록 강요하는 대신, 이 새로운 방법은 여러 개의 독립적인 계산을 하나의 넓은 데이터 줄로 그룹화합니다. 이는 보통 한 번에 한 품목씩 다루는 공장의 조립 라인을 상상해 보십시오. 이 새로운 접근 방식은 작업자들이 한 번에 전체 쟁반을 잡고 한꺼번에 처리할 수 있게 하여, 저장 선반을 향해 왔다 갔다 하는 데 소비되는 시간을 획기적으로 줄여줍니다. 연산의 순서를 재배치하고 데이터를 더 조밀하게 채움으로써, 연구진은 이러한 더 작은 기기들이 특히 짧은 텍스트를 처리하거나 기기가 단어 하나하나를 생성할 때 훨씬 더 빠르게 작동할 수 있음을 발견했습니다.
연구팀은 실제 하드웨어와 상세한 컴퓨터 시뮬레이션을 모두 사용하여 TinyLlama, Llama 3.2, Qwen2.5를 포함한 여러 가지 언어 모델을 테스트했습니다. RISC-V 프로세서를 사용하는 Banana Pi BPI-F3라는 물리적 개발 보드에서, 이 새로운 방법은 엄청난 개선을 보여주었습니다. 기기가 짧은 입력을 읽을 준비를 할 때, 새로운 접근 방식은 표준적인 비최적화 버전보다 12배에서 14배 더 빨랐습니다. 기기가 단어별로 텍스트를 생성할 때는 4배에서 5배 더 빨랐습니다. 또한 연구진은 프로세서가 훨씬 더 넓은 데이터 줄을 갖추어 한 번에 훨씬 더 큰 정보 덩어리를 처리할 수 있도록 구축되었을 경우 어떻게 성능을 발휘할지 확인하기 위해 광범위한 시뮬레이션을 수행했습니다. 이러한 시뮬레이션은 데이터 줄이 넓어질수록 속도 이득이 계속 상승하여, 입력을 준비하는 최상의 시나리오에서는 기본 버전보다 최대 42배까지 빨라진다는 것을 보여주었습니다.
하지만 이 연구는 이러한 기기들이 얼마나 더 빨라질 수 있는지에 대한 명확한 한계도 밝혀냈습니다. 연구진은 프로그램의 어텐션 부분은 더 넓은 데이터 줄로부터 큰 이득을 얻는 반면, 숫자를 예측값으로 변환하는 다른 레이어들은 그렇지 않다는 것을 발견했습니다. 이 레이어들은 공간을 절약하기 위해 데이터를 압축하는 양자화(quantization)라는 특정한 숫자 저장 방식을 사용합니다. 현재 이 데이터가 패킹되는 방식은 넓은 데이터 줄과 구조적 불일치를 일으켜, 프로세서가 숫자를 분리하고 다시 결합하기 위해 추가적이고 비효리적인 작업을 수행하게 만듭니다. 시뮬레이션 결과, 이러한 특정 레이어들의 경우 한 번에 더 많은 데이터를 처리함으로써 얻는 시간 절감 효과가 데이터를 재배치하는 데 드는 시간에 의해 완전히 상쇄되는 것으로 나타났습니다. 이는 어텐션 메커니즘은 믿을 수 없을 정도로 빠르게 만들 수 있지만, 시스템의 전체 속도는 현재 이러한 다른 구성 요소들에 의해 제약받고 있음을 의미하며, 이는 미래의 개선이 단순히 어떻게 프로세싱하느냐가 아니라 어떻게 이 압축된 숫자들을 저장하고 이동시키느냐의 문제를 해결해야 함을 시사합니다.
이러한 결과는 인공지능을 일상적인 기기에서 더 쉽게 사용할 수 있도록 하는 명확한 경로를 제시합니다. 새로운 방법인 FlashAttention-V는 현대 언어 모델의 설계와 확장 가능한 벡터 프로세서의 역량 사이의 간극을 성공적으로 메웠습니다. 이는 컴퓨터가 과업을 생각하는 방식을 단순히 재정렬하고 데이터를 더 효율적으로 채우는 것만으로도 새로운 하드웨어 없이 상당한 성능 향상이 가능하다는 것을 증명합니다. 연구는 짧은 작업과 실시간 텍스트 생성에 있어 이러한 최적화된 프로세서들이 매우 효과적일 수 있음을 확인해 줍니다. 동시에, 이는 현재 기기를 위한 데이터 압축 방식이 한계에 부딪혔을 수도 있다는 경고의 메시지도 담고 있습니다. 즉, 더 넓은 프로세서의 잠재력을 완전히 끌어올리는 것은 이 압축된 숫자들을 더 지능적으로 저장하고 이동시키는 퍼즐을 푸는 것에 달려 있을 것입니다.
연구 분야의 논문에 파묻히고 계신가요?
연구 키워드에 맞는 최신 논문의 일일 다이제스트를 받아보세요 — 기술 요약 포함, 당신의 언어로.