tiny-vllm: C++와 CUDA로 vLLM을 직접 구현하며 배우는 LLM 추론 엔진

[9bow (박정환)]

tiny-vllm 소개

tiny-vllm은 vLLM의 더 작은 형제 격인 LLM 추론 엔진을, C++와 CUDA로 밑바닥부터 직접 구현해 보는 오픈소스 프로젝트입니다. 한 저장소 안에 두 가지가 들어 있습니다. 하나는 실제로 동작하는 추론 서버의 전체 소스 코드이고, 다른 하나는 그 엔진을 한 단계씩 만들어 가도록 안내하는 코스(course) 형태의 글입니다. 만든 사람은 Jędrzej Maczan(GitHub jmaczan) 으로, 학습 도구로 쓰거나 대학 강의 자료로 활용해도 좋다고 밝히고 있습니다.

추론 서버(inference server)는 학습이 끝난 모델 가중치 파일을 실제로 실행해 프롬프트에 대한 응답을 만들어 내는 프로그램입니다. tiny-vllm은 이 과정을 왜 하필 C++와 CUDA로 작성하는지부터 설명합니다. LLM의 연산은 대부분 행렬 곱, 즉 수많은 벡터의 내적으로 환원되기 때문에, 응답을 빠르게 돌려주고 여러 프롬프트를 동시에 처리하려면 GPU에서 직접 코드를 돌리는 것이 유리합니다. CUDA는 그 GPU 코드를 작성하는 언어이자 생태계입니다.

이 프로젝트가 다루는 범위는 학습(training)이 아니라 추론(inference)입니다. 이미 학습된 Llama 3.2 1B Instruct 모델을 가져와, NVIDIA GPU 위에서 여러 요청을 병렬로 빠르게 실행하는 프로그램을 직접 만드는 것이 목표입니다. 본 게시물에서는 tiny-vllm이 구현하는 엔진의 구성 요소, 추론 forward pass의 동작 흐름, CUDA 커널 구현의 한 예, 그리고 설치와 실행 방법을 정리합니다.

tiny-vllm이 구현하는 추론 엔진의 구성 요소

저자가 README 상단에 정리한 구현 체크리스트에 따르면, tiny-vllm의 추론 엔진은 다음 요소들을 포함합니다.

• Safetensors 모델 로딩: Safetensors 포맷으로 저장된 Llama 3.2 1B Instruct 가중치를 읽어 들입니다.
• 전체 forward pass: 프롬프트 전체를 한 번에 처리하는 프리필(prefill)과, 토큰을 하나씩 생성하는 디코드(decode)를 모두 구현합니다.
• 모든 연산의 CUDA 커널화: 임베딩 조회, 정규화, 어텐션 등 핵심 연산을 직접 작성한 CUDA 커널로 수행합니다.
• KV 캐시(KV cache): 이미 계산한 키와 값을 재사용해 디코드 단계의 중복 계산을 줄입니다.
• 정적 배칭(static batching)과 연속 배칭(continuous batching): 여러 요청을 묶어 처리량을 높이는 두 가지 방식입니다.
• 온라인 softmax(online softmax): FlashAttention 계열의 방식으로 softmax를 점진적으로 계산합니다.
• PagedAttention: 운영체제의 페이징(paging) 아이디어를 KV 캐시 메모리 관리에 적용한 PagedAttention 기법입니다.

다만 솔직히 짚어둘 점이 있습니다. 위 기능들은 저자의 구현 체크리스트 기준이며, 코스 본문 글은 핵심 forward pass(임베딩부터 argmax까지)까지는 상세하게 작성되어 있지만, 배칭과 온라인 softmax, PagedAttention을 다루는 뒤쪽 장(章)은 아직 < TODO write more > 나 Incoming! 로 표시된 작성 중 상태입니다. 코스 텍스트를 끝까지 따라 읽으려는 독자라면 이 부분을 감안하는 것이 좋습니다.

tiny-vllm의 forward pass 동작 방식

tiny-vllm이 기준으로 삼는 Llama 3.2 1B는 16개의 트랜스포머 레이어(transformer layer)로 이루어진 모델입니다. 저자는 PyTorch가 출력하는 모델 구조를 그대로 분석의 출발점으로 삼습니다.

LlamaForCausalLM(
  (model): LlamaModel(
    (embed_tokens): Embedding(128256, 2048)
    (layers): ModuleList(
      (0-15): 16 x LlamaDecoderLayer(
        (self_attn): LlamaAttention(
          (q_proj): Linear(in_features=2048, out_features=2048, bias=False)
          (k_proj): Linear(in_features=2048, out_features=512, bias=False)
          (v_proj): Linear(in_features=2048, out_features=512, bias=False)
          (o_proj): Linear(in_features=2048, out_features=2048, bias=False)
        )
        (mlp): LlamaMLP(
          (gate_proj): Linear(in_features=2048, out_features=8192, bias=False)
          (up_proj): Linear(in_features=2048, out_features=8192, bias=False)
          (down_proj): Linear(in_features=8192, out_features=2048, bias=False)
          (act_fn): SiLUActivation()
        )
        (input_layernorm): LlamaRMSNorm((2048,), eps=1e-05)
        (post_attention_layernorm): LlamaRMSNorm((2048,), eps=1e-05)
      )
    )
    (norm): LlamaRMSNorm((2048,), eps=1e-05)
    (rotary_emb): LlamaRotaryEmbedding()
  )
  (lm_head): Linear(in_features=2048, out_features=128256, bias=False)
)

이 구조를 바탕으로, 입력 텍스트가 다음 토큰으로 변환되는 순서는 다음과 같습니다. 먼저 텍스트를 토큰(token)으로 나누고, 각 토큰에 대응하는 임베딩(embedding) 벡터를 조회합니다. 이 프로젝트에서 임베딩 한 개의 길이는 2048이므로, 토큰 5개를 입력하면 (5, 2048) 크기의 행렬이 됩니다. 이어서 16개의 트랜스포머 레이어를 통과하는데, 각 레이어는 RMSNorm, 잔차 연결(residual connection), 마스크드 그룹 쿼리 어텐션(masked grouped-query attention), 피드포워드 네트워크(feed-forward network)로 구성됩니다. 마지막으로 최종 RMSNorm과 선형 출력(lm_head)을 거친 뒤 argmax로 가장 확률이 높은 토큰 하나를 고르면 모델이 생성한 첫 토큰이 나옵니다.

어텐션은 기본 형태가 O(n^2 \cdot d) 의 계산 복잡도를 갖는 비싼 연산입니다. tiny-vllm은 여기에 그룹 쿼리 어텐션(Grouped-Query Attention, GQA)을 적용합니다. 위 구조에서 쿼리 프로젝션(q_proj) 은 출력 차원이 2048인 반면 키와 값 프로젝션(k_proj, v_proj) 은 512로 더 작은데, 이는 여러 쿼리 헤드가 더 적은 수의 키/값 헤드를 공유해 KV 캐시의 크기를 줄이는 GQA의 특징을 그대로 보여줍니다.

tiny-vllm의 CUDA 커널 구현 예: RMSNorm

tiny-vllm의 학습 가치는 추상적인 설명이 아니라 실제로 돌아가는 CUDA 커널을 직접 보여준다는 데 있습니다. 예를 들어 RMSNorm은 GPU에서 병렬 리덕션(parallel reduction)으로 구현되는데, 저자가 공개한 커널은 다음과 같습니다.

__global__ void rmsNormKernel(__nv_bfloat16 *input, __nv_bfloat16 *output, __nv_bfloat16 *norm_weights)
{
    __shared__ float rms_vector[1024];
    int workIndex = threadIdx.x + blockIdx.x * 2048;
    rms_vector[threadIdx.x] = (float)input[workIndex] * (float)input[workIndex] + (float)input[workIndex + 1024] * (float)input[workIndex + 1024];
    __syncthreads();
    // tree reduction
    for (int i = 1; i < 1024; i = i * 2)
    {
        if (threadIdx.x % (i * 2) == 0)
        {
            rms_vector[threadIdx.x] = rms_vector[threadIdx.x] + rms_vector[threadIdx.x + i];
        }
        __syncthreads();
    }
    if (threadIdx.x == 0)
    {
        rms_vector[0] = sqrt(rms_vector[0] / 2048.0 + 1.0e-5);
    }
    __syncthreads();

    output[workIndex] = (__nv_bfloat16)(((float)input[workIndex] / rms_vector[0]) * (float)norm_weights[threadIdx.x]);
    output[workIndex + 1024] = (__nv_bfloat16)(((float)input[workIndex + 1024] / rms_vector[0]) * (float)norm_weights[threadIdx.x + 1024]);
}

공유 메모리(__shared__) 에 제곱 합을 모은 뒤 트리 리덕션으로 합산하고, 그 결과로 입력을 정규화하는 과정이 한 커널 안에 담겨 있습니다. 코스는 이런 식으로 RMSNorm, RoPE, 잔차 연결, SiLU, softmax 등의 커널을 차례로 직접 작성하게 합니다. 가중치는 bfloat16(BF16) 형식인데, BF16은 16비트 안에서 부호 1비트, 지수 8비트, 가수 7비트로 나뉩니다. 16비트 크기를 유지하면서도 32비트 float와 같은 8비트 지수를 가져 오버플로/언더플로 위험이 작기 때문에 추론에서 널리 쓰인다고 저자는 설명합니다.

CUDA로 직접 행렬 곱을 다룰 때 부딪히는 실전 문제도 다룹니다. cuBLAS의 cublasGemmEx 는 열 우선(column-major) 메모리 레이아웃을 가정하는데, C/C++ 배열은 행 우선(row-major)으로 저장됩니다. tiny-vllm은 이 둘 사이의 전치(transposition) 트릭을 그림과 함께 설명합니다.

tiny-vllm 설치 및 실행

저자가 개발하고 테스트한 환경은 다음과 같습니다.

• Linux (커널 6.19.8 x86_64)
• CUDA Toolkit 13.1
• C++ 17, GCC 15.2.1
• NVIDIA GPU (테스트 기준 RTX 5090), AMD CPU (Ryzen 7 9800X3D)
• Llama 3.2 1B Instruct의 model.safetensors 파일

외부 의존성은 단 하나, JSON 파서인 nlohmann/json 3.12.0 뿐이며 이마저도 단일 헤더 파일(include/json.hpp)로 포함되어 있습니다. NVIDIA GPU가 있다면 약간의 경로 조정만으로 다른 플랫폼에서도 빌드할 수 있다고 안내합니다.

먼저 Hugging Face에서 Llama 3.2 1B Instruct의 model.safetensors 파일을 받은 뒤, 빌드와 실행을 한 번에 처리하는 스크립트를 실행합니다.

./test.sh

이 스크립트는 프로젝트를 빌드한 직후 곧바로 실행합니다. 저자는 자신의 환경에 맞게 저장소를 포크해 CMakeLists.txt와 .vscode/c_cpp_properties.json의 CUDA, GCC, NVCC 경로를 조정한 뒤, 개선 사항이 있으면 풀 리퀘스트로 기여해 줄 것을 권하고 있습니다.

tiny-vllm의 라이선스

tiny-vllm은 Apache License 2.0으로 공개되어 있어 개인 및 상업적 목적으로 자유롭게 사용할 수 있습니다.

tiny-vllm 프로젝트 GitHub 저장소

GitHub - jmaczan/tiny-vllm: Build your own high performance LLM inference...

더 읽어보기

• Nano-vLLM: 처음부터 다시 구현한, 1200줄 가량의 경량 vLLM 구현체

• LLM 추론 엔진 이해하기: Nano-vLLM 내부 살펴보기 - 1부 (Understanding LLM Inference Engines: Inside Nano-vLLM (Part 1))

• LLM 추론 엔진 이해하기: Nano-vLLM 내부 살펴보기 - 2부 (Understanding LLM Inference Engines: Inside Nano-vLLM (Part 2))

• Mini-SGLang: SGLang의 핵심을 담은 경량화 고성능 LLM 추론 프레임워크

• TokenSpeed: 에이전트형 워크로드를 위한 빠른 LLM 추론 엔진 (feat. Kimi K2.5, NVIDIA Blackwell)

• FlashKDA: MoonshotAI가 공개한, Kimi Delta Attention을 위한 고성능 CUDA 커널 라이브러리

---

이 글은 GPT 모델로 정리한 글을 바탕으로 한 것으로, 원문의 내용 또는 의도와 다르게 정리된 내용이 있을 수 있습니다. 관심있는 내용이시라면 원문도 함께 참고해주세요! 읽으시면서 어색하거나 잘못된 내용을 발견하시면 덧글로 알려주시기를 부탁드립니다.

파이토치 한국 사용자 모임이 정리한 이 글이 유용하셨나요? 회원으로 가입하시면 주요 글들을 이메일로 보내드립니다! (기본은 Weekly지만 Daily로 변경도 가능합니다.)

아래쪽에 좋아요를 눌러주시면 새로운 소식들을 정리하고 공유하는데 힘이 됩니다~