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，以在產生 token 時最佳化記憶體使用。
  • Batching: 支援靜態與連續批次，能同時處理多個請求。
  • Computational Optimizations: 使用線上 softmax 與類 FlashAttention 機制提升效能。
  • Core Operations: 在 CUDA 中直接實作 RMSNorm、RoPE、SiLU 激活與前饋網路（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 模型。