Tiny-vLLM:用C++和CUDA构建高性能LLM推理引擎
Tiny-vLLM是一个开源项目,提供用C++和CUDA实现的高性能LLM推理引擎的完整源代码和配套课程,旨在帮助开发者从零理解并实现推理引擎的核心技术。
您将使用C++和CUDA构建高性能LLM推理引擎——tiny-vllm,它是vLLM的年轻且更小的兄弟。一路上我们会学到很多东西,犯错误,并从头推导出想法和数学。该仓库包含两部分:1. 推理服务器的完整源代码;
2. 一门课程,我将带领您完成实现引擎的过程。欢迎您将其用作学习工具,如果您是讲师,也欢迎将其用作大学的教学资源。推理引擎包括:- 从Safetensors加载真实的LLM模型(Llama 3.
2 1B Instruct)- 完整的LLM前向传播(预填充+解码)- 所有计算均使用CUDA内核- KV缓存- 静态批处理- 连续批处理- 在线softmax,类似FlashAttention- PagedAttention给自己泡一杯热饮,我们开始吧。
- tiny-vllm- 简介:LLM、vLLM、模型、
推理服务器- 技术先决条件- Safetensors和你的模型- 浮点数的工作原理以及为什么我们使用bfloat16- GPU和CPU内存- 单token推理- 分词- 嵌入- CUDA内核工程——嵌入- CUDA中的RMSNorm和并行归约- RoPE- 残差连接- c
ublasGemmEx- 列主序到行主序的转置技巧- 预填充与解码- 为什么存在KV缓存- 注意力- GQA- SiLU- Softmax- 因果掩码- Argmax- 前馈网络- 缓冲区重用- 静态批处理- 连续批处理- 在线softmax- Paged Attenti
on- 分页KV缓存- Paged Attention CUDA内核近年来发生了这么多事情,很容易迷失。让我们来梳理一下。LLM是一个模型。物理上,LLM是一个包含大量浮点数的文件。从概念上讲,这些数字代表操作的权重。
权重在训练阶段被学习/发现/找到。一些操作使用这些权重。每个操作都是一个函数,它将一些数据作为输入,对其进行处理并产生数据作为输出。操作及其顺序由LLM的架构定义。
每个模型都有自己的架构,由工程师和研究人员设计。从零到LLM写出文本的过程如下:- 设计模型——工程师和研究人员使用Python等高级语言以及PyTorch或tinygrad等张量库来设计模型的架构。
他们训练模型的小版本,用不同的操作、数据和超参数(操作的参数)进行实验。这是确定规范的阶段。- 实现模型——一旦他们决定了最终的模型架构并准备好训练数据,他们就会编写定义最终模型的代码。这也可以用PyTorch或类似工具完成。
- 训练模型——使用虚拟权重初始化所选的模型架构。他们编写一个脚本,再次使用PyTorch或类似工具,在大量硬件(如GPU和TPU)上运行反向传播等学习算法。这个阶段消耗大量的能源、金钱和计算能力。
训练阶段的产物是一个包含模型权重的文件,采用某种格式,例如Safetensors格式。因此,训练阶段是寻找这样一组权重,使得使用给定架构能够生成良好的文本。- 服务模型(我们在这里)——包含权重的文件不能在计算机上直接运行。
它不是可执行文件。它只是一堆数字。架构本身也无法运行——它只是一个计划、一个蓝图、一个计算描述。
为了实际运行模型,我们需要一个程序,将架构及其操作转化为可执行代码,并使用模型权重文件将权重加载到架构中。一旦你编写了实现操作的程序,并且程序加载了权重(权重在程序启动时加载),你最终就可以向模型发送提示并获得有意义的响应。
从模型生成输出称为推理。这就是为什么我们在这里构建的东西被称为推理服务器或推理引擎。了解了需要推理服务器的原因后,让我们思考为什么我们用C++和CUDA来构建它。这是因为我们希望最大限度地高效利用硬件并获得高性能。
这意味着我们希望快速获得响应,并且希望能够同时处理多个提示。CUDA是整个生态系统,但也是一种用于编写在GPU上运行的代码的语言。我们需要在GPU上编写代码,因为LLM中的许多操作都是乘法和加法。
如果你需要做少量的数学运算,CPU就足够了;如果很多,GPU更好。LLM主要涉及矩阵乘法,这归结为计算两个向量的点积,涉及大量数字和大量向量。LLM的数学很简单,我们需要线性代数的基础知识,你可以在编码时学习并随时填补空白。
我发现这种即时学习方式最有效,也许你也会喜欢。我对人工智能和计算之间关系的看法,你可能会觉得有用:智能来自模型的大量参数以及使用这些参数对输入值进行的大量计算。没有一个单一的元素可以让你指着说:“这就是使模型变得智能或有用的原因”。
你可以用不同的部分替换模型的每个部分,并获得不同的权衡作为回报,例如用准确性换取复杂性。我希望稍后当我们触及注意力的数学时不会忘记回到这个话题。因为默认的注意力机制在计算上非常复杂(O(n^2*d))。
这种复杂性可能会受到挑战,事实上人们确实在这样做,并找出替代的注意力机制,例如线性注意力。如果更多人觉得这门课程有用,我会考虑再开一门关于ML编译器(Python或C++中的实用编译器,加上一些SSA理论)或关于替代注意力机制(数学+CUDA内核)的课程。
如果你有兴趣,请告诉我!如果你觉得这门课程有价值,请告诉其他人。范围之外: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_properties.json中的CUDA或GCC,或CMakeLists.txt中的NVCC。
txt我建议你fork这个仓库,进行必要的调整使其在你的机器上运行,然后向jmaczan/tiny-vllm创建pull request,将你的更改上游,以便其他读者受益。
我开发和测试它的确切环境:- Linux (6.19.8 x64_64)- CUDA Toolkit (13.1)- C++ 17- GCC (15.2.1)- 唯一的外部依赖是y
本文为机器翻译辅以 AI 润色,仅供参考。原始事实以原文为准。