7.0
热点
AI SCORE
开源项目2026-05-30 03:38
Tiny-vLLM：C++ 高性能 LLM 推理引擎开源
Hacker News · github.com#开源#LLM推理#C++
Editor brief · 编辑速览
用 C++ 和 CUDA 实现的轻量 LLM 推理引擎，追求低延迟高吞吐，适合本地部署和边缘推理场景。
你将使用 C++ 和 CUDA 构建一个高性能 LLM 推理引擎——tiny-vllm，它是 vLLM 更年轻、更小巧的兄弟。
在此过程中，我们会学到很多东西，也会犯错，并从零开始推导其中的思想和数学原理。
这个仓库包含两部分：1. 推理服务器的完整源代码；2. 一门课程，我会带领你完成引擎的实现过程。欢迎你在自己的学习之路上将它用作学习工具；如果你是讲师，也欢迎将其用作大学里的教学资源。
该推理引擎包含：
给自己准备一杯热饮，让我们开始吧。
tiny-vllm 简介：LLM、vLLM、模型、推理服务器
技术前置要求
Safetensors 与你的模型
浮点数的工作原理，以及为什么使用 bfloat16
GPU 与 CPU 内存
单 token 推理
Tokenization
Embeddings
CUDA kernel 工程——embeddings
RMSNorm 与 CUDA 中的并行归约
RoPE
残差连接
cublasGemmEx
从列主序到行主序的转置技巧
Prefill 与 decode
KV cache 存在的原因
Attention
GQA
SiLU
Softmax
因果掩码
Argmax
前馈网络
缓冲区复用
静态批处理
连续批处理
在线 softmax
Paged Attention
Paged KV cache
Paged Attention CUDA kernel
近几年发生了太多事情，人们很容易迷失其中。让我们逐一梳理。
LLM 是一种模型。从物理形式来看，LLM 是一个包含大量浮点数的文件。从概念上看，这些数字代表各种运算的权重。权重是在训练阶段学习、发现或找到的。部分运算会使用这些权重。每个运算都是一个函数：它接收一些数据作为输入，对其执行某种处理，然后生成数据作为输出。运算及其顺序由 LLM 的架构定义。每个模型都有自己的架构，而架构是由工程师和研究人员设计的。
从零开始，到 LLM 写出文本，整个过程如下：
设计模型——工程师和研究人员使用 Python 这样的高级语言，以及 PyTorch 或 tinygrad 这样的张量库来设计模型架构。他们会训练模型的小型版本，使用不同的运算、数据和超参数（运算的参数）进行实验。这是确定规范的阶段。
实现模型——确定最终的模型架构并准备好训练数据后，他们会编写定义最终模型的代码。这些代码同样可以使用 PyTorch 或类似工具编写。
训练模型——首先使用虚拟权重初始化选定的模型架构。然后，他们编写脚本，再次使用 PyTorch 或类似工具，在 GPU、TPU 等大量硬件上运行反向传播之类的学习算法。这个阶段会消耗大量能源、资金和算力。训练阶段的产物是一个以某种格式保存模型权重的文件，例如 Safetensors 格式。因此，训练阶段就是寻找这样一组权重：在给定架构下，它能够生成优质文本。
部署模型（我们正在这里）——包含权重的文件无法直接在计算机上运行。它不是可执行文件，只是一大堆数字。架构本身也无法运行——它只是一份计划、蓝图，是对计算过程的描述。要真正运行模型，我们需要一个程序，将架构及其中的运算转化为可执行代码，并使用模型权重文件把权重加载到架构中。编写好实现这些运算的程序，并让程序加载权重之后（权重会在程序运行时、启动阶段加载），你终于可以向模型发送提示词并获得有意义的响应。从模型生成输出的过程称为推理。因此，我们在这里构建的程序被称为推理服务器或推理引擎。
理解为什么需要推理服务器后，让我们思考一下为什么要使用 C++ 和 CUDA 来构建它。原因是我们希望最大限度地提高硬件的使用效率并获得高性能。这意味着，我们既希望快速得到响应，也希望能够同时处理多个提示词。CUDA 是一整套生态系统，同时也是一种用于编写 GPU 运行代码的语言。我们需要在 GPU 上执行代码，因为 LLM 内部的许多运算都需要对大量数字进行乘法和加法。如果只需要进行少量数学运算，CPU 就足够了；如果计算量很大，GPU 会更合适。LLM 的核心主要是矩阵乘法，归根结底，就是对大量数字和大量向量计算两个向量的点积。LLM 的数学并不复杂，我们只需要掌握线性代数的基础知识。你可以一边编程一边学习，在过程中补齐知识缺口。我认为这种即时学习方式最有效，也许你也会喜欢。
关于 AI 与计算之间的关系，我有一种理解，也许对你有所帮助：智能来源于模型的大量参数，以及使用这些参数对输入值进行的大量计算。不存在某个单独的元素，让你可以指着它说：“正是这个东西让模型拥有智能或变得有用。”模型的每个部分都可以替换成其他方案，并得到不同的权衡，例如用准确性换取更低的复杂度。我希望之后讲到 Attention 的数学原理时，不会忘记回到这个话题。因为默认的 Attention 机制在计算上非常复杂（O(n^2*d)）。这种复杂度并非不可挑战；事实上，人们一直在这样做，并设计出了线性 Attention 等替代机制。如果更多人认为这门课程有用，我会考虑再制作一门课程，主题可能是 ML 编译器（使用 Python 或 C++ 的实践课程，加上一些 SSA 理论），也可能是替代 Attention 机制（数学原理 + CUDA kernel）。如果你感兴趣，请告诉我！如果你认为这门课程有价值，也请把它分享给其他人。
范围之外：本课程不会涉及 LLM 的训练阶段。我们会使用一个已经训练好的 LLM，并编写一个程序，让这个 LLM 能够在 NVIDIA GPU 上快速运行，并行处理多个请求。如果你想训练自己的 LLM，我强烈推荐 Karpathy 老师的 nanoGPT 和 llm.c 仓库，以及他的 YouTube 频道。同样，我们也不会设计模型，但张量库也是一个引人入胜、值得从零开始理解的主题。George Hotz 的 tinygrad 项目使用极少量代码实现了一个张量库。如果你想从中获得灵感并学习其内部原理，这是一个很好的起点（他们的 Discord 也很不错）！Andrej Karpathy 还有一个年代稍早、规模更小的版本——micrograd。既然提到了 Discord，我还想向你推荐 Mark Saroufim 的 GPU MODE，那里聚集了许多优秀的人！如果你对这里正在发生的事情感到迷茫，而且刚刚踏上 AI/ML 学习之旅，可以从 Jeremy Howard 和 Rachel Thomas 的 fastai 一书开始。我在这里很方便地略过了数据科学与数据工程部分，因为我对此了解不多。Kaggle 也许是一个不错的起点，可以通过动手实践来学习。最后但同样重要的是，我们将使用 C++ 和 CUDA 编程，并在适用的地方使用 cuBLAS。你可以边做边学。NVIDIA 的官方资源很不错，也很有帮助。
只要拥有 NVIDIA GPU，你就可以在任何平台上构建并运行它，只需进行少量修改。你可能需要调整一些路径，例如 c_cpp_propertiesjson 中的 CUDA 或 GCC 路径，或者 CMakeLists.txt 中的 NVCC 路径。
我建议你 fork 这个仓库，进行必要的调整，使其能够在你的机器上运行，然后向 jmaczan/tiny-vllm 创建 pull request，把你的修改贡献回上游，让其他读者也能从中受益。
我进行开发和测试时使用的确切环境如下：
nlohmann/json 3.12.0，它是一个单头文件 include/json.hpp898999bd25b40516fce5a5b8f0948f4c81c650bc），你只需要该仓库中的 model.safetensors 文件安装依赖项，然后使用 ./test.sh 运行程序——它会构建程序并立即执行。
它也可以通过 ROCm/HIP 在 AMD GPU 上运行。向 CMake 传递 -DUSE_HIP=ON，它将使用 hipcc 针对 hipBLAS 构建，而不是 nvcc 和 cuBLAS；CUDA 源代码通过一个轻量级的 src/cuda_to_hip.h 兼容层头文件原样重用。使用 -DCMAKE_HIP_ARCHITECTURES 选择你的 GPU 架构（例如，MI200 用 gfx90a，RDNA3 用 gfx1100，RDNA4 用 gfx1201）—— 它没有硬编码，所以要设置为与你的卡相匹配的值：
cmake -B build -DUSE_HIP=ON -DCMAKE_HIP_ARCHITECTURES=gfx1100 -DCMAKE_PREFIX_PATH=/opt/rocm -G Ninja
cmake --build build
-DCMAKE_PREFIX_PATH=/opt/rocm 让 CMake 找到 hip 和 hipBLAS 包；如果 /opt/rocm/bin 已经在你的 PATH 中，可以省略它，或者如果 ROCm 在其他位置，改为相应路径。我在 gfx90a、gfx1100 和 gfx1201 上测试了 AMD 的构建路径。默认构建（不加 -DUSE_HIP）保持不变，仍然通过 CUDA 针对 NVIDIA。
如果你构建或运行失败，且你选择的 AI 也无法帮助，请在 GitHub 上开一个 Issue —— 我会尽力帮助。请务必提供所有有用的上下文信息。
首先你需要下载一个大语言模型来运行推理。我选择 Llama 3.2 1B Instruct，因为它小巧、易用、针对对话优化，对我们的目的足够了。从我们这些推理服务器构建者的角度来看，模型就是一个包含权重的单一文件。
该模型采用 Safetensors 格式。还有其他格式，如 Pickle 和 Parquet。Safetensors 非常流行，应用广泛，而我们选择的模型也是以 Safetensors 格式托管的。
让我们先暂停一下，在继续之前理解 Safetensors 格式。
Safetensor 文件包含 3 个部分，顺序始终如下：头部大小、头部和张量数据。头部大小总是 8 字节。这 8 字节是一个无符号 64 位整数，表示实际头部占用多少字节。
std::ifstream safetensors_file("model.safetensors", std::ios_base::binary);
uint64_t header_size;
safetensors_file.read(reinterpret_cast<char *>(&header_size), 8);
头部是一个 JSON，包含文件中的所有张量信息。JSON 是一组 <key, value> 对，其中 key 是张量名称的唯一字符串，value 是另一个 JSON 对象，包含该张量的信息。JSON 中的每个 key 都是某个张量的名称，除了一个名为 __metadata__ 的特殊 key（可能用于存储其他必要信息，我们不会用到它，规范说它是"用于存储自由格式文本-文本映射的特殊 key"）。每个 value 都是包含三个 key 的 JSON —— dtype、shape 和 offsets。dtype 表示张量存储的数据类型。shape 表示张量的维度。offsets 表示张量在张量数据部分内的存储位置。每个 shape 是一个长度不确定的整数列表，每个 offsets 值是一个恰好包含两个整数的向量。第一个元素表示张量的起始位置，最后一个元素表示张量的结束位置。
现在你面临第一个设计决策。你想让你的推理服务器与架构无关，使其能够运行任何模型，只要你实现了它所需的操作，还是想从简单开始并专注于我们选择的模型？
无论你选择什么，从单个模型开发然后泛化总是比一开始就试图在不确定代码最终样子的情况下使其灵活要容易得多。你总是可以后来回到这一点并在你选择的时候更新它。
如果你想让你的服务器与模型无关，你需要根据 Safetensors 头部动态分配内存、设置模型数据形状和类型（dtype），以及实现更多操作，确保覆盖所有你想支持的模型使用的操作。你可能仍然需要提供一些模型架构的蓝图，因为 Safetensors 文件不会告诉你应该用哪个操作来处理这些数据、以什么顺序处理等等。我不确定最优的方法是什么，但如果你自己或通过阅读 vLLM/TensorRT 等代码找到了答案，欢迎分享你的发现。
我将假设我们就为 Llama 3.2 1B Instruct 架构编码服务器。下面是加载 meta-llama/Llama-3.2-1B-Instruct 模型后的 Hugging Face Transformers LlamaForCausalLM 对象转储。我们可以检查需要实现哪些操作，以及需要使用什么数据形状和数据类型：
LlamaForCausalLM(
(model): LlamaModel(
(embed_tokens): Embedding(128256, 2048)
(layers): ModuleList(
(0-15): 16 x LlamaDecoderLayer(
(self_attn): LlamaAttention(
(q_proj): Linear(in_features=2048, out_features=2048, bias=False)
(k_proj): Linear(in_features=2048, out_features=512, bias=False)
(v_proj): Linear(in_features=2048, out_features=512, bias=False)
(o_proj): Linear(in_features=2048, out_features=2048, bias=False)
)
(mlp): LlamaMLP(
(gate_proj): Linear(in_features=2048, out_features=8192, bias=False)
(up_proj): Linear(in_features=2048, out_features=8192, bias=False)
(down_proj): Linear(in_features=8192, out_features=2048, bias=False)
(act_fn): SiLUActivation()
)
(input_layernorm): LlamaRMSNorm((2048,), eps=1e-05)
(post_attention_layernorm): LlamaRMSNorm((2048,), eps=1e-05)
)
)
(norm): LlamaRMSNorm((2048,), eps=1e-05)
(rotary_emb): LlamaRotaryEmbedding()
)
(lm_head): Linear(in_features=2048, out_features=128256, bias=False)
)
首先，从这个转储中我们既看不出操作的顺序，也看不出数据类型。但是！Hugging Face 页面上的模型卡告诉我们权重采用 BF16 格式。我们稍后会回到这个格式。
我们需要理解操作的顺序，才能知道如何编码。Sebastian Raschka 有一个大语言模型架构库，很好地展示了操作是如何组织的 —— 看这里（左边的那个）。
通过查看 Sebastian 的图表，我们看到 Llama 3.2 1B 中的操作顺序如下：
向模型发送一些文本
将其转换为 token（一个新概念，我们还没提到）
为每个 token 检索一个嵌入
16 个 Transformer 块（也称为层），包括：
遮蔽分组查询注意力，包括：Q 投影 K 投影 V 投影 RoPE 与 Q 投影 RoPE 与 K 投影 注意力 注意力得分 因果掩码 Softmax 残差连接 注意力得分与 V 投影
RoPE 与 Q 投影
RoPE 与 K 投影
注意力得分与 V 投影
O 投影（输出投影）
残差连接相加
前馈（类似于第一代神经网络中的多层感知器），包括：门投影、第一个线性层 上投影、第二个线性层 SiLU 激活函数、类似于 ReLU 但更像 sigmoid 下投影、第三个线性层 残差连接相加
门投影、第一个线性层
上投影、第二个线性层
SiLU 激活函数、类似于 ReLU 但更像 sigmoid
下投影、第三个线性层
残差连接相加
经过这些步骤后，我们应该获得由在我们服务器上运行的语言模型生成的第一个 token。
如果你对某些操作还不熟悉，不用担心。一旦我们在课程中取得进展，你会从感受上理解它们。我也经常忘记它们是如何工作的，经常需要再查一遍，所以你不必为反复查看或使用聊天机器人或互联网搜索而感到不好。
让我们提醒自己想要实现什么。我们想加载一个模型。我们已经知道了模型文件的结构，即 Safetensors 文件。我们知道了我们的参考模型架构。我们检查了模型权重以 BF16 类型存储。让我们花点时间思考一下这个类型和浮点数的一般情况。
计算机上的一切最终都是二进制的。计算机数字格式也是如此。
在机器学习模型中，权重几乎从不是整数。它们是实数。而计算机是二进制的。这意味着人们不得不想出如何在编程语言中以内存高效的方式表示实数。在这种背景下，这意味着你可以将大量信息打包在少量比特中。当然，你可以构建一个复杂的数据类型（这里的"复杂"指的是"非平凡"，而不是"实部+虚部"），其中小数点左边的部分存储为一个整数，分隔符后的部分（即小于 1 的部分）存储为另一个整数。但你可以看到这有多么低效。它只在你真正需要完整精度时才有用。这就是为什么 Python 中存在像 decimal 这样的东西的原因。
好，现在让我们想想还能采用什么方法。最简单的替代方案是取一个整数，比如 1234，然后规定：“我要在 12 和 34 之间加一个小数点”，这样就得到了 12.34。人们通过缩放因子解决了这个问题，比如对于 1234，直接使用缩放因子 1/100，就能得到 12.34。你仍然需要考虑如何在计算机中表示它。在这种格式中，小数部分（小数点之后的部分）的长度始终固定——小数点恰好位于中间。这种格式称为定点数。与浮点数相比，它的使用范围没那么广。
既然有定点数，就必然会有非定点数。我们先粗略地想一想。非定点可以意味着小数点（point）能够四处移动。小数部分可以更长，也可以更短。听起来这是一种内存效率更高的解决方案。
和其他所有数值类型一样，浮点数也表示为一串比特。不同浮点数的设计选择略有差异，因此我们重点讨论常规的 16 位浮点数（float16、FP16、IEEE 754-2008）。
Float16 和许多其他浮点数的工作方式相似。它们由三个部分组成：符号、指数和小数。总共占用 16 位。
[ sign | exponent | fraction ]
[ 0 | 0 1 0 0 1 | 0 0 1 0 0 0 0 0 0 0 ]
直观理解是这样的：在浮点数中，你可以移动小数点。小数部分是你用来移动小数点的那个数。就像前面的例子一样，你的小数部分可以是 1234。这一次，小数点的位置并不固定，因此你可以把它放在任何位置——1.234、0.0001234、1234000，等等。
符号占 1 位。0 表示正数，1 表示负数。
指数占 5 位。它控制数值的大小。
小数占 10 位。它控制我们用哪个数来移动小数点。它还有一些其他名称：有效数（significand）、尾数（mantissa）。
现在来看公式：$(-1)^{sign} * 2^{exponent-bias} * (1.fraction)$
注意这里出现了两个新东西——小数部分前面的 $1.$，以及 $bias$。这个 1. 是一种用于提高数值精度的设计选择：无须在内存中显式存储，就可以免费获得 1 位。因此，它被称为隐含位（implicit）。关于这一点，可以看看 Stack Overflow 上这个很不错的问答，里面有更好的解释。核心思想是：它并不是什么必不可少的东西，而是一种无需增加内存占用就能存储更多数据的巧妙技巧——浮点数规范直接规定这一位始终存在（零除外）。接下来出现的新东西是 $bias$，再看看它所在的位置：$2^{exponent-bias}$。你能猜到为什么需要它吗？
答案是：为了也能够表示小于 1 的数。更具体地说，是为了让 $2$ 能够拥有负指数。指数始终是一个正数，由若干比特表示（在 float16 中是 5 位）。如果不从中减去偏置量，那么 $2^{exponent}$ 就会始终是正的，而且可能是个相当大的数（二进制 11111 对应十进制 31）。这样一来，我们可以生成任何大于 1 或小于 -1 的数（准确地说几乎是任何数，仍然受浮点数精度限制），却无法表示任何大于 0 的小数，例如 0.0001234。因此，我们需要从指数中减去偏置量。偏置量既不能太大，也不能太小。知道这些之后，你认为 float16 的偏置量是多少？别忘了，指数的最大值是 31。
所以，float16 的偏置量是 15。这意味着你既能表示较小的数，也能表示较大的数。
浮点数规范中还有更多有趣的内容，例如如何编码正无穷和负无穷、如何编码非数值（not-a-number）等。Wiki 中有很多关于指数编码的例子。
指数和小数部分的大小，是不同浮点数类型之间的主要区别。
让我们来看一个例子，把数字 12.34 编码为 float16。12.34 是正数，所以符号位是 0。为了表示 12.34，我们需要先把 12 转换为二进制：
12 / 2 = 6，余数为 0
6 / 2 = 3，余数为 0
3 / 2 = 1，余数为 1
1 / 2 = 0，余数为 1
倒序读取，得到 1100。
所以，十进制的 12 就是二进制的 1100。
现在把小数部分 0.34 转换为二进制：
0.34 * 2 = 0.68，整数部分为 0
0.68 * 2 = 1.36，整数部分为 1，减去 1 后继续
0.36 * 2 = 0.72，整数部分为 0
0.72 * 2 = 1.44，整数部分为 1，减去 1 后继续
0.44 * 2 = 0.88，整数部分为 0
0.88 * 2 = 1.76，整数部分为 1，减去 1 后继续
0.76 * 2 = 1.52，整数部分为 1，减去 1 后继续
0.52 * 2 = 1.04，整数部分为 1，减去 1 后继续
0.04 * 2 = 0.08，整数部分为 0
0.08 * 2 = 0.16，整数部分为 0
此时，我们已经计算出了 10 位，正好等于小数部分／尾数／有效数的长度。乍看之下，我们似乎刚好计算出了所需的位数，但其实并不完全正确。这是因为，我们还需要把被编码数字的第一部分——12（二进制为 1100）——放进尾数中。$1.$ 始终是这个数的隐含部分，因此需要从 $1100$ 中移除第一个 $1$。剩下的 $100$ 必须放进尾数。这意味着我们多计算了 3 位，最终需要舍弃它们，因为尾数中必须留出 3 位来容纳 $12$。所以我们并没有犯错，只是做了一些不必要的工作。如果你有一天决定实现自己的数值类型，这一点值得牢记（如果你真的自己实现了，请务必告诉我！）。
顺便看看最后一项：“0.08 * 2 = 0.16，整数部分为 0”。这个过程还没有结束，因为我们尚未得到 0（而是停在了 0.16）。这意味着在表示这个数时，我们将损失精度。
按照从上到下的顺序读取这些位，十进制的 0.34 近似等于二进制的 0101011100。
把 12 和 0.34 放在一起，就得到了 1100 和 0101011100。再次强调，$1.0$ 是隐含的，因此剩下的是 100 和 0101011100。把它们拼在一起，得到 1000101011100，一共是 13 位。移除最低有效的 3 位（右边的 3 位），最终得到这个数的小数部分的二进制表示：1000101011。
我知道自己有点跑题了，你本来应该是在构建一个推理服务器。但说真的，我们为什么要着急呢？如果你对这些内容不感兴趣，完全可以直接跳到后面。我相信，深入学习事物本身有一种内在的乐趣；如果你也有同感，那就让我们继续。
尾数部分已经完成了。看看我们是如何把 $1100$ 变成 $1.100$ 的。如果仍然身处十进制世界（以 10 为基数），我们会说，要把 $1.100$ 变成 $1100$，我们需