jmaczan/tiny-vllm 来源：https://github.com/jmaczan/tiny-vllm

tiny-vllm

你将用C++和CUDA构建一个高性能LLM推理引擎——tiny-vllm，它是vLLM（https://github.com/vllm-project/vllm）的更年轻、更小巧的兄弟。在这个过程中我们将学到很多东西，会犯错误，并从零推导出各种思路和数学原理。

本仓库包含两个部分：

1. 推理服务器的完整源代码
2. 一份教程，我将带领你逐步实现这个引擎。欢迎你将其作为学习工具，或者如果你是讲师，也欢迎在课堂上作为教学资源使用。

推理引擎包含以下功能：

• 从Safetensors加载真实LLM模型（Llama 3.2 1B Instruct）
• 完整的LLM前向传播（prefill + decode）
• 所有计算均使用CUDA内核
• KV缓存
• 静态批处理
• 连续批处理
• 在线softmax，类似FlashAttention（https://courses.cs.washington.edu/courses/cse599m/23sp/notes/flashattn.pdf）
• PagedAttention（https://arxiv.org/pdf/2309.06180）

给自己泡杯热饮，让我们开始吧。

• tiny-vllm
• 简介：LLM、vLLM、模型、推理服务器
• 技术前置条件
• Safetensors与你的模型
• 浮点数如何工作以及为什么使用bfloat16
• GPU和CPU内存
• 单token推理
• 分词
• 嵌入
• CUDA内核工程——嵌入
• RMSNorm与CUDA中的并行归约
• RoPE
• 残差连接
• cublasGemmEx
• 列主序到行主序的转置技巧
• Prefill与decode
• 为什么存在KV缓存
• 注意力机制
• GQA
• SiLU
• Softmax
• 因果掩码
• Argmax
• 前馈网络
• 缓冲区复用
• 静态批处理
• 连续批处理
• 在线softmax
• Paged Attention
• Paged KV缓存
• Paged Attention CUDA内核

简介：LLM、vLLM、模型、推理服务器

近年来各种事情层出不穷，很容易迷失方向。让我们来梳理一下。

LLM就是一个模型。从物理上讲，LLM是一个包含大量浮点数的文件（https://en.wikipedia.org/wiki/Floating-point_arithmetic）。概念上，这些数字代表操作的权重。权重是在训练阶段学习/发现/找到的。一些操作会用到这些权重。每个操作都是一个函数，接收一些数据作为输入，对其进行处理，然后产生数据作为输出。操作及其顺序由LLM的架构定义。每个模型都有自己独特的架构，由工程师和研究人员设计。

从零开始到LLM写出文本的过程如下： 0. 设计模型 —— 工程师和研究人员使用像PyTorch（https://github.com/pytorch/pytorch）或tinygrad（https://github.com/tinygrad/tinygrad）这样的张量库，用高级语言（如Python）设计模型架构。他们训练小版本模型，用不同的操作、数据和超参数（操作的参数）做实验。这是确定规范（specification）的阶段。

1. 实现模型 —— 一旦他们确定了最终模型架构并准备好训练数据，他们编写定义最终模型的代码。也可以用PyTorch或类似工具。
2. 训练模型 —— 选定的模型架构用随机权重初始化。他们编写一个脚本（再次使用PyTorch等工具），在大量硬件（如GPU（https://en.wikipedia.org/wiki/Graphics_processing_unit）和TPU（https://en.wikipedia.org/wiki/Tensor_Processing_Unit））上运行反向传播等学习算法。这个阶段消耗大量能源、资金和计算能力。训练的产物是一个包含模型权重的文件，采用某种格式，比如Safetensors格式（https://huggingface.co/docs/safetensors/index）。因此，训练阶段就是找到一组权重，使得给定架构能够生成高质量的文本。
3. 部署模型（我们所在阶段） —— 权重文件不能直接在计算机上运行。它不是可执行文件。它只是一堆数字。架构本身也无法运行——它只是一个计划、蓝图、计算描述。要实际运行模型，我们需要一个程序，将架构及其操作转化为可执行代码，并使用权重文件将权重加载到架构中。一旦你编写了一个实现这些操作的程序，并且该程序加载了权重（权重在程序运行时、启动时加载），你终于可以将提示发送给模型并得到有意义的响应。从模型生成输出称为推理。这就是我们这里构建的东西被称为推理服务器或推理引擎的原因。

知道了推理服务器需求背后的原因，我们再来思考为什么用C++和CUDA来构建它。这是因为我们希望最大化硬件利用效率，获得高性能。这意味着我们希望快速得到响应，并能够同时处理多个提示。CUDA是整个生态系统，也是一种用于编写在GPU上运行代码的语言。我们需要编写在GPU上运行的代码，因为LLM内部的许多操作都是对大量数字进行乘法和加法。如果只需要少量计算，CPU就足够了。如果需要大量计算，GPU更合适。LLM主要涉及矩阵乘法，归根结底就是计算两个向量的点积，而且是对大量数字和大量向量进行计算。LLM的数学很简单，我们只需要线性代数的基础知识，你可以在编码过程中学习，并随时填补知识空白。我发现这种即时学习（JIT learning）方式最有效，也许你也会喜欢。

我对AI与计算之间关系的一点看法（可能对你有帮助）是：智能来自于大量模型参数，以及使用这些参数对输入值进行大量计算。没有一个单一的元素可以让你指着它说：“这就是让模型变得智能或有用的东西”。模型中的每个部分都可以用不同的部分替换，从而得到不同的权衡，比如用准确性换取复杂度。我希望以后当我们接触到注意力的数学原理时，我不会忘记回到这个话题。因为——默认的注意力机制在计算上非常复杂（O(n²*d)）。这种复杂性是可以被挑战的，事实上人们已经做到了，并提出了替代的注意力机制，比如线性注意力（https://haileyschoelkopf.github.io/blog/2024/linear-attn/）。如果更多人觉得这个教程有用，我会考虑再写一个：关于ML编译器（使用Python或C++ + 一些SSA理论的实战教程），或者关于替代注意力机制（数学 + CUDA内核）。如果你感兴趣，请告诉我！如果你觉得这个教程有价值，请告诉其他人。

范围外：LLM的训练阶段不在本教程范围内。我们拿一个训练好的LLM，编写一个程序，使其能在NVIDIA GPU上快速并行处理多个请求。如果你想训练自己的LLM，强烈推荐Karpathy老师的仓库，如nanoGPT（https://github.com/karpathy/nanoGPT）和llm.c（https://github.com/karpathy/llm.c），以及他的YouTube频道（https://www.youtube.com/@AndrejKarpathy）。同样，我们也不设计模型，但张量库本身也是一个迷人的话题，值得从零理解。George Hotz的tinygrad（https://github.com/tinygrad/tinygrad）是一个用极少代码实现张量库的项目，如果你想获得灵感并学习内部原理，这是一个好地方（他们的Discord也不错（https://discord.com/invite/ZjZadyC7PK））！还有一个更早更小的版本，由Andrej Karpathy编写——micrograd（https://github.com/karpathy/micrograd）。既然我提到了Discord，我还想向你推荐Mark Saroufim（https://www.marksaroufim.com/）的GPU MODE（https://discord.com/invite/gpumode）。那里有很多优秀的人聚集！如果你在这里感到迷茫，并且你刚踏上AI/ML之旅，可以从Jeremy Howard和Rachel Thomas（https://www.fast.ai/about）的fastai书籍（https://course.fast.ai/Resources/book.html）开始。我在这里方便地略过了数据科学和工程部分，因为我对这方面了解不多。Kaggle（https://www.kaggle.com/）可能是一个开始并实践学习的好地方。

最后但同样重要的是，我们将用C++和CUDA编程，并在适用时使用cuBLAS（https://developer.nvidia.com/cublas）。你可以在实践过程中学习。NVIDIA官方资源（https://docs.nvidia.com/cuda/cuda-programming-guide/）很好也很有帮助。

技术前置条件

你可以在任何平台上构建和运行，只需少量修改，只要你有NVIDIA GPU。你可能需要调整一些路径，比如c_cpp_properties.json中的CUDA或GCC路径，或者CMakeLists.txt中的NVCC路径。

我建议你fork这个仓库，进行必要的调整使其能在你的机器上运行，然后向jmaczan/tiny-vllm（https://github.com/jmaczan/tiny-vllm）创建pull request，将你的更改向上游提交，以帮助其他读者。

我开发和测试的具体环境：

• Linux（6.19.8 x64_64）
• CUDA Toolkit（https://docs.nvidia.com/cuda/cuda-installation-guide-linux/）（13.1）
• C++ 17
• GCC（https://gcc.gnu.org/）（15.2.1）
• 唯一的外部依赖是JSON解析器nlohmann/json（https://github.com/nlohmann/json）3.12.0，它是一个单头文件include/json.hpp
• AMD CPU（Ryzen 7 9800X3D）
• NVIDIA GPU（RTX 5090）
• 我使用了来自Hugging Face的Llama 3.2 1B Instruct（https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct）（提交哈希898999bd25b40516fce5a5b8f0948f4c81c650bc），你只需要这个仓库中的model.safetensors文件。

安装依赖项，然后用./test.sh运行程序——它将构建并立即执行。

如果你构建或运行失败，并且你选择的AI无法帮助，请在GitHub上创建一个Issue——我会尽力帮助。请确保提供所有有用的上下文。

Safetensors与你的模型

第一件事是下载一个LLM，我们将用它来运行推理。我选择Llama 3.2 1B Instruct，因为它简单、小巧、专为对话优化，而且对我们来说足够好。从我们构建推理服务器的工程师角度来看，模型只是一个包含权重的文件。模型采用Safetensors格式（https://huggingface.co/docs/safetensors/index）。还有其他格式，如Pickle（https://docs.python.org/3/library/pickle.html）和Parquet（https://parquet.apache.org/docs/file-format/）。Safetensors非常流行且广泛使用，我们选择的模型正是以Safetensors格式托管的。

让我们停下来稍作理解，在继续之前了解Safetensors格式。

一个safetensor文件包含三个部分，始终按此顺序：头部大小、头部和张量数据。头部大小总是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都是一个张量名称，除了一个特殊的key，叫做__metadata__，可能用于一些附加信息（但我们不会用到，规范说它是“用于存储自由格式文本到文本映射的特殊键”）。每个value都是一个JSON，包含三个键：dtype、shape和offsets。dtype表示张量的数据类型。shape表示张量的维度（https://en.wikipedia.org/wiki/Tensor#As_multidimensional_arrays）。offsets表示张量在张量数据段中的存储位置。每个shape是一个长度未知的整数列表，每个offsets值是一个恰好包含两个整数的向量。第一个元素表示张量开始的位置，最后一个元素表示张量结束的位置。

现在你面临第一个设计决策。你是想让你的推理服务器与架构无关，以便它能运行任意模型（只要你实现了它所需的所有操作），还是想从简单的开始，专注于我们选择的模型？无论你决定什么，先针对单个模型开发，然后再进行泛化，总是比在一开始不确定最终代码会是什么样子时就试图使其灵活更容易。你可以随时回来更新，如果你选择这样做的话。

如果你想使服务器与模型无关，你需要根据Safetensors头部动态分配内存，设置模型数据的shape和类型（dtype），并实现更多操作，确保覆盖你想支持的所有模型使用的所有操作。你可能仍然需要提供模型架构的一些蓝图，因为Safetensors文件不会告诉你用这些数据运行哪个操作、按什么顺序等。我不确定实现这一点的最佳方法，但如果你自己找到了方法，或者通过阅读vLLM/TensorRT（https://github.com/NVIDIA/TensorRT）/其他代码找到了，欢迎分享你的发现。

我将假设我们仅为Llama 3.2 1B Instruct架构编码服务器。以下是加载了meta-llama/Llama-3.2-1B-Instruct模型的Hugging Face Transformers（https://huggingface.co/docs/transformers/index）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页面上的模型卡片（https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct）告诉我们权重是BF16（https://en.wikipedia.org/wiki/Bfloat16_floating-point_format）格式。我们很快就会回到这个格式。我们需要了解操作的顺序才能知道如何编码。Sebastian Raschka有一个LLM架构图库，很好地展示了操作的组织方式。