PagedAttention

PagedAttention은 대규모 언어 모델(LLM) 서빙을 위한 어텐션 알고리즘이자 메모리 관리자입니다. 각 요청의 KV 캐시를 GPU 메모리의 어느 위치에나 배치할 수 있는 작은 고정 크기 블록에 저장하고, 운영 체제가 가상 메모리 페이지를 매핑하는 방식처럼 요청별 블록 테이블로 블록을 추적합니다. UC Berkeley와 협력 대학의 Woosuk Kwon, Zhuohan Li 및 동료들은 vLLM 서빙 엔진과 함께 이 기술을 개발해 SOSP 2023에서 Efficient Memory Management for Large Language Model Serving with PagedAttention 논문으로 발표했습니다. 연속적인 캐시 할당에서 발생하는 단편화를 제거하면 GPU 하나가 동시에 더 많은 시퀀스를 저장할 수 있습니다. 논문은 동일한 지연 시간에서 FasterTransformer 및 Orca보다 처리량이 2~4배 높다고 보고했습니다.

해결하려는 메모리 문제#

텍스트 생성 중 트랜스포머는 이전 토큰의 키 및 값 벡터를 저장해 재계산을 방지합니다. 이 캐시는 계층마다 토큰 하나당 항목 하나씩 늘어나며, 파라미터 130억 개로 구성된 LLaMA 모델의 시퀀스 하나가 최대 1.7GB를 차지할 수 있습니다. 요청이 들어올 때는 출력 길이를 알 수 없으므로, 초기 서빙 시스템은 요청마다 가능한 최대 길이에 맞는 연속 영역을 예약했습니다. 이 영역 대부분은 비어 있었고, 영역 사이의 빈 공간은 너무 작아 재사용할 수 없었습니다. vLLM 팀의 측정에 따르면 이러한 시스템은 KV 캐시 메모리의 60~80%를 낭비해 연산 자원이 부족해지기 훨씬 전에 배치 크기를 제한했습니다.

페이징의 작동 방식#

PagedAttention은 캐시를 고정된 수의 토큰에 대한 키와 값을 저장하는 블록으로 분할합니다. vLLM의 기본값은 블록당 토큰 16개입니다. 요청은 마지막 블록이 가득 찼을 때만 새 블록을 받으며, 빈 블록이 있는 곳이면 어디든 할당되고 블록 테이블에 순서가 기록됩니다. 어텐션 메커니즘 커널은 이 테이블을 따라 블록별로 키와 값을 가져오므로 시퀀스의 캐시를 연속된 영역에 저장할 필요가 없습니다. 낭비되는 공간은 각 시퀀스의 마지막 블록에서 사용되지 않은 부분으로 제한되며, vLLM 측정에서는 4% 미만이었습니다.

이 PyTorch 예시 코드는 공유 풀에서 요청 하나의 키를 분산된 블록에 저장하고, 블록 테이블을 통해 다시 읽습니다.

import torch

block_size, head_dim = 16, 64
key_pool = torch.zeros(64, block_size, head_dim)  # physical blocks shared by every request
free_blocks = [41, 7, 19, 2, 33]  # free blocks are scattered across the pool

# A 40-token request needs ceil(40 / 16) = 3 blocks, wherever they happen to be free
keys = torch.randn(40, head_dim)
num_blocks = -(-len(keys) // block_size)
block_table = [free_blocks.pop(0) for _ in range(num_blocks)]
for pos, key in enumerate(keys):
    key_pool[block_table[pos // block_size], pos % block_size] = key

# The attention kernel reads the request's keys back through its block table
gathered = key_pool[block_table].reshape(-1, head_dim)[: len(keys)]
print(f"block table {block_table}, keys intact: {torch.equal(gathered, keys)}")
print(f"unused slots: {num_blocks * block_size - len(keys)} paged vs {2048 - len(keys)} with a 2,048-token reservation")

블록 테이블 [41, 7, 19]과 keys intact: True을 출력한 다음, 페이징을 사용하는 요청에서는 사용하지 않은 슬롯 8개가 발생하고 토큰 2,048개를 위해 예약된 영역에서는 2,008개가 발생함을 보여줍니다.

시퀀스 간 블록 공유#

블록 테이블을 사용하면 여러 시퀀스가 같은 물리 블록을 가리킬 수도 있습니다. 병렬 샘플링, 빔 서치, 동일한 시스템 프롬프트로 시작하는 요청은 공유 프리픽스의 복사본 하나를 재사용하며, 각 블록에는 참조 횟수가 기록됩니다. 다른 시퀀스가 아직 사용하는 블록에 한 시퀀스가 기록해야 할 경우, 엔진은 먼저 해당 블록을 복사합니다. 이는 운영 체제에서 가져온 쓰기 시 복사 방식입니다.

vLLM 저자들은 이러한 공유가 병렬 샘플링과 빔 서치의 메모리 사용량을 최대 55% 줄여 최대 2.2배 높은 처리량을 제공한다고 보고했습니다. 전체 블록을 해싱하면 이 방식을 요청 간에도 확장할 수 있으며, 이것이 vLLM이 프롬프트 캐싱을 구현하는 방식입니다.

관련 기법#

PagedAttention은 캐시된 키와 값을 저장할 위치를 결정하며, 저장할 데이터의 양이나 읽기 속도를 바꾸는 방법과 함께 사용할 수 있습니다.

• Flash Attention: 어텐션 계산을 타일로 나누어 전체 점수 행렬이 GPU 메인 메모리에 기록되지 않도록 합니다. 최신 엔진은 페이징된 블록에서 읽는 FlashAttention 방식의 커널을 실행합니다.
• 연속 배칭: 다른 요청이 끝날 때마다 새 요청을 받아들입니다. 페이징은 완료된 시퀀스의 블록을 즉시 여유 풀에 반환하므로 새 요청이 사용할 공간을 확보할 수 있습니다.
• 그룹 쿼리 어텐션: 키 및 값 헤드를 쿼리 헤드 간에 공유해 각 블록의 내용을 줄입니다.
• FP8 캐시 양자화: 키와 값을 8비트 정밀도로 저장해 16비트 형식에 비해 캐시 메모리를 대략 절반으로 줄입니다.

Ultralytics LLM 인터페이스는 요청 인수를 변경 없이 전달합니다. 따라서 n=3을 사용하는 Chat Completions 호출을 해당 인터페이스의 base_url를 통해 로컬 vLLM 서버로 보내면, 프롬프트의 캐시 블록을 공유하는 샘플 세 개가 반환되고 서로 달라지는 후속 부분에만 새 메모리가 사용됩니다.