jmaczan/tiny-vllm

Build your own high performance LLM inference engine in C++ and CUDA - a smaller version of vLLM

What it solves

tiny-vllm은 고성능 LLM 추론 엔진으로, vLLM의 더 작고 교육용 버전으로 설계되었습니다. NVIDIA GPU에서 대형 언어 모델(LLM)을 실행하는 데 필요한 복잡한 연산을 실용적으로 구현한 예시를 제공하며, C++와 CUDA를 사용해 처음부터 추론 서버를 구축하는 방법을 배울 수 있습니다.

How it works

이 프로젝트는 CUDA 커널을 사용해 LLM 전체 전방 패스(prefill 및 decode)를 구현합니다. 특히 Llama 3.2 1B Instruct 아키텍처를 목표로 하며, Safetensors 파일에서 모델 가중치를 로드합니다. 주요 기술 구성 요소는 다음과 같습니다.

  • Memory Management: KV 캐시와 PagedAttention을 구현해 토큰 생성 시 메모리 사용을 최적화합니다.
  • Batching: 정적 배치와 연속 배치를 모두 지원해 여러 요청을 병렬 처리합니다.
  • Computational Optimizations: 온라인 softmax와 FlashAttention 유사 메커니즘을 사용해 효율성을 높입니다.
  • Core Operations: CUDA에서 RMSNorm, RoPE, SiLU 활성화, Feed Forward Network(FFN) 등 핵심 LLM 컴포넌트를 직접 구현합니다.

Who it’s for

이 프로젝트는 LLM 추론 엔진의 내부 메커니즘을 이해하고자 하는 개발자, 학생, 강사를 위한 것입니다. 특히 C++와 CUDA 배경을 가지고 최신 LLM 서빙의 수학 및 엔지니어링을 실습을 통해 배우고자 하는 분들에게 유용합니다.

Highlights

  • Full CUDA Implementation: 모든 계산을 맞춤형 CUDA 커널로 수행해 하드웨어 효율성을 극대화합니다.
  • Advanced Inference Techniques: PagedAttention 및 연속 배치를 포함해 vLLM과 같은 프로덕션 급 엔진의 기능을 구현합니다.
  • Educational Focus: 소스 코드 저장소와 가이드 코스로 구성되어, 필요한 수학과 로직을 처음부터 도출합니다.
  • Llama 3.2 Support: Llama 3.2 1B Instruct 모델 실행에 특화되어 최적화되었습니다.