[MLOps, LLMOps] PagedAttention, vLLM 소개

chatGPT 이후로 LLM에 대한 연구가 많이 활성화되었고, 그에 따라 최근에는 LLM 서빙을 다루는 경우가 많아지고 있는 것 같습니다.

vLLM은 LLM을 더 효율적으로 서빙하기 위한 PagedAttention이란 기술을 적용하였다고 하는데, 오늘은 이 부분들을 소개해보도록 하겠습니다.

 

PagedAttention의 원리

vLLM을 개발한 팀은 아래의 논문을 통해 PagedAttention을 먼저 소개했고, 직접 개발하여 vLLM 라이브러리를 출시했습니다.

- 논문 제목: Efficient Memory Management for Large Language Model Serving with PagedAttention

- 링크: https://arxiv.org/pdf/2309.06180.pdf

 

Key-Value Cache

PagedAttention을 이해하기 전에, Key-Value Cache에 대해서 먼저 짚고 넘어가보도록 하겠습니다.

생성형 모델에서는 토큰을 생성하면서 매 request마다 할당하고 해제하여 사용하는 메모리가 많이 생기는데, 현재의 플랫폼들은 아래 이미지와 같이 메모리를 활용하면서 생성 길이에 따른 메모리를 미리 할당 받아놓고 사용하지 않는 경우도 있고, 새로운 메모리를 할당할 때 비워놓는 공간도 생기면서 매우 비효율적으로 메모리를 사용한다고 합니다.

출처: https://arxiv.org/pdf/2309.06180.pdf

실제 메모리를 프로파일링 해보았을 때, Key-Value 어텐션을 위해서 사용하는 메모리가 30%가 넘는 것으로 확인되었다고 하는데 이 부분을 논문에서는 Key-Value Cache라고 하고 있습니다.

 

출처: https://developer.nvidia.com/blog/mastering-llm-techniques-inference-optimization/

위의 이미지에서 Prefill 단계는 흔히 우리가 생각하는 promt에 대한 부분을 계산하는 것이고, Decode 단계가 생성하는 단계인데 Decode 단계에서는 한 토큰씩만 생성할 수 있는데 이전의 Key-Value Attention을 항상 처음부터 새로 계산하는 부분에 대해서 Caching을 해두겠다는 것이 PagedAttention의 main contribution입니다.

PagedAttention

그래서 PagedAttention은 Key-Value Cache를 메모리에 캐싱하여 효율적으로 GPU 메모리를 사용할 수 있도록 했는데 구체적으로 어떻게 캐싱을 하는지 조금 더 알아보도록 하겠습니다.

PagedAttention의 캐싱은 OS(Operating System; 운영체제)에서의 가상 메모리(Virtual Memory)와 페이징(Paging) 방식에서 고안하여 매우 유사하게 만들어졌는데요.

출처: https://arxiv.org/pdf/2309.06180.pdf

그래서 PagedAttention은 위의 이미지와 같이 메모리를 블록으로 나누어 놓고 매핑 테이블을 관리하는 방식으로 설계되어 낭비되는 메모리도 줄이고, 캐싱을 통해 다시 계산하지 않아도 되는 효과를 얻을 수 있었다고 합니다.

캐싱은 GPU 메모리에서도 활용하지만, CPU 메모리도 같이 활용하여 캐싱을 많이 할 수 있도록 실제로 개발되었다고 합니다.

 

그래서 저자들은 PagedAttention을 통해서 메모리를 효율적으로 사용하기 때문에, batch 사이즈를 더 크게 사용할 수 있고 최종적으로 같은 환경에서 throughput을 많이 올릴 수 있었다고 합니다. 

출처: https://arxiv.org/pdf/2309.06180.pdf

특히, 위의 이미지에서 prompt 길이를 1-shot으로 가져가는 것 보다 5-shot으로 길게 가져가게 되면, 캐싱을 많이 활용하기 때문에 훨씬 더 큰 효과를 발휘한다고 합니다.

 

PagedAttention은 메모리 사용을 효율적으로 해주기 때문에, batch를 키워서 throughput은 높였지만 실제 단일 inference 속도가 올라간 것은 아니란 것을 잊지 않으셨으면 좋겠습니다.

vLLM 소개

vLLM은 PagedAttention의 저자들이 직접 구현하여 발표한 라이브러리입니다.

- 링크: https://blog.vllm.ai/2023/06/20/vllm.html

출처: https://blog.vllm.ai/2023/06/20/vllm.html

HF는 HuggingFace, TGI는 HuggingFace Text Generation Inference인데, vLLM을 활용하면 훨씬 더 높은 throughput을 얻을 수 있었다고 합니다.

 

제가 개인적으로 생각했을 때, vLLM의 또 다른 장점은 쉬운 Interface입니다.

설치는 아래 명령어 한 줄로 python package만 설치하면 되고 사용 방법도 매우 쉽습니다.

pip install vllm

from vllm import LLM

prompts = ["Hello, my name is", "The capital of France is"]  # Sample prompts.
llm = LLM(model="lmsys/vicuna-7b-v1.3")  # Create an LLM.
outputs = llm.generate(prompts)  # Generate texts from the prompts.

위와 같이 huggingface 모델들은 import 제외하곤 단 3줄로 vLLM을 사용할 수 있게 구성되어 있습니다.

 

오늘은 PagedAttention, vLLM에 대해서 소개를 해보았는데요.

PagedAttention 자체는 NVIDIA에서 최근에 새로 출시한 TensorRT-LLM에도 탑재되었을 정도로 매우 좋은 알고리즘이라고 생각이 되고, vLLM은 아주 쉽게 쓸 수 있는 인터페이스를 가지고 있어서 매우 좋은 것 같습니다.

PagedAttention은 이제 LLM 서빙에선 빠질 수 없는 기법이 될 것 같은데요.

LLM 서빙에 관심이 있으시다면, 꼭 한 번쯤 공부해보면 좋을 것 같습니다.