Mini vLLM: 从零实现 LLM 推理引擎

Published: January 01, 2024

项目概述

Mini vLLM 是一个从零实现的简化版 vLLM 系统，目标是完整复现 LLM 推理引擎的核心模块，展示对大模型推理系统底层原理的深入理解和工程实现能力。

项目的定位不是”比 vLLM 更快”，而是“理解 vLLM 为什么快，并能从零实现”。

项目实现了七大核心模块：

模块	功能
KV Cache	推理过程中的 Key-Value 缓存管理，避免自回归生成中的重复计算
KV Cache Manager	多个 KVCache 实例的分配与释放管理
Continuous Batching	迭代级动态批处理，请求完成立即退出、新请求立即加入
PagedAttention	分页显存管理，类比操作系统的虚拟内存分页机制，减少内存碎片
Request Scheduler	基于堆的优先级队列调度器，支持 FCFS、优先级调度和抢占机制
Inference Engine	推理引擎，组装所有模块提供统一推理接口
Model Adapter	集成 HuggingFace 的 Qwen2.5 模型，管理 KV Cache 格式转换

项目实现了 PagedAttention 的四个版本，体现从高层到底层的优化路径：

PyTorch 版本 — 用于验证正确性
Triton 版本 — 使用 Triton JIT 编写 GPU kernel，实现 Online Softmax，比 PyTorch 快 2-10 倍
CUDA 基础版本 — 原始 CUDA kernel，每线程处理一个 batch/head
CUDA 优化版本 — 使用 Warp Shuffle 做并行归约，每个线程处理一个 head_dim 维度，比 PyTorch 快 6-9 倍

推理性能对比（Qwen2.5-0.5B, 50 tokens 输入/输出）：

实现	延迟(s)	吞吐量(tokens/s)	显存(MB)
HuggingFace	0.073	683.4	953.2
Mini vLLM	0.062	812.0	1017.8
vLLM	0.046	1086.8	4501.0

Mini vLLM 吞吐量比 HuggingFace 高 19%。

PagedAttention 加速比（相对于 PyTorch）：

系统性从零复现工业级推理系统核心模块。 逐模块实现 KV Cache、Continuous Batching、PagedAttention、Request Scheduler，每一步都有详细的代码注释解释”为什么这样做”。

PagedAttention 四版本渐进式实现。 完整展示 GPU 编程从高层抽象到底层优化的路径，包括 Online Softmax 数值稳定性和 Warp Shuffle 并行归约。

完整的工程闭环。 包含 1772 行测试代码、三篇深入的技术博客、以及系统的性能 benchmark 对比。

真实模型集成。 通过 QwenModelAdapter 集成真实的 Qwen2.5 模型，处理 HuggingFace DynamicCache 与自研 KVCache 之间的格式转换。

项目包含三篇深入的技术博客：