LLM 서빙 메모리 낭비를 제거하는 vLLM PagedAttention과 연속 배칭(Continuous Batching) 아키텍처

1. 기존 LLM 추론 서빙의 메모리 병목과 단편화 한계

대규모 언어 모델(LLM)을 운영 환경에서 서빙할 때 가장 큰 병목은 모델 파라미터 자체의 크기가 아니라, 추론 과정에서 생성되는 Key-Value(KV) 캐시의 기하급수적인 GPU VRAM 점유입니다. 트랜스포머의 자기주의(Self-Attention) 메커니즘은 이전 토큰들의 연산 결과를 반복 계산하지 않기 위해 KV 캐시를 메모리에 유지해야 합니다.

그러나 전통적인 추론 프레임워크(HuggingFace TGI 초기 버전, FasterTransformer 등)는 다음과 같은 심각한 메모리 단편화 문제를 안고 있었습니다.

• 사전 고정 할당으로 인한 내부 단편화(Internal Fragmentation): 각 요청의 최대 출력 토큰 길이(예: 2,048토큰)에 맞춰 연속적인 가상 메모리 공간을 미리 할당합니다. 실제 요청이 100토큰만 생성하고 조기 종료되면 나머지 95%의 공간이 낭비됩니다.
• 정적 배칭(Static Batching)의 유휴 시간 낭비: 배치 내에서 가장 긴 시퀀스가 종료될 때까지 이미 완료된 짧은 시퀀스들이 GPU 메모리를 점유한 채 대기해야 하므로 GPU 연산 유닛(SM)의 가동률이 급락합니다.
• 예약 낭비(Reservation Waste): 아직 생성되지도 않은 미래 토큰을 위해 미리 공간을 잡아둠으로써, 실제로는 60%~80%에 달하는 GPU VRAM이 허수로 낭비되었습니다.

Kwon et al.의 SOSP 2023 논문에 따르면 기존 시스템의 유효 메모리 사용률은 20%~40%에 불과했습니다.

2. vLLM PagedAttention의 핵심 아키텍처 원리

vLLM 공식 연구(SOSP 2023)는 이 문제를 해결하기 위해 현대 운영체제(OS)의 가상 메모리 페이징(Paging) 메커니즘을 최초로 LLM 서빙에 이식했습니다. 물리 메모리를 비연속적인 고정 크기의 페이지 프레임으로 나누고 페이지 테이블을 통해 매핑하듯, KV 캐시를 고정 크기의 물리 블록(Physical Block, 기본 16 또는 32 토큰) 단위로 분할 관리합니다.

PagedAttention의 구조적 혁신은 다음과 같습니다.

• 논리-물리 블록 매핑(Block Table): 논리적 토큰 시퀀스는 연속적으로 보이지만, 실제 GPU VRAM에는 물리적으로 흩어진 블록에 저장됩니다. 블록 테이블(Block Table)이 논리 블록 번호를 물리 블록 주소로 변환합니다.
• 온디맨드 블록 할당(On-Demand Allocation): 토큰이 실제로 생성되는 순간에만 블록을 할당하며, 마지막 블록의 남은 공간을 제외하면 내부 단편화율을 4% 미만으로 극적으로 억제합니다.
• Copy-on-Write 기반 무손실 메모리 공유: 병렬 샘플링(Parallel Sampling)이나 빔 서치(Beam Search) 시 동일한 프롬프트의 KV 캐시 블록을 복사하지 않고 참조 카운트(Ref Count)만 증가시켜 물리 메모리를 100% 공유합니다. 토큰이 갈라지는 분기점에서만 블록을 복제합니다.

3. 반복 레벨 연속 배칭(Continuous Batching) 스케줄링 메커니즘

PagedAttention의 메모리 효율성을 기반으로, vLLM은 요청 단위가 아닌 토큰 생성 반복 단위(Iteration-Level)로 배치를 유동적으로 재구성하는 연속 배칭(Continuous Batching) 스케줄러를 가동합니다.

• 조기 완료 즉시 퇴거: 특정 시퀀스가 EOS 토큰을 출력하거나 최대 길이에 도달하면, 해당 반복(Step)이 끝나는 즉시 KV 캐시 블록을 해제하여 풀(Free Pool)로 반환합니다.
• 신규 요청 즉시 인입: 다른 시퀀스들이 디코딩(Decoding) 단계에 있더라도, 대기열(Queue)에 있는 신규 요청을 프리필(Prefill)하여 동일한 포워드 패스 배치에 즉시 합류시킵니다.
• 선점(Preemption) 및 재계산: 메모리가 일시적으로 고갈될 경우, 우선순위가 낮은 시퀀스를 선점하여 블록을 회수한 뒤 스왑 공간(CPU RAM)으로 이동시키거나 추후 프리필을 재실행합니다.

4. 실전 파이프라인 구현: PagedAttention 블록 할당자 및 연속 배칭 엔진

다음은 vLLM의 핵심 메커니즘인 블록 할당자(Block Allocator), 블록 테이블 매핑, 그리고 반복 레벨 연속 배칭 스케줄러를 파이썬으로 완벽히 재현한 아키텍처 구현체입니다.

import time
import math
import logging
from typing import List, Dict, Optional, Tuple
from dataclasses import dataclass

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] (%(name)s) %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger('PagedAttentionScheduler')

@dataclass
class PhysicalBlock:
    block_id: int
    ref_count: int = 0
    token_ids: List[int] = None

    def __post_init__(self):
        if self.token_ids is None:
            self.token_ids = []

class BlockAllocator:
    def __init__(self, num_blocks: int, block_size: int = 16):
        self.block_size = block_size
        self.free_blocks: List[PhysicalBlock] = [PhysicalBlock(i) for i in range(num_blocks)]
        self.allocated_blocks: Dict[int, PhysicalBlock] = {}

    def allocate(self) -> PhysicalBlock:
        if not self.free_blocks:
            raise MemoryError('GPU KV 캐시 블록 고갈 (Out of Memory)')
        block = self.free_blocks.pop(0)
        block.ref_count = 1
        block.token_ids = []
        self.allocated_blocks[block.block_id] = block
        return block

    def free(self, block_id: int):
        if block_id in self.allocated_blocks:
            block = self.allocated_blocks.pop(block_id)
            block.ref_count = 0
            block.token_ids.clear()
            self.free_blocks.append(block)

class SequenceContext:
    def __init__(self, seq_id: str, prompt_tokens: List[int], max_tokens: int):
        self.seq_id = seq_id
        self.tokens: List[int] = list(prompt_tokens)
        self.max_tokens = max_tokens
        self.block_table: List[int] = []
        self.is_finished: bool = False

class ContinuousBatchingEngine:
    def __init__(self, num_blocks: int = 128, block_size: int = 16):
        self.allocator = BlockAllocator(num_blocks, block_size)
        self.running_seqs: List[SequenceContext] = []
        self.waiting_queue: List[SequenceContext] = []
        self.block_size = block_size

    def add_request(self, seq_id: str, prompt_tokens: List[int], max_tokens: int):
        seq = SequenceContext(seq_id, prompt_tokens, max_tokens)
        self.waiting_queue.append(seq)
        logger.info(f'신규 추론 요청 대기열 인입: {seq_id} (입력 길이: {len(prompt_tokens)})')

    def _allocate_blocks_for_seq(self, seq: SequenceContext, num_tokens_to_add: int = 1):
        total_tokens = len(seq.tokens) + num_tokens_to_add
        needed_blocks = math.ceil(total_tokens / self.block_size)
        current_blocks = len(seq.block_table)

        while current_blocks < needed_blocks:
            new_block = self.allocator.allocate()
            seq.block_table.append(new_block.block_id)
            current_blocks += 1

    def step(self) -> Dict[str, Any]:
        while self.waiting_queue and len(self.allocator.free_blocks) > 4:
            next_seq = self.waiting_queue.pop(0)
            self._allocate_blocks_for_seq(next_seq, num_tokens_to_add=0)
            self.running_seqs.append(next_seq)
            logger.info(f'연속 배칭 풀(Pool) 승격: {next_seq.seq_id}')

        generated_tokens = {}
        finished_ids = []

        for seq in self.running_seqs:
            self._allocate_blocks_for_seq(seq, num_tokens_to_add=1)
            next_token = (seq.tokens[-1] + 1) % 50000
            seq.tokens.append(next_token)
            generated_tokens[seq.seq_id] = next_token

            curr_block_id = seq.block_table[-1]
            self.allocator.allocated_blocks[curr_block_id].token_ids.append(next_token)

            if len(seq.tokens) >= seq.max_tokens:
                seq.is_finished = True
                finished_ids.append(seq.seq_id)

        for fid in finished_ids:
            finished_seq = next(s for s in self.running_seqs if s.seq_id == fid)
            for b_id in finished_seq.block_table:
                self.allocator.free(b_id)
            self.running_seqs.remove(finished_seq)
            logger.info(f'시퀀스 완료 및 KV 캐시 블록 해제: {fid} (최종 길이: {len(finished_seq.tokens)})')

        return {
            'active_batch_size': len(self.running_seqs),
            'waiting_size': len(self.waiting_queue),
            'free_blocks': len(self.allocator.free_blocks),
            'generated_tokens': generated_tokens
        }

if __name__ == '__main__':
    engine = ContinuousBatchingEngine(num_blocks=32, block_size=16)
    engine.add_request('req_01', [101, 102, 103], max_tokens=6)
    engine.add_request('req_02', [201, 202], max_tokens=4)

    for i in range(5):
        step_stat = engine.step()
        logger.info(f'Iteration {i+1}: 활성 배치 {step_stat.get('active_batch_size')}, 잔여 블록 {step_stat.get('free_blocks')}')

5. 성능 벤치마크 및 프로덕션 도입 가이드

Llama 3 70B 모델을 단일 H100 GPU 노드에서 동시 요청 64개 환경으로 벤치마크를 수행한 결과는 다음과 같습니다.

프로덕션 AI 인프라에 vLLM을 도입할 때는 다음 아키텍처 튜닝을 필수적으로 적용해야 합니다.

• 블록 크기(Block Size) 튜닝: 블록 크기가 8~16이면 내부 단편화가 최소화되지만 커널 런처 오버헤드가 증가합니다. 32~64는 긴 문맥(Long Context) 서빙 시 GPU 메모리 대역폭 효율을 극대화하므로, 평균 프롬프트 길이에 맞춰 16 또는 32를 기본값으로 권장합니다.
• GPU 메모리 활용도(gpu_memory_utilization) 상한 설정: 모델 가중치 로드 후 KV 캐시 풀로 할당할 VRAM 비율을 기본 0.90으로 지정하여, 프리필 연산 중 발생하는 활성화 텐서(Activation Tensor)에 의한 OOM을 사전 차단해야 합니다.
• 청크 단위 프리필(Chunked Prefill) 활성화: 매우 긴 프롬프트가 유입될 때 기존 디코딩 배치들의 연산 지연(Inter-token Latency Spike)을 방어하기 위해 --enable-chunked-prefill 옵션을 활성화하여 프리필 토큰을 일정 크기로 분할 연산하도록 설정해야 합니다.