tiny-vllm

여러분은 C++와 CUDA를 사용하여 고성능 LLM 추론 엔진을 구축하게 될 것입니다. 바로 vLLM의 더 젊고 작은 형제인 tiny-vllm입니다.

우리는 이 과정에서 많은 것을 배우고, 실수를 하며, 아이디어와 수학적 원리를 처음부터 유도해 나갈 것입니다.

이 저장소는 두 가지로 구성되어 있습니다: 1. 추론 서버의 전체 소스 코드, 2. 엔진 구현 과정을 안내하는 강의입니다. 여러분의 학습 경로에서 학습 도구로 사용하시거나, 강연자라면 대학에서 교육 자료로 자유롭게 활용하시기 바랍니다.

추론 엔진은 다음을 포함합니다:

• Safetensors로부터 실제 LLM 모델 로드 (Llama 3.2 1B Instruct)
• 전체 LLM 순전파 (forward pass) (prefill + decode)
• 모든 연산을 CUDA 커널로 수행
• KV 캐시 (KV cache)
• 정적 배치 (static batching)
• 연속 배치 (continuous batching)
• 온라인 소프트맥스 (online softmax), FlashAttention 방식
• PagedAttention

따뜻한 음료를 준비하시고 시작해 봅시다.

• tiny-vllm
  • Intro: LLM, vLLM, models, inference servers
  • Technical prerequisities
  • Safetensors and your model
  • How floating-point numbers work and why we use bfloat16
  • GPU and CPU memory
  • Single token inference
  • Tokenization
  • Embeddings
  • CUDA kernel engineering - embeddings
  • RMSNorm and parallel reduction in CUDA
  • RoPE
  • Residual connections
  • cublasGemmEx
  • The column-major to row-major transposition trick
  • Prefill vs decode
  • Why KV cache exists
  • Attention
  • GQA
  • SiLU
  • Softmax
  • Causal mask
  • Argmax
  • Feed forward network
  • Buffer reuse
  • Static batching
  • Continuous batching
  • Online softmax
  • Paged Attention
  • Paged KV cache
  • Paged Attention CUDA kernel

LLM, vLLM, 모델 및 추론 서버 소개

최근 몇 년간 너무 많은 것이 일어나서 혼란스러울 수 있습니다. 하나씩 살펴보겠습니다.

LLM은 모델입니다. 물리적으로 LLM은 수많은 부동 소수점 숫자 (float numbers)를 포함하는 파일입니다. 개념적으로 이 숫자들은 연산의 가중치 (weights)를 나타냅니다. 가중치는 학습 (training) 단계 동안 학습되거나 발견되거나 찾아집니다. 일부 연산들은 이 가중치들을 사용합니다. 모든 연산은 데이터를 입력으로 받아, 무언가를 처리하고, 데이터를 출력으로 생성하는 함수입니다. 연산과 그 순서는 LLM의 아키텍처 (architecture)에 의해 정의됩니다. 모든 모델은 엔지니어와 연구자들에 의해 설계된 고유한 아키텍처를 가집니다.

0에서 시작하여 LLM이 텍스트를 작성하기까지의 과정은 다음과 같습니다:

0. 모델 설계 (Design the model) - 엔지니어와 연구자들은 PyTorch나 tinygrad와 같은 텐서 라이브러리 (tensor library)를 포함한 Python과 같은 고수준 언어를 사용하여 모델의 아키텍처를 설계합니다. 이들은 모델의 작은 버전들을 학습시키고, 다양한 연산, 데이터, 그리고 하이퍼파라미터 (hyperparameters, 연산을 위한 파라미터)를 가지고 실험을 진행합니다. 이는 사양 (specification)을 결정하는 단계입니다.
1. 모델 구현 (Implement the model) - 최종 모델 아키텍처를 결정하고 학습을 위한 데이터를 준비하면, 최종 모델을 정의하는 코드를 작성합니다. 이 또한 PyTorch나 유사한 도구로 작성될 수 있습니다.
2. 모델 학습 (Train the model) - 선택된 모델 아키텍처는 더미 가중치 (dummy weights)로 초기화됩니다. 이들은 다시 PyTorch나 유사한 도구를 사용하여 GPU (Graphics Processing Unit) 및 TPU (Tensor Processing Unit)와 같은 수많은 하드웨어에서 역전파 (backpropagation)와 같은 학습 알고리즘을 실행하는 스크립트를 작성합니다. 이 단계는 엄청난 양의 에너지, 비용, 그리고 컴퓨팅 파워를 소모합니다. 학습 단계의 결과물은 Safetensors 형식 (Safetensors format)과 같은 특정 형식의 모델 가중치가 담긴 파일입니다. 따라서 학습 단계는 주어진 아키텍처를 사용하여 좋은 텍스트를 생성할 수 있는 가중치 집합을 찾아내는 과정입니다.

모델 서빙 (현재 단계) - 가중치가 담긴 파일은 컴퓨터에서 바로 실행할 수 없습니다. 그것은 실행 파일 (executable)이 아닙니다. 그저 수많은 숫자들의 집합일 뿐입니다. 아키텍처 (architecture) 또한 실행할 수 없습니다. 그것은 단지 계획이자 설계도, 즉 연산에 대한 기술일 뿐입니다. 모델을 실제로 실행하려면, 아키텍처와 그 연산들을 실행 가능한 코드 (executable code)로 변환하고, 모델 가중치 파일(file with model weights)을 사용하여 가중치를 아키텍처에 로드하는 프로그램이 필요합니다. 연산을 구현하는 프로그램을 작성하고, 프로그램이 가중치를 로드하고 나면 (가중치는 프로그램의 런타임 (runtime) 중 시작 시점에 로드됩니다), 마침내 모델에 프롬프트 (prompt)를 보내고 의미 있는 응답을 얻을 수 있습니다. 모델로부터 출력을 생성하는 것을 추론 (inference)이라고 부릅니다. 이것이 바로 우리가 여기서 만들고 있는 것을 추론 서버 (inference server) 또는 추론 엔진 (inference engine)이라고 부르는 이유입니다.

추론 서버가 필요한 이유를 알았으니, 왜 이것을 C++ 및 CUDA로 만드는지 생각해 봅시다. 그것은 하드웨어의 효율적인 사용을 극대화하고 높은 성능 (high performance)을 얻고 싶기 때문입니다. 이는 응답을 빠르게 얻고 싶으며, 동시에 여러 개의 프롬프트를 처리할 수 있기를 원한다는 것을 의미합니다. CUDA는 전체 생태계(ecosystem)이기도 하지만, GPU에서 실행되는 코드를 작성하는 데 사용하는 언어이기도 합니다. 우리는 GPU 상에서 코드를 작성해야 하는데, 그 이유는 LLM 내부의 많은 연산이 여러 숫자를 곱하고 더하는 작업이기 때문입니다. 적은 양의 수학 연산이 필요하다면 CPU로 충분합니다. 하지만 양이 많다면 GPU가 더 낫습니다. LLM은 대부분 행렬 곱셈 (multiplying the matrices)에 관한 것이며, 이는 결국 수많은 숫자와 수많은 벡터에 대해 두 벡터의 내적 (dot product)을 계산하는 것으로 귀결됩니다. LLM의 수학은 단순합니다. 선형 대수학 (linear algebra)의 기초가 필요할 것이며, 코딩을 하면서 배우고 부족한 부분을 그때그때 채워나갈 수 있습니다. 저는 이러한 방식의 JIT (Just-In-Time) 학습이 가장 효과적이라고 생각하며, 여러분도 아마 좋아하게 될 것입니다.

AI와 연산 사이의 관계에 대한 저의 견해 중 여러분에게 유용할 수도 있는 점은, 지능은 모델의 수많은 파라미터 (parameters)와 이 파라미터들을 사용한 입력값의 방대한 연산 (computation)으로부터 나온다는 것입니다.

범위 외 (Out of scope): LLM의 학습 (training) 단계는 이 과정에서 다루지 않습니다. 우리는 이미 학습된 LLM을 가져와서, NVIDIA GPU 상에서 여러 요청을 병렬로 빠르게 실행할 수 있는 프로그램을 작성할 것입니다. 만약 직접 LLM을 학습시키고 싶다면, sensei Karpathy의 nanoGPT 및 llm.c 리포지토리와 그의 YouTube 채널을 강력히 추천합니다. 마찬가지로, 우리는 모델을 설계하지는 않지만, 텐서 라이브러리 (tensor libraries) 또한 매우 매력적인 주제이며 밑바닥부터 이해할 가치가 있습니다. George Hotz의 tinygrad는 매우 적은 양의 코드로 텐서 라이브러리를 구현한 프로젝트이므로, 영감을 얻고 내부 동작 원리를 배우고 싶다면 좋은 시작점이 될 것입니다 (또한 그들의 Discord도 좋습니다)! Andrej Karpathy가 만든 조금 더 오래되고 작은 버전인 micrograd도 있습니다. Discord 이야기가 나온 김에, Mark Saroufim의 GPU MODE를 추천하고 싶습니다. 많은 훌륭한 사람들이 그곳에서 활동하고 있습니다! 만약 여기서 무엇이 일어나고 있는지 혼란스럽고 AI/ML 여정이 처음이라면, Jeremy Howard와 Rachel Thomas의 fastai 교재부터 시작하세요. 저는 데이터 과학 (data science) 및 엔지니어링 (engineering) 부분에 대해서는 잘 모르기 때문에 여기서는 편리하게 생략하겠습니다. 아마 Kaggle이 이를 시작하고 직접 배우기에 좋은 장소가 될 수 있을 것입니다. 마지막으로, 우리는 C++ 및 CUDA로 코드를 작성할 것이며, 적용 가능한 경우 cuBLAS를 사용할 것입니다. 진행하면서 배우시면 됩니다. NVIDIA의 공식 리소스는 유용하고 도움이 됩니다.

기술적 요구 사항 (Technical prerequisities)

NVIDIA GPU가 있다는 가정하에, 약간의 변경만 거치면 어떤 플랫폼에서든 빌드하고 실행할 수 있습니다. .vscode/c_cpp_properties.json의 CUDA 또는 GCC 경로와 같이 CMakeLists.txt의 NVCC와 같은 일부 경로를 조정해야 할 수도 있습니다.

이 저장소를 포크(fork)하여 사용자의 환경에서 작동하도록 필요한 조정을 마친 후, jmaczan/tiny-vllm에 풀 리퀘스트(pull request)를 생성하여 다른 독자들을 위해 변경 사항을 업스트림(upstream)할 것을 권장합니다.

제가 개발 및 테스트를 진행한 정확한 설정 환경은 다음과 같습니다:

• Linux (6.19.8 x64_64)
• CUDA Toolkit (13.1)
• C++ 17
• GCC (15.2.1)
• 가져오게 될 유일한 외부 의존성(external dependency)은 단일 헤더 파일인 include/json.hpp 형태의 JSON 파서 nlohmann/json 3.12.0입니다.
• AMD CPU (Ryzen 7 9800X3D)
• NVIDIA GPU (RTX 5090)
• Hugging Face의 Llama 3.2 1B Instruct (commit hash 898999bd25b40516fce5a5b8f0948f4c81c650bc)를 사용했습니다. 이 저장소에서 model.safetensors 파일만 있으면 됩니다.

의존성을 설치하고 ./test.sh로 프로그램을 실행하세요. 빌드 후 즉시 실행됩니다.

빌드나 실행에 실패했고 선택한 AI가 도움을 줄 수 없다면, GitHub에 이슈(Issue)를 생성해 주세요. 도움을 드리도록 노력하겠습니다. 모든 유용한 컨텍스트(context)를 제공해 주시기 바랍니다.

Safetensors와 모델 (Safetensors and your model)

가장 먼저 해야 할 일은 추론 (inference)을 실행하는 데 사용할 LLM을 다운로드하는 것입니다. 저는 Llama 3.2 1B Instruct를 선택했는데, 사용하기 쉽고 크기가 작으며, 대화에 최적화되어 있고, 우리에게 충분히 훌륭하기 때문입니다. 추론 서버를 구축하는 엔지니어의 관점에서 볼 때, 모델은 가중치 (weights)를 포함하는 단일 파일일 뿐입니다.

모델은 Safetensors 형식 (Safetensors format)으로 되어 있습니다. Pickle이나 Parquet과 같은 다른 형식들도 존재합니다. Safetensors는 매우 대중적이고 널리 사용되며, 우리가 선택한 모델은 Safetensors로 호스팅되어 있습니다.

다음으로 넘어가기 전에 잠시 멈춰서 Safetensors 형식을 이해해 봅시다.

Safetensor 파일은 항상 다음 순서대로 3개의 섹션으로 구성됩니다: 헤더 크기 (header size), 헤더 (header), 그리고 텐서 데이터 (tensors data). 헤더 크기는 항상 8바이트입니다. 이 8바이트는 실제 헤더가 차지하는 바이트 수를 나타내는 부호 없는 64비트 정수 (unsigned 64-bit integer)입니다.

std::ifstream safetensors_file("model.safetensors", std::ios_base::binary);
uint64_t header_size;
safetensors_file.read(reinterpret_cast<char *>(&header_size), 8);