Paper review

vAttention: Dynamic Memory Management for Serving LLMs without PagedAttention

임로켓 2025. 3. 11. 21:08
728x90

https://arxiv.org/abs/2405.04437

 

vAttention: Dynamic Memory Management for Serving LLMs without PagedAttention

PagedAttention is a popular approach for dynamic memory allocation in LLM serving systems. It enables on-demand allocation of GPU memory to mitigate KV cache fragmentation -- a phenomenon that crippled the batch size (and consequently throughput) in prior

arxiv.org

 

서론

대규모 언어 모델(LLM)을 서비스할 때 메모리 관리는 시스템 성능과 안정성에 핵심적인 요소입니다. 특히 각 요청(Request)에 대한 KV 캐시는 디코딩 과정 내내 과거 토큰들의 정보를 저장하여 반복 연산을 줄여주지만, 이 KV 캐시가 차지하는 GPU 메모리를 효율적으로 관리하지 못하면 대용량 배치(batch) 처리가 어려워집니다​. 기존에는 모델 최대 컨텍스트 길이만큼 메모리를 미리 확보(Reservation)해 두는 방식이 사용되었으나, 실제 생성되는 토큰 수는 그보다 훨씬 적은 경우가 많아 내부 단편화(internal fragmentation)로 메모리 낭비가 발생했습니다​. 이로 인해 배치 크기를 키우지 못하고 처리량이 저하되는 문제가 있었죠. 이를 개선하기 위해 등장한 기법이 PagedAttention입니다. PagedAttention은 OS의 수요 페이징(demand paging) 아이디어를 차용하여 KV 캐시를 작은 블록 단위로 나누어 필요할 때마다 동적으로 GPU 메모리를 할당하는 방식입니다​. 이 접근법은 미리 메모리를 예약하지 않고도 KV 캐시 단편화를 크게 완화하여, 한정된 GPU 메모리로도 더 많은 요청을 동시에 처리할 수 있게 만들었습니다​. 실제로 현재 HuggingFace TGI, TensorRT-LLM, LightLLM 등 여러 LLM 추론 프레임워크들이 PagedAttention 모델을 도입하여 대용량 배치를 지원하고 있습니다​.

 

그럼에도 불구하고, PagedAttention에는 몇 가지 한계점과 부작용이 존재합니다. 첫째, KV 캐시를 여러 메모리 블록으로 관리하면서 가상 메모리 주소가 불연속적으로 변하기 때문에, 기존 연속 메모리를 가정하고 작성된 어텐션 커널 코드를 수정해야 합니다​. 예를 들어 FlashAttention과 같은 고성능 어텐션 구현에 PagedAttention을 적용하려면, KV 캐시 블록들의 주소를 따라가도록 커널을 변경해야 하며 이는 신규 최적화 기법마다 반복되어야 합니다​. 실제로 vLLM의 PagedAttention 커널은 원본 FlashAttention 커널보다 최대 2.85배까지 느리게 동작하며​, FlashAttention 측에서도 페이징 지원을 추가하려다 유닛 테스트를 통과하지 못한 사례가 보고되었습니다​. 둘째, 사용자 공간에 별도의 메모리 관리자를 구현해야 합니다. 여러 개의 KV 캐시 블록을 할당하고 추적하기 위한 블록 테이블(block table)을 운영해야 하며, 커널에 매 iteration마다 해당 블록들의 위치 정보를 전달하는 등 운영체제가 하는 일을 애플리케이션이 중복 수행하게 됩니다​. 이는 코드 복잡도를 높이고 유지보수를 어렵게 만드는 요인입니다. 셋째, 추론 경로 상에 추가적인 CPU/GPU 오버헤드가 발생합니다. 어텐션 연산 시 불연속적인 여러 메모리 조각을 모아 참조하기 위한 주소 변환 로직이 GPU 커널에 들어가고, 매 iteration마다 블록 테이블을 생성/조회하는 CPU 부담도 생깁니다​. 연구에 따르면 이러한 이유로 PagedAttention을 사용할 때 어텐션 연산 속도가 10% 이상 감소하고, CPU 측에서도 최대 10% 가량의 오버헤드가 추가될 수 있다고 합니다​. 실제 사례로, PagedAttention을 적용한 FlashInfer에서는 첫 토큰 생성 지연(Time-to-first-token)이 45%까지 증가한 보고도 있습니다​. 정리하면 PagedAttention은 메모리 단편화 문제는 해결했지만, 그 대가로 성능 저하와 높은 구현 복잡도라는 과제를 남긴 셈입니다.

vAttention 개요

이러한 한계점을 해결하기 위해 Microsoft Research 등에서 제안한 것이 vAttention입니다​. vAttention의 핵심 아이디어는 KV 캐시의 가상 메모리 연속성은 유지하면서도 물리 메모리는 동적으로 할당함으로써, PagedAttention처럼 사용자 코드에서 페이징을 흉내낼 필요 없이 운영체제/드라이버 수준의 가상 메모리 관리 기능을 활용하자는 것입니다​. 구체적으로, CUDA 가상 메모리 관리(VMM) API를 이용하여 가상 메모리 주소 공간을 대용량으로 미리 연속적으로 예약만 해 두고, 실제 GPU 물리 메모리 할당은 필요할 때마다 수행합니다​. 이렇게 하면 KV 캐시는 논리적으로 하나의 거대한 연속 배열처럼 다룰 수 있어 기존 어텐션 커널 코드를 그대로 재사용할 수 있습니다. 개발자는 더 이상 페이징을 고려한 특별한 커널이나 메모리 관리 코드를 작성하지 않아도 되고, 기존의 최적화된 GPU 커널(Fused kernels, FlashAttention 등)을 그대로 활용하면서 동적 메모리 할당의 이점을 얻을 수 있게 됩니다​. vAttention 논문에서는 이러한 접근법을 통해 소프트웨어 복잡도를 크게 줄이고 이식성(portability)을 높였으며, 성능 또한 향상되었음을 보여줍니다​. 예를 들어, vAttention은 vLLM 대비 최대 1.97배 빠르게 토큰을 생성하고, FlashAttention/FlashInfer의 PagedAttention 버전에 대비해서도 각각 최대 3.92배, 1.45배까지 프롬프트 처리 속도를 높였다고 보고합니다​. vAttention의 기술적 구현을 한마디로 요약하면 “CUDA의 가상 메모리 기능을 활용한 KV 캐시 동적 메모리 관리”입니다​. CUDA 11 이상에서 제공되는 저수준 가상 메모리 API (예: cuMemAddressReservecuMemCreatecuMemMap 등)를 사용하면, 개발자가 가상 메모리 주소 공간 할당과 물리 메모리 할당을 분리해서 제어할 수 있습니다​. vAttention은 이 API들을 활용하여, KV 캐시용 큰 가상 메모리 버퍼를 미리 확보만 해 두고(물리 메모리 비할당), 런타임에 실제 필요한 만큼의 물리 페이지를 온디멘드로 매핑하는 방식을 취합니다​. 이를 통해 KV 메모리 레이아웃은 항상 연속적으로 유지되므로 어텐션 연산시 일반 배열처럼 접근하면 되고, 동시에 메모리 낭비도 없도록 필요한 시점에만 물리 메모리를 붙여주는 것입니다. 요약하면 vAttention은 PagedAttention의 장점(동적 할당)은 취하고 단점(불연속 메모리로 인한 비효율)은 OS 레벨 기능으로 대체한 솔루션이라 볼 수 있습니다. 다음 섹션에서는 PagedAttention과 vAttention의 KV 캐시 관리 방법 및 성능을 자세히 비교해보겠습니다.

PagedAttention vs vAttention 비교

KV 캐시 관리 방식: PagedAttention과 vAttention의 가장 큰 차이는 KV 캐시 메모리의 주소 관리입니다. PagedAttention에서는 각 요청의 KV 캐시를 여러 개의 작은 메모리 블록들로 나누어 관리하며, 이 블록들은 가상 주소상에서도 서로 떨어져 있을 수 있습니다​. 따라서 블록들의 목록(블록 테이블)을 유지하고, 어텐션 연산마다 해당 요청이 보유한 블록들을 일일이 찾아가서 데이터에 접근해야 합니다​. 반면 vAttention에서는 각 요청(reqID) 별로 KV 캐시용 연속 가상 메모리 영역을 통째로 할당합니다. 물리 메모리는 처음엔 붙지 않지만 가상 주소 공간은 비어있는 큰 배열처럼 잡아두고, 거기에 필요한 페이지만 순차적으로 매핑해 나갑니다​. 이때 vAttention은 요청별 KV 캐시의 현재 사용량을 추적하여, 마지막 할당한 페이지가 꽉 찰 때 새로운 물리 메모리 페이지를 할당합니다​. 결국 KV 캐시의 가상 주소는 끝까지 연속성을 가지므로, 어텐션 커널 입장에서는 마치 고정 크기의 거대한 배열에 접근하듯이 동작하면 됩니다. 이는 PagedAttention처럼 특별한 주소 변환 로직이나 블록 테이블을 거칠 필요가 없음을 의미합니다​.

 

어텐션 커널 및 최적화 관점: 앞서 언급했듯 PagedAttention은 GPU 어텐션 커널을 재작성해야 하는 부담이 있습니다​.  이로 인해 새롭게 제안되는 고성능 커널 최적화를 곧바로 적용하기 어렵고, 적용하더라도 추가 오버헤드가 발생할 수 있습니다​. 이는 페이지 처리 과정에서 추가 명령어 실행(7~13% 증가) 및 CPU 단의 블록 테이블 관리 오버헤드가 누적되기 때문입니다​. 반면 vAttention은 어텐션 커널을 수정하지 않고 그대로 사용하므로, FlashAttention이나 FasterTransformer 같은 고성능 커널의 이점을 100% 활용할 수 있습니다​. 또한 PagedAttention처럼 커널 내부에 페이지 접근을 위한 분기나 메모리 간접 참조가 없어 순수 연산 성능도 더 높게 유지됩니다. 요약하면 vAttention은 구현의 용이성과 실행 효율 면에서 모두 PagedAttention보다 우수하며, 특히 최신 최적화 기법과의 호환성 측면에서 큰 장점을 가집니다.

 

성능 및 메모리 효율 비교: 두 기법 모두 메모리 단편화를 해결하여 대용량 배치를 가능케 한다는 공통점이 있지만, 성능 지표에서는 차이를 보입니다. vAttention 논문의 평가에 따르면, 프리필(prefill) 단계와 디코드(decode)단계 모두에서 vAttention이 PagedAttention 대비 높은 처리량을 보여주었습니다. 프리필 단계에서는 입력 프롬프트 토큰들을 병렬로 처리하는데, vAttention은 FlashAttention 및 FlashInfer의 기존 프리필 커널을 그대로 활용하여 PagedAttention 버전 대비 일관되게 더 높은 토큰 처리 속도를 기록했습니다​. 특히 프롬프트 길이가 길어질수록(어텐션 연산 비중이 커질수록) 성능 차이가 커져, 예컨대 Yi-6B 모델에서 컨텍스트 192K 토큰을 프리필할 때 FlashAttention의 Paged 커널보다 처리율이 3.9배 높았다고 보고됩니다​. 한편 짧은 컨텍스트의 경우에는 어텐션 계산 이외의 선형 연산이 지배적이어서 FlashAttention 기반에서는 Paged와 vAttention 간 차이가 적었지만, FlashInfer와 같이 PagedAttention 적용 시 부가 오버헤드가 큰 커널의 경우 짧은 프롬프트에서도 vAttention이 약 17~36% 가량 더 높은 프리필 처리량을 보였습니다​.

 

디코드 단계는 각 요청당 한 토큰씩 순차 생성하기 때문에 메모리 병목이 두드러지며, vAttention의 이점이 더욱 부각됩니다. 긴 컨텍스트(예: 16K 토큰 입력)에 대해 배치 크기를 변화시키며 성능을 측정한 결과, vAttention은 모든 배치 크기에서 vLLM의 PagedAttention 커널보다 높은 디코드 처리량을 기록했습니다. 구체적으로, vLLM 대비 최대 약 1.5~2배 가량 더 많은 출력 토큰을 초당 생성할 수 있었으며, FlashAttention를 PagedAttention으로 개조한 커널(블록 크기 128 사용)과는 유사한 성능을 보였습니다​. 중요한 점은 배치 크기가 커질수록 vAttention의 상대적 이점이 커진다는 것인데​, 이는 배치 내 총 토큰 수가 증가할수록 어텐션 연산이 차지하는 비중이 높아져서 효율적인 커널의 효과가 극대화되기 때문입니다. 예를 들어 Yi-6B 모델에서 배치 1→8로 늘리면 vLLM 대비 처리량 향상이 1.x배에서 1.7배 이상으로 커졌습니다​. 흥미로운 점은 FlashAttention에 큰 블록 크기(128)를 사용한 Paged 커널은 vAttention과 비슷한 성능을 낼 수 있었는데, 이는 블록 크기를 키우면 GPU상의 L1 캐시 히트율이 높아져 일부 오버헤드를 상쇄하기 때문이라고 합니다​. 그러나 이 경우에도 여전히 별도의 페이징 지원 커널을 구현해야 한다는 점에서, 동일한 성능을 달성하면서도 구현 부담이 없는 vAttention의 우수성이 부각됩니다​.

기술적 세부 사항

이 절에서는 vAttention이 실제로 KV 캐시 메모리를 관리하는 방법과, 성능 최적화를 위한 기술적 디테일을 살펴보겠습니다. 핵심 구성요소는 가상/물리 메모리 분리 할당요청 단위의 KV 캐시 인덱싱, 그리고 내부 단편화 해소를 위한 최적화입니다.

 

가상 메모리 사전 예약 & 물리 메모리 온디멘드 할당: vAttention 초기화 시, 서비스 프레임워크는 모델의 레이어별로 K캐시와 V캐시용 대용량 가상 메모리 버퍼를 미리 확보합니다​. 각 버퍼의 크기는 최대 배치 크기 B × 최대 컨텍스트 길이 L × (레이어당 토큰당 KV 캐시 메모리)로 산정되며, 이를 통해 최악의 경우에도 모든 요청의 KV를 담을 수 있는 가상 주소 공간이 마련됩니다​. 예를 들어, 60억 매개변수짜리 모델(Yi-6B)의 경우 한 토큰의 KV 캐시가 대략 64KB 정도이고 레이어 수 32개라 하면, 컨텍스트 16K 토큰, 배치 8 기준으로 약 수백 MB의 가상 공간이 필요합니다 (현대 64비트 시스템에서는 수 테라바이트 수준의 가상 주소 공간이 사용 가능하므로 문제없음). 이렇게 확보된 연속 가상 메모리는 아직 물리 메모리가 결합되지 않은 빈 공간이지만, vAttention은 이를 각 요청별 KV 캐시 “슬롯”으로 활용합니다. 새로운 요청이 들어오면 vAttention에서 고유 reqId를 할당하고 해당 요청의 KV용 가상 메모리 범위를 지정해줍니다​. 이제 추론이 진행되면서 그 요청의 KV 캐시가 커질 때마다 vAttention은 필요한 만큼의 GPU 메모리 페이지를 실제 할당하여 해당 가상 주소 영역에 매핑합니다​.  이때 한 iteration(한 토큰 디코드)에서 각 요청이 새로 필요로 하는 메모리는 많아야 1페이지에 불과하기 때문에, vAttention은 “현재 페이지가 꽉 찼을 때만” 다음 물리 페이지를 매핑하는 식으로 동작합니다​. 덕분에 물리 메모리는 최소한으로 사용되고, 사용되지 않는 부분은 여전히 가상 주소로만 존재하게 됩니다.

 

요청 수준 KV 캐시 인덱싱: vAttention에서 각 요청은 reqId로 식별되며, KV 캐시도 reqId별로 논리적으로 분리됩니다​. 서비스 프레임워크는 매 스텝(step)마다 현재 진행 중인 요청들의 reqId와 해당 현재 컨텍스트 길이를 vAttention에 전달하고, vAttention은 이 정보를 바탕으로 각 요청에 필요한 물리 메모리가 충분히 매핑되었는지 확인합니다​. 만약 어떤 요청의 KV 캐시가 다음 토큰을 위해 추가 메모리가 필요하면, 즉시 새로운 페이지를 할당하여 매핑한 뒤 실행을 진행합니다. 반대로 요청이 완료되면(free_reqid 호출) 해당 요청의 KV 캐시에 매핑되었던 물리 메모리를 바로 회수하지 않고 일시 보류할 수 있습니다​. 이렇게 “deferred reclamation”, 즉 지연 해제를 하는 이유는, 곧 이어서 새로운 요청이 들어올 때 그 메모리를 재활용하기 위함입니다. 실제로 R1 요청이 끝나고 바로 R2 요청이 들어오는 경우, vAttention은 R1의 reqId를 R2에 재할당하여 이미 물리 메모리가 붙어있는 KV 캐시 영역을 그대로 활용합니다​. R2의 프롬프트 길이가 R1이 쓰던 것보다 길지만 않다면 추가 할당 없이 바로 사용이 가능하고, 필요 시에만 추가 페이지를 붙입니다. 이처럼 reqId 재사용을 통해 메모리 할당 호출 자체를 줄이는 기법이 vAttention에 내장되어 있습니다. 더 나아가 “즉각 할당(eager allocation)” 최적화로, 항상 예비용 가상 KV 영역 하나를 선정해 약간의 물리 메모리를 미리 매핑해두는 방법도 사용합니다​. 이렇게 하면 새로운 요청이 들어올 때 이미 할당된 페이지를 가진 reqId를 바로 배정할 수 있어, 초반 프리필 단계의 메모리 할당 지연을 최소화할 수 있습니다. 이러한 동적 메모리 관리 작업(배경 페이지 할당, 메모리 회수 등)은 vAttention 내부의 백그라운드 스레드가 전담하여 수행하며, GPU 연산과 최대한 병렬로 진행되도록 설계되었습니다​.

 

내부 단편화 및 메모리 효율: vAttention은 물리 메모리를 페이지 단위로 할당하는데, NVIDIA CUDA 드라이버의 기본 최소 페이지 크기가 2MB라면 작은 요청의 경우에도 2MB를 통째로 할당해야 해 낭비되는 공간이 생길 수 있습니다​

. 예를 들어 2MB 페이지는 모델 설정에 따라 대략 1024 토큰 분량의 KV를 담을 수 있는데, 만약 어떤 요청이 100토큰으로 끝나버리면 마지막 할당된 2MB 중 상당 부분이 사용되지 않고 남습니다. 이를 해결하기 위해 연구진은 오픈소스 NVIDIA UVM 드라이버를 수정하여 더 작은 페이지 크기(64KB, 128KB, 256KB)를 지원하도록 했습니다​. vAttention은 설정에 따라 이 중 하나의 페이지 크기를 사용할 수 있으며, 기본값으로 256KB를 권장합니다​. 256KB 페이지만 해도 2MB 대비 1/8 크기이므로 단편화로 인한 낭비를 크게 줄일 수 있습니다. 논문의 분석에 따르면, 64KB 페이지를 사용할 경우 요청당 최대 낭비 메모리가 몇 MB 수준으로 떨어져 2MB 페이지 대비 최대 95% 이상 절약되는 것으로 나타났습니다​. 예컨대 Yi-34B 모델을 2-way 텐서 병렬로 구동할 때 2MB 페이지는 요청 하나가 최대 240MB까지 쓰지 않은 메모리를 남길 수 있었지만, 64KB 페이지에서는 약 15MB 이하로 감소합니다​.

 

즉 vAttention은 가상 메모리 낭비는 개의치 않되, 물리 메모리 낭비는 최소화하는 방향으로 설계되어 있습니다. 또한 CUDA VMM API 호출 자체의 지연도 간과할 수 없는데, 2MB 페이지 1개 할당에는 수 마이크로초 단위 시간이 걸리므로 레이어가 많은 모델에서는 한 번 새로운 토큰을 받을 때 수십~수백 마이크로초 지연이 누적될 수 있습니다​. 이를 해결하기 위해 vAttention은 메모리 할당을 GPU 연산과 겹쳐 수행(overlap)하는 최적화를 도입했습니다​. 디코드 단계에서는 다음 토큰이 생성되기 전에 어느 요청이 새 페이지가 필요할지 예측할 수 있으므로, 이전 토큰을 생성하는 동안 백그라운드로 미리 cuMemMap을 실행하는 방식입니다​. 이렇게 하면 정작 새로운 토큰을 쓸 때는 이미 메모리 준비가 완료되어 있어서, 추론 경로 상의 지연 없이 메모리가 확장됩니다. 실험 결과, 이러한 최적화가 없을 때는 일정 간격으로 디코드 지연이 10ms 이상씩 튀는 현상이 보였지만​, 최적화를 적용한 후에는 CUDA 메모리 할당으로 인한 지연이 거의 0에 수렴함을 확인했습니다​. 결국 vAttention은 페이지 크기 최적화와 할당 지연 숨기기 기법을 통해, 운영체제 페이징의 편리함을 취하면서도 성능 희생을 최소화하는 데 성공했습니다. 특히 GPU 메모리 할당 대역폭 측면에서 보면, vAttention이 제공하는 할당 처리량은 초당 최대 ~7.6GB로, 디코드 단계에서 실제 필요한 메모리 증분(수백 MB/s 수준)을 한창 초과하여 충분한 여유를 보입니다​.

실험 및 평가

vAttention 논문에서는 다양한 벤치마크 실험을 통해 앞서 설명한 설계의 효과를 검증하였습니다​. 실험은 프리필 단계 성능디코드 단계 성능메모리 할당 효율 등의 측면에서 PagedAttention 기반 구현들과 비교가 이뤄졌습니다. 테스트에는 80GB 메모리를 가진 NVIDIA A100 GPU 한 대(또는 두 대 연결)가 사용되었고, Yi-6B, Llama2-8B(논문 표기 Llama-3-8B), Yi-34B 세 가지 모델이 포함되었습니다​. 각 실험은 동일한 추론 프레임워크(vLLM 기반)에 PagedAttention과 vAttention을 각각 적용한 뒤 수행되었습니다.

 

프리필(prefill) 성능: 프리필 단계에서는 주어진 프롬프트 토큰들을 병렬로 처리하여 KV 캐시를 채우는데, 이때의 처리량을 초당 처리 토큰 수로 측정했습니다. FlashAttention 2와 FlashInfer 라이브러리가 제공하는 프리필 전용 커널들을 대상으로, 해당 커널에 PagedAttention을 적용한 경우(FA_Paged, FI_Paged)와 vAttention을 적용한 경우(FA_vAttn, FI_vAttn)를 비교했습니다​. 실험에는 긴 프롬프트를 여러 조각으로 쪼개어 순차 처리하는 “청크 프리필(chunked-prefill)” 최적화도 사용되었는데, 모든 비교 대상에 동일하게 적용되었습니다​. 그 결과, vAttention을 사용한 경우가 모든 상황에서 PagedAttention 대비 높은 프리필 처리량을 보여주었습니다​. 예를 들어 Yi-6B 모델의 경우, 중간 길이 컨텍스트(16K)에서 FlashAttention 커널 기반 vAttention이 PagedAttention 대비 약 10% 이상 빠르게 프리필을 완료했고, 컨텍스트가 길어질수록 그 차이가 벌어져 192K 토큰 길이에서는 3.92배까지 속도 차이가 벌어졌습니다​. 이는 앞서 설명했듯이 긴 컨텍스트에서는 어텐션 연산 비중이 커지는데, vAttention이 연속 메모리 기반의 빠른 어텐션 커널을 사용함으로써 큰 이점을 얻었기 때문입니다​. 반면 짧은 컨텍스트(예: 수천 토큰 이하)에서는 FlashAttention의 경우 Paged vs vAttn 차이가 미미했고, FlashInfer의 경우에도 vAttention이 소폭 우세하긴 하나 격차가 크진 않았습니다​. 이는 짧은 프롬프트에서는 어텐션 이외에도 MLP 연산 등이 성능을 좌우하기 때문에, 어텐션 커널 최적화의 효과가 상대적으로 제한적이기 때문입니다​. 하지만 전반적인 추세는 프리필 단계에서도 vAttention이 안정적으로 더 높은 처리량을 제공한다는 것이었습니다​. 또한 PagedAttention 적용으로 인해 FlashInfer의 경우 첫 토큰 출력 지연(TTFT)이 커지는 현상이 관찰되었는데, vAttention을 사용하면 이러한 TTFT가 크게 감소하여 긴 프롬프트일수록 응답 시작 시간도 빨라지는 장점이 있었습니다​.

 

디코드(decode) 성능: 디코드 단계에서는 각 요청마다 순차적으로 토큰을 한 개씩 생성하므로, 메모리 효율과 커널 속도가 곧바로 출력 토큰 처리량에 영향을 줍니다​. 실험에서는 일정 길이(예: 16K) 프롬프트를 가진 다수의 요청을 배치로 투입한 뒤, 요청당 256 토큰씩 생성하는 시나리오를 가정하여 배치 크기를 1부터 8까지 늘려가며 초당 출력 토큰 수를 측정했습니다​. 비교 대상은 (a) vLLM의 기본 PagedAttention 커널(vLLM), (b) FlashAttention 커널에 PagedAttention을 적용한 커널(FA_Paged, 블록 크기 16과 128 각각), 그리고 (c) FlashAttention 커널에 vAttention을 적용한 경우(FA_vAttn)였습니다​. 그 결과, vAttention 적용 시가 모든 경우에서 vLLM 대비 높은 처리량을 기록했고, 특히 배치 크기가 클수록 그 격차가 벌어졌습니다​. 예를 들어 Yi-6B 모델에서 배치=1일 때 vLLM 대비 약 1.3배, 배치=8일 때는 약 1.97배의 출력 속도 향상이 있었다고 보고됩니다​. Llama-2-8B와 Yi-34B와 같은 큰 모델에서도 vAttention의 이점은 뚜렷하여, 배치 8 기준 vLLM 대비 각각 ~1.3배, ~1.6배 가량 더 많은 토큰을 생성했습니다​. 한편 FlashAttention 기반 Paged 커널(FA_Paged, 블록128 사용)은 vAttention과 근접한 성능을 보였지만, 블록 16과 같이 작은 페이지 블록을 사용한 경우 현저히 느려지는 모습을 보였습니다​. 이는 블록 크기가 작을수록 매 iteration마다 CPU에서 블록 테이블을 설정하는 오버헤드가 커지기 때문이며, 논문에서도 vLLM의 작은 블록 크기(16) 설정이 큰 성능 저하를 유발함을 확인했습니다​. 결과적으로 vAttention은 별도 최적화 없이도 FlashAttention Paged 커널의 최고 성능을 달성하거나 웃도는 수준을 보였고​, 배치 규모가 커질수록 이러한 성능상의 이점이 두드러졌습니다.

 

메모리 활용 및 안정성: vAttention의 궁극적 목표는 메모리 단편화 없이 높은 처리량을 얻는 것입니다. 실험에서 관찰된 바에 따르면, vAttention은 GPU 메모리 용량을 매우 효율적으로 활용하여 대용량 배치 추론을 안정적으로 수행했습니다. 우선 메모리 할당 처리량 측면에서, vAttention이 CUDA VMM을 통해 물리 메모리를 할당/해제할 수 있는 속도는 GPU 한 개당 최소 수 GB/s 이상으로, 디코드 단계에서 요구되는 메모리 증가 속도(수백 MB/s)를 충분히 상회했습니다​. 이는 동적 할당으로 인한 병목이 없었음을 의미합니다. 또한 내부 단편화율을 보면, 작은 페이지 크기 설정 덕분에 실제 낭비되는 물리 메모리가 미미했습니다. 예컨대 256KB 페이지를 사용하면, 각 요청마다 못 쓰고 남는 메모리가 최대 몇십 MB 수준에 불過하지 않았고, 대부분의 요청이 그보다 훨씬 적은 여유 메모리만을 남긴 채 완료되었습니다​. 반면 2MB 페이지였다면 요청당 최대 수백 MB가 낭비될 수도 있는 것을 감안하면, 이는 크게 개선된 수치입니다. 끝으로, 안정성 측면에서도 vAttention은 유리한데, 이유는 GPU 메모리를 탄력적으로 운영하면서도 커널 수정이 없으므로 예기치 않은 버그나 비호환성이 적기 때문입니다. 논문에서는 서로 다른 두 어텐션 커널 (FlashAttention ↔ FlashInfer)을 동일한 vAttention 환경에서 손쉽게 교체하여 사용하는 예제를 보여주며, PagedAttention과 달리 커널 교체를 위한 추가 코드 수정이 거의 필요 없음을 강조합니다​. 이는 곧 새로운 모델이나 최적화 기법을 도입할 때 vAttention이 가진 실용적 이점이라 할 수 있습니다.

결론

LLM 서비스에서 동적 메모리 관리는 고성능과 메모리 효율을 모두 충족시키는 데 필수적인 요소입니다. 기존 PagedAttention 기법은 수요 페이징을 흉내내어 KV 캐시 단편화 문제를 해결하고 대용량 배치를 가능케 했지만, 커널 수정 필요성, 높은 소프트웨어 복잡도, 런타임 오버헤드라는 트레이드오프가 있었습니다​. Microsoft Research의 vAttention은 CUDA 가상 메모리 기능을 통해 이러한 문제를 우아하게 풀어낸 접근법으로, KV 캐시를 가상 메모리 상에서는 연속적으로 유지하면서 물리 메모리를 동적으로 할당함으로써 PagedAttention의 함정들을 피해갔습니다​. 그 결과 추가적인 GPU 커널 개발이나 사용자 수준 메모리 관리 없이도 동적인 KV 캐시 운용이 가능해졌고, 실험을 통해 성능 향상과 구현 단순화를 모두 입증했습니다​

728x90