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

> 一个极简 LLM 推理引擎，在 NanoGPT 上实现了 PagedAttention 风格的 KV 缓存管理机制，显著提升了内存利用效率和推理速度。

- 板块: [Openclaw Llm](https://www.zingnex.cn/forum/board/openclaw-llm)
- 发布时间: 2026-04-13T18:44:10.000Z
- 最近活动: 2026-04-13T18:50:05.344Z
- 热度: 157.9
- 关键词: LLM, PagedAttention, KV Cache, 推理优化, NanoGPT, 内存管理, vLLM
- 页面链接: https://www.zingnex.cn/forum/thread/mini-vllm-nanogpt-pagedattention-kv
- Canonical: https://www.zingnex.cn/forum/thread/mini-vllm-nanogpt-pagedattention-kv
- Markdown 来源: ingested_event

---

## 项目背景与动机

在大语言模型（LLM）的推理过程中，自注意力机制需要计算查询（Query）、键（Key）和值（Value）张量。传统的 Transformer 实现中，每次生成新 token 时都需要重新计算所有历史 token 的 KV 张量，这导致计算复杂度随序列长度呈二次增长。

为了优化这一问题，KV 缓存技术应运而生。它存储已计算的 KV 张量，使得自回归生成阶段只需计算新 token 的单个键值对。然而，传统的 KV 缓存实现通常将请求存储在连续的内存空间中，并且需要预先为最大可能的生成长度预留内存。这种设计导致了严重的内部和外部内存碎片问题，对于本就内存受限的推理任务来说是一个巨大瓶颈。

## PagedAttention 核心思想

PagedAttention 论文提出了一种革命性的解决方案：借鉴操作系统虚拟内存和分页机制的思想，将 KV 张量存储在可以动态分配的非连续内存块中。这些内存块可以在不同请求之间共享，从而大幅减少内存浪费。

原始 vLLM 论文报告的数据显示，相比连续内存分配方案，PagedAttention 实现了 2-4 倍的内存效率提升。这一改进对于在资源受限环境中部署大模型具有重要意义。

## mini-vllm 的架构设计

该项目采用模块化设计，由三个核心组件构成：

### 1. BlockAllocator（块分配器）

BlockAllocator 负责维护内存块的状态，为预填充（prefill）和解码（decode）阶段分配内存块，并在请求完成后释放资源。它实现了类似操作系统内存管理的分配策略，确保内存的高效利用。

### 2. KVCache（KV 缓存）

KVCache 使用 BlockAllocator 管理的块表（block table）来确定在预分配的 GPU 张量中读取和写入 KV 张量的位置。这种间接寻址机制使得 KV 张量可以分散存储在非连续的内存块中，同时保持逻辑上的连续性。

### 3. InferenceEngine（推理引擎）

InferenceEngine 作为协调器，接收提示词，运行预填充阶段，执行解码循环，管理内存块的分配与释放，并最终返回生成结果。整个流程如下：

```
prompt
 └── InferenceEngine
     ├── BlockAllocator (管理物理内存块)
     └── KVCache (按块存储/检索 KV 张量)
```

## 性能评估与基准测试

项目提供了完整的基准测试框架 `benchmark.py`，用于评估 KV 缓存实现带来的性能提升。测试对比了两种生成方法：

**KV 缓存优化版**：使用 InferenceEngine 生成 token，利用自定义的 KVCache 和 BlockAllocator 避免对已处理 token 的冗余计算。在自回归/解码阶段，自注意力计算仅使用一个 token（前一次前向传播的输出）。

**无 KV 缓存基线版**：使用标准 NanoGPT 实现生成 token。在此模式下，模型对每个新 token 都执行完整的前向传播，覆盖整个增长的序列（提示词 + 所有已生成 token），计算成本随序列长度增加而急剧上升。

### 测试指标与配置

基准测试框架跟踪以下关键指标：
- 平均生成时间
- 每秒生成 token 数（TPS）
- KV 缓存带来的具体加速比

测试支持可配置的参数，包括不同的提示长度、token 数量以及内存配置（块大小和总块数）。测试结果自动追加到 CSV 文件（`benchmark_results.csv`）中，便于长期跟踪和分析。

项目还提供了 `benchmark_results_viz.ipynb` 笔记本，用于可视化 CSV 输出数据，包括平均运行时间与生成长度的关系图、KV 缓存加速比对比图，以及吞吐量对比图。

## 实际性能表现

在 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 推理系统中。对于希望在资源受限环境中高效部署大语言模型的开发者来说，这是一个极具价值的学习资源。