본 논문은 장문맥스트를 활용하는 고급 AI 시스템에서 확장성 문제를 야기하는 어텐션 계산의 효율성을 높이는 방법을 제시합니다. 기존의 scaled dot-product attention (SDPA)는 일부 중요 토큰만 출력에 크게 기여하는 토큰 스파스성을 가지지만, 이를 효과적으로 활용하는 것은 어려웠습니다. 기존 방법들은 품질 저하 또는 추가 자원 소모라는 문제가 있었습니다. 본 논문에서는 중요 토큰 식별을 Maximum Inner Product Search (MIPS) 문제로 정의하고, GPU 친화적이지 않고 query와 key 분포의 차이로 성능이 저하되는 기존 MIPS 솔루션의 문제점을 지적합니다. 대신, 중요 토큰 식별을 추천 문제로 재구성하는 HashAttention을 제안합니다. HashAttention은 학습된 매핑 함수를 사용하여 해밍 공간에서 키와 쿼리를 인코딩하여 의미적 유사성을 포착하고, 비트 연산을 통해 효율적으로 중요 토큰을 식별하여 어텐션을 계산합니다. 일반 데이터로 학습된 HashAttention은 최소한의 품질 손실로 최대 16배의 토큰 사용량을 줄이며, 토큰당 32비트의 보조 메모리만 필요합니다. 특정 작업에 대한 미세 조정을 통해 스파스성을 32배까지 향상시킬 수 있으며, A100 GPU에서 32배 스파스성을 달성하여 GPT-FAST와 FlashDecode의 어텐션 지연 시간을 각각 최대 4.3배 및 2.54배 줄이고, GPT-FAST의 처리량을 최대 3.12배 향상시킵니다.
