C++/CUDA 手写 LLM 推理引擎：拆解 vLLM 核心 PagedAttention 与连续批处理

一、项目释义

本工程是一套C++ + CUDA 从零实现的轻量化大模型推理服务，对标vLLM核心设计思想，剥离工业级框架复杂封装，专门用于学习现代LLM推理底层原理。

整套工程由两大部分构成：完整推理服务源码 + 配套学习教程。模型选用Llama 3.2 1B Instruct，权重采用Safetensors格式存储，全部计算逻辑由手写CUDA Kernel结合cuBLAS完成。

项目覆盖现代推理引擎必备全套能力：模型权重加载、Prefill预填充+Decode解码完整前向链路、KV缓存、静态批处理、连续批处理 、Online Softmax、类FlashAttention实现、PagedAttention分页KV缓存。

项目定位不是面向线上大规模生产部署，而是教学型工程。开发者可以顺着代码，从头推导数学公式、调试CUDA核函数，亲手实现vLLM最核心的分页缓存与连续批调度逻辑；高校讲师也可以把这套工程作为大模型推理课程实验素材。

推理流程一句话概括：读取Safetensors权重文件，加载BF16模型参数到GPU显存；接收多用户prompt，分词后送入Transformer网络；Prefill一次性处理全部输入token并构建分页KV缓存；Decode循环利用PagedAttention，配合连续批调度并发生成token，直到输出终止符。

二、行业技术知识点

通读这套工程，你会完整掌握以下LLM推理硬核知识点：

1. Safetensors权重文件解析：文件结构、Header解析、张量偏移映射，从二进制文件读取BF16模型权重，区别于传统pickle格式。
2. BF16浮点数据原理：16bit，8bit指数位、7bit尾数，和FP16的差异；大模型选用BF16做推理的原因，解决溢出问题。
3. CUDA内存模型 ：Host主机内存（CPU内存）、Device显存、Shared共享内存；cudaMalloc、cudaMemcpy主机设备数据拷贝，内存复用与生命周期管理。
4. Transformer Decoder架构：Llama3.2结构，RMSNorm、残差连接、RoPE旋转位置编码、GQA分组查询注意力、SiLU激活的MLP。
5. Prefill & Decode两阶段推理：Prefill一次性处理整段prompt，一次性生成KV；Decode每次仅输入单个新token，复用历史KV缓存，大幅降低算力。
6. KV缓存与PagedAttention分页缓存：vLLM最核心创新，把KV缓存拆分为固定大小显存页，按需分配、复用空闲显存页，解决传统连续KV缓存显存碎片问题。
7. 批处理策略：静态批处理 vs 连续批处理
   • 静态批：一批请求全部完成Prefill之后，统一解码；短板：长prompt阻塞整批请求，长尾延迟高。
   • 连续批：维护多个batch slot，某个slot生成结束后立刻填充新请求，新请求单独跑Prefill，其余slot继续解码，提升GPU利用率。
8. CUDA并行规约（Tree Reduction）：在Shared内存上并行求和、求最大值，RMSNorm、Softmax的底层加速手段。
9. cuBLAS矩阵乘法 ：cublasGemmEx、列主序/行主序转换技巧，模型线性层QKV投影、MLP的矩阵运算。
10. Online Softmax：在线Softmax，优化Attention分数计算，减少中间显存占用；带Causal掩码，保证token只能看见上文，看不到未来token。
11. CUDA Kernel编写 ：SIMT线程模型，block/thread/warp，__shared__共享内存，__syncthreads()线程同步。

三、架构设计思路

整体架构分为四大模块：权重加载模块、Token预处理模块、CUDA计算内核模块、批调度推理服务模块。

3.1 权重加载模块

模型权重存储在Safetensors文件，文件分为3段：8字节Header长度、JSON头部、张量二进制数据。

JSON头部记录每一个张量的名称、数据类型、shape、在二进制段内的偏移起止位置。

设计思路：一次性分配一块连续GPU显存，读取所有张量，根据JSON内offset，把每一层权重映射到大显存块的对应位置。

模型权重全部为__nv_bfloat16，包括Embedding、RMSNorm权重、QKV/MLP线性层权重。

3.2 Token预处理模块

CPU侧完成文本分词，将字符串转为token ID数组；token数组拷贝至GPU显存，送入Embedding核函数。分词器使用外部Llama分词实现，C++内核不内置分词逻辑，简化GPU端代码。

3.3 CUDA计算内核模块

所有基础算子全部拆分为独立CUDA Kernel，算子之间通过GPU显存buffer传递数据，配合cuBLAS完成大矩阵乘法：

• Embedding收集Kernel：根据token index读取embedding权重
• RMSNorm Kernel：并行规约计算RMS，归一化hidden state
• RoPE旋转位置编码Kernel：对Q、K向量做旋转变换注入位置信息
• Residual残差相加Kernel
• SiLU激活Kernel
• Softmax、Online Softmax Kernel
• Causal掩码Kernel
• Attention计算Kernel（GQA分组查询注意力）
• MLP前向计算（cuBLAS矩阵乘法 + SiLU）

3.4 批调度推理服务模块

实现静态批 + 连续批两套调度器，基于slot机制管理并发请求。

• Slot：GPU上预留的推理槽位，每个slot维护独立Paged KV缓存；
• 连续批调度逻辑：slot完成生成、输出终止token后，标记为空闲；调度队列取出新请求，分配空闲slot，单独执行Prefill；其余活跃slot继续并行Decode解码。

流程原理示意图

四、核心代码实现原理 & 源码逐段讲解

4.1 GPU状态检测工具函数（C++）

用于启动时读取GPU信息，算力、显存大小，提前判断硬件是否满足运行条件

#define B_TO_MB (1024*1024)
#define B_TO_GB (1024*1024*1024)
int checkGPUStatus()
{
    int device_count = 0;
    cudaGetDeviceCount(&device_count);
    if (device_count == 0)
    {
        std::cerr << "No CUDA devices found\n";
        return 1;
    }
    cudaDeviceProp prop;
    cudaGetDeviceProperties(&prop, 0);
    std::cout << "Device: " << prop.name << "\n";
    std::cout << "Compute capability: " << prop.major << "." << prop.minor << "\n";
    std::cout << "Global memory: " << prop.totalGlobalMem / B_TO_MB << " MB\n";
    std::cout << "SM count: " << prop.multiProcessorCount << "\n";
    std::cout << "Max threads per block: " << prop.maxThreadsPerBlock << std::endl;
    size_t free_mem;
    size_t total_mem;
    cudaMemGetInfo(&free_mem, &total_mem);
    std::cout << "Free memory: " << free_mem / B_TO_GB << "GB, total memory: " << total_mem / B_TO_GB << "GB\n";
    return 0;
}

解读：cudaGetDeviceProperties读取GPU硬件属性；cudaMemGetInfo查询当前空闲/总显存。工程启动第一时间调用，防止显存不足直接跑推理导致CUDA报错。

4.2 Embedding收集CUDA Kernel

输入token id，并行取出对应embedding向量。受限于GPU单Block最大1024线程限制，每个线程处理2个embedding元素。

__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];
}

void embeddingGather(int *gpu_input_tokens, __nv_bfloat16 *gpu_input_embeds, __nv_bfloat16 *embed_tokens, int num_input_tokens)
{
    embeddingGatherKernel<<<num_input_tokens, 1024>>>(gpu_input_tokens, gpu_input_embeds, embed_tokens);
#ifdef DEBUG
    cudaError error = cudaGetLastError();
    if (error != cudaError::cudaSuccess)
    {
        std::cout << "CUDA last error: " << cudaGetLastError() << std::endl;
    }
#endif
}

解读：

• blockIdx.x 代表当前token编号，一个block处理1个token；block内1024个thread，每个线程读取2个元素，覆盖2048维embedding。
• gpu_input_tokens[blockIdx.x]拿到token id，乘以2048得到该token embedding在权重数组上的起始位置。

4.3 RMSNorm CUDA Kernel（Tree Reduction并行规约）

RMSNorm公式：

RMS(a)=1n∑i=0n−1ai2+ϵ\text{RMS}(a)=\sqrt{\frac{1}{n}\sum_{i=0}^{n-1}a_i^2+\epsilon}RMS(a)=n1i=0∑n−1ai2+ϵ

normalizedi=aiRMS(a)⋅wi\text{normalized}_i=\frac{a_i}{\text{RMS}(a)} \cdot w_inormalizedi=RMS(a)ai⋅wi

__global__ void rmsNormKernel(__nv_bfloat16 *input, __nv_bfloat16 *output, __nv_bfloat16 *norm_weights)
{
    __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();
    // tree reduction 树规约求和
    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]);
}

解读：

1. 每个线程读取2个元素，计算平方存入shared内存数组；
2. 循环树规约，不断对半合并求和，__syncthreads()保证所有线程同步；
3. block第0号线程计算RMS，增加1e-5防止除0；
4. 全部线程读取RMS值，完成归一化，乘上RMSNorm可学习权重，写回输出。

4.4 RoPE旋转位置编码Kernel

RoPE通过旋转Q、K向量，把token相对位置信息注入向量，注意力分数只和token相对距离有关。

#define HEAD_DIM 64
__global__ void ropeKernel(__nv_bfloat16 *input, int num_tokens, int proj_dim)
{
    if (2 * threadIdx.x + 1 + blockIdx.x * proj_dim < num_tokens * proj_dim)
    {
        int double_i = 2 * (threadIdx.x % 32);
        float theta = 1.0 / (pow(500000.0, ((float)double_i / HEAD_DIM)));
        float angle = blockIdx.x * theta;
        __nv_bfloat16 prev_2i = input[2 * threadIdx.x + blockIdx.x * proj_dim];
        __nv_bfloat16 prev_2i_1 = input[2 * threadIdx.x + 1 + blockIdx.x * proj_dim];
        input[2 * threadIdx.x + blockIdx.x * proj_dim] = (__nv_bfloat16)((float)prev_2i * cos(angle) - (float)prev_2i_1 * sin(angle));
        input[2 * threadIdx.x + 1 + blockIdx.x * proj_dim] = (__nv_bfloat16)((float)prev_2i * sin(angle) + (float)prev_2i_1 * cos(angle));
    }
}
void rope(__nv_bfloat16 *input, int num_tokens, int proj_dim)
{
    int num_threads = proj_dim / 2;
    if (num_threads > 1024)
    {
        std::cout << "Can't launch more than 1024 threads, RoPE kernel not launched";
        return;
    }
    ropeKernel<<<num_tokens, num_threads>>>(input, num_tokens, proj_dim);
#ifdef DEBUG
    cudaError error = cudaGetLastError();
    if (error != cudaError::cudaSuccess)
    {
        std::cout << "CUDA last error: " << cudaGetLastError() << std::endl;
    }
#endif
}

4.5 SiLU激活Kernel（MLP）

SiLU公式：SiLU(x)=x⋅σ(x)\text{SiLU}(x)=x\cdot\sigma(x)SiLU(x)=x⋅σ(x)，σ\sigmaσ是sigmoid。MLP的gate支路经过SiLU，再与up支路逐元素相乘。

__global__ void siluKernel(__nv_bfloat16 *a, __nv_bfloat16 *b)
{
    int workIndex = threadIdx.x + blockIdx.x * 8192;
    for (int i = 0; i < 8192; i += 1024)
    {
        a[workIndex + i] = (__nv_bfloat16)((float)a[workIndex + i] * (1 / (1 + expf(-(float)a[workIndex + i]))) * (float)b[workIndex + i]);
    }
}
void silu(__nv_bfloat16 *a, __nv_bfloat16 *b, int num_tokens)
{
    siluKernel<<<num_tokens, 1024>>>(a, b);
#ifdef DEBUG
    cudaError error = cudaGetLastError();
    if (error != cudaError::cudaSuccess)
    {
        std::cout << "CUDA last error: " << cudaGetLastError() << std::endl;
    }
#endif
}

4.6 Causal Mask Kernel

因果掩码：保证token只能看到上文，未来位置的Attention分数设置为负无穷，softmax后概率趋近0。

__global__ void causalMaskKernel(__nv_bfloat16 *input, int num_tokens)
{
    if (threadIdx.x + blockIdx.x * blockDim.x >= num_tokens * num_tokens * 32)
    {
        return;
    }
    int column = threadIdx.x;
    int row = blockIdx.x % num_tokens;
    if (column > row)
    {
        input[blockIdx.x * num_tokens + threadIdx.x] = -HUGE_VALF;
    }
}
void causalMask(__nv_bfloat16 *input, int num_tokens)
{
    if (num_tokens > 1024)
    {
        std::cout << "Can't launch more than 1024 threads, Causal mask kernel not launched";
        return;
    }
    causalMaskKernel<<<num_tokens * 32, num_tokens>>>(input, num_tokens);
#ifdef DEBUG
    cudaError error = cudaGetLastError();
    if (error != cudaError::cudaSuccess)
    {
        std::cout << "CUDA last error: " << cudaGetLastError() << std::endl;
    }
#endif
}

4.7 cuBLAS矩阵乘法调用示例（K投影线性层，行主序/列主序转换技巧）

cuBLAS原生使用列主序存储，而模型权重是行主序，通过调换转置标记，不需要原地转置数据，节省显存开销。

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

五、环境配置与完整运行测试教程

硬件需求

• NVIDIA显卡，支持CUDA，支持BF16；推荐RTX 4090、RTX5090；
• 也支持AMD GPU，使用ROCm/HIP编译，替换nvcc为hipcc，hipBLAS替代cuBLAS。

软件依赖

• Linux x86_64
• C++17，GCC编译器
• CUDA Toolkit（推荐13.1）
• CMake + Ninja
• 仅外部依赖：单头文件JSON库 nlohmann/json 3.12.0

5.1 编译NVIDIA CUDA版本

# 1. 准备工程目录，创建build编译文件夹
rm -fr build
cmake -B build -G Ninja
cmake --build build

5.2 编译AMD ROCm HIP版本

cmake -B build -DUSE_HIP=ON -DCMAKE_HIP_ARCHITECTURES=gfx1100 -DCMAKE_PREFIX_PATH=/opt/rocm -G Ninja
cmake --build build

5.3 准备模型权重

1. 获取Llama3.2 1B Instruct Safetensors权重文件，只需要model.safetensors；
2. 将权重文件放到工程weights目录。

5.4 启动推理服务

# 一键启动脚本，自动加载权重并运行推理
./test.sh

程序启动后，首先执行GPU状态检测，打印显卡名称、算力、显存；加载Safetensors权重，解析Header，把BF16权重映射到GPU显存；进入交互推理模式，输入prompt即可生成回答。

5.5 批处理模式测试

修改调度配置，可以切换静态批 / 连续批模式：

• 静态批：设置固定batch size，一次性提交多条prompt，等待全部请求完成解码；
• 连续批：开启slot调度，多请求并发，请求完成后slot复用，持续接收新prompt。

5.6 调试手段

开启DEBUG宏，会打印每一个CUDA Kernel执行错误，定位核函数报错。可以打印中间张量，查看Embedding、RMSNorm输出、Attention分数用于调试。

六、落地用途

1. vLLM原理学习首选工程
   工业vLLM代码几十万行，阅读门槛极高。这套轻量化实现保留PagedAttention、连续批这些最核心创新，代码量精简，可逐行断点调试，亲手验证KV分页缓存、批调度的工作流程，理解现代LLM高吞吐推理底层原理。
2. CUDA大模型算子学习实验台
   可以修改RMSNorm、RoPE、Softmax、Attention Kernel，尝试不同并行规约策略、Online Softmax变体，快速验证算子性能，不需要维护重型推理框架。
3. 课程教学实验
   高校AI/大模型课程实验，让学生手写CUDA核函数，实现Prefill/Decode、PagedAttention，直观理解LLM推理硬件调度。
4. 端侧GPU推理原型验证
   基于这套C++/CUDA基础代码，可以裁剪修改，移植到边缘NVIDIA设备，快速验证模型推理吞吐、显存占用。
5. 自定义推理引擎二次开发底座
   在此基础上扩展：增加量化（GPTQ/AWQ）、更多模型架构、流式输出、服务接口，构建自研推理服务。

If you need the complete source code, please add the WeChat number (c17865354792)

七、总结

这套C++/CUDA推理工程，完整复刻vLLM核心能力，用精简代码实现PagedAttention分页KV缓存、连续批调度、GQA、RoPE、RMSNorm、Online Softmax。

全部基础算子手写CUDA Kernel，结合cuBLAS加速矩阵运算，同时兼顾代码可读性。工程最核心价值不是拿来直接上线部署，而是让开发者跳出PyTorch高层API，深入GPU硬件视角，看懂大模型从权重加载、两阶段推理、显存管理到请求调度的全链路。

Welcome to follow WeChat official account【程序猿编码】