본문으로 건너뛰기

[vLLM] PagedAttention: OS 페이징 기법으로 KV 캐시를 관리하는 방법

들어가며

LLM 서빙에서 가장 큰 병목 중 하나는 KV 캐시 메모리 관리이다. Transformer의 어텐션 연산은 이전 토큰들의 Key, Value 텐서를 보관해야 하는데, 요청마다 시퀀스 길이가 다르기 때문에 메모리를 사전에 연속 할당하면 심각한 내부 단편화(internal fragmentation)가 발생한다. PagedAttention은 운영체제의 가상 메모리 페이징을 KV 캐시에 적용해 이 문제를 해결한다.

공식 문서

vLLM 공식 문서: Paged Attention

핵심 구조/코드 분석

KVCacheManager: 블록 단위 할당의 중심

vLLM의 KV 캐시 관리 핵심은 vllm/v1/core/kv_cache_manager.py의 KVCacheManager 클래스이다. OS의 페이지 테이블처럼 요청별로 블록 매핑을 관리한다.

class KVCacheManager:
    def __init__(
        self,
        kv_cache_config: KVCacheConfig,
        max_model_len: int,
        hash_block_size: int,
        enable_caching: bool = True,
        use_eagle: bool = False,
        ...
    ) -> None:
        self.coordinator = get_kv_cache_coordinator(
            kv_cache_config=kv_cache_config,
            max_model_len=self.max_model_len,
            use_eagle=self.use_eagle,
            enable_caching=self.enable_caching,
            ...
        )
        self.block_pool = self.coordinator.block_pool

coordinator가 실제 블록 할당/해제/캐시 히트 탐색을 담당하며, block_pool은 GPU 메모리에 미리 할당된 고정 크기 블록들의 풀이다. OS의 물리 페이지 프레임 풀과 동일한 역할이다.

블록 할당: allocate_slots

요청이 스케줄링되면 allocate_slots가 호출된다. 이 메서드의 블록 레이아웃을 보면 OS 페이지 할당과의 유사성이 명확하다:

def allocate_slots(self, request, num_new_tokens, ...):
    # 슬라이딩 윈도우 밖의 블록 해제 (OS의 페이지 아웃과 유사)
    self.coordinator.remove_skipped_blocks(
        request.request_id, total_computed_tokens
    )
    # 필요한 블록 수 계산
    num_blocks_to_allocate = self.coordinator.get_num_blocks_to_allocate(...)
    # 여유 블록 확인 (OS의 free frame list 확인)
    if num_blocks_to_allocate > self.block_pool.get_num_free_blocks():
        return None  # OOM - 스케줄러가 preemption 수행
    # 새 블록 할당
    new_blocks = self.coordinator.allocate_new_blocks(...)
    return self.create_kv_cache_blocks(new_blocks)

PagedAttention 커널: 비연속 메모리 접근

실제 어텐션 연산은 vllm/v1/attention/ops/paged_attn.py에서 수행된다:

class PagedAttention:
    @staticmethod
    def write_to_paged_cache(
        key, value,
        key_cache, value_cache,
        slot_mapping,        # 가상 → 물리 블록 매핑
        kv_cache_dtype,
        k_scale, v_scale,
    ) -> None:
        ops.reshape_and_cache(
            key, value, key_cache, value_cache,
            slot_mapping.flatten(), kv_cache_dtype, k_scale, v_scale,
        )

slot_mapping은 OS의 페이지 테이블에 해당한다. 각 토큰 위치가 실제 물리 블록의 어느 슬롯에 매핑되는지를 나타내며, CUDA 커널이 이 매핑을 따라 비연속 메모리에 흩어진 KV를 읽고 쓴다.

KVCacheBlocks: 블록 ID 관리

@dataclass
class KVCacheBlocks:
    blocks: tuple[Sequence[KVCacheBlock], ...]

    def get_block_ids(self, allow_none=False):
        return tuple(
            [blk.block_id for blk in group] for group in self.blocks
        )

요청마다 KVCacheBlocks를 들고 있으며, 이것이 OS 프로세스의 페이지 테이블에 해당한다. block_id가 물리 페이지 프레임 번호이다.

왜 이 설계인가

  1. 외부 단편화 제거: 고정 크기 블록 단위 할당으로 메모리 단편화가 원천 차단된다. 논문에 따르면 기존 방식 대비 메모리 낭비가 최대 4% 미만으로 줄어든다.

  2. Prefix Caching 자연 지원: 블록 단위로 해시를 계산하면 동일 프리픽스를 공유하는 요청들이 같은 물리 블록을 참조할 수 있다. get_computed_blocks가 이 캐시 히트를 탐색한다.

  3. 동적 메모리 공유: Copy-on-Write 시맨틱으로 빔 서치 등에서 KV 캐시를 효율적으로 공유할 수 있다. block_pool의 reference counting이 이를 지원한다.

  4. 스케줄러와의 긴밀한 연동: allocate_slots가 None을 반환하면 스케줄러가 preemption(OS의 스왑 아웃)을 수행한다. 이 설계 덕분에 메모리 부족 시 graceful degradation이 가능하다.

PagedAttention은 단순히 "블록 단위로 나누자"는 아이디어가 아니라, OS의 수십 년간 검증된 메모리 관리 패턴을 GPU KV 캐시에 충실하게 이식한 설계이다. vLLM이 높은 처리량을 달성할 수 있는 가장 근본적인 이유이다.

논문 핵심 내용

PagedAttention 논문의 핵심 기여는 LLM 서빙에서 KV 캐시의 메모리 낭비를 거의 제로 수준으로 줄인 것이다. 기존 시스템들은 요청의 최대 시퀀스 길이에 맞춰 연속 메모리를 사전 할당하기 때문에, 실제 사용량과의 차이에서 오는 내부 단편화가 심각했다. 논문에서는 기존 시스템의 KV 캐시 메모리 중 60-80%가 단편화와 예약으로 인해 낭비된다고 분석했다. PagedAttention은 OS의 가상 메모리 페이징을 적용하여 이 낭비를 4% 미만으로 줄였다.

벤치마크 결과를 보면 vLLM은 동일 지연시간 조건에서 기존 최고 성능 시스템(FasterTransformer, Orca) 대비 2-4배 높은 처리량을 달성했다. 특히 시퀀스가 길어질수록, 모델이 커질수록, 빔 서치 같은 복잡한 디코딩 알고리즘을 사용할수록 개선 폭이 더 커진다. Copy-on-Write 메커니즘 덕분에 빔 서치에서 빔 간 KV 캐시 공유가 가능해져, 빔 크기에 비례하던 메모리 사용량이 크게 줄었다.

비교 항목 기존 시스템 vLLM (PagedAttention)
KV 캐시 메모리 낭비 60-80% 4% 미만
처리량 (vs FasterTransformer) 1x 2-4x
메모리 공유 (빔 서치) 불가 Copy-on-Write 지원
프리픽스 캐싱 불가 블록 단위 자연 지원

댓글

관련 포스트

vLLM 의 다른글