Tiny-vLLM 实战：用 C++/CUDA 手写高性能 LLM 推理引擎

大模型推理的性能瓶颈往往不在计算本身，而在内存管理与算子调度。当 PyTorch 的抽象层成为性能桎梏时，直接用 C++/CUDA 重写推理引擎成为极致优化的必由之路。tiny-vLLM 项目提供了一个从底层构建的完整实现，本文聚焦其三大核心技术：内存池管理、算子融合与 Kernel 级优化。

内存池管理：从操作系统借来的智慧

KV Cache 是 LLM 推理中最关键的内存消耗源。传统实现为每个请求预分配连续内存，导致严重的内存碎片和浪费。tiny-vLLM 借鉴操作系统虚拟内存管理思想，实现了 PagedAttention 机制。

核心思想是将 KV Cache 分割为固定大小的页（page），每个页存储固定数量的 token 的 Key 和 Value 向量。当请求生成新 token 时，动态分配新的页；请求结束时，回收页到空闲池。这种设计带来三个直接收益：

• 消除内存碎片：页大小固定，空闲页可立即被其他请求复用
• 支持动态扩展：长序列请求按需分配页，无需预先分配最大长度
• 实现内存共享：不同请求可共享相同的页（适用于前缀缓存场景）

在 tiny-vLLM 的实现中，内存池的核心数据结构是一个空闲页链表。每个页存储 page_size × head_dim × num_kv_heads 个 BF16 数值。对于 Llama 3.2 1B 模型，KV head 维度为 512，若设置页大小为 16 tokens，则每页占用 16 × 512 × 2 = 16KB 显存。

可落地参数：

• 页大小建议设置为 16-32 tokens，平衡内存碎片与分配开销
• 内存池初始容量按最大并发请求数 × 平均序列长度 × 1.5 倍系数预分配
• 启用内存预分配池，避免推理过程中的 cudaMalloc 延迟

算子融合：减少 Kernel 启动开销

LLM 推理涉及大量小算子，每个算子启动都伴随 CUDA Kernel 调用开销。tiny-vLLM 通过算子融合显著降低这一开销。

以 Softmax 为例，标准实现需要三次遍历：求最大值、求指数和、归一化。tiny-vLLM 实现了 Online Softmax，将三次遍历压缩为单次遍历，通过维护运行中的最大值和指数和累加器，在读取数据的同时完成计算。

另一个关键融合点是 FlashAttention-like 实现。传统 Attention 需要显式构造完整的注意力矩阵（seq_len × seq_len），内存复杂度为 O (n²)。FlashAttention 通过分块计算和在线 Softmax，将内存复杂度降至 O (n)，同时保持计算强度，充分利用 GPU 算力。

tiny-vLLM 中的 RMSNorm 实现展示了典型的融合模式：

// 并行归约计算 RMS
__shared__ float rms_vector[1024];
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();
}

工程 Checklist：

• 识别高频小算子（RMSNorm、SiLU、残差连接）进行纵向融合
• Attention 计算采用分块策略，块大小匹配 SM 共享内存容量
• 使用 CUDA Graph 捕获静态计算图，消除动态调度开销

CUDA Kernel 优化：榨干硬件算力

tiny-vLLM 的 Kernel 优化围绕三个核心原则展开：

1. 共享内存最大化利用

共享内存（Shared Memory）比全局内存快两个数量级。在 RMSNorm 和 Softmax 中，使用 __shared__ 声明的数组存储中间结果，避免线程间的全局内存竞争。每个 Block 的共享内存使用应控制在 48KB 以内（RTX 4090/5090 的 SM 限制），以允许单个 SM 同时驻留多个 Block。

2. 并行归约模式

树形归约（Tree Reduction）是 CUDA 中累加操作的标准模式。每次迭代将活跃线程数减半，直到 thread 0 持有最终结果。注意每次归约步骤后必须调用 __syncthreads() 确保数据可见性。

3. 数据类型精度权衡

tiny-vLLM 采用 BF16 存储权重和激活值，但在 Kernel 内部使用 FP32 进行累加计算。BF16 的 8 位指数与 FP32 相同，避免了 FP16 的数值溢出问题；而 7 位尾数在累加计算中通过 FP32 中间值补偿精度损失。

关键参数配置：

• 线程块大小：最大 1024 线程 / 块（受限于硬件），实际配置通常为 256-512 以优化占用率
• Grid 维度：通常与 batch size 或序列长度对齐，确保每个元素有对应线程
• 寄存器使用：通过 __launch_bounds__ 限制每个线程的寄存器使用量，提高 SM 并行度

连续批处理：吞吐与延迟的平衡

tiny-vLLM 实现了连续批处理（Continuous Batching）以提升 GPU 利用率。与静态批处理等待所有请求完成不同，连续批处理允许在批次中动态替换已完成的请求。

实现要点：

• 将批次划分为固定数量的槽位（slots），每个槽位对应一个请求
• 槽位状态机管理：等待 → 预填充 → 解码 → 完成
• 当槽位进入 "完成" 状态时，从队列拉取新请求进行预填充
• 预填充阶段独占执行，解码阶段所有活跃槽位并行

配置建议：

• 槽位数量设置为 GPU 显存可容纳的最大并发请求数
• 预填充与解码阶段分离，避免长预填充阻塞短请求的解码
• 实现请求优先级队列，支持抢占和重排序

总结

tiny-vLLM 展示了从头构建高性能推理引擎的核心技术路径：通过 PagedAttention 解决内存管理难题，通过算子融合减少 Kernel 开销，通过精细的 CUDA Kernel 优化榨取硬件性能。这些技术不仅适用于教育场景，更是生产级推理引擎（如 vLLM、TensorRT-LLM）的核心构建块。

对于希望深入推理优化的工程师，建议从 tiny-vLLM 的 RMSNorm 和 Softmax Kernel 入手，逐步理解共享内存、线程同步和并行归约模式，再扩展到更复杂的 Attention 和 FeedForward 优化。

参考资源：

• tiny-vLLM GitHub 仓库 - 完整源码与教程
• vLLM: Efficient Memory Management for LLM Serving - PagedAttention 论文
• FlashAttention: Fast and Memory-Efficient Exact Attention - 在线 Softmax 与分块 Attention

ai-systems