ARTICLE · 1301187
《Nano-vLLM源码逐行解析》第2篇 sequence.py、scheduler.py
《Nano-vLLM源码逐行解析》第2篇 sequence.py、scheduler.py
前文图文版讲清了nano-vllm的架构和概念，但参数为什么这么设、张量形状怎么变、流程如何串起来，只有逐行读源码才能理解。这个系列贴出逐行中文注释的源码，建议和图文版对照阅读。
注释版源码：https://github.com/llm-infer/nano-vllm-annotated
nano-vllm 代码结构
• nanovllm/： • llm.py：对外接口LLM，直接继承LLMEngine • config.py：引擎级配置（调度上限、显存比例、KV block等） • sampling_params.py：请求级采样参数（temperature、最大生成长度、是否忽略 EOS） • engine/：推理引擎核心，负责调度、kv cache管理和模型执行 • llm_engine.py：推理入口，自己不做计算，把 tokenizer、scheduler、model_runner 串起来 • scheduler.py：调度器，决定每个step计算哪些序列、各算多少token👈 本篇 • sequence.py：一个请求对应一条序列，记录token、状态和 block table👈 本篇 • block_manager.py：在 CPU 上管理 KV block 的分配、回收和前缀缓存 • model_runner.py：每张 GPU 一个，负责加载模型、分配 KV cache、执行前向和采样 • layers/：模型的基础组件，大多支持张量并行 • attention.py：用 Triton kernel把K/V写入KV cache，再调用 FlashAttention • linear.py：按列、按行切分的并行线性层 • embed_head.py：词表并行的Embedding和LM Head • rotary_embedding.py：旋转位置编码RoPE • layernorm.py：RMSNorm • activation.py：SwiGLU 激活 • sampler.py：按温度从 logits 中采样下一个token • models/：具体的模型结构 • qwen3.py：Qwen3 模型定义 • utils/：工具 • context.py：一次前向所需的注意力元数据，避免层层传参 • loader.py：从 safetensors 文件加载权重到模型
sequence.py
from copy import copyfrom enum import Enum, autofrom itertools import countfrom nanovllm.sampling_params import SamplingParams# 序列的生命周期状态# auto()：自动给枚举成员分配值（1、2、3...），只关心名字不关心值时使用# 状态流转图：## add_request# │# ▼# ┌─────────┐ prefill算完整个prompt ┌─────────┐ 生成EOS或达到max_tokens ┌──────────┐# │ WAITING │ ─────────────────────────> │ RUNNING │ ──────────────────────────> │ FINISHED │# └─────────┘ └─────────┘ └──────────┘# ▲ │ chunked prefill没算完， │# │ └─ 留在WAITING，下一步继续 │# └──────────────────────────────────────────┘# KV cache不够，被抢占(preempt)，释放所有block，回到WAITING队头，之后重新prefillclassSequenceStatus(Enum):# 在等待队列中，还没做完prefill（请求刚进来、chunked prefill做了一半、资源被抢占） WAITING = auto()# prefill已完成，在运行队列中逐token decode RUNNING = auto()# 遇到EOS或达到max_tokens，KV block已释放 FINISHED = auto()# 一条请求在引擎内部的表示：保存token、状态、KV block映射和采样参数# 各计数字段变化示例（block_size=256, prompt长600，前缀缓存命中1个block，num_max_batched_tokens=200）# 初始化： num_tokens=600，num_prompt_tokens=600，num_cached_tokens=0，block_table=[]# 第1次调度： num_cached_tokens=256（命中1个block），block_table=[5,9,12]（三个block分别存放256+256+88=600个token）# num_scheduled_tokens=200（预算只有200，chunked prefill）# 第1次postprocess： num_cached_tokens=456，num_scheduled_tokens=0，还没prefill完成，丢弃采样结果# 第2次调度： num_scheduled_tokens=144（剩余600-456），状态变为RUNNING# 第2次postprocess： num_cached_tokens=600，追加第1个生成token -> num_tokens=601# 之后每次decode： num_scheduled_tokens=1，num_cached_tokens+1，num_tokens+1classSequence:# 类变量，所有实例共享# 读取self.block_size时，实例上没有，就去类上找 block_size = 256# 全局自增器，为每条序列分配唯一的id counter = count()# 默认SamplingParams()只在函数定义时创建一次，所有未传参的调用共享一个SamplingParams()# 因为是只读sampling_params的字段，所以共享没有问题def__init__(self, token_ids: list[int], sampling_params = SamplingParams()):# 从0开始自增，也代表提交顺序self.seq_id = next(Sequence.counter)# 新序列状态为waitingself.status = SequenceStatus.WAITING# 全部的token id：prompt+已生成的token# 浅拷贝：后续append_token不影响调用方传进来的listself.token_ids = copy(token_ids)# 最后一个token，decode时作为输入，计算KVself.last_token = token_ids[-1]# 当前总token数（prompt + 已生成）self.num_tokens = len(self.token_ids)# prompt的token数，区分prompt和生成部分self.num_prompt_tokens = len(token_ids)# 已算好kv并写入kv cache的token数# 来源：前缀缓存命中的部分+本序列step计算过的部分（prefill/decode）self.num_cached_tokens = 0# 本次step需要计算的token数# prefill：本次处理的prompt长度（chunked prefill时可能小于prompt剩余长度）# decode：固定为1self.num_scheduled_tokens = 0# 是否处在prefill阶段self.is_prefill = True# kv 逻辑block -> kv 物理block映射表# block_table[i]代表第i个逻辑block对应的物理block id# 类似操作系统的页表，物理block在显存中不必连续# 例如：block_table=[5,9,12]，表示第0~255个kv在物理block 5,第256~511个kv在物理block 9self.block_table = []# 请求可配置采样参数，将其拷贝到序列上，便于ModelRunner批量读取self.temperature = sampling_params.temperature# 此序列允许生成token的上限，达到上限后直接结束self.max_tokens = sampling_params.max_tokens# 是否忽略EOS继续生成self.ignore_eos = sampling_params.ignore_eos# 以双下划钱开头和结尾的方法叫‘魔法方法’，会在特定语法下自动调用# len(seq) 等价于 seq.__len__()，返回token总数def__len__(self):returnself.num_tokens# 获取某个或某段token，例如：seq[3] -> self.token_ids[3]；seq[0:5] -> self.token_ids[0:5]# key可以是下标或切片对象def__getitem__(self, key):returnself.token_ids[key]# @property把方法变成‘只读属性’# 适用于由其他字段计算得到的值，每次访问都重新计算 @propertydefis_finished(self):returnself.status == SequenceStatus.FINISHED# 已生成的token数，判断是否达到max_tokens @propertydefnum_completion_tokens(self):returnself.num_tokens - self.num_prompt_tokens @propertydefprompt_token_ids(self):returnself.token_ids[:self.num_prompt_tokens]# 已生成的token，用于detokenizer生成文本，返回客户端 @propertydefcompletion_token_ids(self):returnself.token_ids[self.num_prompt_tokens:]# 当前token数需要占用多少kv block，向上取整# 新请求需占用block = num_blocks - 命中缓存的block数 @propertydefnum_blocks(self):return (self.num_tokens + self.block_size - 1) // self.block_size# 最后一个block存了多少token = 总token数 - 前num_blocks-1 个满block存的token数# 用于计算当前token存储在物理block上的偏移量 @propertydeflast_block_num_tokens(self):returnself.num_tokens - (self.num_blocks - 1) * self.block_size# 取第i个逻辑block对应的token ids，用于计算前缀缓存的hash# 超出长度时自动截断，例如block_size=256，num_tokens=600# block(2) -> token_ids[512,768]，实际只有88个元素，即token_ids[512,599]defblock(self, i):assert0 <= i < self.num_blocksreturnself.token_ids[i*self.block_size: (i+1)*self.block_size]# decode每生成一个token，就追加到序列末尾defappend_token(self, token_id: int):self.token_ids.append(token_id)self.last_token = token_idself.num_tokens += 1# 自定义序列化字段# python序列化库pickle默认序列化所有字段（self.__dict__）# 定义了__getstate__后，pickle只序列化它的返回值；还原时把这个值传给__setstate__# 传输路径：rank 0负责调度指挥，把seq派发给其他rank# rank 0：pickle.dumps(seq) -> 调用 seq.__getstate__() -> 字符串 -> 共享内存# rank 1：共享内存 -> 字符串 -> pickle.loads -> 调用 __setstate__ -> seqdef__getstate__(self):# prefill：需要完整的token_ids来计算下一个token# decode：历史token的kv存在kv cache里，只需要最后一个token，避免每个step都传输整条序列，节省序列化开销 last_state = self.last_token ifnotself.is_prefill elseself.token_ids# rank只做前向计算，所以只发送前向需要的数据即可return (self.num_tokens, self.num_prompt_tokens, self.num_cached_tokens, self.num_scheduled_tokens, self.block_table, last_state)# 反序列化def__setstate__(self, state):self.num_tokens, self.num_prompt_tokens, self.num_cached_tokens, self.num_scheduled_tokens, self.block_table, last_state = state# prefill阶段last_state是完整的token_idsifisinstance(last_state, list):self.token_ids = last_stateself.last_token = self.token_ids[-1]else:# decode阶段last_state是上一个step生成的token idself.token_ids = []self.last_token = last_statescheduler.py
from collections import dequefrom nanovllm.config import Configfrom nanovllm.engine.sequence importSequence, SequenceStatusfrom nanovllm.engine.block_manager import BlockManager# 调度器，决定每step计算哪些seq，各计算多少token# 策略：# 1. prefill优先：只要等待队列里有能调度的请求，本step就只做prefill# 2. 没有prefill可做时，对所有running的序列做一次decode，各生成1个token# 3. 同一个step里不会混合prefill和decode# 4. KV cache不够时，抢占最晚加入running的序列，释放其block，放回waiting队列等待重新prefill# 流程图：# add() prefill算完 结束(EOS/max_tokens)# 新请求 ───────> ┌─────────┐ ────────────────────> ┌─────────┐ ──────────────────> 释放block，移出# │ waiting │ │ running │# └─────────┘ <──────────────────── └─────────┘# ▲ │ preempt：放回waiting队头（appendleft）# └──────┘# chunked prefill未完成：留在waiting队头，下一步继续# 决策流程：# ┌────────────────────── (1) prefill 阶段 ──────────────────────┐# │ while waiting非空 且 未达max_num_seqs: │# │ seq = waiting队头 │# │ 剩余token预算为0？是->跳出循环 │# │ 首次调度 且 KV block不够？是->跳出循环 │# │ 预算放不下整个seq 且 不是第一条？是->跳出循环 │# │ 分配block（仅首次） │# │ 本次计算min(剩余要算的token数，剩余预算)个token │# │ prompt全部算完？是->移入running队列 │# │ 加入scheduled_seqs │# └──────────────────────────────────────────────────────────────┘# │# scheduled_seqs非空？ ── 是 ──> return (scheduled_seqs, True)# │ 否（没有能做的prefill）# ▼# ┌────────────────────── (2) decode 阶段 ───────────────────────┐# │ while running非空 且 未达max_num_seqs: │# │ seq = 从running队头取出 │# │ while KV block不够追加1个token: │# │ running还有序列？ 是 -> 抢占running队尾的序列 │# │ 否 -> 抢占seq自己，跳过它 │# │ （block够了）加入scheduled_seqs │# └─────────────────────────────────────────────────────────────┘# │# ▼# scheduled_seqs放回running队头，return (scheduled_seqs, False)classScheduler:def__init__(self, config: Config):# 单step最多调度的序列数，主要影响decode性能# 相较于prefill，decode持续时间长，并发度很容易超过max_num_seqs# decode step阶段读取一次权重，增加max_num_seqs，可以把权重读取成本摊薄# max_num_seqs存在上限，由可用的kv block决定self.max_num_seqs = config.max_num_seqs# 单step最多计算的token数（主要约束prefill；decode每条序列1个token）self.max_num_batched_tokens = config.max_num_batched_tokens# 结束符的token id，由tokenizer提供self.eos = config.eos# 每个kv block可容纳的token数self.block_size = config.kvcache_block_size# num_kvcache_blocks在ModelRunner初始化时，根据剩余显存自动计算得到self.block_manager = BlockManager(config.num_kvcache_blocks, config.kvcache_block_size)# 等待prefill 的序列（新请求、chunked prefill未完成的、被抢占的）self.waiting: deque[Sequence] = deque()# 已完成prefill，处于decode阶段的序列self.running: deque[Sequence] = deque()# 空容器在条件判断中被当做Falsedefis_finished(self):returnnotself.waiting andnotself.runningdefadd(self, seq: Sequence):self.waiting.append(seq)# 返回本step需要计算的seq，以及是否为prefilldefschedule(self) -> tuple[list[Sequence], bool]: scheduled_seqs = []# 记录本step消耗的token预算 num_batched_tokens = 0# prefill# 从waiting中取序列# prefill优先调度：每次调度优先看waiting队列whileself.waiting andlen(scheduled_seqs) < self.max_num_seqs:# 只看队头，不取出；能算完时才真正取出 seq = self.waiting[0]# 本step剩余的token预算 remaining = self.max_num_batched_tokens - num_batched_tokensif remaining == 0:break# 第一次被调度，或者被抢占后重新调度，此时还没申请block_tableifnot seq.block_table:# 检查block是否足够，并查询缓存命中情况 num_cached_blocks = self.block_manager.can_allocate(seq)# 返回-1，代表空闲block不够；队列后面的seq也不调度了，保持FIFO，防止队头饿死if num_cached_blocks == -1:break# 本次step需要计算的token = 总数 - 前缀缓存命中的token数 num_tokens = seq.num_tokens - num_cached_blocks * self.block_sizeelse:# 已经分配过block，说明是chunked prefill的后续块，继续计算剩余部分 num_tokens = seq.num_tokens - seq.num_cached_tokens# 剩余token预算不够用来调度此seq，而且前面已经有序列被选中，则把此seq留到下一个step# 为啥不用chunked prefill调度此seq？scheduled_seqs不为空，if remaining < num_tokens and scheduled_seqs: # only allow chunked prefill for the first seqbreak# 一次性为整条序列分配所有block（包括本次还没算到的chunk）# 传入num_cached_blocks，不再为缓存的token申请blockifnot seq.block_table:self.block_manager.allocate(seq, num_cached_blocks)# 本次实际计算的token数，预算不足时只算一部分 seq.num_scheduled_tokens = min(num_tokens, remaining) num_batched_tokens += seq.num_scheduled_tokens# 本次能把prompt全部算完，则移入running队列，下一个step开始decode# 否则继续留在waiting队头，下一个step继续算下一个chunkif seq.num_cached_tokens + seq.num_scheduled_tokens == seq.num_tokens: seq.status = SequenceStatus.RUNNINGself.waiting.popleft()self.running.append(seq) scheduled_seqs.append(seq)# 如果有可调度的prefill请求，则先执行prefill。prefill优先体现在这里if scheduled_seqs:return scheduled_seqs, True# decode# 每个running的序列各计算1个token# 抢占示例（running=[A,B,C,D]，空闲block为0，每条序列刚好都需要新block）# 取A：block不够 -> 抢占队尾D（释放D的block）-> block够了 -> A被调度# 取B：block不够 -> 抢占队尾C -> 够了 -> B被调度# running为空，退出循环，本轮调度A、B；C和D回到waiting队头等待重新prefillwhileself.running andlen(scheduled_seqs) < self.max_num_seqs:# 从队头中取出 seq = self.running.popleft()whilenotself.block_manager.can_append(seq):ifself.running:# 空闲block不够，不断抢占running队尾的seq，直到block够用self.preempt(self.running.pop())else:# 没有别的seq可以抢了，只能抢占自己self.preempt(seq)break# block足够，本序列可以参与本次decodeelse: seq.num_scheduled_tokens = 1 seq.is_prefill = False# 检查是否需要申请新的block，追加到block_table# 当前block存满时，才需要追加self.block_manager.may_append(seq) scheduled_seqs.append(seq)assert scheduled_seqs# 放回running队列，默认它们还会继续decode# 当seq遇到EOS或达到max_tokens时，会调用self.running.remove(seq)把它移除self.running.extendleft(reversed(scheduled_seqs))return scheduled_seqs, False# 抢占running状态的seq，释放kv block，放回waiting队列队头，下次schedule优先考虑它# 这里是重计算式抢占，而不是把KV换到CPU# 由于block释放后，hash仍保留，如果这些block没有被覆盖，重新prefill时仍有可能命中缓存defpreempt(self, seq: Sequence): seq.status = SequenceStatus.WAITING seq.is_prefill = Trueself.block_manager.deallocate(seq)self.waiting.appendleft(seq)# 模型forward和sampling之后的后处理：更新seq状态，追加新token，判断是否结束# seqs：本step调度的序列；token_ids：本step生成的token，与seqs一一对应defpostprocess(self, seqs: list[Sequence], token_ids: list[int], is_prefill: bool):# zip把两个列表按位置配对：(seqs[0], token_ids[0]), (seq[1], token_ids[1]), ...for seq, token_id inzip(seqs, token_ids):# 把本次计算填满的block登记到前缀缓存self.block_manager.hash_blocks(seq)# 本次计算的token的KV已写入cache seq.num_cached_tokens += seq.num_scheduled_tokens seq.num_scheduled_tokens = 0# chunked prefill还没算完整个prompt，此时计算出的token没有意义，直接丢弃if is_prefill and seq.num_cached_tokens < seq.num_tokens:continue seq.append_token(token_id)# 结束条件：生成了EOS，或者达到最大长度# 释放kv block、从running队列中移除if (not seq.ignore_eos and token_id == self.eos) or seq.num_completion_tokens == seq.max_tokens: seq.status = SequenceStatus.FINISHEDself.block_manager.deallocate(seq)self.running.remove(seq)
