DuoAttention 프레임워크

DuoAttention은 매사추세츠 공과대학(MIT)의 한 송 팀이 2024년에 제안한 새로운 프레임워크로, 긴 텍스트 맥락을 처리할 때 대규모 언어 모델(LLM)의 추론 효율성을 개선하는 것을 목표로 합니다. 관련 논문 결과는 다음과 같습니다.DuoAttention: 검색 및 스트리밍 헤드를 사용한 효율적인 장문맥 LLM 추론".

이 프레임워크는 "검색 헤드"와 "스트리밍 헤드"라는 두 가지 유형의 어텐션 헤드를 똑똑하게 구별하여 모델의 메모리 사용량과 컴퓨팅 속도를 최적화합니다. 검색 헤드는 장거리 종속성 처리에 중점을 두고 전체 키-값(KV) 캐시가 필요한 반면, 스트리밍 헤드는 최근 토큰과 주의 수렴 지점에 중점을 두고 고정 길이의 KV 캐시만 필요합니다. 이 디자인은 디코딩 및 사전 채우기 중에 모델의 메모리 사용량과 대기 시간을 크게 줄이는 동시에 모델이 긴 텍스트 컨텍스트를 처리하는 기능을 유지합니다.

DuoAttention은 검색 헤더에 전체 KV 캐시를 적용하고 스트리밍 헤더에 가볍고 고정 길이의 KV 캐시를 적용하여 메모리와 컴퓨팅 리소스를 최적화합니다. 이러한 개선을 통해 모델 디코딩 속도와 사전 채우기 효율성이 향상될 뿐만 아니라 긴 텍스트를 처리할 때 지연 시간도 단축됩니다. 멀티헤드 어텐션(MHA) 모델의 경우 최대 2.55배까지 줄일 수 있으며, 그룹 쿼리 어텐션(GQA) 모델의 경우 최대 1.67배까지 줄일 수 있습니다. 동시에 디코딩 속도 측면에서 멀티 헤드 어텐션(MHA) 모델은 최대 2.18배까지 향상될 수 있으며, 그룹 쿼리 어텐션(GQA) 모델은 최대 1.50배까지 향상될 수 있습니다. 사전 채우기 속도 측면에서 멀티 헤드 어텐션(MHA) 모델은 최대 1.73배까지 증가할 수 있고, 그룹 쿼리 어텐션(GQA) 모델은 최대 1.63배까지 증가할 수 있으며, 전체 어텐션 모드와 비교했을 때 정확도 손실이 최소화됩니다. 특히, 양자화 기술과 결합된 듀얼 어텐션 프레임워크는 단일 A100 GPU에서 330만 개의 텍스트 길이를 가진 Llama-3-8B 모델을 디코딩할 수 있게 해줍니다.