使用 Tiny-vLLM 从零开始构建高性能 LLM 推理引擎

现代大语言模型 (LLM) 的复杂性往往掩盖了实际运行它们所需的底层工程实现。虽然像 PyTorch 这样高层库让模型设计变得触手可及，但从权重文件到响应式推理服务器的转变，需要深入研究内存管理、GPU kernel 工程以及线性代数。

Tiny-vLLM 是一个旨在揭开这一过程神秘面纱的教育项目。它不是提供一个黑盒，而是作为一个引导式课程，使用 Llama 3.2 1B Instruct 模型作为参考，教你如何使用 C++ 和 CUDA 构建高性能推理引擎。它弥合了理论架构与硬件加速这一艰巨现实之间的鸿沟。

推理服务器的解剖结构

要构建一个推理引擎，首先必须理解 LLM 在物理层面是什么：一个包含数百万个浮点数（权重）的文件，以及一个定义这些数字如何在一系列操作中使用的蓝图（架构）。

推理服务器的主要任务是将该蓝图转化为可执行代码。为了实现高性能，这通常使用 C++ 和 CUDA 来完成，以最大限度地提高硬件利用率。核心目标是在最小化延迟的同时处理多个 prompt，这需要将计算从 CPU 转移到 GPU，因为矩阵乘法——LLM 的核心——可以在数千个核心上进行并行化。

加载模型：Safetensors 与 BF16

大多数现代模型都以 Safetensors 格式分发。一个 Safetensors 文件由头部大小、包含张量元数据（dtype、shape 和 offsets）的 JSON 头部以及原始张量数据组成。

Bfloat16 的选择

LLM 推理中的一个关键技术决策是数值格式。虽然标准的 16 位浮点数 (FP16) 很常见，但许多模型使用 Bfloat16 (BF16)。

BF16 以精度换取范围。它使用 8 位用于指数（与 32 位浮点数相同），仅使用 7 位用于小数部分。这使得它在 transformer 层所需的巨量求和过程中，显著降低了溢出和欠拟合问题的发生概率；经验证据表明，精度损失对于其提供的稳定性来说是一个可以接受的权衡。

GPU 内存管理

高效推理需要一套严格的数据在 Host (CPU/DRAM) 与 Device (GPU/VRAM) 之间移动的策略。 由于 GPU 无法直接访问系统 DRAM，数据必须通过 cudaMalloc 进行显式分配，并通过 cudaMemcpy 进行传输。

为了优化性能，引擎遵循以下几条黄金法则：

  1. 最小化分配： 预先分配大型缓冲区。
  2. 复用内存： 使用生命周期分析来识别一旦数据不再需要即可重新利用的缓冲区。
  3. 减少传输： 尽可能减少将数据复制到 GPU 的次数。

工程化 CUDA Kernels

实现 transformer 的前向传播需要编写自定义 CUDA kernels。这涉及以 SIMT (Single Instruction, Multiple Threads) 的思维进行思考，即在被分组为 blocks 和 warps 的数千个线程中执行相同的函数。

Embedding Gather

第一步是将输入 token (整数) 映射到其对应的 embedding 向量。由于大多数 GPU 每个 block 的线程限制为 1024，但 Llama 3.2 的 embedding 长度为 2048，一种常见的优化方法是让每个线程处理 embedding 的两个元素，从而有效地将每个 block 的吞吐量翻倍。

RMSNorm 与 Parallel Reduction

均方根层归一化 (RMSNorm) 需要计算整个 embedding 的平方值的平均值的平方根。这产生了一个同步挑战：如何跨不同线程对数值进行求和？

Tiny-vLLM 使用 Parallel Reduction (Tree Reduction) 来实现这一点。线程不使用单个共享变量（这会导致竞态条件），而是将它们的局部和写入一个 __shared__ 内存向量。然后，它们以树状结构迭代地对这些值进行求和，每一步将活动线程数减半，直到在索引 0 处保留一个总和。

为了确保数值稳定性并防止除以零（这会导致 NaN 并导致推理崩溃），在 RMS 计算中加入了一个 epsilon 值（例如 1e-05）。

矩阵乘法挑战

矩阵乘法是引擎中最耗费计算资源的环节。虽然可以编写自定义 kernel，但行业标准是 cuBLAS，NVIDIA 的高度优化线性代数库。

然而，存在一个显著的摩擦点：cuBLAS 期望矩阵采用 column-major format（列主序格式），而 LLM 权重通常以 row-major format（行主序格式）存储。Tiny-vLLM 采用了一种数学转置技巧来避免在内存中物理重新排列数据。通过操作转置标志 (CUBLAS_OP_T 和 CUBLAS_OP_N) 以及 GEMM (General Matrix Multiply) 调用中的操作顺序，引擎可以将行主序数据视为列主序，从而显著降低开销。

优化吞吐量：Batching 与 Caching

Prefill 与 Decode

推理分为两个截然不同的阶段：

  • Prefill： 引擎处理整个输入 prompt 以生成第一个 token。这在计算上很重，但具有高度并行性。
  • Decode： 引擎逐个生成后续 token。由于它只处理最后一个生成的 token，因此受限于内存带宽。

KV Cache

为了避免在 decode 阶段为每个之前的 token 重新计算 Key (K) 和 Value (V) 的投影，引擎实现了 KV Cache。该缓存存储了所有已处理 token 的 K 和 V 向量，使得模型只需追加新 token 的投影并对现有历史进行 attention 运算。

高级 Batching

为了扩展服务器，Tiny-vLLM 探索了两种类型的 batching：

  • Static Batching： 一次处理多个请求。虽然这增加了吞吐量，但它引入了延迟，因为整个 batch 必须等待最长的序列完成。
  • Continuous Batching： 一种更复杂的方法，新请求可以在任何空闲槽位出现时立即进入 batch，从而消除了与静态填充 (static padding) 相关的浪费。

结论

从零开始构建推理引擎揭示了，LLM 的“智能”并非源于单行代码，而是源于参数的规模以及计算的效率。通过实现这些底层原语——从 RMSNorm 中的 tree reductions 到 cuBLAS 的转置技巧——开发者可以超越“将 LLM 作为服务使用”的阶段，开始理解使现代 AI 成为可能的软硬件协同设计。

Sources