从零用 C++ 和 CUDA 实现 Llama 3.2 1B 推理引擎:6 小时视频与 85 页手写笔记
6 Hours of YouTube and 85 handtyped pages of notes on LLM Inference – Start here
作者开源了 samir-vllm,一个用 C++17 和 CUDA 从零构建的 Llama 3.2 1B Instruct 推理引擎,覆盖权重加载、RMSNorm、RoPE、GQA、SwiGLU MLP、分页 KV cache 与连续批处理。
通过用 C++ 和 CUDA 构建一个小型引擎,从第一性原理讲解 LLM 推理。
本系列讲述从提示词到生成 token 之间发生的过程,从加载模型权重到在 GPU 上运行 transformer 的每个部分。
我的推理笔记
阅读 Inference from Scratch (PDF),其中包含 80 多页我的手写笔记和图表。本 PDF 中的所有文字均由我本人撰写,没有 AI 写作。
构建系列 | 6 小时以上视频
每个视频的配套笔记都在 学习笔记索引 中,另有单独的 Markdown 和 PDF 文件夹。你也可以阅读 完整 PDF 合集。
目前已完成的功能
- 一个 C++17 / CUDA 项目,使用 CMake 构建,并链接 ICU 以支持 Unicode 感知的分词。
- 直接
model.safetensors加载:解析元数据、校验张量偏移,并将原始 BF16 权重复制到 GPU 内存。此外还有一个内存映射加载器实现,以避免额外的 CPU 侧拷贝。 - 固定的 Llama 3.2 1B Instruct 权重布局,直接指向全部 16 层的嵌入、归一化、注意力和 MLP 张量。
- 一个本地 BPE 分词器,支持 Unicode、特殊 token、编码/解码以及 Instruct 聊天提示格式;其预期输出由专门的测试可执行文件覆盖。
- 一个 CUDA 嵌入收集内核和预填充路径,将提示词 token ID 复制到 GPU 并分配激活值。
- 一个用于 2,048 宽隐藏状态的 CUDA RMSNorm 内核:它以 FP32 进行归约,并在整个前向传播中应用学习到的 BF16 权重。
- 稠密因果预填充现在运行全部 16 层:Q/K/V 投影、RoPE、分组查询注意力、输出投影、注意力残差、注意力后 RMSNorm、SwiGLU MLP 以及 MLP 残差。
- 最终的 RMSNorm 和绑定的嵌入/LM 头投影产生 FP32 logits;GPU argmax 选择每个下一个 token。
- 一个共享的 BF16 KV 页池,每个序列和层都有块表。每个物理页保存 16 个 token 的 K 和 V;预填充填充提示词页,解码在其绝对 RoPE 位置追加一个新 token。
- 使用 Instruct 聊天格式进行贪婪响应生成,支持回合结束/文本结束停止,并可在 2,048 token 上下文内配置输出长度。
- 持久化暂存缓冲区和由预填充与解码共享的一个 cuBLAS 句柄。Q 存储复用于输出/MLP 投影,归一化存储复用于注意力输出。
- 批处理解码,具有运行时大小的缓冲区、打包投影,以及每个序列单独的 KV 页映射和注意力长度。请求数量不会扩大 KV 池。
- 静态批处理与连续调度器根据可用的页预算接纳排队中的提示,并在完成后归还其页和预算。
- 直接分页解码注意力,采用 FP32 在线 softmax。每个查询头流式读取其序列的缓存页并累加加权值,无需收集连续前缀或存储注意力分数。
第 8 部分讲解 Prefill、Decode 与 KV Cache。第 9 部分接着介绍 静态批处理与连续批处理,包括当前的分页缓存和在线 softmax 解码注意力。
当前范围
Llama 3.2 1B 推理路径包括稠密因果预填充、分页解码注意力,以及静态或连续批处理。模型维度和 2,048 token 的上下文在此学习引擎中保持固定。KV 池默认为 128 MiB,可通过 --kv-cache-mib N 设置;活动请求容量取决于工作负载的页预算,而非固定的两请求上限。
准入机制为每个请求的最大已处理前缀预留空间:prompt_tokens + allowed_new_tokens - 1,在 16 层中每层向上取整到 16 token 的页。允许的输出预算被限制在上下文上限内,最终预测的 token 直接返回而不再处理。这些准入额度与物理页分离,物理页按需惰性分配。提前完成会同时释放已分配的页和剩余额度。
当预算无法与活动请求共存时,请求按 FIFO 顺序等待。活动请求有空间完成,因此未来的页边界增长不会耗尽池或导致生成死锁。预算超过整个池的请求会被拒绝并给出明确错误。这种保守策略避免了抢占和重计算;即使模型最终提前停止,较大的请求输出预算也可能限制并发。
--stats 报告池大小、峰值活动请求数、最大解码批大小、峰值已分配页数、准入等待次数和最终空闲页数。预填充在解码迭代之间仍一次处理一个提示。
BF16 计算在不同批大小下可能产生不同的舍入结果。当词表 logits 接近并列时,贪心输出 token 因此在单次运行和批处理运行之间可能不同。不同批大小之间的逐位输出相等无法保证。
构建
要求
- 支持 CUDA 的 NVIDIA GPU 和 CUDA Toolkit
- CMake 3.24+
- C++17 编译器
- ICU 开发库
nlohmann/json头文件- 位于
models/llama-3.2-1b-instruct/model.safetensors的模型权重 - 位于
models/llama-3.2-1b-instruct/tokenizer.model的匹配分词器(从models/Llama-3.2-1B-Instruct/original/tokenizer.model复制)
mkdir -p models/llama-3.2-1b-instruct cp models/Llama-3.2-1B-Instruct/original/tokenizer.model models/llama-3.2-1b-instruct/ cmake -S . -B build cmake --build build -j ./build/samir-vllm "Explain the KV cache in two sentences." 128 # Continuous batching (default): admit requests whose page budgets fit. ./build/samir-vllm "Say hello." 64 "Explain the KV cache." "What is RoPE?" # Static batching: wait for the current group before admitting more. ./build/samir-vllm --static "Say hello." 64 "Explain the KV cache." "What is RoPE?" # More than two requests share the same default 128 MiB pool. ./build/samir-vllm --stats "Explain paging." 64 "Explain GQA." "Explain RoPE." "Explain softmax." # A smaller pool queues requests when their declared budgets cannot fit together. ./build/samir-vllm --kv-cache-mib 8 --stats "Explain paging." 64 "Explain GQA." "Explain RoPE." ./build/tokenizer-test models/llama-3.2-1b-instruct/tokenizer.model
在可选的 --static、--continuous、--kv-cache-mib N 和 --stats 标志之后,第一个参数是提示,第二个是 max_new_tokens(默认值:128)。其余参数是额外提示;token 预算适用于每个请求。每个提示都格式化为 Instruct 用户消息。生成在 EOS、EOT 或 EOM,或输出/上下文上限处独立停止,每个完成的响应都会打印其 token ID 和请求编号。
提示一次预填充一个,期间暂停解码。每个解码步骤将每个活动序列的一个 token 打包到共享缓冲区中。注意力只读取该序列的缓存;请求可以有不同的提示和输出长度。连续批处理在页额度可用时接纳等待中的请求,而静态批处理在当前组完成后接纳下一组。
项目布局
src/ Engine entry point, CUDA kernels, and tokenizer implementation
include/ Public kernel and tokenizer headers
tests/ Tokenizer tests
python/ Tokenizer reference script
proper_notes/ Study-notes index, markdown/ sources, and pdf/ reading editions
scripts/ Rebuild the study-note PDFs
notes.md Working notes from the build
许可证
src/、include/、tests/ 和 python/ 中的引擎源代码,以及 CMakeLists.txt、.gitignore 和本 README,均根据 MIT 许可证 授权。
proper_notes/、notes.md 和 podLibraryInstall.md 中的笔记,以及本项目随附的手写笔记、扫描件、插图、笔记照片和 PDF,均保留所有权利,不在 MIT 许可范围内。允许个人阅读和学习。无论免费还是付费,再分发、再出版、重新托管、转售或分发改编版本均需事先获得书面许可。完整条款和例外情况请参阅 LICENSE-NOTES。
公开仓库仍可在 GitHub 上依据 GitHub 服务条款 查看和 fork。
模型资产与依赖项
使用 Llama 构建。 随附的 Llama 3.2 分词器以及任何单独下载的 Llama 3.2 权重均保留 Meta 的 Llama 3.2 社区许可 和 可接受使用政策。它们不在 MIT 许可范围内。分发任何 Llama 材料时,请一并保留协议、政策和 署名声明。
第三方依赖项,包括 CUDA、cuBLAS、ICU 和 nlohmann/json,保留各自的许可证。
来源:Hacker News · github.com