PagedAttention. 정말 최적인가? 진짜?

추론 과정에서 GPU 메모리의 대부분을 차지하는 KV 캐시를 어떻게 관리하느냐가 서비스 성능에 큰 영향을 미치게 되는데, 이전 시스템들은 모델이 지원하는 최대 컨텍스트 길이를 기준으로 각 리퀘스트별로 메모리를 할당하는 "naive"한 방법을 사용했습니다. 하지만 이렇게 하면 모든 리퀘스트가 항상 모델의 최대 컨텍스트 길이만큼 디코드를 하는 것은 아니므로 메모리 낭비가 심하고, 낭비되는 메모리만큼 서비스를 더 할 수 없게 됩니다.

Inference Engine의 (거의) 표준이 되가고 있는 vLLM은 OS 기반 가상 메모리 시스템의 demand-paging 방식에서 영감을 받아, 할당된 블록이 모두 사용되었으면, 필요에 따라 작은 GPU 메모리 블록을 할당하는 PagedAttention 방식을 제안했습니다. 이 방법으로 fragmentation을 완화하기 위해 대부분의 inference engine이 이 방식을 사용하고 있습니다. (TGI, TensorRT-LLM 등)

최근 자체 개발 중인 엔진에서 paged attention이 정말 최적인가?하는 이야기들을 하고 있었습니다. 성능적인 측면에서의 문제와, 메모리를 reservation 함으로써 여러 모델/데이터 타입을 유연하게 지원하기 어렵다는 점, 새로운 어텐션 커널 등을 지원하거나, 새로운 feature 도입 시 이 paged 방식을 사용하는 것 때문에 수정 사항이 많이 발생하는 등의 문제점이 있다고 생각했거든요. 그럼에도 불구하고 memory fragmentation의 솔루션으로 (거의) 표준처럼 사용하고 있다는 점에서, 우리도 쓸 수 밖에 없나? 하는 고민들이 있었습니다. paged attention을 비롯해서 여러 opensource library, framework 들이 다 엉켜있어서 하나를 잘 못 도입하게 되면 전체 구조에 영향을 줄 것 같아서 불안한 마음이.. 

검색을 좀 해보니, 이미 먼저 고민하고 솔루션을 도출해 낸 논문들이 있었고, 그 중에 vAttention 방식이 간결하고 효과적인것 같아 공유합니다.

기존 PagedAttention 방식의 문제점은 사용자 레벨에서 할당하는 텐서 객체로 페이지를 만들기 때문에 이 객체들의 실제 물리 메모리 주소는 연속적이지 않게 됩니다. 따라서 KV cache의 레이아웃이 non-contiguous하게 변경되므로 attention 연산 시 kernel에서 동적으로 할당된 가상 메모리 블록을 연결하기 위한 메모리 관리자를 통해 메모리를 참조해서 연산을 해야 하는 문제점이 생기게 됩니다.

vAttention은 CUDA VMM (Virtual Memory Management) API를 사용해서 가상 메모리와 물리적 메모리의 할당을 아예 분리해버림으로써, 이 문제를 해결하죠.
문제를 없애면 문제를 풀지 않아도 돼! 의 긍정적인 적용 케이스..

평가 결과, vAttention은 attention 커널의 성능과 prefill 및 decode 단계의 전체 성능을 향상시키는 것으로 나타났습니다. 특히 긴 문맥의 프리필 단계에서 FlashAttention-2 백엔드를 사용할 때 vAttention은 페이지 처리 버전보다 더 우수한 성능을 보였습니다. 또한 디코드 처리량에서도 vAttention은 최신 PagedAttention과 동등하게 성능을 발휘하며, PagedAttention을 초월하는 경우도 있었습니다.

논문이 나온지 조금 시간이 있었고 사용한 버전들이 약간.. 구버전이라서 물론 지금 시점에서 평가를 해보면 약간 수치들이 다른 부분이 있을 수는 있겠습니다만, 논문에서 말하는 것처럼 새로운 어텐션 방법들이 나오거나 했을 때 일정 시간/노력이 걸리게 되는 것을 입증한 것 같기도 하고요.

https://arxiv.org/abs/2405.04437

끝.