LLM 서빙에서는 GPU 연산 성능뿐 아니라 메모리 용량과 대역폭이 처리량을 제한한다. 특히 요청마다 커지는 KV 캐시를 어떻게 할당하고 공유하는지가 동시 처리 가능한 요청 수에 영향을 준다.

이 글은 PagedAttention 논문에서 제안한 vLLM의 설계와 실험을 정리한다. 구현 규모와 성능 수치는 논문 당시의 조건에 해당하며, 현재의 모든 vLLM 버전과 백엔드에 그대로 적용되는 값은 아니다.

자기회귀 생성과 KV 캐시

자기회귀 언어 모델은 이전 토큰을 조건으로 다음 토큰의 확률을 계산한다. causal self-attention에서는 현재 위치의 쿼리(Q)가 현재와 과거 위치의 키(K)를 참조하고, 어텐션 점수로 값(V)의 가중 합을 구한다. 점수는 q_iᵀ k_j / √d_k를 유효한 과거 위치에 대해 softmax 정규화한 값이다.

생성 단계마다 과거 토큰의 K·V를 다시 계산하는 대신 저장해 재사용하는 것이 KV 캐시다. 같은 토큰도 위치와 문맥에 따라 값이 달라지므로, 동일한 단어라는 이유만으로 캐시를 공유할 수는 없다.

LLM 추론은 두 단계로 나눌 수 있다.

  • Prefill: 입력 프롬프트를 처리하고 초기 KV 캐시를 만든다. 주어진 토큰들을 병렬로 처리할 수 있으며, 큰 행렬 연산으로 GPU를 활용한다.
  • Decode: 이전 결과에 의존해 토큰을 하나씩 생성한다. 과거 KV 캐시를 읽고 새 토큰의 K·V를 추가한다. 한 시퀀스의 생성 단계를 동시에 진행할 수는 없지만, 단계 내부와 서로 다른 요청 사이에서는 병렬 처리가 가능하다.

Decode는 작은 배치에서 연산량에 비해 가중치와 KV 캐시를 읽는 양이 많아 메모리 대역폭의 영향을 크게 받는다. Prefill과 Decode의 병목은 모델, 배치 크기, 시퀀스 길이에 따라 달라진다.

배칭과 메모리 용량

여러 요청을 묶으면 가중치 접근 비용을 공유하고 GPU 연산 활용도를 높일 수 있다. 그러나 요청 도착 시점과 입출력 길이가 서로 달라 정적 배칭에서는 대기와 패딩 비용이 생긴다. 요청 단위 스케줄링을 개선하더라도 동시에 유지할 KV 캐시의 용량은 여전히 제약이다.

논문은 OPT-13B의 긴 시퀀스에서 요청 하나의 KV 캐시가 최대 약 1.6GB에 이를 수 있음을 설명한다. 이 크기는 모델 구조, 시퀀스 길이, 데이터 정밀도에 따라 달라진다.

출력 길이를 알 수 없는 상태에서 최대 길이만큼 연속 메모리를 예약하면 다음 비용이 생긴다.

  • 아직 생성되지 않은 토큰을 위한 예약 공간이 필요하다.
  • 실제 출력이 짧으면 할당한 공간 일부를 사용하지 못한다.
  • 요청별 할당 사이에 남은 틈을 다른 요청이 활용하기 어려울 수 있다.

논문에서 분석한 기존 시스템은 KV 캐시용 공간 중 실제 토큰에 사용하는 비율이 약 20–40%였다. 이 결과는 논문의 비교 시스템과 워크로드에 대한 관찰이다.

PagedAttention의 설계

PagedAttention은 운영체제의 페이징에서 아이디어를 얻어 KV 캐시를 고정 크기 블록으로 나눈다. 요청의 논리적 토큰 순서와 GPU의 물리적 저장 위치를 분리하고, 블록 테이블로 둘을 연결한다.

스케줄러는 실행할 요청을 선택하고, 블록 매니저는 필요한 메모리 블록을 할당하고 회수한다. 요청이 늘어날 때 실제로 필요한 블록을 추가하므로 최대 길이 전체를 미리 확보할 필요가 줄어든다.

어텐션 커널은 블록 테이블을 따라 K·V를 가져온다. 예를 들어 한 시퀀스의 앞부분과 뒷부분이 서로 다른 물리 블록에 있어도 논리적인 순서대로 참조할 수 있다. 모든 유효 토큰에 걸친 softmax 정규화를 유지하며 결과를 결합한다. 블록마다 독립적으로 정규화한 값을 단순히 더하는 방식은 아니다.

이 설계는 어텐션의 의미를 유지하면서 저장 위치의 연속성 제약을 완화한다. 다만 마지막 블록에는 사용하지 않는 슬롯이 남을 수 있고, 블록 테이블 조회 비용도 추가된다.

블록 공유

병렬 샘플링과 빔 서치에서는 여러 출력 시퀀스가 같은 프롬프트에서 시작한다. 이 경우 공통 부분의 KV 블록을 공유하고, 수정이 필요한 시점에 복사하는 방식으로 메모리를 절약할 수 있다.

공유 가능성은 토큰과 문맥이 같은지에 달려 있다. 메모리 공유는 결과가 같은 부분에만 적용해야 하며, 단순히 문자열 일부가 같다는 이유로 임의의 KV 캐시를 재사용할 수는 없다.

구현과 실험 환경

논문 당시 엔진은 약 8,500줄의 Python 코드와 2,000줄의 C++/CUDA 코드로 구현되었다. 스케줄러와 블록 관리는 Python에서, 핵심 어텐션 연산은 CUDA 커널에서 처리했다.

평가는 NVIDIA A100 GPU에서 OPT-13B·66B·175B와 LLaMA-13B를 사용했다. ShareGPT와 Alpaca를 바탕으로 입출력 길이가 다른 워크로드를 구성하고, 요청 도착은 포아송 분포로 모델링했다.

비교 대상은 FasterTransformer와 논문 저자들이 재구현한 Orca다. 주요 지표는 요청률을 높일 때의 정규화된 지연 시간과 처리량이다. 아래 결과는 동일한 모델과 비교 설정에서 읽어야 한다.

주요 결과

  • 기본 샘플링의 ShareGPT 실험에서는 비슷한 지연 수준에서 Orca(Oracle) 대비 1.7–2.7배, Orca(Max) 대비 2.7–8배 높은 요청률을 처리했다.
  • 병렬 샘플링과 빔 서치에서는 공통 프롬프트의 블록 공유로 메모리 사용을 줄였다.
  • 공유 프리픽스 실험에서는 Orca 대비 1.67–3.58배 높은 처리량을 보고했다.
  • 챗봇 시나리오에서는 메모리 낭비를 줄여 Orca 대비 약 2배 높은 요청률을 처리했다.

이 수치를 단일 어텐션 커널의 속도 향상으로 해석하면 안 된다. 메모리 활용이 좋아지면서 더 많은 요청을 함께 처리할 수 있게 된 시스템 수준의 이점이다.

커널 비용과 블록 크기

PagedAttention 커널은 비교 대상인 FasterTransformer 커널보다 약 20–26% 높은 지연을 보였다. 블록 조회 등 추가 작업이 있기 때문이다. 그럼에도 더 큰 배치로 요청을 처리하면서 전체 서빙 처리량에서는 이점을 얻었다.

논문에서는 블록 크기 16이 메모리 낭비와 GPU 활용 사이에서 좋은 균형을 보였다. 블록이 너무 크면 마지막 블록의 낭비가 커지고, 너무 작으면 관리와 접근 비용이 증가할 수 있다. 최적 크기는 워크로드와 커널 구현에 따라 달라진다.

적용 범위와 관련 연구

이 접근은 길이가 변하는 시퀀스를 동시에 처리하고 KV 캐시 용량이 병목이 되는 상황에 적합하다. 텐서 크기가 정적이거나 연산량이 주요 병목인 다른 GPU 작업에서도 동일한 이점을 얻는다고 보장할 수는 없다.

논문은 메모리가 부족할 때 CPU 메모리로 시퀀스를 이동하는 스왑과, 나중에 KV 캐시를 다시 계산하는 재연산을 함께 다룬다. 데이터를 다시 가져오는 비용과 연산 비용을 비교해 복구 방식을 선택한다.

Orca의 스케줄링과 vLLM의 메모리 관리는 서로 다른 병목을 다룬다. FlexGen은 제한된 GPU에서 오프로딩을 활용한 배치 추론에, FlashAttention은 어텐션 연산의 메모리 접근 비용에 초점을 둔다. 어떤 기술이 항상 우수한지보다 해결하려는 문제와 평가 조건을 비교해야 한다.

정리

vLLM의 핵심은 KV 캐시의 저장 위치를 논리적인 시퀀스와 분리해 필요한 만큼 할당하고 공유하는 데 있다. 커널 하나의 지연을 줄이는 것보다 메모리 활용을 개선해 동시에 처리할 수 있는 요청을 늘린다는 관점에서 이해할 수 있다.

참고 자료