KV 캐시를 가상 메모리처럼 다루는 PagedAttention

LLM 서빙 처리량은 GPU 연산력보다 KV 캐시 메모리 관리가 좌우한다는 사실을 정면으로 다룬 연구입니다. SOSP 2023에 발표된 이 논문은 오늘날 사실상 표준 서빙 엔진이 된 vLLM의 출발점이기도 합니다.

문제 배경

각 요청은 생성한 토큰의 키와 값을 KV 캐시로 GPU 메모리에 유지하는데, 이 캐시는 크기가 크고 생성이 진행되며 동적으로 늘었다 줄었다 합니다. 기존 시스템은 요청마다 최대 길이만큼 연속 메모리를 미리 할당했기 때문에 내부 단편화와 예약 낭비가 심했고, 이것이 배치 크기를 제한해 GPU 연산력을 놀리게 만들었습니다. 논문의 프로파일링에 따르면 기존 시스템에서 실제 토큰 상태 저장에 쓰인 KV 캐시 메모리는 20.4%에서 38.2%에 불과했습니다.

방법

PagedAttention은 운영체제의 가상 메모리와 페이징에서 착안했습니다. KV 캐시를 고정 크기 블록으로 나누어 비연속적으로 할당하고, 블록 테이블이 논리 블록과 물리 블록을 매핑합니다. 시퀀스가 실제로 길어지는 만큼만 메모리를 쓰므로 예약 낭비가 사라지고, 병렬 샘플링이나 빔 서치처럼 프리픽스를 공유하는 요청 간에는 같은 물리 블록을 참조해 중복 저장을 없앱니다. 이 알고리즘 위에 스케줄러와 분산 실행까지 갖춘 서빙 시스템으로 구현한 것이 vLLM입니다.

결과

vLLM은 KV 캐시 낭비를 거의 0으로 만들었고 FasterTransformer와 Orca 등 당시 최신 시스템 대비 같은 지연 수준에서 인기 LLM의 처리량을 2배에서 4배 개선했습니다. 논문에 따르면 시퀀스가 길수록, 모델이 클수록, 디코딩 알고리즘이 복잡할수록 개선 폭이 더 커집니다.

한계와 주의점

블록 단위 관리는 공짜가 아닙니다. 어텐션 커널이 블록 테이블을 따라가도록 다시 작성돼야 하고 블록 관리 자체의 오버헤드도 있습니다. 프리픽스 공유가 없는 짧은 단발 요청 위주의 워크로드라면 개선 폭은 논문 수치의 하단에 가까워집니다. 결국 자체 서빙 코드가 아니라 vLLM 같은 검증된 런타임을 채택하는 결정이 됩니다.

기업 적용 관점

온프레미스 서빙에서 GPU 대수를 결정하는 것은 모델 크기보다 동시 처리량인 경우가 많습니다. 사내 문서 질의응답처럼 점심시간에 트래픽이 몰리는 워크로드에서 같은 장비로 수 배의 요청을 받아낼 수 있으므로, 하드웨어 견적 단계에서 PagedAttention 기반 서빙을 전제로 GPU 수량을 산정하면 도입 비용 자체가 달라집니다.