PagedAttention

PagedAttention은 대형 언어 모델(LLM)을 효율적으로 서빙하기 위한 어텐션 알고리즘이다. 이 알고리즘은 2023년 우석 권(Woosuk Kwon)과 동료들이 논문 《PagedAttention을 이용한 대형 언어 모델 서빙을 위한 효율적인 메모리 관리》에서 VLLM 서빙 엔진과 함께 처음 소개하였다. 이 방법은 자기회귀 디코딩 중에 사용되는 키-값 캐시를 비연속적인 물리 메모리에 매핑할 수 있는 고정 크기 블록으로 저장하며, 가상 메모리, 페이징, 운영체제 설계의 개념을 차용한다.

배경

트랜스포머 추론에서 키-값 캐시는 시퀀스 길이와 동시 요청 수에 따라 증가한다. 권 등은 기존 서빙 시스템이 일반적으로 연속적인 캐시 영역을 미리 예약하며, 이로 인해 예약된 공간 낭비, 내부 단편화 및 외부 단편화가 발생한다고 주장하였다. 논문은 실험을 통해 이전 시스템의 유효 메모리 활용률이 20.4%까지 떨어질 수 있음을 보고하였다.

설명

PagedAttention은 각 시퀀스의 캐시를 고정 크기의 KV 블록으로 분할한다. 요청의 캐시는 논리적 블록의 시퀀스로 표현되며, 블록 테이블은 이러한 논리적 블록을 물리적 GPU 메모리 블록으로 매핑한다. 결과적으로 인접한 논리적 블록이 물리 메모리에서 연속적일 필요가 없으며, 생성이 진행됨에 따라 필요할 때 새로운 블록을 할당할 수 있다.

이 설계는 또한 관련 디코딩 경로 간에 캐시 상태를 공유하기 쉽게 만든다. VLLM에서 물리적 블록은 참조 카운트될 수 있으며 요청이나 분기 간에 공유될 수 있고, 공유 블록을 수정해야 할 때는 블록 단위의 카피 온 라이트가 사용된다. 원 논문은 이 설계를 병렬 샘플링, 빔 탐색, 공유 접두사가 있는 프롬프트에 적용하였다.

수학적 공식

인과적 셀프 어텐션(causal self-attention)에서 쿼리 토큰 ${\displaystyle i}$에 대한 표준 어텐션 출력은 다음과 같이 작성할 수 있다. $${\displaystyle a_{ij}={\frac {\exp(\mathbf {q} _{i}^{\top }\mathbf {k} _{j}/{\sqrt {d}})}{\sum _{t=1}^{i}\exp(\mathbf {q} _{i}^{\top }\mathbf {k} _{t}/{\sqrt {d}})}},\;\mathbf {o} _{i}=\sum _{j=1}^{i}a_{ij}\mathbf {v} _{j}}$$ 여기서 ${\displaystyle \mathbf {q} _{i}}$, ${\displaystyle \mathbf {k} _{j}}$, ${\displaystyle \mathbf {v} _{j}}$는 각각 쿼리, 키, 값 벡터이며 ${\displaystyle d}$는 어텐션 차원이다.

캐시가 ${\displaystyle B}$ 크기의 블록으로 분할되면 키 및 값 블록은 다음과 같이 작성될 수 있다. $${\displaystyle \mathbf {K} _{j}=(\mathbf {k} _{(j-1)B+1},\ldots ,\mathbf {k} _{jB}),\;\mathbf {V} _{j}=(\mathbf {v} _{(j-1)B+1},\ldots ,\mathbf {v} _{jB})}$$ PagedAttention은 이후 블록 단위로 계산을 수행한다. $${\displaystyle \mathbf {A} _{ij}={\frac {\exp(\mathbf {q} _{i}^{\top }\mathbf {K} _{j}/{\sqrt {d}})}{\sum _{t=1}^{\lceil i/B\rceil }\exp(\mathbf {q} _{i}^{\top }\mathbf {K} _{t}/{\sqrt {d}})}},\;\mathbf {o} _{i}=\sum _{j=1}^{\lceil i/B\rceil }\mathbf {V} _{j}\mathbf {A} _{ij}^{\top }}$$ 여기서 ${\displaystyle \mathbf {A} _{ij}}$는 ${\displaystyle j}$번째 KV 블록에 대한 어텐션 점수 벡터이다. 권 등이 제시한 공식에서 이는 인과적 어텐션 계산을 보존하면서도 키 및 값 블록이 비연속적인 물리 메모리에 위치할 수 있도록 한다.

성능 및 사용

VLLM 논문은 평가된 작업 부하에서 PagedAttention과 관련 메모리 관리 설계를 사용한 결과, FasterTransformer 및 Orca를 포함한 비교 대상 대비 서빙 처리량이 2~4배 향상되었으며 모델 출력 결과는 동일하게 유지되었다고 보고하였다. Alpaca 트레이스를 사용한 OPT-13B 실험에서 이 논문은 KV 블록 공유를 통해 병렬 샘플링에서는 6.1~9.8%, 빔 탐색에서는 37.6~55.2%의 메모리 절감 효과를 보였다고 보고하였다.

2024년 LLM 서빙 시스템 조사에 따르면 PagedAttention은 TGI, VLLM, TensorRT-LLM에 대한 지원을 바탕으로 LLM 서빙 프레임워크의 산업 표준이 되었다고 기술하였다.

한계 및 대안

후속 연구들은 이 접근 방식의 장단점을 기술하였다. 2025년 vAttention 논문은 PagedAttention이 페이징을 지원하기 위해 어텐션 커널을 재작성해야 하며, 소프트웨어 복잡성, 이식성 문제, 중복성 및 실행 오버헤드를 증가시킨다고 주장하였다. 대신 가상 메모리에서 캐시를 연속적으로 유지하면서 물리적 할당을 위해 요구 페이징(demand paging)에 의존하는 메모리 관리자를 제안하였다.

같이 보기

• AlphaDev 및 AlphaEvolve — 컴퓨터 과학 알고리즘을 발견하고 최적화하기 위한 구글 딥마인드의 AI 시스템
• 어텐션 메커니즘
• 빔 탐색
• 캐시 교체 정책
• 인공지능 알고리즘 목록

각주