DuoAttention: Efficient Long-Context LLM Inference with Retrieval and Streaming Heads

Deploying long-context large language models (LLMs) is essential but poses significant computational and memory challenges. Caching all Key and Value (KV) states across all attention heads consumes substantial memory. Existing KV cache pruning methods either damage the long-context capabilities of LLMs or offer only limited efficiency improvements. In this paper, we identify that only a fraction of attention heads, a.k.a, Retrieval Heads, are critical for processing long contexts and require full attention across all tokens. In contrast, all other heads, which primarily focus on recent tokens and attention sinks--referred to as Streaming Heads--do not require full attention. Based on this insight, we introduce DuoAttention, a framework that only applies a full KV cache to retrieval heads while using a light-weight, constant-length KV cache for streaming heads, which reduces both LLM's decoding and pre-filling memory and latency without compromising its long-context abilities. DuoAttention uses a lightweight, optimization-based algorithm with synthetic data to identify retrieval heads accurately. Our method significantly reduces long-context inference memory by up to 2.55x for MHA and 1.67x for GQA models while speeding up decoding by up to 2.18x and 1.50x and accelerating pre-filling by up to 1.73x and 1.63x for MHA and GQA models, respectively, with minimal accuracy loss compared to full attention. Notably, combined with quantization, DuoAttention enables Llama-3-8B decoding with 3.3 million context length on a single A100 GPU. Code is provided in https://github.com/mit-han-lab/duo-attention.

한 줄 요약

DuoAttention은 LLM의 장문맥 추론에서 메모리와 계산 효율성을 극대화하기 위해 Retrieval Heads와 Streaming Heads를 구분하는 새로운 인퍼런스 프레임워크이다.

핵심 기여도

Retrieval Heads와 Streaming Heads의 구분을 통해 KV 캐시를 선택적으로 적용함으로써 메모리 사용량을 MHA 기준 2.55×, GQA 기준 1.67× 줄임.
최적화 기반 알고리즘으로 Retrieval Heads를 정확히 식별하여, 인공 데이터를 사용한 학습 없이도 정확도 손실 최소화 (최대 0.1% 이하).
디코딩 속도는 MHA 기준 2.18×, GQA 기준 1.50× 향상, 프리필링 속도는 각각 1.73×, 1.63× 향상.
8-bit weight + 4-bit KV 캐시 양자화와 결합 시, Llama-3-8B 모델이 단일 A100 GPU에서 3.3M 토큰 처리 가능.

핵심 아이디어

기존 KV 캐시 최적화 방법은 전체 어텐션 헤드에 동일한 처리를 적용하여 장문맥 능력을 저하시키거나 효율성 향상이 제한적이었다. 본 연구는 어텐션 헤드가 두 가지로 분류될 수 있음을 발견했다: Retrieval Heads는 전체 토큰에 대한 전체 어텐션을 필요로 하며, Streaming Heads는 최근 토큰과 attention sink에만 집중할 수 있다. 이 통찰을 바탕으로, DuoAttention은 Retrieval Heads에만 전체 KV 캐시를 적용하고, Streaming Heads에는 고정 길이의 경량 KV 캐시를 사용함으로써 메모리와 계산 비용을 줄인다. Retrieval Heads 식별은 최적화 기반 알고리즘과 인공 데이터를 활용하여 이루어지며, 기존의 어텐션 패턴 분석 기반 방법과 달리 토큰 드롭 시 출력 편차를 직접 측정하여 정확도를 높인다.

기술적 접근법

**DuoAttention 프레임워크**: 각 Transformer 레이어에 두 가지 KV 캐시를 적용 — Retrieval Heads에는 전체 KV 캐시, Streaming Heads에는 고정 길이 KV 캐시 (최근 토큰과 attention sink만 포함).
**Retrieval Heads 식별 알고리즘**: 최적화 기반, 인공 데이터를 사용한 토큰 드롭 기반 출력 편차 측정.
**지원 모델 아키텍처**: MHA (Multi-Head Attention), GQA (Grouped-Query Attention).
**양자화 호환**: 8-bit weight + 4-bit KV 캐시 양자화와 결합 가능.
**하드웨어**: A100 GPU 기준, 3.3M 토큰 처리 가능.

주요 결과

**메모리 절감**: MHA 기준 2.55×, GQA 기준 1.67×.
**디코딩 속도 향상**: MHA 기준 2.18×, GQA 기준 1.50×.
**프리필링 속도 향상**: MHA 기준 1.73×, GQA 기준 1.63×.
**Llama-3-8B 성능**: 8-bit weight + 4-bit KV 캐시 양자화 시, 단일 A100 GPU에서 3.3M 토큰 처리 가능 (기존 FP16 대비 6.4× 증가).
**정확도 손실**: 최대 0.1% 이하.

의의 및 한계

DuoAttention은 장문맥 추론에서 메모리와 계산 효율성을 극대화하면서도 모델의 장문맥 능력을 유지하는 데 성공했다. 특히, 기존 KV 캐시 최적화 방법과 달리 Retrieval Heads 식별을 통해 선택적 캐시 적용이 가능하다는 점에서 혁신적이다. 또한, GQA와 양자화 기법과의 호환성을 통해 실용적 적용 가능성을 높였다. 그러나 본 연구는 특정 모델 아키텍처 (예: Llama-3-8B)에 기반한 실험 결과를 제시했으며, 다른 모델 아키텍처나 데이터셋에 대한 일반화 가능성은 추가 연구가 필요하다. 또한, Retrieval Heads 식별 알고리즘의 복잡도와 인공 데이터 생성 과정의 비용도 고려해야 한다.

실용적 활용

DuoAttention은 대규모 텍스트, 영상, 비디오 처리가 필요한 산업 분야 (예: 미디어 분석, 고객 지원, 법률 문서 요약)에서 LLM의 장문맥 추론을 효율적으로 수행할 수 있도록 지원한다. 특히, GPU 메모리 제한이 있는 환경에서 장문맥 처리가 필요한 클라우드 인프라 및 엣지 기기에서 유용하게 활용될 수 있다.