배경 — PagedAttention과 KV-캐시 단편화 문제

배경 — PagedAttention과 KV-캐시 단편화 문제

고처리량 LLM 추론에서 GPU 메모리를 효율적으로 쓰는 것은 필수예요. 이 글은 vAttention 논문의 도입부를 바탕으로, 왜 기존 접근이 메모리를 낭비했고 PagedAttention이 어떤 문제를 가져왔는지 살펴봅니다.

출처: https://arxiv.org/html/2405.04437v1

사전 할당 방식의 문제

이전 시스템들은 KV-캐시용 메모리를 미리(ahead-of-time) 예약해 뒀어요. 그 결과 내부 단편화(internal fragmentation) 때문에 용량이 낭비됐습니다. 예를 들어 Orca는 KV-캐시 메모리를 사전에 예약하는데, 정확히 필요한 만큼만 쓰지 못하고 통째로 예약하다 보니 실제 배치를 키우기 어렵고 메모리가 남는 상황이 벌어졌어요.

PagedAttention의 등장

OS 기반 가상 메모리 시스템에서 영감을 받아, vLLM은 PagedAttention 을 제안했어요. KV-캐시를 고정 크기의 블록으로 나누고 한 번에 한 블록씩 메모리를 할당하는 방식이에요. 이렇게 하면 요청이 필요로 하는 만큼만, 그리고 필요할 때에만 할당하므로 단편화를 없애고 더 큰 배치로 고처리량 서빙이 가능해집니다.

PagedAttention이 만든 복잡성

하지만 물리 메모리를 동적으로 할당하려면 PagedAttention은 KV-캐시의 레이아웃을 연속 가상 메모리에서 비연속 가상 메모리로 바꿔야 해요. 이 변경 때문에 두 가지 부담이 생겼습니다.

  • 페이지드를 지원하도록 어텐션 커널을 재작성해야 한다.
  • 서빙 프레임워크가 메모리 매니저를 구현해야 한다.

논문의 표(Tabel 2)를 보면 각 시스템의 메모리 관리 방식을 비교할 수 있어요. vLLM은 Block-Table(들)로 관리하고, FlashInfer는 압축된 Block-Table을 쓰며, FlashAttention·FasterTransformer는 페이지드가 아닌(비페이지드) 커널입니다. vLLM의 구현은 페이지드 어텐션이 비페이지드 FasterTransformer 커널보다 20~26% 느리다고 인정하는데, 주로 Block-Table 조회 오버헤드와 추가 분기 실행 때문이에요. 벤더마크상 vLLM의 페이지드 커널은 FlashAttention-2 커널보다 최대 2.8배 느린 경우도 있어요.

또한 실제로는 동적 할당된 객체가 반드시 연속적이라는 보장이 없어요. 이 지점이 vAttention이 파고드는 핵심 문제입니다.

더 알아보기