PagedAttention + vLLM

Transformer Serving의 문제점

• Transformer의 generation process는 memory-bound임
  • GPUs의 computation power를 제대로 사용하지 못함
  • 이로 인해 serving throughput을 제한함

• Throughput을 개선하기 위해선 multi-requests를 모아 처리해야 함
  • 하지만 많은 requests를 배치 처리하기 위해서는 각 requests에 대한 메모리 공간을 효율적으로 관리해야 함
    • 모델 weight는 일정하며 KV cache를 관리하는 방식이 최대 batch size를 결정함
    • KV cache를 비효율적으로 관리하면 batch size가 크게 되므로 LLM의 throughput이 제한됨
• 기존 LLM serving system은 KV cache 메모리를 효율적으로 관리하지 못함
  • KV cache는 기존 딥러닝 워크로드와 달리 고유한 특성을 가짐
    • LLM이 새 토큰을 생성함에 따라 동적으로 늘어나며 수명과 길이를 미리 알 수 없음
    • 이와 같은 KV cache 특성으로 인해 연속적인 공간에 저장할 경우 비효율성이 발생함
  • 기존 LLM serving 시스템의 문제점
    • Internal fragmentation & external fragmentation 
    • 메모리 공유 기회를 활용할 수 없음
      • 디코딩 알고리즘 중에서 병렬 샘플링 & beam 검색을 이용하는 경우, KV cache를 부분적으로 공유할 수 있는 여러 시퀀스로 구성됨
      • 기존 시스템은 KV catche가 별도의 연속된 공간에 저장되기 때문에 메모리 공유가 불가능함
• PagedAttention 제안
  • Requests의 KV cache를 블록으로 나누고 OS의 가상 메모리처럼 KV cache를 유연한 방식으로 관리함
    • Block=page, token=byte, request=process
  • 상대적으로 작은 블록을 사용하고 필요에 따라 함당함으로써 internal fragmentation을 완화
  • 모든 블록의 크기를 동일하게 하여 external fragmentation을 제거함
  • 동일한 requests과 다른 requests간 조차도 블록 단위로 메모리 공유가 가능하게 함

LLM serving의 메모리 문제

• 배치 처리를 할 수 있는 requests는 여전히 GPU 메모리 용량, 특히 KV cache를 저장하기 위해 할당된 공간에 의해 제한
• LLM serving system의 throughput은 memory-bound임. 다음과 같은 문제를 해결해야 함
  • 큰 KV cache
    • KV cache 크기는 requests에 따라 빠르게 증가함
    • 사용 가능한 메모리를 모두 KV cache에 할당하더라도 수용할 수 있는 requests은 수십개에 불과함
    • A100에서 H100으로 세대가 변경되었을 때 FLOPS는 2배 이상 증가하지만 GPU 메모리는 최대 80GB로 유지됨 (메모리 병목이 심화됨
  • 복잡한 디코딩 알고리즘
    • 디코딩 알고리즘에 따라 프롬프트 부분의 KV cache를 공유하여 메모리 사용량을 최소화할 수 있음
    • 생성 단계 중 KV cache는 서로 다른 샘플 결과와 컨텍스트 및 위치에 대한 의존성으로 인해 공유되지 않은 상태로 유지되어야 함
    • KV cache 공유 범위는 사용된 특정 디코딩 알고리즘에 따라 다름   
  • 알 수 없는 입력과 출력 길이에 대한 스케줄링
    • 입력: 광범위한 프롬프트 길이를 수용할 수 있는 메모리 관리 시스템이 필요함
    • 디코딩 시 requests의 출력 길이가 늘어나면 KV cache에 필요한 메모리 공간이 커짐
    • Incoming requests이나 기존 프로프트에 대한 지속적인 생성에 사용 가능한 메모리가 소진됨
• 기존 시스템의 메모리 관리
  • 현 딥러닝 프레임워크는 텐서를 대부분은 연속 메모리에 저장함
  • 예측할 수 없는 입력과 출력 길이로 인해 최대 시퀀스 길이를 기반으로 요청에 대한 메모리 청크를 정적으로 할당함
  • 기존 시스템의 pre-allocation 방법의 메모리 낭비 요소
    • Reserved slots for future tokens
    • Internal fragmentation: 최대 시퀀스 길이와 관계
    • External fragmentation: buddy allocator와 같은 allocator의 영향

해결 방법: Paged Attention

• PagedAttention
  • 고전적인 paging에 영감을 받아 연속적인 KV을 연속되지 않은 메모리 공간에 저장함
  • 어텐션 계산 중에 PagedAttention 커널은 다양한 KV 블럭을 별도로 식별하고 가져옴
• KV Cache Manager
  • VLLM은 가상 메모리기반의 아이디어를 사용하여 LLM 서비스에서 KV cache를 관리함
  • PagedAttention을 사용하여 KV cache를 가상 메모리의 페이지와 같은 고정 크기 KV 블록으로 구성함
  • Requests의 KV caches는 논리적 KV cache의 series로 표현되며 새로운 토큰과 해당 KV cache가 생성될 때 왼쪽에서 오른쪽으로 채워짐. 마지막 채워지지 않은 KV 블럭은 미래 생성을 위해 예약됨
  • GPU 워커의 block engine은 GPU DRAM의 연속적인 chunk를 할당하고 그것을 물리적인 KV blocks으로 분할함
  • KV block 매니저는 또한 각 requests의 논리적 & 물리적 KV 블럭간의 매핑을 포함한 block tables를 유지함
  • 논리적 KV block과 물리적 KV block을 분리하면 vLLM이 KV cache 메모리를 동적으로 늘릴 수 있으므로 기존 시스템에서 대부분의 메모리 낭비가 제거됨
• PagedAttention & vLLM으로 decoding하기
  • ① OS의 가상 메모리에서와 마찬가지로 vLLM은 처음에 생성 가능한 최대 시퀀스 길이를 위해 메모리를 예약할 필요가 없음. 대신 프롬프트 계산 중에 생성된 KV 캐시를 수용하기 위해 필요한 KV 블록만 예약함.이 경우 프롬프트에는 7개의 토큰이 있으므로 vLLM은 처음 2개의 논리적 KV 블록(논리적 Block 0 및 1)을 2개의 물리적 KV 블록(각각 물리적 Block 7 및 1)에 매핑함. Prefill 단계에서 vLLM은 프롬프트와 첫 번째 출력 토큰의 KV 캐시를 기존의 self-attention 알고리즘으로 생성함. 그런 다음 vLLM은 처음 4개 토큰의 KV 캐시를 논리 블록 0에, 다음 3개 토큰을 논리 블록 1에 저장함. 나머지 슬롯은 후속 자동 회귀 생성 단계를 위해 예약됨
  • ② 첫 번째 auto-regressive 디코딩 단계에서 vLLM은 물리적 블록 7과 1에서 PagedAttention 알고리즘을 사용하여 새 토큰을 생성함. 마지막 논리 블록에 사용 가능한 슬롯이 하나 남아 있으므로 새로 생성된 KV 캐시가 여기에 저장되고 Block Table의 # filled 레코드가 업데이트됨.
  • ③ 두 번째 디코딩 단계에서 마지막 논리 블록이 가득 차면 vLLM은 새로 생성된 KV 캐시를 새 논리 블록에 저장하고, 이를 위해 새 물리적 블록(물리적 Block 3)을 할당하고 이 매핑을 블록 테이블에 저장함

  • 

  • 전역적으로, 각 디코딩 반복에 대해 vLLM은 먼저 일괄 처리할 후보 시퀀스 집합을 선택하고, 새로 필요한 논리 블록을 위한 물리적 블록을 할당함. 그런 다음 vLLM은 현재 반복의 모든 입력 토큰(즉, 프롬프트 단계 요청에 대한 모든 토큰과 생성 단계 요청에 대한 최신 토큰)을 하나의 시퀀스로 연결하고 이를 LLM에 공급함. LLM이 계산하는 동안 vLLM은 PagedAttention 커널을 사용하여 논리적 KV 블록의 형태로 저장된 이전 KV 캐시에 액세스하고 새로 생성된 KV 캐시를 물리적 KV 블록에 저장함.
  • KV 블록(Block Size > 1)에 여러 개의 토큰을 저장하면 PagedAttention 커널이 더 많은 위치에서 KV 캐시를 병렬로 처리할 수 있으므로 하드웨어 사용률을 높이고 지연 시간을 줄일 수 있음. 그러나 블록 크기가 커지면 메모리 조각화도 증가함
  • 다시 말하지만, vLLM은 더 많은 토큰과 해당 토큰의 KV 캐시가 생성됨에 따라 새로운 물리적 블록을 논리적 블록에 동적으로 할당함. 모든 블록이 왼쪽에서 오른쪽으로 채워지고 이전 블록이 모두 채워질 때만 새로운 물리적 블록이 할당되기 때문에, vLLM은 한 블록 내에서 요청이 낭비되는 메모리를 모두 제한하여 모든 메모리를 효과적으로 활용할 수 있음. 이를 통해 더 많은 요청을 메모리에 일괄 처리할 수 있으므로 처리량이 향상됨. 요청이 생성을 완료하면 해당 KV 블록은 다른 요청의 KV 캐시를 저장하기 위해 해제될 수 있음.
  • 그림 7에서는 두 개의 시퀀스에 대한 메모리를 관리하는 vLLM의 예를 보여줌. 두 시퀀스의 논리 블록은 GPU 워커의 블록 엔진이 예약한 공간 내에서 서로 다른 물리적 블록에 매핑됨. 두 시퀀스의 인접한 논리 블록은 물리적 GPU 메모리에서 인접할 필요가 없으며 물리적 블록의 공간은 두 시퀀스 모두에서 효과적으로 활용할 수 있음.

  • 

    Fig 7

• 다른 디코딩 시나리오에 적용하기
  • 많은 성공적인 LLM 애플리케이션에서 LLM 서비스는 복잡한 액세스 패턴과 더 많은 메모리 공유 기회를 보여주는 더 복잡한 디코딩 시나리오를 제공해야 함. 이 섹션에서는 이들에 대한 vLLM의 일반적인 적용 가능성을 보여줌
  • Parallel sampling
    •  LLM 기반 프로그램 어시스턴트에서 LLM은 하나의 입력 프롬프트에 대해 여러 개의 샘플링된 출력을 생성하며, 사용자는 다양한 후보 중에서 원하는 출력을 선택할 수 있음. 지금까지는 요청이 단일 시퀀스를 생성한다고 암묵적으로 가정함. 이 백서의 나머지 부분에서는 요청이 여러 개의 시퀀스를 생성하는 보다 일반적인 경우를 가정함. Parallel sampling에서는 하나의 요청에 동일한 입력 프롬프트를 공유하는 여러 샘플이 포함되므로 프롬프트의 KV 캐시도 공유할 수 있음. vLLM은 PagedAttention 및 페이징 메모리 관리를 통해 이러한 공유를 쉽게 실현하고 메모리를 절약할 수 있음.
    • 그림 8은 두 출력에 대한 병렬 디코딩의 예임. 두 출력 모두 동일한 프롬프트를 공유하기 때문에 프롬프트 단계에서 프롬프트 상태의 복사본을 위한 공간을 하나만 예약하고, 두 시퀀스의 프롬프트에 대한 논리 블록은 동일한 물리적 블록에 매핑됨(두 시퀀스의 논리 블록 0과 1은 각각 물리적 블록 7과 1에 매핑됨).
    • 하나의 물리적 블록이 여러 논리 블록에 매핑될 수 있으므로 각 물리적 블록에 대한 참조 횟수를 도입합니다. 이 경우 물리 블록 7과 1의 참조 카운트는 모두 2임. 생성 단계에서 두 출력은 서로 다른 출력 토큰을 샘플링하므로 KV 캐시를 위한 별도의 스토리지가 필요함. vLLM은 여러 시퀀스에 의해 수정이 필요한 물리 블록에 대해 블록 단위로 copy on write 메커니즘을 구현하며, 이는 OS 가상 메모리의 copy on write 기법과 유사함(예: 프로세스 fork 시. Copy on Write는 리소스가 복제되었지만 수정되지 않은 경우에 새 리소스를 만들 필요 없이 복사본과 원본이 리소스를 공유하고, 복사본이 수정되었을 때만 새 리소스를 만드는 리소스 관리 기법). 
    • 구체적으로 그림 8에서 샘플 A1이 마지막 논리 블록(논리 블록 1)에 쓰기를 해야 할 때 vLLM은 해당 물리 블록(물리 블록 1)의 참조 카운트가 1보다 크다는 것을 인식하고 새로운 물리 블록(물리 블록 3)을 할당하고 블록 엔진에 물리 블록 1의 정보를 복사하도록 지시한 후 참조 카운트를 1로 줄임. 다음으로, 샘플 A2가 물리적 블록 1에 기록할 때 이미 참조 카운트가 1로 줄어들었으므로 A2는 새로 생성된 KV 캐시를 물리적 블록 1에 직접 기록함.
    • 요약하자면, vLLM을 사용하면 여러 출력 샘플에 걸쳐 프롬프트의 KV 캐시를 저장하는 데 사용되는 대부분의 공간을 공유할 수 있음(단, 최종 논리 블록은 카피 온 쓰기 메커니즘으로 관리됨). 여러 샘플에 걸쳐 물리적 블록을 공유하면 특히 긴 입력 프롬프트의 경우 메모리 사용량을 크게 줄일 수 있음

    • 

      Fig. 8

  • Beam Search
    • Parallel Sampling과 달리 Beam Search 기능은 초기 프롬프트 블록뿐만 아니라 서로 다른 후보 간에 다른 블록도 공유하며, 공유 패턴은 디코딩 프로세스가 진행됨에 따라 동적으로 변경됨. 이는 복합 포크로 생성된 OS의 프로세스 트리와 유사함. 그림  9는 vLLM이 빔 검색 예를 위해 KV 블록을 관리하는 방법을 보여줌.(k=4) 
    • 점선으로 표시된 반복 이전에는 각 후보 시퀀스가 ​​4개의 전체 논리 블록을 사용함. 모든 빔 후보는 첫 번째 블록 0(즉, 프롬프트)을 공유합니다. 후보 3은 두 번째 블록의 다른 후보와 다릅니다. 후보 0~2는 처음 3개 블록을 공유하고 네 번째 블록에서 분기됩니다. 후속 반복에서 상위 4개의 가능한 후보는 모두 후보 1과 2에서 유래합니다. 원래 후보 0과 3은 더 이상 상위 후보에 속하지 않으므로 해당 논리 블록이 해제되고 해당 물리 블록의 참조 카운트가 감소합니다. vLLM은 참조 횟수가 0에 도달한 모든 물리적 블록(블록 2, 4, 5, 8)을 해제합니다. 그런 다음 vLLM은 새 물리적 블록(블록 9-12)을 할당하여 새 후보의 새 KV 캐시를 저장합니다. 이제 모든 후보는 블록 0, 1, 3을 공유합니다. 후보 0과 1은 블록 6을 공유하고, 후보 2와 3은 블록 7을 추가로 공유합니다.
    • 이전 LLM 서비스 시스템에는 빔 후보 전체에 걸쳐 KV 캐시의 메모리 복사본이 자주 필요했습니다. 예를 들어, 그림 9 의 경우  점선 이후에 후보 3은 생성을 계속하려면 후보 2의 KV 캐시의 상당 부분을 복사해야 합니다. 이러한 빈번한 메모리 복사 오버헤드는 vLLM의 물리적 블록 공유를 통해 크게 줄어듭니다. vLLM에서는 서로 다른 빔 후보의 대부분의 블록을 공유할 수 있습니다. 쓰기 중 복사 메커니즘은 병렬 디코딩에서와 같이 새로 생성된 토큰이 이전 공유 블록 내에 있는 경우에만 적용됩니다. 여기에는 하나의 데이터 블록만 복사하는 작업이 포함됩니다.

    • 

• 스케줄링 &  Preemption
  • 요청 트래픽이 시스템 용량을 초과하면 vLLM은 요청의 하위 집합에 우선순위를 지정해야 합니다. vLLM에서는 모든 요청에 대해 선착순(FCFS) 스케줄링 정책을 채택하여 공정성을 보장하고 요청 고갈을 방지합니다. 요청을 선점해야 하는 경우 vLLM은 가장 먼저 도착한 요청이 먼저 제공되고 가장 늦게 도착한 요청이 먼저 선점되도록 합니다.
  • LLM 서비스는 입력 프롬프트의 길이가 크게 달라질 수 있고, 입력 프롬프트와 모델에 따라 결과 출력 길이를 선험적으로 알 수 없다는 독특한 문제에 직면해 있습니다. 요청과 그 출력의 수가 증가함에 따라 vLLM은 새로 생성된 KV 캐시를 저장하기 위한 GPU의 물리적 블록이 부족해질 수 있습니다. 이러한 상황에서 vLLM이 답해야 하는 두 가지 전형적인 질문이 있습니다.
    • (1) 어떤 블록을 퇴거시켜야 하는가?
    • (2) 필요한 경우 퇴거된 블록을 다시 복구하는 방법은 무엇인가? 
  • 일반적으로 퇴거 정책은 휴리스틱을 사용하여 향후 가장 멀리 접근될 블록을 예측하고 해당 블록을 퇴거시킵니다. 우리의 경우 시퀀스의 모든 블록이 함께 액세스된다는 것을 알기 때문에, 시퀀스의 모든 블록을 퇴거하거나 아무것도 퇴거하지 않는 all-or-nothing 퇴거 정책을 구현합니다.
  • 또한 하나의 요청 내에 있는 여러 시퀀스(예: 하나의 빔 검색 요청에 있는 빔 후보)는 시퀀스 그룹으로 갱 스케줄링됩니다. 한 시퀀스 그룹 내의 시퀀스는 해당 시퀀스 간의 메모리 공유 가능성으로 인해 항상 함께 선점되거나 재조정됩니다. 퇴거된 블록을 복구하는 방법에 대한 두 번째 질문에 답하기 위해 두 가지 기술을 고려합니다
  • Swapping
    • 이는 퇴거된 페이지를 디스크의 스왑 공간에 복사하는 대부분의 가상 메모리 구현에서 사용하는 고전적인 기법입니다. 우리의 경우, 퇴거된 블록을 CPU 메모리에 복사합니다. 
    • vLLM에는 GPU block allocator 외에도 CPU RAM으로 스왑된 물리적 블록을 관리하기 위한 CPU block allocator가 포함되어 있습니다. 새로운 토큰을 위한 물리적 블록의 여유 공간이 소진되면 vLLM은 일련의 시퀀스를 선택하여 해당 KV 캐시를 퇴거하고 CPU로 전송합니다. 
    • 시퀀스를 선점하고 해당 블록을 퇴거시키면 vLLM은 선점한 모든 시퀀스가 완료될 때까지 새 요청 수락을 중지합니다. 요청이 완료되면 해당 블록이 메모리에서 해제되고 선점된 시퀀스의 블록이 다시 가져와서 해당 시퀀스의 처리를 계속합니다. 
    • 이 설계에서는 CPU RAM으로 스왑되는 블록의 수가 GPU RAM의 총 물리적 블록 수를 초과하지 않으므로 CPU RAM의 스왑 공간은 KV 캐시에 할당된 GPU 메모리에 의해 제한된다는 점에 유의하세요.
  • Recomputation
    • 이 경우, 선점된 시퀀스가 다시 예약될 때 KV 캐시를 다시 계산하기만 하면 됩니다. 디코딩 시 생성된 토큰을 원래 사용자 프롬프트와 새 프롬프트로 연결하여 모든 위치의 KV 캐시를 한 번의 프롬프트 단계 반복으로 생성할 수 있으므로 재계산 지연 시간이 원래 지연 시간보다 훨씬 짧아질 수 있습니다.
  • Swapping 및 Recomputation의 성능은 CPU RAM과 GPU 메모리 사이의 대역폭과 GPU의 연산 능력에 따라 달라짐
• 분산 실행
  • 많은 LLM의 매개변수 크기가 단일 GPU의 용량을 초과합니다. 따라서 이를 분산된 GPU에 분할하여 모델 병렬 방식으로 실행해야 합니다. 이를 위해서는 분산 메모리를 처리할 수 있는 메모리 관리자가 필요합니다. vLLM은 트랜스포머에서 널리 사용되는 메가트론-LM 스타일 텐서 모델 병렬 처리 전략을 지원함으로써 분산 환경에서 효과적입니다. 이 전략은 선형 레이어를 분할하여 블록 단위 행렬 곱셈을 수행하고, GPU가 올 리듀스 연산을 통해 중간 결과를 지속적으로 동기화하는 SPMD(단일 프로그램 다중 데이터) 실행 스케줄을 준수하는 방식입니다. 특히 attention 연산자는 attention head 차원으로 분할되며, 각 SPMD 프로세스는 MHA에서 attention head의 하위 집합을 처리합니다.
  • 모델 병렬 실행을 사용하더라도 각 모델 샤드는 여전히 동일한 입력 토큰 세트를 처리하므로 동일한 위치에 대한 KV 캐시가 필요하다는 것을 관찰할 수 있습니다. 따라서 vLLM은 그림 4와 같이 중앙 집중식 스케줄러 내에 단일 KV 캐시 관리자를 제공합니다. 여러 GPU 워커가 이 매니저와 논리적 블록에서 물리적 블록으로의 매핑을 공유합니다. 이러한 공통 매핑을 통해 GPU 워커는 각 입력 요청에 대해 스케줄러가 제공하는 물리적 블록으로 모델을 실행할 수 있습니다. 각 GPU 워커는 동일한 물리적 블록 ID를 갖지만, 워커는 해당 관심 헤드에 대해 KV 캐시의 일부만 저장합니다.
  • 각 단계에서 스케줄러는 먼저 일괄 처리의 각 요청에 대한 입력 토큰 ID와 각 요청에 대한 블록 테이블이 포함된 메시지를 준비합니다. 다음으로 스케줄러는 이 제어 메시지를 GPU 워커에 브로드캐스트합니다. 그러면 GPU 워커가 입력 토큰 ID로 모델을 실행하기 시작합니다. 관심 레이어에서 GPU 워커는 제어 메시지의 블록 테이블에 따라 KV 캐시를 읽습니다. 실행 중에 GPU 워커는 스케줄러의 조정 없이 중간 결과를 올 리듀스 통신 프리미티브와 동기화합니다. 결국 GPU 워커는 이 반복의 샘플링된 토큰을 스케줄러로 다시 전송합니다. 요약하면, GPU 워커는 각 디코딩 반복이 시작될 때 단계 입력과 함께 모든 메모리 관리 정보를 받기만 하면 되므로 메모리 관리에서 동기화할 필요가 없습니다.
• 구현
  • 커널수준 최적화
    • 페이징어텐션은 기존 시스템에서 효율적으로 지원하지 않는 메모리 액세스 패턴을 도입했기 때문에 이를 최적화하기 위해 몇 가지 GPU 커널을 개발했습니다. 
      • (1) Fused reshape and block write: 모든 Transformer 레이어에서 새로운 KV 캐시는 블록으로 분할되고, 블록 읽기에 최적화된 메모리 레이아웃으로 재형성된 다음 블록 테이블에 지정된 위치에 저장됩니다. 커널 실행 오버헤드를 최소화하기 위해 이를 단일 커널로 융합합니다. 
      • (2) Fusing block read & attention:  블록 테이블에 따라 KV 캐시를 읽고 즉시 어텐션 연산을 수행할 수 있도록 FasterTransformer(NVIDIA, 2023a)의 어텐션 커널을 조정합니다. 통합 메모리 액세스를 보장하기 위해 각 블록을 읽기 위해 GPU 워프를 할당합니다. 또한 요청 배치 내에서 가변 시퀀스 길이에 대한 지원을 추가합니다. 
      • (3) Fused block copy: copy-on-write 메커니즘에 의해 실행되는 블록 복사 작업은 불연속적인 블록에서 작동할 수 있습니다. cudaMemcpyAsync API를 사용하면 작은 데이터 이동이 여러 번 호출될 수 있습니다. 이러한 오버헤드를 완화하기 위해 여러 블록에 대한 복사 작업을 한 번의 커널 실행으로 일괄 처리하는 커널을 구현했습니다.

    • 

      Fig. 11

  • 다양한 디코딩 알고리즘 지원
    • vLLM은 fork, append, free의 세 가지 주요 방법을 사용하여 다양한 디코딩 알고리즘을 구현합니다.
      • fork 메서드는 기존 시퀀스에서 새 시퀀스를 생성합니다.
      • append 메서드는 시퀀스에 새 토큰을 추가합니다.
      • 마지막으로 free 메서드는 시퀀스를 삭제합니다.
    • 예를 들어, Parallel Sampling에서 vLLM은 fork 메서드를 사용하여 단일 입력 시퀀스에서 여러 개의 출력 시퀀스를 생성합니다. 그런 다음 모든 반복에서 추가를 사용하여 이러한 시퀀스에 새 토큰을 추가하고 자유를 사용하여 중지 조건을 충족하는 시퀀스를 삭제합니다. 동일한 전략이 vLLM의 빔 검색과 접두사 공유에도 적용됩니다. 향후 디코딩 알고리즘도 이러한 방법을 결합하여 지원할 수 있다고 생각합니다.
• Evaluation
  • 실험 셋업
    • Model & Server Configuration
      • 평가에는 13B, 66B, 175B 파라미터를 가진 OPT(Zhang et al., 2022) 모델과 13B 파라미터를 가진 LLaMA(Touvron et al., 2023)를 사용합니다. 13B와 66B는 LLM 리더보드(ORG, 2023)에서 볼 수 있듯이 LLM에 널리 사용되는 크기이며, 175B는 유명한 GPT-3(Brown et al., 2020) 모델에 사용되는 크기입니다. 모든 실험에서는 구글 클라우드 플랫폼에서 NVIDIA A100 GPU가 탑재된 A2 인스턴스를 사용했습니다. 자세한 모델 크기와 서버 구성은 표 1에 나와 있습니다. 
      •  
    • Workloads
      • 실제 LLM 서비스의 입력 및 출력 텍스트를 포함하는 ShareGPT(Team, 2023) 및 Alpaca(Taori 외, 2023) 데이터셋을 기반으로 워크로드를 합성합니다. ShareGPT 데이터 세트는 사용자가 ChatGPT(OpenAI, 2022)와 공유한 대화의 모음입니다. Alpaca 데이터 세트는 GPT-3.5에서 자체 인스트럭션으로 생성된 명령어 데이터 세트입니다(Wang et al., 2022a).
      • 우리는 데이터 세트를 토큰화하고 입력 및 출력 길이를 사용하여 클라이언트 요청을 합성합니다. 그림 11에서 볼 수 있듯이 ShareGPT 데이터셋은 알파카 데이터셋보다  8.4× 더 긴 입력 프롬프트와 5.8× 더 긴 출력을 가지고 있으며, 분산도 더 큽니다. 이러한 데이터 세트에는 타임스탬프가 포함되어 있지 않기 때문에, 요청 속도가 다른 푸아송 분포를 사용하여 요청 도착 시간을 생성합니다.
    • Baseline 1: FasterTransformer

      • (엔비디아,2023a) 은 지연 시간에 고도로 최적화된 분산 추론 엔진입니다. FasterTransformer에는 자체 스케줄러가 없으므로 Triton  (NVIDIA,[N. 디.]) . 구체적으로 최대 배치 크기를 설정합니다.비GPU 메모리 용량에 따라 각 실험마다 최대한 큰 값을 사용합니다. 스케줄러는 다음을 수행합니다.비가장 먼저 도착한 요청 수를 확인하고 처리를 위해 일괄 처리를 FasterTransformer로 보냅니다.

    • Baseline 1: Orca
      • Orca(Yu et al. ,2022년)는 처리량에 최적화된 최첨단 LLM 서비스 시스템입니다. Orca는 공개적으로 사용할 수 없으므로 자체 버전의 Orca를 구현합니다. Orca는 버디 할당 알고리즘을 사용하여 KV 캐시를 저장할 메모리 주소를 결정한다고 가정합니다. 요청 출력을 위해 공간을 얼마나 과도하게 예약하는지에 따라 Orca의 세 가지 버전을 구현합니다.
      • Orca (Oracle): 우리는 시스템이 요청에 대해 실제로 생성될 출력의 길이를 알고 있다고 가정합니다. 이는 실제로 달성하기 어려운 Orca의 상한 성능을 보여줍니다.
      • Orca (Pow2): 우리는 시스템이 출력을 위한 공간을 최대 2만큼 초과 예약한다고 가정합니다.×. 예를 들어 실제 출력 길이가 25이면 출력용으로 32개의 위치를 ​​예약합니다.
      • Orca (Max): 우리는 시스템이 항상 모델의 최대 시퀀스 길이, 즉 2048 토큰까지 공간을 예약한다고 가정합니다.
    • Key metrics
      • 우리는 서비스 처리량에 중점을 둡니다. 특히, 요청 속도가 다른 워크로드를 사용하여 Orca에서와 같이 모든 요청의 엔드투엔드 지연 시간을 출력 길이로 나눈 평균인 시스템의 정규화 지연 시간을 측정합니다(Yu et al., 2022). 처리량이 많은 서빙 시스템은 높은 요청 속도에 대해 낮은 정규화 지연 시간을 유지해야 합니다. 대부분의 실험에서는 1시간 추적 데이터를 사용해 시스템을 평가합니다. 예외적으로, 비용 제한으로 인해 OPT-175B 모델의 경우 15분 트레이스를 사용합니다.
  • Basic Sampling
    • 세 가지 모델과 두 가지 데이터 세트에 대해 기본 샘플링(요청당 하나의 샘플)을 사용하여 vLLM의 성능을 평가합니다. 그림 12의 첫 번째 행은 ShareGPT 데이터 세트에 대한 결과를 보여줍니다. 곡선은 요청 속도가 증가함에 따라 지연 시간이 처음에는 점진적인 속도로 증가하다가 갑자기 폭발적으로 증가하는 것을 보여줍니다. 이는 요청 속도가 서버 시스템의 용량을 초과하면 대기열 길이가 무한히 증가하고 요청의 지연 시간도 증가하기 때문일 수 있습니다.
    • ShareGPT 데이터셋에서 vLLM은 동일한 latency를 유지하면서 Orca(오라클)에 비해 1.7x - 2.7x Orca(Max)에 비해 2.7x-8x 더  높은 request rates을 유지할 수 있습니다. 이는 vLLM의 PagedAttention이 메모리 사용량을 효율적으로 관리할 수 있어 Orca보다 더 많은 요청을 일괄 처리할 수 있기 때문입니다. 예를 들어, 그림 13(a)에서 볼 수 있듯이 OPT-13B vLLM는  Orca(Oracle)에 비해 동시에 2.2x 더 많은 requests를 처리할 수 있으며 Orca(MaX)에 비해 4.3x 더 많은 requests를 처리할 수 있다. FasterTransformer와 비교했을 때, vLLM은 최대 22× 더 높은 requests rates를 견딜 수 있는데, 이는 FasterTransformer가 세분화된 스케줄링 메커니즘을 사용하지 않고 Orca(Max)처럼 메모리를 비효율적으로 관리하기 때문입니다.
    • 그림 12와 그림 13(b)의 두 번째 행은 ShareGPT 데이터 세트와 유사한 추세를 따르는 Alpaca 데이터 세트의 결과를 보여줍니다. 한 가지 예외가 있는데, 그림 12 (f)에서는 Orca(Oracle) 및 Orca(Pow2)에 비해 vLLM의 우위가 덜 두드러집니다. 이는 OPT-175B의 모델 및 서버 구성(표 1)이 KV 캐시를 저장하는 데 사용할 수 있는 큰 GPU 메모리 공간을 허용하는 반면, 알파카 데이터세트는 시퀀스가 짧기 때문입니다. 이 설정에서 Orca(Oracle)와 Orca(Pow2)는 메모리 관리의 비효율성에도 불구하고 많은 수의 요청을 일괄 처리할 수 있습니다. 그 결과, 시스템의 성능은 메모리 제한이 아닌 컴퓨팅 제한이 됩니다.
    •  
  • Parallel Sampling & Beam Search
    • Parallel Sampling과 Beam Search이라는 두 가지 널리 사용되는 샘플링 방법을 사용하여 PagedAttention에서 메모리 공유의 효과를 평가합니다. 병렬 샘플링에서는 요청의 모든 병렬 시퀀스가 프롬프트에 대한 KV 캐시를 공유할 수 있습니다. 그림 14의 첫 번째 행에서 볼 수 있듯이, 샘플링할 시퀀스 수가 많을수록 vLLM은 Orca 기준선보다 더 많은 개선을 가져옵니다. 마찬가지로 그림 14의 두 번째 행은 빔 폭이 다른 빔 검색의 결과를 보여줍니다. 빔 검색을 통해 더 많은 공유가 가능하기 때문에 vLLM은 훨씬 더 큰 성능 이점을 보여줍니다. OPT-13B 및 알파카 데이터 세트에서 vLLM의 개선은 basic sampling에서 1.3x부터 6의 width를 가진 beam search에서 2.3x로 간다. 
  • Shared Prefix
  • Chatbot
    • 챗봇(OpenAI, 2022; Google, 2023; Chiang et al., 2023)은 LLM의 가장 중요한 응용 분야 중 하나입니다. 챗봇을 구현하려면 모델이 채팅 기록과 마지막 사용자 쿼리를 프롬프트에 연결하여 응답을 생성하도록 합니다. 채팅 기록과 사용자 쿼리는 ShareGPT 데이터 세트를 사용하여 합성합니다. OPT-13B 모델의 컨텍스트 길이가 제한되어 있기 때문에 프롬프트를 마지막 1024개 토큰으로 잘라내고 모델이 최대 1024개의 토큰을 생성하도록 합니다. 서로 다른 대화 라운드 사이에 KV 캐시를 저장하면 대화 라운드 사이에 다른 요청을 위한 공간을 차지하게 되므로 저장하지 않습니다.
    • 그림 17은 vLLM이 세가지 Orca baseline 대비 2× 더 높은 요청 속도를 유지할 수 있음을 보여줍니다. ShareGPT 데이터 세트에는 긴 대화가 많이 포함되어 있기 때문에 대부분의 요청에 대한 입력 프롬프트에는 1024개의 토큰이 있습니다. 버디 할당 알고리즘으로 인해, Orca 기준선은 출력 길이를 예측하는 방식에 관계없이 요청 출력을 위해 1024개의 토큰을 위한 공간을 예약합니다. 이러한 이유로 세 가지 Orca 기준선은 비슷하게 작동합니다. 이와는 대조적으로 vLLM은 메모리 조각화 및 예약 문제를 해결하기 때문에 긴 프롬프트를 효과적으로 처리할 수 있습니다.
• Ablation Studies
  • Kernel Microbenchmark
    • PagedAttention의 동적 블록 매핑은 저장된 KV 캐시와 관련된 GPU 작업, 즉 블록 읽기/쓰기 및 주의의 성능에 영향을 줍니다. 기존 시스템과 비교했을 때, 우리의 GPU 커널(§5)은 블록 테이블 액세스, 추가 분기 실행, 가변 시퀀스 길이 처리와 같은 추가 오버헤드를 수반합니다. 그림 18(a)에서 볼 수 있듯이, 이는 고도로 최적화된 FasterTransformer 구현에 비해 20~26% 더 높은 주의 커널 지연 시간을 초래합니다. 이 오버헤드는 주의 연산자에게만 영향을 미치고 리니어와 같은 모델의 다른 연산자에게는 영향을 미치지 않기 때문에 크지 않다고 생각합니다. 이러한 오버헤드에도 불구하고 PagedAttention은 엔드투엔드 성능에서 vLLM이 FasterTransformer보다 훨씬 뛰어난 성능을 발휘하도록 합니다(§6).
  • Impact of Block Size
    • 블록 크기 선택은 vLLM의 성능에 상당한 영향을 미칠 수 있습니다. 블록 크기가 너무 작으면 vLLM이 KV 캐시를 읽고 처리하는 데 GPU의 병렬 처리 기능을 충분히 활용하지 못할 수 있습니다. 블록 크기가 너무 크면 내부 조각화가 증가하여 공유 확률이 감소합니다.
    • 그림 18(b)는 고정된 요청 속도에서 기본 샘플링이 포함된 ShareGPT 및 알파카 트레이스를 사용하여 블록 크기가 다른 vLLM의 성능을 평가한 것입니다. ShareGPT 추적에서는 16~128의 블록 크기가 가장 좋은 성능을 보입니다. 알파카 추적에서는 블록 크기 16과 32가 잘 작동하지만, 블록 크기가 클수록 시퀀스가 블록 크기보다 짧아지기 때문에 성능이 크게 저하됩니다. 실제로 블록 크기 16은 GPU를 효율적으로 활용하기에 충분히 크고 대부분의 워크로드에서 심각한 내부 조각화를 피하기에 충분히 작은 것으로 나타났습니다. 따라서 vLLM은 기본 블록 크기를 16으로 설정합니다.
  • Recomputation & Swapping 비교하기
    • vLLM은 복구 메커니즘으로 재계산과 스와핑을 모두 지원합니다. 두 방법 간의 장단점을 이해하기 위해 그림 19와 같이 엔드투엔드 성능을 평가하고 오버헤드를 마이크로벤치마킹했습니다. 그 결과, 스와핑은 블록 크기가 작을 때 과도한 오버헤드를 발생시킨다는 것을 알 수 있었습니다. 이는 블록 크기가 작으면 CPU와 GPU 간에 수많은 작은 데이터 전송이 발생하여 유효 PCIe 대역폭이 제한되기 때문입니다. 반면, 재계산은 KV 블록을 활용하지 않기 때문에 재계산 오버헤드는 다양한 블록 크기에 걸쳐 일정하게 유지됩니다. 따라서 블록 크기가 작을 때는 재계산이 더 효율적이고, 블록 크기가 클 때는 스와핑이 더 효율적이지만 재계산 오버헤드는 스와핑 지연 시간의 20%를 넘지 않습니다. 16에서 64 사이의 중간 블록 크기에서는 두 가지 방법이 비슷한 엔드 투 엔드 성능을 보입니다.
• Discussion
  • 가상 메모리 및 페이징 기술을 다른 GPU 워크로드에 적용하기. 
    • 가상 메모리와 페이징의 개념은 워크로드에 동적 메모리 할당(출력 길이를 선험적으로 알 수 없으므로)이 필요하고 그 성능이 GPU 메모리 용량에 의해 제한되기 때문에 LLM 서빙에서 KV 캐시를 관리하는 데 효과적입니다. 그러나 이는 일반적으로 모든 GPU 워크로드에 적용되는 것은 아닙니다. 예를 들어, DNN 훈련에서 텐서 모양은 일반적으로 정적이므로 메모리 할당을 미리 최적화할 수 있습니다. 또 다른 예로, LLM이 아닌 DNN을 서비스하는 경우, 성능은 주로 컴퓨팅에 제한되어 있기 때문에 메모리 효율이 증가해도 성능 개선이 이루어지지 않을 수 있습니다. 이러한 시나리오에서는 메모리 방향성 및 비연속 블록 메모리의 추가 오버헤드로 인해 vLLM의 기술을 도입하면 오히려 성능이 저하될 수 있습니다. 하지만 LLM 서비스와 유사한 속성을 가진 다른 워크로드에도 vLLM의 기술이 적용될 수 있기를 기대합니다.
  • 가상 메모리 및 페이징을 적용할 때 LLM에 특화된 최적화
    • vLLM은 애플리케이션별 시맨틱을 활용하여 가상 메모리 및 페이징의 개념을 재해석하고 보강합니다. 한 가지 예로 요청을 처리하려면 해당 토큰 상태가 모두 GPU 메모리에 저장되어야 한다는 사실을 활용하는 vLLM의 전부 또는 전무 스왑아웃 정책을 들 수 있습니다. 또 다른 예로는 퇴거된 블록을 복구하기 위한 재계산 방법이 있는데, 이는 OS에서는 불가능합니다. 또한 vLLM은 메모리 액세스 연산을 위한 GPU 커널과 주의와 같은 다른 연산을 위한 커널을 융합하여 페이징에서 메모리 간접 처리의 오버헤드를 완화합니다.
• 결론
  • 이 백서에서는 주의 키와 값을 비연속적인 페이징 메모리에 저장할 수 있는 새로운 주의 알고리즘인 PagedAttention을 제안하고, PagedAttention을 통해 효율적인 메모리 관리가 가능한 고처리량 LLM 서빙 시스템인 vLLM을 소개합니다.
  • 운영 체제에서 영감을 받아 가상 메모리와 copy-on-write와 같은 기존 기술을 어떻게 적용하여 KV 캐시를 효율적으로 관리하고 다양한 디코딩 알고리즘을 처리할 수 있는지 보여줍니다. 실험 결과, vLLM은 SOTA 시스템을 넘어 2~4×의 처리량 향상을 달성하는 것으로 나타났습니다.