nano-vllm解读

一个简化版的vllm，实现了以下的内容flashattntriton。

qq_43491590

236人浏览 · 2026-05-26 21:26:55

qq_43491590 · 2026-05-26 21:26:55 发布

nano-vllm是什么

nano-vllm一个简化版的vllm，实现了以下的内容

一个简单但完整的推理引擎
1. 多请求调度
2. 动态请求切换
3. 前缀缓存
4. 分块预填充
完整的Qwen3模型
1. 使用flashattn 库实现关键的注意力机制
2. 使用triton 编写部分算子，包括kvcache的写入
3. 使用张量并行支持多卡推理
PagedAttention
其他
1. CUDA Graph
2. Gumbel-Max采样等

项目的整体流程

入口在LLMEngine.generate ，流程如下

添加请求
1. 使用调度器来处理，其实就是简单加到等待列表中
循环直到结束
1. 也就是调度器结束，等待列表和运行列表都为空
2. 执行step
3. 保存最新输出内容
所有输出进行解码

step

流程如下

调度器调度多个请求
准备，将请求转为模型的输入
运行模型
1. 模型是nn.Module
2. 得到logits后采样
后处理
1. 把输出的结果保存到请求中，更新请求的数据
2. 如果结束（遇到eos或者长度超了）则释放请求

PagedAttention是如何实现的

PageAttention原理

PagedAttention就是把KVCache给分成多个页，每个页比如长256个词元，这样可以按照256为最小单位分配显存（直接从预分配好的大空间中取256，而不是临时开辟），相比以往直接按照最大长度分配的粒度小很多，可以减少内存空洞和浪费

做法就是在计算注意力的时候（也只有这里用得到KVCache）额外传入每个词元对应的块编号，然后GPU计算的时候读取KVCache间接查询得到地址

原来的过程：

给一个连续的地址addr
根据索引计算得到要读取的KV值地址addr+i*BLOCK_SIZE_N
使用计算

现在：

给一个连续的地址addr
根据索引计算得到所在的块和块内偏移block_id, block_offset
计算得到KV值地址addr+block_table[block_id] * block_size + block_offset
使用计算

实际在GPU上计算的时候，一次读取N长度的KV，只要BlockSize可以被N整除（比如256可以被64整除），就可以保证读取不会跨块，每次读取都是连续内存

Nano-vllm中的使用

# attention.py::Attention::forward
o = flash_attn_with_kvcache(q.unsqueeze(1), k_cache, v_cache,
                            cache_seqlens=context.context_lens, 
                            block_table=context.block_tables, 
                            softmax_scale=self.scale, causal=True)

这是flashattn 库的函数，这个k_cache 和v_cache 是大段连续空间，供所有请求使用，而每个请求的块表放在block_table 中，cache_seqlens 则表示每个请求的缓存长度

详细的参数如下

q
- 尺寸为sum_seq_lens x H x d
- 经过unsqueeze 后变成reqs x 1 x H x d
k_cache和v_cache
- 尺寸为max_blocks x block_size x H_kv x d
cache_seqlens
- 表示每个请求的缓存长度
- 尺寸为 reqs ，为int32
block_table
- 表示每个请求的多个块编号
- 尺寸为 reqs x max_len_blocks ，类型为int32
- 按照最长块数在右侧补齐长度，填充的部分写-1就行，这里只有一个序号，没有额外显存开销
softmax_scale
- 就是 1/sqrt(HEAD_DIM)

KVCache的写入

这个部分使用triton 自己实现了一个核

block_table 只有块的序号，这里新增了一个slot_mapping 的状态，直接是每个kv对应整个K/V缓存的偏移（除以D，D=HxHeadDim），这样直接可以根据slot_mapping 得到偏移，省的额外的计算，这个步骤在准备请求转输入的位置实现

Page的分配

初始分配

计算一个块的显存大小，然后根据总可用显存，得到最多能用多少块

一个块的大小如此计算

$2\times L\times \text{BlockSize}\times H_{kv} \times D_{\text{head}}\times \text{sizeof}(\text{dtype})$

直接使用一个empty 分配得到最大块数乘以每个块的空间，然后分配给每层

# model_runner.py::ModelRunner::allocate_kv_cache

for module in self.model.modules():
    if hasattr(module, "k_cache") and hasattr(module, "v_cache"):
        module.k_cache = self.kv_cache[0, layer_id]
        module.v_cache = self.kv_cache[1, layer_id]
        layer_id += 1

不过这样只能支持一个模型

单个请求分配

这里的分配都是分配一个块编号，一个块编号唯一对应一块显存

分两种情况：

预填充
- 需要把所有输入都缓存起来
- 长度就是num_tokens
- 这里涉及前缀缓存，前缀尽量使用，复用的不用申请，剩下的申请
- 复用的需要增加引用
解码
- 只需要额外缓存一个
- 如果 num_tokens % block_size == 1 就额外获得一个块序号

释放

简单减少引用，如果为0就释放，所谓释放也就是把这个编号返回到空闲编号队列中，不涉及显存释放

所以这个引用仅用于前缀缓存

前缀缓存

前缀缓存是分块的，如何判别两个块相同，最简单的就是两个块的词元id完全相同，用散列即可。另外输出是前向依赖的，所以光判别本块散列相同还不够，还需要判别前面的所有块都相同。

做法就是从一个固定值开始，把前一个块的散列和本块一起参加散列，得到的值再作为下一个块的散列

整理得到预填充时的分配：

将提示词分词
将词元列表分块
对于每个块，计算一个前缀散列
- 如果散列值已经存在，则复用这个块，引用+1
- 否则新开辟一个，引用+1

SGLang在这里改进为使用前缀树来处理，可以应对任意长度的前缀，不再限制必须整块匹配。对于不完整匹配可以复制一份出来，不匹配的清空即可

BlockManager

负责管理块索引，不涉及显存

维护两个队列

used
free

申请就从free中移除，放到used中，释放则反过来

如何支持多请求

每个请求包含哪些内容

所有的词元id
1. 数量统计
  1. 已缓存
  2. 待缓存
  3. 总数
分块表
请求标识，id
请求状态
1. 等待
2. 运行
3. 完成等
阶段
1. 预填充还是解码

如何适配不同请求不同序列长度

整个Qwen3，可以分为以下几种操作

嵌入
- 只涉及最后一维
矩阵乘法
- 只涉及最后两维
单点操作和Norm
- 只涉及最后一维
注意力
- 只涉及最后两维（如果算上头部就是最后三维）

所以实际上并没有明确依赖B维度，并且在预填充中，不同请求的长度不同，如果使用B维度就需要统一长度，造成浪费

做法就是把不同请求的输入在序列长度维度堆叠起来，组成 sum_of_seq_lens x D ，对于非注意力部分没有任何影响

注意力则需要拆分不同请求，因为因果遮罩只能在每个请求的QKV内生效

做法就是

# attention.py::Attention::forward
o = flash_attn_varlen_func(q, k, v,
                           max_seqlen_q=context.max_seqlen_q, cu_seqlens_q=context.cu_seqlens_q,
                           max_seqlen_k=context.max_seqlen_k, cu_seqlens_k=context.cu_seqlens_k,
                           softmax_scale=self.scale, causal=True, block_table=context.block_tables)

这也是flashattn 的库实现，参数如下

q
- 尺寸为sum_seq_lens x H x d
kv
- 尺寸为sum_seq_lens x H_kv x d
max_seqlen_q，max_seqlen_k
- 猜测是内部按照最大长度分块，然后不同请求之间并行
- 两个值在有前缀缓存的情况下可以不同
cu_seqlens_q，cu_seqlens_q
- 累加长度，而不是单独的长度

多请求是如何调度的

在Scheduler 这个类中，schedule 这个函数实现

维护两个队列，一个等待，一个正在运行，添加一个请求直接加到等待队列中

一次前向传播执行完毕，即触发调度

一、优先执行预填充调度（Prefill）

循环条件：存在等待队列（waiting） 且已调度序列数小于最大并发序列数（max_num_seqs）
取出等待队列第一个序列，计算剩余可批处理 token 数，若为 0 则终止预填充
序列无块表（block_table）→ 全新序列：
- 向块管理器申请可分配的块数量（支持块复用）
- 若分配失败（返回 - 1），说明显存不足，直接终止预填充
- 本次计算 token 数 = 序列总 token 数 - 可分配块总大小
序列有块表（block_table）→ 已有缓存：
- 本次计算 token 数 = 序列总 token 数 - 已缓存 token 数
非首个调度序列时，若剩余可批 token 数不足本次计算量，停止添加，本请求也不加入
1. 也就是预填充最多两次算完，不能再多了
全新序列：执行块分配，写入块表
最终调度 token 数 = 本次计算 token 数与剩余可批 token 数的较小值
累加批处理 token 数，更新序列状态：
- 完成全部预填充 → 标记为运行中，从等待队列移入运行队列
- 未完成 → 保留在等待队列
将序列加入调度列表，继续循环

二、预填充无可用序列 → 执行解码调度（Decode）

循环条件：存在运行队列（running） 且已调度序列数小于最大并发序列数
取出运行队列第一个序列，检查是否可追加块：
- 可追加：直接设置调度 token 数 = 1，标记为解码，执行块追加，加入调度列表
- 不可追加：
  - 若有其他运行中序列 → 抢占末尾序列释放资源
  - 若无其他运行中序列 → 直接抢占当前序列释放资源
  - 重新尝试块追加，直到成功或抢占完成
调度完成后，将调度列表序列放回运行队列头部

CUDAGraph优化

实际上，nanovllm只有在预填充才会直接调用self.model ，而解码则会使用cuda graph，代码如下

bs = input_ids.size(0)
context = get_context()
graph = self.graphs[next(x for x in self.graph_bs if x >= bs)]
graph_vars = self.graph_vars
graph_vars["input_ids"][:bs] = input_ids
graph_vars["positions"][:bs] = positions
graph_vars["slot_mapping"].fill_(-1)
graph_vars["slot_mapping"][:bs] = context.slot_mapping
graph_vars["context_lens"].zero_()
graph_vars["context_lens"][:bs] = context.context_lens
graph_vars["block_tables"][:bs, :context.block_tables.size(1)] = context.block_tables
graph.replay()
return self.model.compute_logits(graph_vars["outputs"][:bs])

如果整个流程的核函数、核函数启动参数、核函数之间的依赖关系完全一致，cuda可以将其保存为一个计算图，进行一些优化，然后一次计算一个计算图

对于整个LLM而言，唯一的变数就是序列长度N，N会影响核函数启动参数，预填充的序列长度更加多变，它不仅有灵活的输入长度，而且Q的长度和KV长度完全可能不同。而解码的序列长度就是请求数量，较为固定

nano-vllm选择一些预置长度，对每种长度先算一遍，保存一下计算图，使用时找一个不小于实际序列长度最小长度的图，而这些图变量长度都是预置长度，只填充需要的部分

然后每次修改一些上下文参数以及输入参数，直接重放计算图。

图创建如下

input_ids = torch.zeros(max_bs, dtype=torch.int64)
positions = torch.zeros(max_bs, dtype=torch.int64)
slot_mapping = torch.zeros(max_bs, dtype=torch.int32)
context_lens = torch.zeros(max_bs, dtype=torch.int32)
block_tables = torch.zeros(max_bs, max_num_blocks, dtype=torch.int32)

self.graph_bs = [1, 2, 4, 8] + list(range(16, max_bs + 1, 16))
for bs in reversed(self.graph_bs):
    graph = torch.cuda.CUDAGraph()
    set_context(False, slot_mapping=slot_mapping[:bs], context_lens=context_lens[:bs], block_tables=block_tables[:bs])
    outputs[:bs] = self.model(input_ids[:bs], positions[:bs])    # warmup
    with torch.cuda.graph(graph, self.graph_pool):
        outputs[:bs] = self.model(input_ids[:bs], positions[:bs])    # capture
    if self.graph_pool is None:
        self.graph_pool = graph.pool()
    self.graphs[bs] = graph
    torch.cuda.synchronize()
    reset_context()

self.graph_vars = dict(
    input_ids=input_ids,
    positions=positions,
    slot_mapping=slot_mapping,
    context_lens=context_lens,
    block_tables=block_tables,
    outputs=outputs,
)

预置的图变量直接使用最大长度，它的显存占用相比KVCache而言非常有限

Gumbel-Max采样

一般得到logits后除以温度，然后进行softmax，然后根据概率分布随机采样一个

对一个长度为N的分类进行采样，做法就是求累积概率，然后随机采样一个 $u\sim U(0,1)$ ，找到累计概率不小于u的最小的类别，这个过程不便进行并行

Gumbel-Max是对其的并行加速：

Gumbel分布是

$G=-\ln(-\ln(u)), u \sim U(0,1)$

得到

$F_G(g) = \exp(-\exp(-g)), f_G(g)=\exp(-g-\exp(-g))$

令

$X_i = \ln p_i + G_i, I = \argmax X_i$

于是有

$P(I=i)=P(X_j \le X_i, \forall j \ne i)=P(G_j \le \ln(\frac{p_i}{p_j})+G_i, \forall j\ne i)$

这些 $G_j$ 都是独立同分布采样出来的，从而有

$P(I=i)=\prod_{j\ne i} \int_{-\infin}^\infin f_G(g) F_G(\ln(\frac{p_i}{p_j}+g)\mathrm{d}g$

展开就有

$P(I=i)=\int_{-\infin}^\infin f_G(g) \prod_{j\ne i} \exp(-\frac{p_j}{p_i}\exp(-g))\mathrm{d}g=\\ \int_{-\infin}^\infin f_G(g) \exp(-\frac{1-p_i}{p_i}\exp(-g))\mathrm{d}g$