제목

Efficient Memory Management for Large Language Model Serving with PagedAttention
- 즉, LLM을 서빙할 때 PagedAttention을 사용하여 효율적인 메모리 관리하는 방법론에 대해서 다루는 논문이다.

소개

LLM 서빙에서 throughput(단위 시간당 처리할 수 있는 업무량)을 늘려 request 당 비용을 낮추는 것이 점점 중요해지고 있다.
왜 throughput을 늘리는 게 어렵나?
- 토큰을 하나 새로 생성할 때, 이전에 생성한 모든 토큰을 알고 있어야 그걸 계산할 수 있다.
- 이런 sequential 생성 방식이 CPU/GPU의 성능보다 메모리 대역폭이 병목이 되게끔 만든다. (makes the workload memory-bound)
위 도표는 A100 40GB에 13B LLM을 올렸을 때, 전체의 몇 퍼센트가 어떤 부분을 차지한 것인지 보이는 도표다.
- Parameters: 모델의 가중치. 전체에서 65%를 차지하나 상수이므로 조절하기 어렵다.
- Others: activation과 같은 짜투리. 이 부분을 최적화해봤자 차지하는 크기가 작으므로 큰 의미가 없다.
- KV Cache: key와 value 텐서들을 저장해놓는 캐싱. 최적화할 여지가 있다!
  - 어떻게 도움이 될까?
    - 원래 hidden state에 가중치 $W_{K}$ 나 $W_{V}$ 를 곱하여 key와 value 텐서를 얻을 수 있는데, 매번 계산을 하는 것보다 Next Token Prediction에 필요한 key와 value 텐서만 저장하고 있으면 더 빠르게 생성이 가능하다.

tensor가 연속된 메모리 공간 상에 있어야, 대부분의 연산을 적용할 수 있다.
근데 출력을 얼마나 해야할지 알 수 없기 때문에, 가능한 최대한의 길이(maximum possible sequence length)를 일단 allocate한다.
그러다보니 아래와 같은 비효율이 발생한다.
1. reserved: 앞으로 생성될 토큰을 위해 예약(reserved)해둔 메모리 공간
2. internal fragmentation: 실제 입출력과 무관하게 가능한 최대로 메모리 공간을 예약하다보니, 내부(internal)에서 영원히 사용되지 않은 메모리 공간(fragmentation)이 존재하게 된다.
3. external fragmentation: 위에 figure처럼 어떤 request도 할당받지 않은 slot이지만, 연속된 slot이 짧기 때문에 할당해줄 수 없는 메모리 공간이 존재하게 된다.
아래는 LLM 서빙 시스템별로 얼마나 메모리 공간이 낭비되고 있는지 측정한 도표다.

Paging을 이용해서 KV Cache를 관리한다. 이름하여 PagedAttention.
$j$ 번째 블록에 저장된 key와 value가 각각 $K_{j} = (k_{(j - 1) B + 1}, \dots, k_{j B})$ , $V_{j} = (v_{(j - 1) B + 1}, \dots, v_{j B})$ 일 때 아래 식으로 다음 토큰을 예측할 수 있다. ( $B$ 는 블록 사이즈)

A_{ij} = \frac{exp ( q _{i}^{⊤} K _{j} / d )}{\sum _{t = 1}^{⌈ i / B ⌉} exp ( q _{i}^{⊤} K _{t} 1 / d )}, o_{i} = t = 1 \sum ⌈ i / B ⌉ V_{j} A_{ij}^{⊤}

즉, 비연속적인 메모리 공간에서도 어텐션 스코어를 계산할 수 있다.
- 장점: physical memory를 미리 한꺼번에 예약할 필요없이, 추가적으로 블록을 할당해야 할 때만 할당하면 된다. → 당장 사용할 수 있는 메모리 공간이 늘어나므로 request를 더 많이 받을 수 있다. → throughput 상승
블록 사이즈가 크면 클수록 좋을까?
- parallel한 연산을 한꺼번에 할 수 있어서 좋지만, internal fragmentation을 감당해야 함.

block engine이 GPU DRAM의 연속적인 청크를 할당하고 이를 물리적인 메모리 공간 상에서 KV Cache를 나눈다.
- 블록 테이블을 이용해서 logical 메모리와 physical 메모리 사이를 mapping한다.

copy-on-write를 이용한다.
- 똑같이 sampling된다면 계속 같은 physical block을 공유하다가, 변화가 생기면 그 block만 다른 physical block에 복사한 다음은 mapping을 바꿔준다.