tiny-vLLM으로 처음부터 고성능 LLM 추론 엔진 구축하기

The complexity of modern Large Language Model (LLM) inference engines often hides the fundamental mathematics and hardware interactions that make them work. While high-level libraries like PyTorch make model design accessible, the actual process of serving these models at scale requires a deep understanding of GPU memory management, CUDA kernel engineering, and linear algebra.

tiny-vLLM은 이 과정을 명확히 설명하기 위해 만든 교육 프로젝트입니다. 블랙박스 구현을 제공하는 대신, C++와 CUDA에 대한 강좌 역할을 하며 개발자들이 고성능 추론 엔진을 구현하도록 안내합니다—산업 표준 vLLM의 “작은 형제”라고 할 수 있습니다. 엔진을 처음부터 직접 구축함으로써, 개발자는 원시 가중치 파일에서 작동하는 서버까지 필요한 수학과 로직을 스크래치에서 도출할 수 있습니다.

추론 서버의 구조

본질적으로 LLM은 방대한 부동소수점 숫자(가중치) 배열을 담은 파일일 뿐입니다. 이 가중치는 비용이 많이 드는 학습 단계의 산물입니다. 그러나 가중치 파일은 실행 파일이 아니라 설계도입니다. 이를 기능하게 만들기 위해 inference server는 다음을 수행해야 합니다:

  1. Load the weights를 Safetensors와 같은 형식에서 로드한다.
  2. Implement the architecture(수학 연산 순서)를 모델이 정의한 대로 구현한다(예: Llama 3.2).
  3. Execute these operations를 하드웨어에서 효율적으로 실행한다. 일반적으로 대규모 행렬 곱셈을 처리하기 위해 GPU를 사용한다.

하드웨어 마스터하기: CUDA와 메모리

고성능을 달성하기 위해 tiny-vLLM은 NVIDIA의 병렬 컴퓨팅 플랫폼인 CUDA를 활용합니다. 이 과정에서 중요한 부분은 **Host(CPU)**와 Device(GPU) 메모리의 차이를 이해하는 것입니다.

GPU는 시스템 DRAM에 직접 접근할 수 없기 때문에 데이터를 명시적으로 이동시켜야 합니다. 일반적인 워크플로는 cudaMalloc으로 GPU에 메모리를 할당하고 cudaMemcpy로 데이터를 전송하는 것입니다. 모든 고성능 엔진의 목표는 이러한 전송을 최소화하고 가능한 한 할당된 버퍼를 재사용하여 빈번한 할당에 따른 오버헤드를 피하는 것입니다.

심층 탐구: 추론 파이프라인 구현

1. Safetensors 로드

Safetensors는 빠르고 안전하기 때문에 널리 사용되는 형식입니다. Safetensors 파일은 헤더 크기, 텐서 이름·형태·오프셋을 포함하는 JSON 헤더, 그리고 원시 텐서 데이터로 구성됩니다. tiny-vLLM은 이러한 오프셋을 GPU 메모리 포인터에 직접 매핑하는 방법을 보여주어 엔진이 특정 가중치(예: 특정 레이어의 K-projection)를 최소 오버헤드로 가져올 수 있게 합니다.

2. 정밀도 트레이드오프: BF16

대부분의 최신 LLM은 **bfloat16 (BF16)**을 사용합니다. 표준 float16이 지수와 가수 비트를 균형 있게 배분하는 반면, BF16은 8비트 지수( float32와 동일)와 7비트 가수를 사용합니다. 이 트레이드오프는 LLM에 있어 매우 중요합니다. 범위(오버플로·언더플로) 문제를 방지하면서 약간의 정밀도 손실을 허용하는데, 실험적 증거에 따르면 이는 모델 정확도에 거의 영향을 주지 않습니다.

3. CUDA 커널 엔지니어링

tiny-vLLM의 가장 가치 있는 측면 중 하나는 맞춤형 CUDA 커널 작성에 중점을 둔다는 점입니다.

  • Embedding Gather: 첫 번째 단계는 토큰 ID를 벡터에 매핑하는 것입니다. NVIDIA GPU는 일반적으로 블록당 1024 스레드로 제한하지만, 임베딩(예: Llama)은 2048 요소이므로, 프로젝트는 각 스레드가 두 요소를 처리하도록 하여 하드웨어 제한 내에 머무는 일반적인 최적화를 가르칩니다.
  • RMSNorm and Parallel Reduction: Root Mean Square Layer Normalization은 벡터 전체에 대한 제곱합을 계산해야 합니다. 레이스 컨디션 없이 효율적으로 수행하기 위해 tiny-vLLM은 __shared__ 메모리와 __syncthreads()를 사용한 **parallel reduction (tree reduction)**을 구현하여 로그 형태로 값을 합칩니다.
  • RoPE (Rotary Positional Embeddings): 모델에 토큰 순서를 인식시키기 위해 RoPE는 은닉 상태에 회전을 적용합니다. 이는 CUDA 커널에서 사인·코사인 변환을 통해 구현된 복소수 수학을 포함합니다.

4. 행렬 곱셈과 cuBLAS 트릭

행렬 곱셈은 엔진에서 가장 무거운 작업입니다. 맞춤형 커널을 작성할 수도 있지만, 산업 표준은 cuBLAS입니다. 그러나 cuBLAS는 column-major 형식을 기대하는 반면, 대부분의 LLM 가중치는 row-major 형식으로 배포됩니다.

메모리에서 데이터를 물리적으로 전치하는 비용이 큰 과정을 피하기 위해, tiny-vLLM은 수학적 트릭을 사용합니다: 전치 플래그(CUBLAS_OP_T와 CUBLAS_OP_N)를 조작하고 곱셈 순서를 재배열($C^T = B \times A^T$)함으로써, 엔진은 행렬 데이터를 한 바이트도 이동시키지 않고 row-major 데이터를 column-major처럼 취급할 수 있습니다.

고급 추론 개념

엔진이 단일 토큰 생성에서 프로덕션 수준 서버로 발전함에 따라 여러 고급 개념이 도입됩니다:

  • Prefill vs. Decode: 첫 번째 토큰(prefill)은 전체 프롬프트를 처리해야 합니다. 이후 토큰(decode)은 가장 최근에 생성된 토큰만 처리하면 됩니다.
  • KV Cache: 디코드 단계에서 이전 토큰마다 Key(K)와 Value(V) 프로젝션을 다시 계산하는 것을 피하기 위해 엔진은 이를 캐시에 저장합니다. 이는 새로운 토큰을 생성하는 복잡도를 $O(n^2)$에서 $O(n)$으로 변환합니다.
  • Continuous Batching: 처리량을 최대화하기 위해 엔진은 전체 요청 배치를 기다리지 않습니다. 대신, 이전 요청이 완료되는 즉시 새로운 요청을 "slot"에 채워 넣습니다.
  • PagedAttention: 운영 체제의 가상 메모리에서 영감을 받아, PagedAttention은 KV 캐시를 비연속 블록(페이지)으로 관리하여 메모리 단편화를 크게 줄이고 더 큰 배치 크기를 허용합니다.

결론

LLM 추론의 "마법"을 구체적이고 구현 가능한 레슨으로 분해함으로써, tiny-vLLM은 AI를 사용하는 단계에서 이를 구동하는 시스템을 구축하고자 하는 모든 사람에게 로드맵을 제공합니다. 한 커뮤니티 멤버가 언급했듯이, 레슨 형식 접근법은 CUDA를 한 번도 다뤄본 적 없는 사람에게도 코드베이스를 친숙하게 만들며, 복잡한 엔지니어링 과제를 구조화된 학습 경험으로 전환합니다.

Sources