vAttention 접근 방식 — CUDA 가상 메모리 관리

vAttention은 물리 메모리의 단편화를 완화하면서 KV-캐시의 가상 메모리 연속성을 유지하는 접근이에요. 핵심은 CUDA 가상 메모리 관리(CUDA VMM) API로 가상 메모리와 물리 메모리 할당을 분리하는 것입니다.

출처: https://www.microsoft.com/en-us/research/publication/vattention-dynamic-memory-management-for-serving-llms-without-pagedattention

가상과 물리를 분리하다

vAttention은 물리 메모리를 사전에 커밋하지 않고도 KV-캐시를 연속적인 가상 메모리에 저장해요. 가상 메모리 주소는 그대로 연속적으로 잡아두되, 실제 물리 메모리는 요청 페이징(demand paging) 을 통해 필요할 때 할당합니다. 이때 저수준 시스템 지원인 CUDA VMM API를 이용합니다.

이렇게 하면 기존에 이미 존재하는 시스템 지원을 활용하므로, 어텐션 커널 개발자가 페이지드를 명시적으로 지원할 필요가 없어요. 서빙 프레임워크에서 메모리 관리 로직을 다시 구현할 필요도 없어집니다.

수정하지 않은 커널에서도 동작

vAttention은 수정하지 않은(unmodified) 어텐션 커널에 동적 메모리 할당을 원활하게 추가할 수 있음을 보여줘요. FlashAttention과 FlashInfer의 기존 커널 구현을 그대로 사용하면서도 동적 메모리 관리가 가능하다는 뜻입니다.

이 특징은 곧 포트폴리오 측면에서도 장점이에요. 이미 만들어져 있는 GPU 커널을 재사용할 수 있고, 서빙 시스템에 메모리 매니저 구현이라는 부담이 없어지니까요. 특히 prefill 바운드 워크로드에서 성능이 개선되는 경우가 많습니다.

정리

PagedAttention 방식은 애플리케이션이 동적 할당된 물리 메모리를 직접 관리해야 하고, 어텐션 커널까지 재작성해야 해요. vAttention은 이와 달리 가상 메모리와 물리 메모리의 할당을 분리해, 가상 연속성을 유지하면서 물리 메모리만 필요할 때 할당합니다. 그 결과 커널 재작성·메모리 매니저 구현이라는 소프트웨어 복잡도와 중복을 피할 수 있습니다.

더 알아보기

배경(PagedAttention 문제)은 PagedAttention과 단편화 참고
구현 상세는 Implementation 참고
논문: "vAttention: Dynamic Memory Management for Serving LLMs without PagedAttention" (arXiv:2405.04437)