vLLM의 PagedAttention: GPU 메모리에 빈자리가 있는데 왜 요청을 못 받을까

언어 모델 서버를 만들 때 가장 먼저 계산하는 숫자는 모델 가중치가 GPU에 들어가는지 여부다. 그런데 모델을 올리는 데 성공한 뒤에도 요청이 조금만 몰리면 처리량이 막힌다. 생성 중인 문장마다 과거 토큰의 정보를 보관해야 하기 때문이다. vLLM의 PagedAttention 논문은 이 두 번째 메모리를 정면으로 다룬다. 신경망 수식을 바꾸기보다 운영체제의 오래된 아이디어를 가져온다는 점이 특히 재미있다.

모델보다 먼저 늘어나는 작업 공간

자기회귀 생성에서는 지금까지 입력된 토큰들의 키와 값 벡터를 KV 캐시에 저장한다. 다음 토큰을 만들 때 과거 토큰들의 K와 V를 처음부터 다시 계산하지 않도록 하는 장치다. 가중치는 여러 요청이 공유하지만, KV 캐시는 요청마다 길이가 다르고 생성이 진행될수록 자란다.

한 사용자가 짧게 답을 받고 끝날 수도 있고, 다른 사용자가 긴 코드를 요청할 수도 있다. 시작할 때 최종 길이를 정확히 알기 어렵다. 이때 요청마다 최대 길이만큼 연속 공간을 미리 확보하면 실제로 쓰지 않는 부분이 생긴다. 필요한 만큼 할당하더라도, 요청의 시작과 종료가 엇갈리면 작은 빈 공간들이 흩어질 수 있다. 논문은 이런 예약 낭비와 단편화, 중복 보관을 배치 크기를 제한하는 원인으로 분석한다. PagedAttention 논문

가상의 서버를 떠올려 보자. 사용 가능한 캐시 용량이 토큰 1,000개분이고, 각 요청은 최대 200개까지 생성할 수 있다고 하자. 실제 평균 생성 길이가 40개여도 요청마다 200개분을 예약하면 동시에 다섯 요청밖에 넣지 못한다. 숫자는 논문 실험이 아니라 설명을 위한 예시다. 요점은 GPU에 실제 정보가 꽉 찬 상태와, 필요할지도 모르는 자리를 비워 둔 상태가 모두 '메모리 부족'으로 나타날 수 있다는 데 있다.

논리적인 순서와 물리적인 위치를 떼어 놓기

PagedAttention은 KV 캐시를 고정 크기의 블록으로 나눈다. 문장 안에서는 첫 블록, 두 번째 블록, 세 번째 블록이라는 순서가 있지만, 실제 GPU 메모리에서는 서로 떨어진 곳에 놓여도 된다. 블록 테이블이 논리적인 번호를 실제 위치로 연결해 준다. 어텐션 커널은 이 표를 따라 필요한 블록을 읽어 계산한다.

서랍에 책을 꽂는 장면으로 바꿔 보면 감이 온다. 한 문서의 페이지를 반드시 연속된 서랍에 넣을 필요 없이 빈 서랍에 나누어 넣고, 목차에 위치를 기록하는 것이다. 새 페이지가 생겨도 옆 서랍이 비어 있는지 기다릴 필요가 없다. 다른 문서가 끝나서 반환한 서랍도 바로 활용할 수 있다.

블록이 꽉 차면 새 블록을 추가하고, 요청이 끝나면 블록들을 반납한다. 남는 자리는 주로 마지막 블록 안으로 제한된다. 블록 크기가 아주 작으면 낭비는 줄지만 관리와 접근 비용이 늘 수 있고, 크면 반대의 문제가 생긴다. 따라서 블록 크기는 '작을수록 좋다'는 단일 규칙보다 커널 효율과 메모리 사용 사이의 선택으로 보는 편이 맞다. 원문 메모리 관리와 블록 크기 분석

여기서 페이지라는 표현 때문에 자동으로 CPU 메모리까지 확장되는 마법을 상상하기 쉽다. 핵심 기여는 GPU 안에서도 논리적 연속성과 물리적 연속성을 분리할 수 있게 한 데 있다. 전체 GPU 용량은 여전히 유한하다. 공간이 부족할 때 어떤 요청을 중단하고 상태를 옮기거나 재계산할지는 스케줄링 문제로 남는다.

같은 프롬프트를 네 번 보관할 이유

하나의 프롬프트에서 답변 후보를 여러 개 생성하는 경우를 생각해 보자. 갈라지기 전까지의 입력은 동일하므로 그 구간의 KV 캐시도 동일하다. 그런데 후보마다 별도의 연속 버퍼를 만들면 프롬프트 부분까지 복제할 수 있다.

vLLM은 여러 시퀀스의 논리 블록이 같은 물리 블록을 가리키게 해서 이 부분을 공유한다. 공유한 블록에 서로 다른 새 값을 써야 하는 순간에는 복사한 뒤 분리한다. 운영체제에서 익숙한 copy-on-write와 닮은 방식이다. 이미 채워진 공통 구간은 유지하고, 분기 이후의 상태만 각자 가진다. 참조 수를 관리해야 안전하게 공유와 반환을 처리할 수 있다. 공식 프로젝트의 공유 메모리 설명

내가 만든 작은 예에서는 800토큰 프롬프트로 200토큰짜리 후보 네 개를 생성한다. 프롬프트를 네 번 보관하면 3,200토큰분의 공통 캐시가 필요하지만, 공유하면 공통 구간은 대체로 800토큰분으로 줄일 수 있다. 마지막 블록 처리와 추가 상태를 무시한 단순한 계산이다. 프롬프트가 길고 분기가 많을수록 공유의 의미가 커진다는 것을 보여 준다.

다만 단어가 겹친다는 이유만으로 캐시를 공유할 수는 없다. 같은 토큰도 앞선 문맥이 달라지면 표현이 달라질 수 있다. 재사용 가능한 접두부는 동일한 조건에서 계산된 동일한 입력 구간이어야 한다. 실제 서비스에서 모델 버전과 입력 구성까지 함께 관리하고 싶은 이유다. 이것은 논문 아이디어를 운영 문제로 옮기며 얻은 내 판단이다.

처리량이 좋아졌다는 말의 정확한 의미

원문은 당시 비교 시스템인 FasterTransformer, Orca 등과 여러 모델 및 요청 길이 조건에서 평가한다. 초록은 비슷한 지연 조건에서 처리량이 2~4배 개선됐다고 요약한다. 이 수치는 논문이 사용한 시스템과 작업 부하에 대한 결과다. 현재 라이브러리들 사이의 보편적인 성능 서열로 읽을 수는 없다. 원문 평가

이 결과에서 눈여겨볼 것은 한 요청의 어텐션 곱셈이 무조건 몇 배 빨라졌다는 주장이 아니라, 같은 GPU 메모리에서 더 많은 요청을 함께 처리할 수 있다는 연결이다. 공간을 효율적으로 쓰면 배치에 들어갈 수 있는 시퀀스가 늘고, 그 결과 장치 활용률과 전체 처리량이 올라갈 여지가 생긴다. 논문은 커널 수준의 비용도 따로 분석한다.

처리량과 응답성은 함께 재야 한다. 요청을 더 모아서 처리하면 초당 생성 토큰 수는 좋게 보이지만 첫 토큰을 받는 시간이 길어질 수 있다. 긴 요청이 섞이면 마지막까지 걸리는 시간도 달라진다. 내가 서비스를 평가한다면 평균만 보기보다 첫 토큰 지연, 토큰 사이 지연, 긴 꼬리의 응답 시간을 같이 보겠다. 사용자가 기다리는 경험은 하나의 토큰 처리량 숫자로 설명되지 않는다.

메모리 효율은 스케줄러의 자유도다

이 논문을 읽으며 가장 오래 남은 부분은 남는 메모리가 단순한 여유 공간 이상의 의미를 가진다는 점이었다. 요청을 하나 더 받을 수 있고, 여러 후보를 비교할 수 있고, 긴 입력도 기존 요청들과 함께 처리할 수 있다. 메모리 관리가 좋아지면 서버가 선택할 수 있는 운영 방식이 넓어진다.

실무에서는 요청 길이 분포와 동시 사용자 수를 먼저 살펴보고 싶다. 짧은 요청 하나만 반복하는 벤치마크에서는 단편화와 공유의 효과가 잘 드러나지 않을 수 있다. 반대로 긴 프롬프트, 다양한 종료 시점, 여러 출력 후보가 섞이면 캐시 관리가 병목인지 분명해진다. 실제 트래픽의 길이 분포를 보존한 재생 실험이 필요한 이유다.

PagedAttention은 언어 모델 서빙을 신경망 추론만의 문제로 바라보지 않게 만든다. 같은 모델도 메모리를 어떻게 나누고 돌려주느냐에 따라 서버의 성격이 바뀐다. GPU를 더 사기 전에, 이미 산 메모리 안에 실제로 무엇이 들어 있는지 들여다볼 가치가 있다.

읽은 논문: Woosuk Kwon 외, Efficient Memory Management for Large Language Model Serving with PagedAttention, SOSP 2023. 논문 · 공식 소개