![[LLM] Paged Attention](https://img1.daumcdn.net/thumb/R750x0/?scode=mtistory2&fname=https%3A%2F%2Fblog.kakaocdn.net%2Fdn%2FA0UWR%2FbtsD3bbFtCF%2FKN7IzZ1sgMwz5xHCk8rhd1%2Fimg.webp)
Background
1. Query, Key, Value
Encoder에서의 의미
분석하고자 하는 단어(Query)를 그 외의 다른 단어들에 대한 문맥상의 의미(내용, 위치 정보등)를 나타내는 임베딩 값(Key)과 비교하면서 우리의 Query가 문장 내 다른 단어들의 Key값 얼마나 유사한지를 평가합니다. 그렇게 생성된 유사도 점수는 문장에서 각 단어들의 의미(Value)에 대한 가중치로 적용이 됩니다. 그렇게 유사도 점수를 Value에 적용하여 구해낸 값이 Attention Score입니다. 이는 우리가 분석하고자 하는 단어가 문맥상의 의미를 고려해서 어떤 의미인지 계산해내는 방법입니다. 이를 Encoder에서는 모든 단어들에 대해서 구하게 됨으로써 입력 시퀀스 내의 각 단어가 다른 단어와 어떻게 관련되어 있는 지를 계산해내게 됩니다.
Decoder에서의 의미
Encoder와 비슷하지만 우리가 비교하는 Key는 이전에 생성해낸 단어들에 대해서만 구하고 이를 기반하여 Attention값을 구하게 됩니다. 이 Attention값은 입력 시퀀스에 대한 정보를 활용하여 현재 생성중인 단어에 대한 컨텍스트를 얻는데 활용됩니다.
2. KV cache
앞써 설명드린 Decoder 부분을 좀 더 이야기 해보고자 합니다. Decoder에서는 Attention을 이용해서 현재 생성 중인 단어에 대한 컨텍스트를 얻는다고 말씀 드렸습니다. 생성 모델 중 Auto Regressive 모형을 가지는 경우 Decoder는 자신이 생성해낸 토큰을 다시 이용하여 다음 토큰을 생성해냅니다. 이 과정에서 Decoder가 매번 자신이 생성해놓은 문장 전체에 대한 Key, Value를 다시 계산해내는 비효율이 발생합니다. KV cache는 이 Key, Value에 대한 값을 caching하여 불필요한 연산과정을 줄이자는 아이디어 입니다.
Problem
KV cache의 아이디어는 참 좋습니다. 다만 Key, Value를 어떻게 memory에 저장할지가 문제입니다. 기존 HF나 TGI같은 Transformer Toolkit의 경우 13B모델이 이를 이용해서 Inference를 제공할 때 1회 inference시 1.7GB의 저장공간이 필요했습니다. 이러한 상황에서는 다음과 같은 2가지 문제가 발생합니다.
1. Key, Value의 캐시 사이즈는 Prompt에 의존적입니다. 즉, 예측할 수 없습니다. 따라서 메모리에 우리가 cache를 저장할 때 미리 공간을 할당하여 사용할 수 없습니다. 이로 인해 메모리 공간의 fragmentation이나 over-commitment가 발생하게 됩니다.
2. 여러명이 inference를 계속 요청하는 경우 몇몇 토큰들에 key, value값은 겹칠 수 있습니다. 이러한 경우 같은 값을 메모리에 여러번 write하는 문제가 있습니다.
Paged Attention
page attention의 해결책은 가상메모리를 통한 OS의 메모리 관리로 부터 영감을 받아 KV cache를 다음과 같이 처리합니다.
1. 각 입력 토근 길이를 일정 길이로 분할 취급한다.
2. 이들을 가상 메모리에 연속적으로 저장하고 Block Table을 통해 물리적 메모리에는 불연속적으로 저장을 하게 됩니다.
위의 방식은 다음과 같은 이점을 남깁니다.
- 기존에 가변적인 KV cache size로 인해 발생하는 fragmentation이나 over-commitment가 가상 메모리 활용으로 없어지게 됩니다.
- Block Table을 통해 병렬 처리 과정에서 중복으로 사용되는 토큰들에 대해서 물리적 메모리에는 1번 저장하고 이를 사용하는 여러 프로세스는 Block Table을 참고하여 사용하기 때문에 불필요한 중복 write를 줄일 수 있게 됩니다.
이렇게 OS의 가상 메모리 관리 방식으로 KV cache에 저장하는 Attention 방식을 Page Attention이라고 하며 이를 통해 기존 HF를 이용해서 만든 Inference server대비 30배 높은 처리량을 처리할 수 있는 vLLM을 탄생시키게 되었습니다.
참고
https://arxiv.org/abs/2309.06180
Efficient Memory Management for Large Language Model Serving with PagedAttention
High throughput serving of large language models (LLMs) requires batching sufficiently many requests at a time. However, existing systems struggle because the key-value cache (KV cache) memory for each request is huge and grows and shrinks dynamically. Whe
arxiv.org
https://blog.vllm.ai/2023/06/20/vllm.html
vLLM: Easy, Fast, and Cheap LLM Serving with PagedAttention
GitHub | Documentation | Paper
blog.vllm.ai
https://www.axelera.ai/decoding-transformers-on-edge-devices
긴 글 읽어주셔서 감사합니다.
틀린 부분이 있으면 댓글을 달아주시면 감사하겠습니다.
📧 : realhwan1202@gmail.com
🔗 : https://github.com/RicardoKim
'인공지능 > LLM' 카테고리의 다른 글
| [RAG] Knowledge Graph Prompting for Multi-Document Question Answering (0) | 2024.04.07 |
|---|---|
| [LLM] SELF-RAG: Learning to Retrieve, Generate and Critique through Self-reflect (1) | 2024.02.13 |
| [LLM] Seven Failure Points When Engineering a Retrieval AugmentedGeneration System (0) | 2024.01.29 |
| [LLM] Flash Attention (0) | 2024.01.22 |