최대 24배 빠른 vLLM의 비밀 파헤치기

HuggingFace 대비 최대 24배 빠른 LLM 추론 엔진 vLLM(v0.1.2)의 내부 구현을 파이썬 스케줄링부터 CUDA 커널까지 파헤친다. KV 캐시를 블록 단위로 관리하고 어텐션 커널을 직접 최적화한 것이 속도의 핵심이다.

핵심 포인트

파이썬 계층은 LLM·LLMEngine·BlockManager·Worker로 나뉘고, 요청을 Waiting·Running·Swapped 3상태로 스케줄링한다.
메모리가 부족하면 우선순위 낮은 시퀀스를 선점(preempt)해 KV 캐시를 재계산하거나 CPU로 swap한다.
CUDA 계층은 reshape-and-cache 커널로 KV 값을 블록 구조에 저장하고, 최신 토큰 쿼리만으로 캐시된 KV에 어텐션하는 single query attention을 쓴다.
그리드를 [num_heads, num_sequences]로 잡고 128 스레드 블록에서 butterfly reduction으로 max·sum을 병렬 계산한다.
float16x2를 uint32_t로 패킹하는 Vec 구조로 SIMD 연산을 활용한다.

상세 정리

LLM 클래스: generate()가 프롬프트와 SamplingParams를 받아 추론을 시작한다.
LLMEngine: 요청 큐와 KV 캐시 초기화를 관리하며 스케줄링을 담당한다.
스케줄링: Waiting·Running·Swapped 3상태로 시퀀스를 관리한다. 실행 중 시퀀스를 우선하고, 메모리 부족 시 낮은 우선순위를 선점하며, 상태에 따라 KV 캐시를 재계산하거나 CPU로 스왑한다.
BlockManager: 참조 카운팅으로 메모리 블록을 할당한다. 블록 할당은 block_size x num_heads x head_size / x 형태로 구성한다.
Worker: 모델 추론과 입력 데이터 준비를 실행한다.
Continuous batching: 자원을 효율적으로 쓰도록 요청을 연속적으로 배치한다.
reshape-and-cache 커널: KV 값을 전용 블록 구조에 저장해 이후 재사용한다.
single query attention: 과거 KV는 캐시에서 읽고 최신 토큰 쿼리 하나만 계산해 중복 연산을 제거한다.
워프 수준 병렬화: 각 스레드 그룹이 한 토큰을 처리하고 블록 간 협력하며, butterfly reduction으로 병렬 max·sum을 효율적으로 수행한다.
Vec·SIMD: float16x2를 uint32_t로 패킹해 벡터 연산으로 처리량을 높인다.
범위: 분석 대상은 vLLM v0.1.2로, 이후 버전은 구현이 크게 다를 수 있음을 명시한다.