mini-vllm：基于 NanoGPT 的 PagedAttention 风格 KV 缓存管理实现

摘要：一个极简 LLM 推理引擎，在 NanoGPT 上实现了 PagedAttention 风格的 KV 缓存管理机制，显著提升了内存利用效率和推理速度。 关键词：LLM, PagedAttention, KV Cache, 推理优化, NanoGPT, 内存管理, vLLM

本帖将详细介绍 mini-vllm 项目的背景、核心技术、架构设计、性能表现及未来规划，帮助大家理解 PagedAttention 风格 KV 缓存优化的实现与价值。

在大语言模型（LLM）推理中，传统 Transformer 实现每次生成新 token 需重新计算所有历史 token 的 KV 张量，复杂度随序列长度二次增长。KV 缓存技术虽解决此问题，但传统实现用连续内存空间并预预留最大长度内存，导致严重的内存碎片问题，成为推理任务的瓶颈。

PagedAttention 借鉴操作系统虚拟内存分页机制，将 KV 张量存储在动态分配的非连续内存块中，可在不同请求间共享，大幅减少内存浪费。原始 vLLM 论文显示，相比连续内存分配，PagedAttention 实现 2-4 倍内存效率提升，对资源受限环境部署大模型意义重大。

mini-vllm 采用模块化设计，核心组件包括：

1. BlockAllocator：维护内存块状态，为预填充/解码阶段分配释放内存，类似操作系统内存管理策略。
2. KVCache：使用 BlockAllocator 的块表确定 GPU 张量中 KV 张量的读写位置，实现非连续存储与逻辑连续。
3. InferenceEngine：协调器，接收提示词，运行预填充、解码循环，管理内存块分配释放，返回结果。

流程：prompt → InferenceEngine → BlockAllocator（管理物理块）+ KVCache（按块存储/检索 KV 张量）。

项目提供 benchmark.py 评估性能，对比两种生成方法：

• KV 缓存优化版：用 InferenceEngine，自定义 KVCache/BlockAllocator，解码阶段仅计算新 token。
• 无 KV 缓存基线版：标准 NanoGPT，每次生成新 token 执行完整前向传播，成本随序列增长急剧上升。

测试指标：平均生成时间、TPS、加速比；支持配置提示长度、token 数、内存参数；结果存 CSV（benchmark_results.csv），并提供 benchmark_results_viz.ipynb 可视化（运行时间-长度关系、加速比对比、吞吐量对比）。

CPU 环境测试显示，KV 缓存实现对中等生成长度（200-500 token）提供 2-10 倍加速，序列增长时延迟显著降低。即使在资源受限环境，PagedAttention 风格内存管理也带来实质性性能提升。

未来规划：

1. BlockAllocator 的 C++ 重写，提升内存分配性能；
2. 前缀共享（prefix-sharing），支持块在不同请求间共享，进一步提升内存效率。

相关资源：

• PagedAttention 论文：https://arxiv.org/abs/2309.06180
• vLLM GitHub：https://github.com/vllm-project/vllm
• NanoGPT 项目：https://github.com/karpathy/nanoGPT

mini-vllm 为理解和实现 PagedAttention 技术提供简洁完整的参考实现。通过模块化架构和详尽基准测试，开发者可深入理解 KV 缓存优化原理，并应用到自己的 LLM 推理系统。对资源受限环境高效部署大模型的开发者，是极具价值的学习资源。