跳到正文
Hacker News· samirswr·· 7 小时前AI 评分58

从零用 C++ 和 CUDA 实现 Llama 3.2 1B 推理引擎:6 小时视频与 85 页手写笔记

6 Hours of YouTube and 85 handtyped pages of notes on LLM Inference – Start here

AI 导读

作者开源了 samir-vllm,一个用 C++17 和 CUDA 从零构建的 Llama 3.2 1B Instruct 推理引擎,覆盖权重加载、RMSNorm、RoPE、GQA、SwiGLU MLP、分页 KV cache 与连续批处理。

正文 · AI 翻译

通过用 C++ 和 CUDA 构建一个小型引擎,从第一性原理讲解 LLM 推理。

本系列讲述从提示词到生成 token 之间发生的过程,从加载模型权重到在 GPU 上运行 transformer 的每个部分。

我的推理笔记

阅读 Inference from Scratch (PDF),其中包含 80 多页我的手写笔记和图表。本 PDF 中的所有文字均由我本人撰写,没有 AI 写作。

构建系列 | 6 小时以上视频

Part 0: AI inference explained
第 0 部分 · AI 推理详解
Part 1: Loading SafeTensor weights
第 1 部分 · 加载模型权重
Part 2: Embedding gather kernel
第 2 部分 · 嵌入收集
Part 3: RMSNorm kernel
第 3 部分 · RMSNorm
Part 4: Rotary position embeddings
第 4 部分 · RoPE
Part 5: Q, K, and V projections with cuBLAS
第 5 部分 · 使用 cuBLAS 的 Q/K/V
Part 6: Multi-head self-attention
第 6 部分 · 自注意力
Part 7: Feed-forward network and SwiGLU
第 7 部分 · 前馈网络
Part 8: Decode and KV cache
第 8 部分 · 解码与 KV 缓存
Part 9: Static and continuous batching
第 9 部分 · 静态批处理与连续批处理
Part 10: Paged KV cache
第 10 部分 · 分页 KV 缓存
Finale: Paged attention
终章 · 分页注意力

每个视频的配套笔记都在 学习笔记索引 中,另有单独的 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