用C++和CUDA构建你自己的高性能LLM推理引擎

用C++和CUDA构建你自己的高性能LLM推理引擎

你是否曾好奇，当你向大型语言模型发送提示时，幕后发生了什么？魔法不仅仅在于模型权重——还在于编排计算的推理引擎。本文带你用C++和CUDA构建自己的高性能LLM推理引擎，灵感来自vLLM的架构，但设计为学习工具。

为什么要构建自己的推理引擎？

大多数开发者通过高级API（如Hugging Face Transformers或OpenAI）与LLM交互。但如果你想了解真正的性能特征——延迟、吞吐量、内存使用——或针对特定硬件进行优化，你需要更深入。从头构建推理引擎教会你：

• LLM实际上如何计算预测
• CUDA内核在加速矩阵运算中的作用
• GPU工作负载的内存管理策略
• 像PagedAttention和连续批处理这样的高级技术

tiny-vllm的架构

项目tiny-vllm实现了Llama 3.2 1B Instruct的完整推理服务器。它涵盖：

• 从Safetensors格式加载真实LLM
• 完整前向传播（预填充+解码）
• 所有计算使用自定义CUDA内核
• KV缓存管理
• 静态和连续批处理
• 在线softmax和类似FlashAttention的技术
• PagedAttention

入门：技术先决条件

要在你的机器上运行，你需要：

• Linux（在6.19.8 x86_64上测试）
• CUDA工具包（13.1）
• C++17编译器（GCC 15.2.1）
• NVIDIA GPU（在RTX 5090上测试）
• 模型文件：来自Llama 3.2 1B Instruct的model.safetensors

唯一的外部依赖是nlohmann/json（单头文件）。

理解Safetensors格式

在编写任何代码之前，你需要了解模型权重是如何存储的。Safetensors文件有三个部分：

1. 头部大小（8字节）：一个64位无符号整数，表示头部长度
2. 头部（JSON）：包含每个张量的元数据——名称、数据类型、形状和偏移量
3. 张量数据：实际的权重值

以下是读取头部大小的方法：

std::ifstream safetensors_file("model.safetensors", std::ios_base::binary);
uint64_t header_size;
safetensors_file.read(reinterpret_cast<char *>(&header_size), 8);

头部JSON告诉你每个张量在数据部分中的起始和结束位置。这允许你将张量直接映射到GPU内存，而无需复制不必要的数据。

模型架构：Llama 3.2 1B

从Hugging Face模型卡中，我们可以检查架构：

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)
)

前向传播遵循以下顺序：

1. 对输入文本进行分词
2. 检索每个标记的嵌入
3. 对于16个Transformer层中的每一层：
   • RMSNorm
   • 分组查询注意力（带RoPE）
   • 残差连接
   • RMSNorm
   • 前馈网络（带SiLU）
   • 残差连接
4. 最终RMSNorm
5. 线性投影到词汇表大小
6. Argmax选择下一个标记

你的第一个CUDA内核：嵌入收集

让我们实现嵌入查找。对于每个输入标记，我们需要从嵌入表中检索一个2048元素的向量。

__global__ void embeddingGatherKernel(int *gpu_input_tokens, __nv_bfloat16 *input_embeddings, __nv_bfloat16 *embed_tokens)
{
    int workIndex = threadIdx.x + blockIdx.x * 2048;
    input_embeddings[workIndex] = embed_tokens[gpu_input_tokens[blockIdx.x] * 2048 + threadIdx.x];
    input_embeddings[workIndex + 1024] = embed_tokens[gpu_input_tokens[blockIdx.x] * 2048 + threadIdx.x + 1024];
}

注意技巧：由于大多数GPU每个块最多支持1024个线程，但每个嵌入有2048个元素，我们每个线程处理两个元素。这避免了启动不必要的块。

使用并行归约的RMSNorm

RMSNorm对每个嵌入向量进行归一化。公式为：

$$\text{RMS}(a) = \sqrt{\frac{1}{n}\sum_{i=1}^{n}a^2_i}$$

为了在GPU上高效计算，我们在共享内存中使用并行归约（树归约）：

__shared__ float rms_vector[1024];
int workIndex = threadIdx.x + blockIdx.x * 2048;
rms_vector[threadIdx.x] = (float)input[workIndex] * (float)input[workIndex] 
                        + (float)input[workIndex + 1024] * (float)input[workIndex + 1024];
__syncthreads();

// 树归约
for (int i = 1; i < 1024; i = i * 2)
{
    if (threadIdx.x % (i * 2) == 0)
    {
        rms_vector[threadIdx.x] = rms_vector[threadIdx.x] + rms_vector[threadIdx.x + i];
    }
    __syncthreads();
}

if (threadIdx.x == 0)
{
    rms_vector[0] = sqrt(rms_vector[0] / 2048.0 + 1.0e-5);
}
__syncthreads();

output[workIndex] = (__nv_bfloat16)(((float)input[workIndex] / rms_vector[0]) * (float)norm_weights[threadIdx.x]);
output[workIndex + 1024] = (__nv_bfloat16)(((float)input[workIndex + 1024] / rms_vector[0]) * (float)norm_weights[threadIdx.x + 1024]);

cuBLAS转置技巧

cuBLAS期望矩阵按列主序格式，但模型通常按行主序存储。我们可以利用矩阵转置性质，而不是转换数据：

$$C = A \times B^T \implies C^T = B \times A^T$$

由于cuBLAS将行主序数据视为列主序（实际上相当于转置），我们相应地设置标志：

cublasGemmEx(cublas_handle, 
             CUBLAS_OP_T, CUBLAS_OP_N, 
             KV_DIM, num_active_slots, EMBEDDING_LENGTH, 
             &k_proj_alpha, weights.w_k[layer], CUDA_R_16BF, EMBEDDING_LENGTH, 
             rms_norms, CUDA_R_16BF, EMBEDDING_LENGTH, 
             &k_proj_beta, k_proj_batched_buffer, CUDA_R_16BF, KV_DIM, 
             CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT);

为什么存在KV缓存

在推理过程中，只有每个标记的K和V投影需要持久化。Q投影、注意力分数和前馈输出在每一步之后都会被丢弃。通过缓存K和V，我们避免了为每个新标记重新计算它们——这是一个巨大的性能提升。

预填充与解码

• 预填充：处理所有输入标记以生成第一个输出标记。这是计算密集型的，因为你一次处理许多标记。
• 解码：一次生成一个后续标记。这是内存密集型的，因为你正在获取缓存的K和V值。

分组查询注意力（GQA）

在GQA中，多个查询头共享相同的键和值头。对于Llama 3.2 1B，4个查询头共享1个键/值头。这减少了内存使用并提高了吞吐量，而不会显著降低质量。

for (int i = 0; i < NUM_Q_HEADS; ++i)
{
    int k_head_idx = i / GQA_Q_TO_K_RATIO; // 4个Q头使用相同的1个K头
    __nv_bfloat16 *q_head = q_proj + i * HEAD_DIM;
    __nv_bfloat16 *k_head = k_proj[layer] + k_head_idx * HEAD_DIM;
    // ... 计算注意力分数
}

连续批处理

静态批处理将N个请求一起处理，但等待最慢的一个。连续批处理通过使用槽来解决这个问题：当一个请求完成时，它的槽立即被队列中的新请求重用。这最大限度地提高了GPU利用率。

下一步是什么？

该存储库正在积极开发中。即将推出的部分包括：

• 在线softmax推导和实现
• 用于高效内存管理的PagedAttention
• 分页KV缓存

结论

从头构建LLM推理引擎是深入理解现代AI系统工作原理的最佳方式之一。你将学习GPU编程、内存管理、线性代数和系统设计——同时构建一个真正运行真实模型的东西。

如果你是讲师，欢迎将此作为教学资源。如果你是学习者，请fork存储库，在你的机器上运行，并贡献你的改进。

祝你编码愉快！

来源

jmaczan/tiny-vllm: 用C++和CUDA构建你自己的高性能LLM推理引擎 - vLLM的简化版本