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；未推断任何额外功能。