Ray Summit 2023 - Fast LLM Serving with vLLM and PagedAttention

Fast LLM Serving with vLLM and PagedAttention

https://youtu.be/5ZlavKF_98U?si=sjZ31eHwRAClmx7a

 

권우석(woosuk.kwon). Berkley Ph.D Student

 

LLM 기반 Application이 많이 등장함. Chat, Programming, CopyWriting...

• 도메인은 다양하지만 결국 LLM Serving이 핵심인 서비스들. LLM Speed / Cost 영향을 크게 받는다.
• 따라서 Serve Fast / Cost effcient 가 중요한 과제가 되었음

 

지금 LLM은 굉장히 비싼 몸.

• High-End GPU를 '많이' 필요함.
• 각각의 GPU가 handling할 수 있는 request는 매우 적음
  • NVidia A100 1개가 TPS 1도 안 됨.

KV Cache in LLM Inference Process

 

Inference Process를 되짚어보자.

• 사용자가 prompt 포함된 query를 하면, Possible Next Token 결과를 받아서 계속 모델에 입력으로 넣는 식.
  • Max Token 제한에 걸리거나, <|end of sentense|> 같은 문장 종결토큰이 모델 결과로 등장할 때까지 반복됨

 

Inference Process의 unique component로 KV Cache라는 게 있다.

• 'Next token Prediction'을 수행할 때, 단순히 '직전 토큰'만 쓰는 게 아니라 All Previous Token 정보가 필요함.
• Previous Token 정보를 저장하는 공간이 KV Cache.

따라서 KV Cache는 Sequence가 길어질수록 많은 공간을 차지하고, 연산이 끝나면 줄어든다.

vLLM: Efficiently Managing KV Cache

 

즉, KV Cache를 효과적으로 관리하는 것이 LLM Serving Throughput의 핵심이다.

• 예컨대 A100 40GB 메모리에 LLM을 올린다고 가정하면
  • 파라미터가 26GB, KV Cache에 12GB를 쓴다.
• KV Cache를 효율적으로 쓰지 못하는 LLM 모델이 많았음. 그래서 request가 증가할 때 Memeory Usage 상승률이 높았다.

vLLM은 KV Cache 사용량을 최적화해서, GPU의 메모리를 보다 효율적으로 사용할 수 있게끔 고안되었다.

• 더 높은 Throughput을 지원한다.

KV Cache에서 정확히 뭐가 문제였나

 

기존 LLM serving에서 메모리 효율성을 떨어뜨리는 세 가지 영역

• Internal fragmentation: sequence 연산 때문에 할당되었지만, 실제로는 안 쓰인 메모리 공간.
  • 모델이 token을 얼마나 생성할지 알 수 없기 때문에 발생함
• Reservation: 지금은 sequence 연산에 안쓰이지만, 곧 쓰이게 될 메모리 공간
• External fragmentation: 요청마다 들어오는 sequence 길이가 달라서 발생하는 메모리 공간

프로파일링 해본 바로는, 20~40% 정도의 KV Cache가 실제로 token state 저장을 위해 쓰임.

vLLM은 어떻게 했나: PagedAttention

 

OS에서 쓰이는 virtual Memory + paging 개념을 KV Cache에 적용

• OS는 fragmentation을 낮추기 위해 paging 기법을 쓰고 있고
• Process 간 efficient / elegant space multplexing을 위해 virtual memory 개념을 사용함

이걸 가능하게 한 것이 PagedAttention

 

KV Cache 공간을 paging한 것처럼, KV block으로 구분한다.

• KV Block: fized-size contiguous chunk of memory that stores KV cache from left to right.
  • block size: 하나의 Block에 얼마나 많은 토큰이 들어갈 수 있는지 결정하는 값

 

즉, KV Cache를 쓰기 위해 'contiguous space'가 필요했던 구조적 한계를 개선하는 것.

• input과 output이 static할 경우 contiguous space가 효율적이지만, LLM은 그렇지 않다.
• KV Cache는 block 단위로 나뉘어서 임의의 메모리 공간에 저장하고, 필요할 때 가져오는 방식.

 

Virtualize KV Cache into logical / physical KB Block도 가능하다.

예컨대 사용자가 input으로 'Alan Turing is a computer scientist" 라고 넣으면

• logical KV Block: order 정보를 포함하는 consecutive block
• physical KV Block: adjacent할 필요 없이, KB Block 특정 공간에 저장됨

consecutive 정보는 일종의 Block Table에서 관리한다. physical block number 정보를 저장한 mapping table을 관리하는 식.

 

logical block에 new token이 생성되면 physical block에도 저장되고

block 공간이 차면, 다른 physical block을 매핑해서 진행하는 식.

• 즉, 기존 방식과 달리 "On-demand" 로 메모리 공간을 할당한다.

 

serving multiple request도 똑같다. os가 메모리 관리하는 것과 근본적으로는 똑같기 때문

How much Efficient?

 

Minimal Internal Fragmentation

• last block of sequence에서만 발생. wasted token 개수는 block size를 넘지 않는다.
  • 보통 token 개수는 100 ~ 1000까지 다양한데, block size는 16 or 32 정도로 설정함.
• waste space를 4% 수준으로 낮춤 -> increase throughput.

vLLM: Enables input prompt sharing

 

위 예시는 Parallel Sampling. multiple output을 응답하는 형태.

• input으로 들어온 문장은 여러 Output process가 공유해서 쓸 수 있다.
• Chat보다는 github copilot처럼 여러 개의 웅답 후보군을 준비해야 하는 경우에 유용함

 

구체적인 동작 방식

• physical block은, 본인은 reference하는 logical block 개수를 저장한다.
• new token을 입력할 때, physical block의 reference count를 보고
  • 1보다 클 경우, copy-on-write 방식으로 다른 physical block에 복사한 후 new token 정보를 추가한다. ref count는 1 줄어든다.
  • 1일 경우, physical block에 직접 쓴다.

vLLM: complicated decoding like beam search

 

beam search: 가장 정확도 높은 machine translation을 찾아주는 기능.

• 동작 방식이 parallel sharing과 똑같다. 가장 나은 translation 문장을 생성하는 방식이 똑같기 때문.
• OS에서 process tree (fork / kill) 하는 것과 비슷함.

Summarize

PagedAttention 기능을 사용해서 메모리 사용량을 효율화

• Memory Fragmentation 비율을 낮췄옴.
• Memory Sharing 활성화.

 

huggingface.generate 방식으로 모델 가져와서 쓴 것보다 성능이 월등히 낫다.

• PagedAttention 없던 시절의 TGI (text generation inference)보다도 나음.

 

우린 지원되는 모델도 많고, 주목받는 오픈소스 중 하나다

 

System Architecture / Implementation

 

• Centralized vLLM Engine: deals with Block table.
  • iteration마다 GPU workers에 memory command 전달.
• gpu block allocator가 physical block 관리

vLLM 자체는 python project이고, PagedAttention 로직 일부는 cuda code.

• Megatron-LM for distributed Execution
• huggingface / pytorch
• Ray for cluster management (control plane communication)

Powering Vicuna & Chatbot Arena

 

오픈소스 공개 이전 시점에서 진행했던 것

• LLaMA-based Model 사용한 Serving
  • 40K conversation per Day
  • GPU 개수 50% 줄이고도 서빙 가능
  • serving 속도는 2x ~ 3x 빠름

Adopters

 

많은 곳에서 쓰고 있다.

PagedAttention 자체가 업계 표준이 되고 있는 추세.

How to use?

 

offlien batch inference 예시

• prompt list를 넣어주면, 알아서 batching process 진행한다.

 

for online serving 예시: with FastAPI

• openAI의 api format과 동일한 interface 지원

---

Q. recurrent 방식의 Neural Network도 PagedAttention 쓸 수 있는가?

• 된다고 생각함. 결국 memory management는 필요하고, output length를 생성 시점에서 알 수 없는 건 똑같기 때문.
• 물론, Attention 방식을 안 쓴 모델은 적용 안 된다 ㅎㅎ

Q. memory problem을 어떻게 인지했나 (원문: How do you identify the memory problem)

• 오픈소스 활성화가 덜 된 시점에서 LLM Serving 시스템 직접 만들어야 했고, Reduce latency보다는 high throughput을 목표로 했었다. (그게 더 비용효율적이었음) batch size를 늘려야 했고, 그러다보니 문제를 발견했다.
• 그 당시의 memory management 방식은 전부 비효율적이었다.
• Hands-on experience였음 ㅎㅎ

Q. Cache misses 같은 버그? 가 존재하는가?

• 불가능하다. Attention 동작방식 자체가 'all previous words'를 받아서 computation 하는 구조. 연산 중 특정 구간의 missing이 발생하지 않는다

Q. PagedAttention 기법이 TGI 같은 기존 serving platform에도 적용되거나 그럴 예정인 듯한데, 그럼에도 vLLM을 써야 하는 이유가 있다면?

• 우리는 Memory Sharing도 있다 (TGI에서는 아직 이게 없다)
• Optimization 잘되어 있다 (TGI 대비 성능표)
• 우리 Apache2 라이센스다 ㅎㅎ

Q. Sharing Prefix across different request 같은, 다른 기능을 테스트하거나 계획중인 게 있는지?

• sharing prefix는 된다. 가능하다. experiment in paper에 있음.
• 아직 engineering issue 때문에 branch merge는 안 되어 있지만, 될 거다

Q. Throughput에 집중했는데, Latency는?

• Throughput을 늘리는 게 비용 효율적이어서 우선순위가 높았고
• Latency의 경우 크게 두 가지 카테고리로 개선점을 잡을 수 있는데
  • Pure Optimization. i.e. Kernel Fusion? 같은 거... 언젠간 반영될 거다.
  • Tradeoff btwn Latency / Throughput. -> 조심스럽게 접근할 이슈라고 본다
    • speculative decoding 같은 거.
• 일단 우리는 첫 번째 step에 집중하고 있다.