Command Palette
Search for a command to run...
블록 어텐션의 혼합 최적화
블록 어텐션의 혼합 최적화
Guangxuan Xiao Junxian Guo Kasra Mazaheri Song Han
초록
블록 어텐션의 혼합(Mixture of Block Attention, MoBA)은 쿼리가 키-밸류 블록 중 일부 작은 서브셋에만 희소하게 어텐션을 수행하도록 함으로써, 대규모 언어 모델(LLM)에서 장거리 컨텍스트를 효율적으로 처리할 수 있는 유망한 구성 요소이다. 이를 통해 계산 비용을 극적으로 감소시킬 수 있다. 그러나 MoBA의 성능을 결정짓는 설계 원리에 대한 이해는 여전히 부족하며, 효율적인 GPU 구현체가 부족하여 실용적 적용이 지연되고 있다. 본 논문에서는 먼저 MoBA의 내부 메커니즘을 분석하기 위한 통계 모델을 제안한다. 제안된 모델은 성능이 라우터(router)가 쿼리-키 유사도 기반으로 관련 블록과 비관련 블록을 정확히 구분할 수 있는 능력에 크게 의존함을 밝혀냈다. 또한, 아키텍처 파라미터와 이 검색 정확도 사이를 공식적으로 연결하는 신호 대 잡음비(signal-to-noise ratio)를 도출하였다. 본 분석을 바탕으로, 성능 향상을 위한 두 가지 핵심 방향을 제시한다. 첫째, 더 작은 블록 크기를 사용하고, 둘째, 키에 짧은 컨볼루션(short convolution)을 적용하여 관련 신호를 군집화함으로써 라우팅 정확도를 향상시키는 것이다. 이론적으로는 더 우수한 성능을 기대할 수 있으나, 작은 블록 크기는 GPU 상에서 효율적이지 못하다. 이러한 격차를 메우기 위해, 본 연구에서는 이론적으로 권장하는 작은 블록 크기에서도 효율적인 MoBA 실행을 가능하게 하는 하드웨어 인식형 CUDA 커널인 FlashMoBA를 제안한다. 우리는 초기부터 LLM을 훈련시켜 본 연구의 통찰을 검증하였으며, 개선된 MoBA 모델이 밀도 있는 어텐션 기준 모델과 동등한 성능을 달성함을 확인하였다. FlashMoBA는 작은 블록에 대해 FlashAttention-2 대비 최대 14.7배의 속도 향상을 달성함으로써, 본 연구에서 제안한 이론 기반 개선 방안이 실용적으로 가능함을 입증하였다.