jmaczan/tiny-vllm

tiny‑vllm – 一個實作導向的 C++/CUDA LLM 推理引擎（及教學）

它是什麼 – tiny‑vllm 是一個以 C++ 編寫的完整開源 LLM 推理伺服器實作，支援 CUDA（並可選支援 AMD GPU 的 HIP）。它不僅提供引擎的原始碼，還附帶一個逐步教學課程，引導讀者理解現代基於 Transformer 的語言模型的每個元件，從載入權重到產生 token。

為何重要 – 多數高效率推論伺服器（例如 vLLM、TensorRT‑LLM）都是大型生產級程式碼庫。tiny‑vllm 故意將問題簡化為單一 10 億參數的 Llama 3.2 模型，讓學習者能清楚看見使推論加速的 CUDA 內核、記憶體佈局與批次處理邏輯。因此，它在理論（注意力機制、RMSNorm、FlashAttention 風格的 softmax、PagedAttention）與可自行編譯執行的實際系統之間架起一座橋樑。

主要功能（如 README 所列）

• 從 Safetensors 檔案載入真實 LLM（已使用 Llama 3.2 1B Instruct、BF16 權重測試）。
• 支援完整的前向傳播，包含 prefill（提示處理）與 decode（單 token 生成）。
• 所有重計算皆透過自訂 CUDA 內核完成，包括：
  • 嵌入查找
  • 帶並行歸約的 RMSNorm
  • 旋轉位置編碼（RoPE）
  • 分組查詢注意力（GQA）
  • Flash 風格的線上 softmax
  • PagedAttention 與分頁 KV 快取，實現長上下文的記憶體高效處理。
• 兩種批次處理策略：
  • 靜態批次 – 適用於簡單工作負載的固定大小批次。
  • 連續批次 – 支援動態新增/移除請求，模擬生產伺服器行為。
• 可選的 AMD GPU 支援透過 HIP/hipBLAS 實現，透過輕量級相容標頭複用相同 CUDA 原始碼。

技術堆疊

• 語言：C++ 17
• GPU API：CUDA 13.1（nvcc）與 cuBLAS；AMD GPU 可選 HIP。
• 模型格式：Safetensors（BF16 權重）。
• 相依性：單頭 JSON 解析器 nlohmann/json（v3.12.0）。
• 建置系統：CMake（支援 Ninja，可切換 -DUSE_HIP=ON）。

適合誰使用

• 學生 / 自學者：希望從核心層級理解 LLM 推論運作原理的人。
• 教師：尋找程式碼與敘事解釋結合的教學資源的人。
• 工程師：對底層效能技巧（FlashAttention、KV 快取分頁）感興趣，並希望獲得最小、可讀性高的參考實作的人。

專案狀態

• README 將所有主要引擎元件標記為已實作（複選框已勾選）。
• test.sh 指令碼可建置專案並執行快速推論示範，顯示程式碼在作者的 Linux + RTX 5090 環境中可編譯並產生輸出。
• 此專案為教育目的積極維護；歡迎提出問題以獲取建置協助。

如何開始

1. 克隆倉儲，如有需要請調整 CUDA/GCC 路徑。
2. 安裝 CUDA Toolkit（AMD 用 ROCm），並確保有相容 GPU。
3. 從 Hugging Face 下載 Llama 3.2 1B Instruct 的 model.safetensors。
4. 執行 ./test.sh – 指令碼將建置引擎並執行最小推論測試。
5. 按照 Markdown 教學章節（如 Tokenization, Attention, Paged KV cache）閱讀理論並檢視對應原始碼。

命名由來 – 它被定位為流行專案 vLLM 的「年輕而小巧的兄弟」，專注於清晰度與學習，而非生產級功能廣度。

---

以上所有細節均直接取自倉儲的 README；未推斷任何額外功能。