nano-vllm源码分析
在大语言模型(LLM)的时代,推理效率成为了制约应用落地的关键瓶颈。想象一下,当你向 ChatGPT 提出问题时,它需要在几秒内处理数千个token,同时服务数万名用户。这背后的技术挑战是什么?答案就是高效的推理引擎。
vLLM 作为当前最高效的 LLM 推理引擎之一,通过创新的 PagedAttention、KV Cache 管理、连续批处理等技术,将推理吞吐量提升了 24倍。然而,vLLM 的生产代码复杂庞大,包含大量工程优化和边界情况处理,对于初学者来说犹如迷宫。
cleanvllm 项目应运而生——它承继了 CleanRL 的教育理念,用单文件(1500行Python代码)实现了 vLLM 的核心概念,去除了工程复杂性,让我们能够专注于理解架构本质。正如 Linus Torvalds 所说:”Talk is cheap. Show me the code.”
本文将带你从零开始,逐步构建一个完整的 vLLM 推理引擎,揭示其背后的设计哲学和技术细节。读完本文,你将能够:
- 理解 vLLM 的核心架构和关键创新
- 掌握 PagedAttention 的工作原理
- 学会如何实现高效的 KV Cache 管理
- 了解连续批处理的调度策略
- 运行自己的 vLLM 实现
vLLM 的核心挑战与设计思路
传统 LLM 推理的痛点
让我们以一个具体的例子来说明:假设你有一台 A100 GPU(80GB显存),要同时处理 100 个用户请求,每个请求平均生成 100 个token。
传统的推理方式会遇到以下问题: 1. 内存碎片化:每个序列的 KV Cache 需要连续内存,但序列长度不同导致严重的内存浪费。实际可用内存只有 60%! 2. 批处理困难:序列A需要10个token,序列B需要100个token,如何高效批处理?传统方法只能等最短序列完成。 3. 动态调度复杂:你无法预知用户会问多长的问题,资源调度变得极其困难。 这些问题导致 GPU 利用率通常只有 30-40%,大量算力被浪费。
vLLM 的解决方案
vLLM 通过三个核心创新解决了这些问题:
- PagedAttention:将 KV Cache 分页管理,类似操作系统的虚拟内存
- 连续批处理:动态添加/移除序列,最大化 GPU 利用率
- 前缀缓存:复用相同前缀的 KV Cache,减少重复计算 让我们在 cleanvllm 中看看这些概念是如何实现的。
数据结构设计 - vLLM 的基石
序列管理:Sequence 类
在 cleanvllm 中,每个用户请求被抽象为一个 Sequence 对象:
class Sequence:
"""
序列类
表示一个文本生成请求的序列,包含prompt tokens和生成的completion tokens。
管理序列的状态、token列表、块表(block table)等信息。
"""
block_size = 256 # KV缓存的块大小
counter = count() # 序列ID计数器
def __init__(self, token_ids: list[int], sampling_params = SamplingParams()):
"""
初始化序列
Args:
token_ids: prompt的token ID列表
sampling_params: 采样参数
"""
self.seq_id = next(Sequence.counter) # 序列的唯一ID
self.status = SequenceStatus.WAITING # 初始状态为等待
self.token_ids = copy(token_ids) # token ID列表的副本
self.last_token = token_ids[-1] # 最后一个token
self.num_tokens = len(self.token_ids) # token总数
self.num_prompt_tokens = len(token_ids) # prompt token数量
self.num_cached_tokens = 0 # 已缓存的token数量(用于前缀缓存)
self.block_table = [] # 关键:**块表**,存储KV缓存块的索引, 映射逻辑块到物理块
self.temperature = sampling_params.temperature # 采样温度
self.max_tokens = sampling_params.max_tokens # 最大生成token数
self.ignore_eos = sampling_params.ignore_eos # 是否忽略EOS token
def __len__(self):
"""返回序列的token数量"""
return self.num_tokens
def __getitem__(self, key):
"""支持索引访问token列表"""
return self.token_ids[key]
#@property 的核心作用: 把“计算得出的值”封装成只读属性,让接口更自然、更安全
# 调用方式是:seq.is_finished # ✅ 像属性一样访问
# 而不是: seq.is_finished() # ❌ 不需要加括号
@property
def is_finished(self):
"""检查序列是否已完成"""
return self.status == SequenceStatus.FINISHED
@property
def num_completion_tokens(self):
"""返回已生成的completion token数量"""
return self.num_tokens - self.num_prompt_tokens
@property
def prompt_token_ids(self):
"""返回prompt部分的token ID列表"""
return self.token_ids[:self.num_prompt_tokens]
@property
def completion_token_ids(self):
"""返回completion部分的token ID列表"""
return self.token_ids[self.num_prompt_tokens:]
@property
def num_cached_blocks(self):
"""返回已缓存的块数量"""
return self.num_cached_tokens // self.block_size
@property
def num_blocks(self):
"""返回序列需要的总块数量"""
return (self.num_tokens + self.block_size - 1) // self.block_size
@property
def last_block_num_tokens(self):
"""返回最后一个块中的token数量"""
return self.num_tokens - (self.num_blocks - 1) * self.block_size
def block(self, i):
"""
获取第i个块的token列表
Args:
i: 块索引
Returns:
该块的token ID列表
"""
assert 0 <= i < self.num_blocks
return self.token_ids[i*self.block_size: (i+1)*self.block_size]
def append_token(self, token_id: int):
"""
向序列追加一个新生成的token
Args:
token_id: 新生成的token ID
"""
self.token_ids.append(token_id)
self.last_token = token_id
self.num_tokens += 1
def __getstate__(self):
"""
序列化函数,用于pickle
为了节省内存,如果序列还没有生成completion tokens,则保存完整的token列表;
否则只保存最后一个token(因为prompt tokens可以从其他地方恢复)
"""
return (self.num_tokens, self.num_prompt_tokens, self.num_cached_tokens, self.block_table,
self.token_ids if self.num_completion_tokens == 0 else self.last_token)
def __setstate__(self, state):
"""
反序列化函数,用于pickle
从序列化状态恢复序列对象
"""
self.num_tokens, self.num_prompt_tokens, self.num_cached_tokens, self.block_table = state[:-1]
if self.num_completion_tokens == 0:
self.token_ids = state[-1]
else:
self.last_token = state[-1]
关键洞察:block_table 是 PagedAttention 的核心数据结构,类似操作系统的页表。它将序列的逻辑内存映射到物理内存块,实现了内存的灵活管理。
执行上下文:Context 类
"""
上下文管理模块
实现了全局上下文管理,用于在模型推理过程中传递序列信息、KV缓存位置等。
使用全局变量存储上下文,避免在函数间传递大量参数。
"""
from dataclasses import dataclass
import torch
#@dataclass: 少写样板代码(boilerplate),让“装数据的类”更干净
@dataclass
class Context:
"""
上下文数据类
存储当前推理步骤的上下文信息,包括:
- 是否为prefill阶段
- 序列长度信息
- KV缓存的slot映射
- 块表信息
"""
is_prefill: bool = False # 是否为prefill阶段
cu_seqlens_q: torch.Tensor | None = None # query的累积序列长度
cu_seqlens_k: torch.Tensor | None = None # key的累积序列长度
max_seqlen_q: int = 0 # query的最大序列长度
max_seqlen_k: int = 0 # key的最大序列长度
slot_mapping: torch.Tensor | None = None # KV缓存的slot映射
context_lens: torch.Tensor | None = None # 每个序列的上下文长度(decode阶段使用)
block_tables: torch.Tensor | None = None # 块表(用于前缀缓存和decode阶段)
_CONTEXT = Context() # 全局上下文对象
def get_context():
"""
获取全局上下文
Returns:
当前上下文对象
"""
return _CONTEXT
def set_context(is_prefill, cu_seqlens_q=None, cu_seqlens_k=None, max_seqlen_q=0, max_seqlen_k=0, slot_mapping=None, context_lens=None, block_tables=None):
"""
设置全局上下文
Args:
is_prefill: 是否为prefill阶段
cu_seqlens_q: query的累积序列长度
cu_seqlens_k: key的累积序列长度
max_seqlen_q: query的最大序列长度
max_seqlen_k: key的最大序列长度
slot_mapping: KV缓存的slot映射
context_lens: 每个序列的上下文长度
block_tables: 块表
"""
global _CONTEXT
_CONTEXT = Context(is_prefill, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, slot_mapping, context_lens, block_tables)
def reset_context():
"""重置全局上下文"""
global _CONTEXT
_CONTEXT = Context()
Context 存储当前推理步骤的全局状态,区分 Prefill(首次处理提示)和 Decode(逐token生成)两个阶段。
配置管理:Config 类
"""
配置管理模块
定义了nano-vllm系统的配置类,包含模型路径、批处理参数、GPU内存使用等配置项
"""
import os
from dataclasses import dataclass
from transformers import AutoConfig
@dataclass
class Config:
"""
系统配置类
包含nano-vllm运行所需的所有配置参数,包括模型路径、批处理大小、
GPU内存使用率、张量并行大小等。
"""
model: str # 模型路径
max_num_batched_tokens: int = 16384 # 批处理的最大token数量
max_num_seqs: int = 512 # 最大序列数量
max_model_len: int = 4096 # 模型支持的最大序列长度
gpu_memory_utilization: float = 0.9 # GPU内存使用率(0-1之间)
tensor_parallel_size: int = 1 # 张量并行大小,支持1-8
enforce_eager: bool = False # 是否强制使用eager模式(不使用CUDA graph)
hf_config: AutoConfig | None = None # HuggingFace配置对象
eos: int = -1 # 结束符token ID
kvcache_block_size: int = 256 # KV缓存块大小,必须是256的倍数
num_kvcache_blocks: int = -1 # KV缓存块数量,-1表示自动计算
def __post_init__(self):
"""
初始化后处理函数
验证配置参数的有效性,并加载HuggingFace配置
"""
# 验证模型路径是否存在
assert os.path.isdir(self.model)
# 验证KV缓存块大小是256的倍数
assert self.kvcache_block_size % 256 == 0
# 验证张量并行大小在有效范围内
assert 1 <= self.tensor_parallel_size <= 8
# 加载HuggingFace配置
self.hf_config = AutoConfig.from_pretrained(self.model)
# 确保最大模型长度不超过配置中的最大位置嵌入
self.max_model_len = min(self.max_model_len, self.hf_config.max_position_embeddings)
# 确保批处理token数量至少等于最大模型长度
assert self.max_num_batched_tokens >= self.max_model_len
这些参数控制着系统的内存使用、并发能力和性能特征。
内存管理 - PagedAttention 的实现
内存块抽象
"""
块管理器模块
实现了KV缓存的块管理功能,包括块的分配、释放、哈希缓存等。
支持前缀缓存(prefix caching)以提高性能。
"""
from collections import deque
import xxhash
import numpy as np
from nanovllm.engine.sequence import Sequence
class Block:
"""
KV缓存块类
表示一个KV缓存块,包含块ID、引用计数、哈希值和token列表。
"""
def __init__(self, block_id):
"""
初始化块
Args:
block_id: 块的唯一ID
"""
self.block_id = block_id # 块的唯一ID
self.ref_count = 0 # 引用计数,表示有多少个序列在使用这个块
self.hash = -1 # 块的哈希值,用于前缀缓存查找
self.token_ids = [] # 块中存储的token ID列表
def update(self, hash: int, token_ids: list[int]):
"""
更新块的内容和哈希值
Args:
hash: 块的哈希值
token_ids: token ID列表
"""
self.hash = hash
self.token_ids = token_ids
def reset(self):
"""重置块的状态"""
self.ref_count = 1
self.hash = -1
self.token_ids = []
## 块管理器:PagedAttention 的核心
class BlockManager:
"""
块管理器类
管理KV缓存块的分配和释放,支持前缀缓存以提高性能。
通过哈希表快速查找已缓存的块,避免重复计算。
"""
def __init__(self, num_blocks: int, block_size: int):
"""
初始化块管理器
Args:
num_blocks: 总块数量
block_size: 每个块的大小(token数量)
"""
self.block_size = block_size # 块大小
self.blocks: list[Block] = [Block(i) for i in range(num_blocks)] # 所有块的列表
self.hash_to_block_id: dict[int, int] = dict() # 哈希值到块ID的映射,用于前缀缓存
self.free_block_ids: deque[int] = deque(range(num_blocks)) # 空闲块ID队列
self.used_block_ids: set[int] = set() # 已使用块ID集合
@classmethod
def compute_hash(cls, token_ids: list[int], prefix: int = -1):
"""
计算token列表的哈希值
Args:
token_ids: token ID列表
prefix: 前缀哈希值,用于链式哈希计算
Returns:
计算得到的哈希值
"""
h = xxhash.xxh64()
if prefix != -1:
h.update(prefix.to_bytes(8, "little"))
h.update(np.array(token_ids).tobytes())
return h.intdigest()
def _allocate_block(self, block_id: int) -> Block:
"""
分配一个块(内部方法)
Args:
block_id: 要分配的块ID
Returns:
分配得到的块对象
"""
block = self.blocks[block_id]
assert block.ref_count == 0
block.reset()
self.free_block_ids.remove(block_id)
self.used_block_ids.add(block_id)
return self.blocks[block_id]
def _deallocate_block(self, block_id: int) -> Block:
"""
释放一个块(内部方法)
Args:
block_id: 要释放的块ID
"""
assert self.blocks[block_id].ref_count == 0
self.used_block_ids.remove(block_id)
self.free_block_ids.append(block_id)
def can_allocate(self, seq: Sequence) -> bool:
"""
检查是否有足够的空闲块来分配序列
Args:
seq: 要检查的序列
Returns:
如果有足够块则返回True,否则返回False
"""
return len(self.free_block_ids) >= seq.num_blocks
def allocate(self, seq: Sequence):
"""
为序列分配KV缓存块
支持前缀缓存:如果块的token序列已经在缓存中,则复用该块,避免重复计算。
Args:
seq: 要分配块的序列
"""
assert not seq.block_table
h = -1 # 当前块的哈希值
cache_miss = False # 是否发生缓存未命中
for i in range(seq.num_blocks):
token_ids = seq.block(i)
# 只有当块是完整大小(block_size)时才计算哈希值
h = self.compute_hash(token_ids, h) if len(token_ids) == self.block_size else -1
# 尝试从哈希表中查找已存在的块
block_id = self.hash_to_block_id.get(h, -1)
# 检查是否真的匹配(哈希碰撞检查)
if block_id == -1 or self.blocks[block_id].token_ids != token_ids:
cache_miss = True
if cache_miss:
# 缓存未命中,分配新块
block_id = self.free_block_ids[0]
block = self._allocate_block(block_id)
else:
# 缓存命中,复用已有块
seq.num_cached_tokens += self.block_size
if block_id in self.used_block_ids:
# 块已被使用,增加引用计数
block = self.blocks[block_id]
block.ref_count += 1
else:
# 块未被使用,分配它
block = self._allocate_block(block_id)
# 如果计算了哈希值,更新块并记录到哈希表
if h != -1:
block.update(h, token_ids)
self.hash_to_block_id[h] = block_id
seq.block_table.append(block_id)
def deallocate(self, seq: Sequence):
"""
释放序列占用的所有块
Args:
seq: 要释放块的序列
"""
# 从后往前遍历块表,减少引用计数
for block_id in reversed(seq.block_table):
block = self.blocks[block_id]
block.ref_count -= 1
# 如果引用计数为0,释放该块
if block.ref_count == 0:
self._deallocate_block(block_id)
seq.num_cached_tokens = 0
seq.block_table.clear()
def can_append(self, seq: Sequence) -> bool:
"""
检查是否可以追加新token到序列(是否需要新块)
Args:
seq: 要检查的序列
Returns:
如果可以追加则返回True
"""
# 如果序列长度模block_size等于1,说明需要新块
return len(self.free_block_ids) >= (len(seq) % self.block_size == 1)
def may_append(self, seq: Sequence):
"""
为序列追加新token做准备(如果需要新块则分配)
当序列追加新token时,可能需要:
1. 分配新块(如果当前块已满)
2. 更新最后一个块的哈希值(如果块刚满)
Args:
seq: 要追加token的序列
"""
block_table = seq.block_table
last_block = self.blocks[block_table[-1]]
if len(seq) % self.block_size == 1:
# 需要新块:当前块已满,需要分配新块
assert last_block.hash != -1
block_id = self.free_block_ids[0]
self._allocate_block(block_id)
block_table.append(block_id)
elif len(seq) % self.block_size == 0:
# 块刚满:更新最后一个块的哈希值
assert last_block.hash == -1
token_ids = seq.block(seq.num_blocks-1)
# 使用前一个块的哈希值作为前缀
prefix = self.blocks[block_table[-2]].hash if len(block_table) > 1 else -1
h = self.compute_hash(token_ids, prefix)
last_block.update(h, token_ids)
self.hash_to_block_id[h] = last_block.block_id
else:
# 块未满:不需要操作
assert last_block.hash == -1
每个 Block 代表一个固定大小(通常256个token)的内存块,可以存储 KV Cache。hash 字段支持前缀缓存优化。
作者:mospan
微信关注:墨斯潘園
本文出处:http://mospany.github.io/2025/12/14/nano-vllm-code-analysis/
文章版权归本人所有,欢迎转载,但未经作者同意必须保留此段声明,且在文章页面明显位置给出原文连接,否则保留追究法律责任的权利。
