Efficient Memory Management for Large Language Model Serving with PagedAttention
본 논문은 대형 언어 모델(LLM) 서빙 환경에서 가장 큰 병목 중 하나인 메모리 관리 문제를 해결하기 위해 PagedAttention이라는 혁신적인 방법을 제안한다. 이 기법은 특히 KV 캐시(Key-Value Cache) 메모리 사용 최적화에 초점을 맞추며, 운영체제의 가상 메모리 및 페이징 기법에서 영감을 받았다.
배경과 문제점
- LLM은 거대한 모델 크기와 긴 시퀀스 입력 처리 때문에 GPU 메모리 소비가 매우 크다.
- 특히 트랜스포머 모델의 attention 계산에서는 이전 토큰들의 key와 value 값을 저장하는 KV 캐시가 큰 비중을 차지한다.
- 기존 서빙 시스템들은 KV 캐시를 물리적으로 연속된 메모리 공간에 할당하는 방식으로, 메모리 과다 할당(over-reservation)과 단편화(fragmentation) 문제를 겪었다.
- 이로 인해 GPU 메모리 활용률이 낮아지고, 처리량 향상에 제한이 있었다.
PagedAttention의 핵심 아이디어
- KV 캐시를 크기가 고정된 여러 개의 블록(block) 단위로 분할하여 관리한다.
- 물리적 메모리 상에서 연속적이지 않은 여러 블록에 KV 데이터를 분산 저장한다.
- 이는 마치 컴퓨터 운영체제의 가상 메모리 시스템이 논리적 페이지를 물리적 페이지에 매핑하는 방식과 유사하다.
- 각 요청 시퀀스에 대해 필요한 만큼의 KV 블록만 동적으로 할당한다. 블록이 꽉 차면 추가 블록을 할당하는 “Pay-as-you-go” 방식을 채택해 메모리 낭비를 줄인다.
- 논리적 KV 블록과 물리적 KV 블록 간의 매핑 정보를 유지하는 페이지 테이블을 도입해 블록들의 순서와 연결성을 관리한다.
PagedAttention 동작 과정
- 입력 토큰 시퀀스를 고정 크기 KV 블록 단위로 나눈다.
- 각 KV 블록은 key, value 벡터를 담고 있으며, 물리적 메모리 상에 분산 저장된다.
- 어텐션 계산 시 필요한 KV 블록만을 선택적으로 불러와 계산한다.
- 계산된 결과를 종합하여 최종 어텐션 출력을 생성한다.
- 동적 할당 및 해제를 통해 메모리 사용을 최적화한다.
기술적 특징과 장점
- 메모리 효율성 극대화: KV 캐시를 필요 최소한의 단위로 분할하여 과다 할당을 방지한다.
- 단편화 감소: 논리적-물리적 메모리 매핑을 통해 메모리 단편화 문제를 완화한다.
- 성능 향상: GPU 활용률이 높아져 추론 처리량(throughput)이 기존보다 2배 이상 증가한다는 실험 결과를 보임.
- 유연한 메모리 관리: 요청 수·길이에 따라 적절한 메모리 블록을 할당하는 유연함.
- 실시간 서빙 최적화: 지속적 배치(Continuous Batching) 전략과 결합해 처리 지연을 최소화한다.
실서비스 적용과 향후 전망
- PagedAttention은 vLLM과 같은 LLM 서빙 프레임워크에서 채택되어 실제 대규모 모델 서비스에서 높은 처리량과 자원 효율성을 입증했다.
- 향후 동적 블록 크기 조정, 멀티모달 모델 확장, 더 정교한 압축 기법과의 결합 등으로 발전 가능성이 크다.