VLLM 소개

September 29, 2023

PagedAttention은 vLLM(very Large Language Model)에서 사용되는 기술로, LLM(언어 모델)의 성능을 향상시키기 위한 것입니다. 이 기술의 주요 내용을 간단하게 설명하면 다음과 같습니다

문제점: 기존 LLM에서 성능이 메모리에 의해 제한되는 문제가 있었습니다. 특히, 자기회귀 디코딩 과정에서 LLM에 입력되는 토큰은 모두 어텐션 키(key)와 값(value) 텐서를 생성하며, 이러한 텐서는 GPU 메모리에 보관되어 다음 토큰을 생성하는 데 사용됩니다. 이러한 키와 값 텐서의 캐시를 KV 캐시라고 합니다.

KV 캐시의 문제점

대용량: LLaMA-13B 모델에서 하나의 시퀀스에 대해 최대 1.7GB의 메모리를 사용합니다.
동적 크기: 시퀀스 길이에 따라 크기가 달라지며, 이는 매우 다양하고 예측하기 어렵습니다. 이러한 KV 캐시를 효율적으로 관리하는 것은 큰 도전입니다. 기존 시스템에서는 메모리의 60% - 80%가 조각화와 과다 예약으로 인해 낭비되는 것으로 나타났습니다.

PagedAttention 소개

PagedAttention은 운영 체제의 가상 메모리와 페이지 기법에서 영감을 얻은 어텐션 알고리즘입니다.
기존의 전통적인 어텐션 알고리즘과 달리, PagedAttention은 연속된 키와 값들을 연속되지 않은 메모리 공간에 저장할 수 있게 합니다.
각 시퀀스의 KV 캐시를 블록으로 분할하고, 각 블록은 고정된 수의 토큰에 대한 키와 값들을 포함합니다. 어텐션 계산 중에 PagedAttention 커널은 이러한 블록을 효율적으로 식별하고 가져옵니다.

메모리 관리 및 공유

블록이 메모리 상에서 연속적이지 않아도 되므로 가상 메모리의 개념을 활용하여 키와 값들을 더 유연하게 관리할 수 있습니다. 블록은 페이지, 토큰은 바이트, 시퀀스는 프로세스로 비유할 수 있습니다.
시퀀스의 연속적인 논리적 블록은 블록 테이블을 통해 연속되지 않은 물리적 블록에 매핑됩니다. 새로운 토큰이 생성될 때마다 물리적 블록이 필요에 따라 할당됩니다.

성능 및 이점

PagedAttention은 메모리 낭비를 거의 최소화하며, 마지막 블록에서만 약 4%의 메모리 낭비가 발생합니다.
메모리 효율성 향상으로 인해 시스템은 더 많은 시퀀스를 함께 배치하고 GPU 활용도를 높이며, 결과적으로 처리량을 크게 향상시킬 수 있습니다.
PagedAttention은 메모리 공유를 통해 병렬 샘플링 및 빔 검색과 같은 복잡한 샘플링 알고리즘의 메모리 오버헤드를 크게 줄일 수 있습니다. 이로 인해 처리량이 최대 2.2배 향상될 수 있습니다.

vLLM과의 연관

PagedAttention은 vLLM의 핵심 기술로 사용되며, 다양한 모델을 지원하면서도 높은 성능과 사용하기 쉬운 인터페이스를 제공하는 LLM 추론 및 서빙 엔진입니다.

from vllm import LLM, SamplingParams

llm = LLM(model="lmsys/vicuna-7b-v1.5")

prompts = [
    "Hello, my name is",
    "The president of the United States is",
    "The capital of France is",
    "The future of AI is",
    "John F. Kennedy International Airport is",
    "In a galaxy far, far away, the Jedi fought against",
    "When life gives you lemons, make",
    "Once upon a time in a land filled with magic,",
    "Behind the mountains, there lies a hidden",
    "Exploring the depths of the ocean, scientists discovered",
    "Walking through the enchanted forest, I stumbled upon",
    "Amidst the bustling city streets, a lone street performer played",
    "As the sun set over the horizon, the sky turned into shades of",
    "Lost in a world of books, I found myself",
    "With a backpack full of dreams, I embarked on a journey to",
    "The sound of laughter echoed through the park as children",
    "Underneath the starry night sky, two lovers sat on a bench and",
    "Beneath the mask, a superhero grappled with",
    "In the laboratory, a team of scientists worked tirelessly to unlock",
    "As the first snowflake fell, a sense of wonder filled",
    "In the heart of the rainforest, a rare species of bird with vibrant feathers",
    ]

sampling_params = SamplingParams(temperature=0)

outputs = llm.generate(user_input, sampling_params)

# Print the outputs.
for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r} {generated_text!r}")

간단히 실행할 수 있지만 현재까지는 양자화된 모델을 사용이 어려운걸로 알고 있습니다.

그리고 실제 테스트 결과 input 데이터가 엄청 길때는 뭔가 버벅이는게 있는 거 같습니다. 확실하지 않지만 개인적인 경험입니다.