⚡ nano-vLLM：轻量级、低延迟 LLM 推理从零开始
引言：LLM 中的推理是什么？
当你听到“ChatGPT 响应”或“LLM 生成文本”时，你正在目睹 推理。
推理是使用训练好的模型进行预测或生成输出的过程。
在 LLM 中，推理意味着
- 接收你的提示词
- 通过数十亿权重进行运行
- 获得智能且相关的输出
但问题在于：推理缓慢、资源消耗大，并且通常未针对边缘或个人设备进行优化。
这就是 vLLM —— 以及现在的 nano-vLLM —— 等优化工具发挥作用的地方。
为什么推理优化很重要
大型模型（即使是 1B+ 参数）往往会
- 消耗大量 VRAM
- 引入延迟，尤其是在长时间生成时
- 需要大规模基础设施才能投入生产
我们想要实现
- 快速令牌生成
- 低内存占用
- 并行请求处理
- 在笔记本电脑、Colab 和边缘设备上可行
vLLM：推理巨头
vLLM 是由加州大学伯克利分校和 Meta 的研究人员构建的生产级推理引擎。
主要优势：
- PagedAttention 用于虚拟内存高效的 KV 缓存
- 连续批处理 用于并行处理提示词
- 预填充 + 解码并行
- 张量并行 用于多 GPU 推理
- 在 HuggingFace 推理 API 中使用
挑战：
- 复杂、庞大的代码库（~10K+ 行代码）
- 使用 C++、CUDA 扩展
- 修改和学习难度较高
- 不适合初学者或 Colab 使用
介绍 nano-vLLM：轻量级重构
nano-vLLM 是 vLLM 的最小重新实现——仅有 ~1200 行简洁的 Python 代码，专为以下目的构建：
- 理解
- 修改
- 在有限硬件上运行
把它看作是 vLLM 的小巧、可读的兄弟——却异常快速和实用。
nano-vLLM 的亮点
| 特性 | nano-vLLM |
|---|---|
| 小巧的代码库 | 约 1.2k 行 |
| 纯 Python 和 Triton | 易于修改 |
| CUDA Graph + torch.compile | 用于更快解码 |
| 支持 Flash Attention | 可选 |
| 可在笔记本电脑/Colab 上运行 | |
| 支持张量并行 | 基础 |
| 无 C++/CUDA 扩展 | 更简单的安装 |
| 可用于研究的修改 | 非常适合学习 |
nano-vLLM 的工作原理（内部）
让我们分解一下 nano-vLLM 引擎的核心部分
1. 提示词分词
- 使用 HuggingFace 分词器
- 支持多提示词（批处理）
- 分为
prefill和decode阶段
2. KV 缓存管理
- 实现 Triton 内核：
store_kvcache_kernel - 高效的键/值注意力内存存储
- 支持前缀缓存
3. Flash Attention
- 使用
flash-attn以提高速度（如果已安装） - 也可回退到默认注意力机制
4. 解码引擎
- 重用缓存值以实现快速的步进式生成
- 如果可能，将解码循环封装在
torch.cuda.graph()中
5. 采样参数
- 实现
temperature、top_k、max_tokens - 没有不必要的抽象——开箱即用
6. 张量并行
- 轻量级
torch.distributed封装器 - 跨多个 GPU 拆分模型（可选）
nano-vLLM 的底层工作原理（面向 ML 开发者的深入探究）
nano-vLLM 简化了 vLLM 的许多高级概念，同时保留了性能关键组件。以下是其内部结构的分解：
1. 提示词分词和输入格式
nano-vLLM 使用 Hugging Face 分词器预处理输入文本。在分词过程中：
- 输入通过
cu_seqlens进行批处理和填充，以支持变长序列。 - 引擎区分：
- 预填充阶段：KV 缓存正在初始化时
- 解码阶段：逐令牌生成时
这种分离使得多轮或流式生成能够更高效地处理。
2. KV 缓存：自定义内存管理
KV（键-值）缓存存储注意力层的隐藏状态，使模型能够
- 在解码阶段重用之前的上下文
- 避免冗余计算
在 nano-vLLM 中
- 使用 Triton 内核 (
store_kvcache_kernel) 高效地将键和值写入预分配的缓存中。 - 缓存槽使用
slot_mapping张量进行映射，避免 Python 级别的索引。 - 缓存布局：
[batch_size, num_heads, head_dim] → [total_slots, head_dim]（为性能而展平）
此设计模仿了 vLLM 的 PagedAttention，但保持了可读性和可修改性。
3. ⚡ Flash Attention（v2 兼容）
如果安装了 flash-attn，nano-vLLM 会
- 在预填充阶段调用
flash_attn_varlen_func - 在解码阶段调用
flash_attn_with_kvcache
FlashAttention v2
- 通过避免注意力矩阵的实例化来减少内存使用
- 在融合的 Triton 内核中计算 softmax 注意力
- 使用块稀疏布局以更好地利用 GPU
对于未安装 flash-attn 的环境，也支持回退到标准注意力机制。
4. torch.compile + CUDA Graphs
对于解码阶段（逐个令牌生成）
- nano-vLLM 将生成循环封装在 CUDA Graph 中（如果支持）
- 还使用
torch.compile()来融合操作并减少 Python 开销
这带来了
- 稳定的内存分配
- 更高的内核启动效率
- 更低的单令牌解码延迟
这在 T4 或 RTX 30/40 系列等消费级 GPU 上尤其有用。
5. SamplingParams：简洁、极简的采样 API
SamplingParams 类支持
temperature：控制随机性top_k：Top-k 过滤max_tokens：每个请求的令牌预算stop_tokens：可选的停止序列强制执行
采样逻辑使用 PyTorch 张量操作高效实现
- 使用 top-k 过滤 logits
- 缩放后应用 Softmax
- 使用
torch.multinomial采样下一个令牌
6. 轻量级张量并行
nano-vLLM 支持使用 torch.distributed 进行基本的张量并行
- 将模型权重分布到多个 GPU 上
- 每个 GPU 包含注意力/MLP 块中线性投影的一部分
- 最终输出在 GPU 之间收集
虽然不如 DeepSpeed 或 vLLM 的 NCCL 分片功能丰富，但它在研究或 Colab 环境中对中小型模型效果良好。
7. 模块化设计和简洁性
该架构被分解为简洁的模块
llm_engine.py：主推理协调器layers/：自定义注意力、MLP、旋转嵌入utils/context.py：全局推理状态管理（例如，块表、缓存长度）
不涉及 C++/CUDA 自定义扩展——一切都是纯 Python + Triton，使其高度可读和可定制。
🛠 总结
| 组件 | 使用的技术/模块 |
|---|---|
| 分词 | HuggingFace 分词器 + cu_seqlens 批处理 |
| KV 缓存 | 带有槽映射的 Triton 内核 |
| 注意力机制 | Flash Attention v2 或自定义回退 |
| 采样 | 基于 PyTorch 的多项式采样 |
| 速度优化 | torch.compile + CUDA Graph（解码阶段） |
| 并行化 | 使用 torch.distributed 的基本张量并行 |
| 代码简洁性 | 约 1.2k 行代码，无 C++ 或不透明抽象 |
这种模块化且具有深度教育意义的结构使 nano-vLLM 成为以下人员的绝佳选择：
- 探索推理的 LLM 工程师
- 尝试新解码算法的研究人员
- 学习系统级机器学习的学生
📊 基准测试（RTX 4070 笔记本电脑）
| 引擎 | 生成的令牌数 | 时间（秒） | 吞吐量（令牌/秒） |
|---|---|---|---|
| vLLM | 133,966 | 98.37 | 1,361.84 |
| nano-vLLM | 133,966 | 93.41 | 1,434.13 |
🚀 在相同设置下比 vLLM 更快。
并非魔法——仅仅是智能缓存、Triton 和精简架构。
🔄 nano-vLLM 与 vLLM：主要区别
| 特性 | vLLM | nano-vLLM |
|---|---|---|
| 代码库大小 | ~10k+ 行代码 | ~1.2k 行代码 |
| 语言 | C++ + Python + CUDA | Python + Triton |
| Flash Attention | 内置 | 可选 |
| 张量并行 | 高级 | 基础 |
| KV 缓存 | PagedAttention | 手动 Triton 内核 |
| 兼容性 | 完全 HF | HF（部分，已测试） |
| 量化 | 仅外部 | 即将通过社区实现 |
| 理想用例 | 生产服务器 | 研究、Colab、设备端 |
实际用例：在 Google Colab（T4 GPU）上运行
pip install git+https://github.com/GeeeekExplorer/nano-vllm.git
huggingface-cli download Qwen/Qwen3-0.6B --local-dir ./Qwen3-0.6B
from nanovllm import LLM, SamplingParams
llm = LLM("./Qwen3-0.6B", enforce_eager=True, tensor_parallel_size=1)
sampling_params = SamplingParams(temperature=0.7, max_tokens=128)
output = llm.generate(["Hello nano-vLLM!"], sampling_params)
print(output[0]['text'])
最终思考
nano-vLLM 不仅仅是一个迷你推理引擎——它还是一个强大的工具，可以完全按照你的意愿学习、适应和部署大型语言模型。
无论你是
- 深入 LLM 内部的研究人员
- 为边缘设备打造工具的黑客
- 在小型 GPU 上追求性能的工程师
nano-vLLM 都是你开源的盔甲。它灵活、敏捷，旨在赋能你的想法。
🌐 有用链接
- 🔗 GitHub: nano-vLLM 仓库
