nano-vLLM 源码解读：1200 行 Python 实现一个 LLM 推理引擎

自我介绍

大家好，我是Aaron。过去一直从事算法领域相关工作，自工作以来，一直持续在关注算法相关技术的发展，当前除了在工作中会进行一些AI的落地实践，工作之余也在通过AI辅助编程做一些有趣的探索。我会持续更新AI实践过程中的心得体会，以及AI知识分享。欢迎大家关注我关注公众号，以便第一时间收到更新推送。

• 微博: https://weibo.com/u/1801679460
• Github: https://github.com/aaronshan

nano-vLLM 是一个从零实现的轻量级 vLLM：约 1200 行干净的 Python（外加少量 Triton kernel），剥掉了生产版 vLLM 的样板代码与工程复杂度，却完整保留了现代推理引擎的核心机制——PagedAttention 式显存管理、prefix caching、连续批处理、张量并行、torch.compile 与 CUDA Graph。它基于 vLLM 早期的 V0 架构（vLLM 主线已演进到 V1+），但这些核心原理在今天的高性能引擎里基本没变。

本文按「一次推理的生命周期」走读 nano-vLLM 的当前源码（main 分支）：先建立全景，再沿 example.py → LLMEngine → Scheduler / BlockManager → ModelRunner 一路向下，最后补齐模型与算子层的关键实现。代码引用写作 文件:行号，均可直接在仓库中核对。

阅读门槛：会读 Python 即可。 token、attention、KV cache、prefill/decode 这些概念不需要提前知道——第 0 章用纯 Python 视角把它们全部补齐，正文首次出现的工程名词也都能在第 0.6 节的名词速查表里找到一句话解释。

零、先修课：只用 Python 也够用的背景包

本章把正文用到的背景知识压缩到「会 Python 就能懂」的颗粒度。每一节都不追求严格，只追求够正文使用。

0.1 大模型推理在算什么

token 与 tokenizer。模型不读文字，只读整数。tokenizer 把文本切成一个个 token 并查表换成 id：

"Write a travel plan"  --tokenizer-->  [1, 2, 3, 4]   # 示意值

词表（vocab）就是「id ↔ 文本片段」的对照表，Qwen3 的词表约 15 万个条目。

一步推理。模型拿到一串 id，对词表里每个词算一个分数（logits），然后按分数挑一个作为下一个 token：

logits = model(input_ids)        # 长度 = 词表大小的一列分数，比如 15 万个浮点数next_id = sample(logits)         # 按分数挑下一个词的 id

softmax 把这列分数变成和为 1 的概率分布；temperature（温度）调节分布的平与尖——温度高，低分词也有机会被选中，输出更随机；温度低，几乎总选最高分，输出更确定。「采样」就是按这个概率分布抽一个 id。

生成 = 循环。每次只产一个 token，把它拼回输入再算下一个——这叫「自回归」：

ids = tokenize(prompt)while not finished:    logits = model(ids)    ids.append(sample(logits))

「写一个旅行计划」这样几百 token 的回答，就是把这个循环跑几百圈。推理引擎存在的意义，就是让这个循环转得尽量快、尽量省。

0.2 Transformer 十分钟版

model(ids) 内部做的事，用极简伪代码写出来是：

x = embedding[ids]                  # 查表：每个 id → 一个向量（比如 1024 维）for layer in layers:                # Qwen3-0.6B 有 28 层，每层结构相同    x = x + attention(norm(x))      # 残差连接：原向量 + 子层输出    x = x + mlp(norm(x))logits = x @ lm_head_weight.T       # 最后一层向量映射回「词表大小」个分数

逐个拆开：

• embedding：一张 [词表大小, hidden_dim] 的大表，按 id 查行，把整数变成向量。hidden_dim 是模型内部的向量宽度（Qwen3-0.6B 是 1024）。
• attention（注意力）：让位置 i 的向量「看到」它前面的所有位置。每个位置算出三个向量：query（我想找什么）、key（我是什么）、value（我提供什么）。位置 i 的 query 与每个历史位置的 key 做点积得到「相关度分数」，softmax 归一化后对所有 value 加权求和——这就是该位置的新表示。因果掩码保证位置 i 只能看不晚于 i 的位置（生成时未来还不存在）。
• 多头（multi-head）：把向量切成 h 份，每份独立做一次 attention（一个「头」），再拼回去。不同头可以关注不同的关系。GQA（分组查询注意力）是让几个 q 头共享同一对 k/v 头的变体——k/v 头数少于 q 头数，为的是省 KV cache 显存（0.4 节和正文 2.9 会用到）。
• norm（RMSNorm）：把向量除以它自己的均方根、再乘一组可学习的缩放系数，把数值稳定在合适尺度。每个子层入口都有一次。
• MLP：两层全连接加激活函数，对每个位置独立地做。
• lm_head：最后把每个位置的向量投影回词表大小的分数列。

位置从哪来：attention 本身看不出顺序——「狗咬人」和「人咬狗」对它的打分一样。所以要把位置信息注入进去。nano-vLLM 用的 RoPE（旋转位置编码）做法是：按位置把 q、k 向量两两一组旋转一个与位置成正比的角度，旋转后两个位置的点积天然携带它们的相对距离。正文 3.4 有完整实现。

0.3 为什么推理慢：FLOPs 与字节

GPU 做两件事的速度差着数量级（以 H100 为例，bf16 精度）：

• 算力 ≈ 1000 TFLOP/s（每秒千万亿次浮点运算）
• 显存带宽 ≈ 3 TB/s（每秒从显存读三万亿字节）

换算一下：每从显存读 1 字节数据，GPU 大约「顺手」能做 300 次浮点运算。「每读一字节能摊到多少次运算」这个比值有个名字——算术强度（arithmetic intensity），正文会反复用到它：算术强度够高才算力打满，不够就只能干等显存。而模型权重必须常驻显存，且每产一个 token 都要把全部权重读一遍。于是同一个模型有两种截然不同的处境：

• prefill（预填充）：处理 prompt 的阶段。一次读权重、并行算几百上千个 token，计算量大到把算力打满——计算受限（compute-bound）。
• decode（解码）：逐 token 生成的阶段。每步读的权重字节数不变，却只算 1 个 token，算力大量闲置，带宽先被打满——显存受限（memory-bound）。

由此得到推理引擎最重要的一条推论：decode 时把 N 个请求捆成一个 batch，权重只读一遍就能为每个请求各产一个 token——batch 越大，每 token 越便宜。正文里几乎所有机制（分页、共享、抢占、chunked prefill、CUDA Graph）说到底都在服务这件事：安全地把 batch 做大、把空转做小。

0.4 KV cache：用显存换重复计算

decode 第 t 步的 attention 需要位置 0 到 t−1 的所有 key/value。不缓存的话，每步都要把整段历史重算一遍，越生成越贵：

# 无缓存：第 t 步重算 t 个位置的 K/V，总开销 O(T²)for t in range(T):    k, v = compute_kv(ids[:t+1])# 有缓存：旧位置的 K/V 只依赖过去、永远不会变，存起来；每步只算新 token 的一份for t in range(T):    cache[t] = compute_kv(ids[t])

这就是 KV cache：拿显存换计算，把 O(T²) 变成 O(T)。代价是显存占用——每个 token 要在每一层存下它的 k 和 v 两份向量：

每 token 每层 = 2(k 和 v) × kv 头数 × 头维 × 每元素字节数

代入 Qwen3-0.6B（28 层、8 个 kv 头、头维 128、bf16 两字节）：2 × 28 × 8 × 128 × 2 ≈ 112 KB/token。一张 24 GB 显卡扣掉权重和杂项后留给 KV cache 的部分，也就十几万 token 的量级。显存因此是推理系统里最稀缺的资源——正文的分页、共享、抢占，全都是围绕「怎么省着用、抢着用这块显存」展开的。

0.5 GPU 工程常识包

正文涉及的几样工程工具，各给一句话：

• kernel 与发射开销：GPU 上的每个操作（矩阵乘、归一化……）是一个 kernel，由 CPU 逐个「发射」过去执行。发射本身有微秒级开销；decode 一步要发射几十上百个小 kernel，CPU 的发射开销可能比 GPU 的实际计算还贵。
• Triton：用 Python 语法写 GPU kernel 的框架。kernel[(N,)](args) 表示启动 N 个并行 program；kernel 内 tl.program_id(0) 拿到自己的编号，tl.arange(0, D) 生成 [0, 1, ..., D-1] 下标向量，tl.load / tl.store 按「指针 + 偏移」批量读写显存。正文 2.11 用它实现「把每个 token 的 K/V 存到指定槽位」。
• CUDA Graph：把一连串 kernel 的执行录制成一张「图」，之后一次 replay() 整体重放，省掉逐个发射的 CPU 开销。
• torch.compile：把一串 PyTorch 操作融合、编译成更少的 kernel，减少发射次数和显存往返。
• 多卡通信（NCCL / all_reduce）：NCCL 是 NVIDIA 的多 GPU 通信库；all_reduce 把各卡上的同名张量按元素求和、结果人人有份。张量并行（TP）把每层的计算切到多卡上，各卡只算部分结果，靠 all_reduce 合并。
• SharedMemory + pickle：Python 多进程之间传数据的朴素方式——一块各进程都能读写的共享内存，配上 pickle 把对象序列化成字节写进去。正文 2.8 主进程指挥 worker 用的就是它。

0.6 名词速查表

正文首次出现的名词都能在这里查到一句话解释：

一、引擎全景

1.1 推理引擎到底在优化什么

LLM 推理引擎的唯一目标是高性能地运行推理，而「性能」最终落到两件事上^[2]：

1. 更低的每 token 成本。GPU 按小时计费，空转和满载花一样的钱——所以成本问题本质是利用率问题。
2. 更低的 token 延迟。聊天、编程 agent 等交互场景对延迟尤其敏感。

理解引擎之前，先要理解推理的两个阶段，因为它们的瓶颈完全不同：

• Prefill：一次性并行处理整个 prompt，是计算受限（compute-bound）的——GPU 算力是瓶颈。
• Decode：逐 token 串行生成，每步都要把全部模型权重从 HBM 读一遍，却只算一个 token，是显存带宽受限（memory-bound）的。

一句话概括：prefill 是计算问题，decode 是显存问题。这两个阶段分别对应两个关键延迟指标：

• TTFT（Time-To-First-Token）：收到请求到产出第一个 token 的时间，主要反映 prefill（含排队）。
• TPOT（Time-Per-Output-Token）：首 token 之后每个输出 token 的平均延迟，主要反映 decode。

端到端延迟近似为 Latency ≈ TTFT + TPOT × (输出 token 数 − 1)。吞吐侧常用 RPS（每秒请求数）和 TPS（每秒 token 数）；更严格的指标是 goodput——满足延迟 SLO 的请求速率，例如系统跑 100 RPS 但只有 70 个满足 TTFT/TPOT 目标，goodput 就是 70。

这组指标解释了推理引擎几乎所有设计决策的动机。比如为什么引擎如此执着于把 batch 做大：decode 阶段每步都要完整搬运一遍权重，单请求时这次昂贵的搬运只换来一个 token；把 N 个请求捆在一个 batch 里，同一次权重搬运就能为每个请求各产一个 token——算术强度上去了，每 token 成本就下来了。PagedAttention、prefix caching、chunked prefill、CUDA Graph，说到底都是为「在安全的前提下把 batch 做大、把空转做小」服务的。

1.2 四个组件，一条流水线

nano-vLLM 由四个职责分离的组件构成：

                 ┌────────────────────────────────────────────┐                 │                 LLMEngine                  │                 │          （编排器，跑在主进程）              │                 └────────────────────────────────────────────┘   prompt ──tokenize──►  ┌───────────┐  can_allocate  ┌────────────────┐                         │ Scheduler │ ◄────────────► │  BlockManager  │                         │  waiting  │   allocate /   │ free_block_ids │                         │  running  │   deallocate   │ hash_to_block  │                         └─────┬─────┘                └────────────────┘                               │ seqs, is_prefill                               ▼                    ┌─────────────────────┐   run   ┌──────────────────┐                    │  ModelRunner rank 0 │ ──────► │ Qwen3 模型（GPU） │                    └──────────┬──────────┘         └──────────────────┘                               │ SharedMemory + Event（多卡时）                               ▼                    ┌─────────────────────┐                    │ ModelRunner rank 1+ │   worker：loop() 等待指令                    └─────────────────────┘

• LLMEngine（编排器）：接收用户 prompt，驱动各子组件，输出最终结果。
• Scheduler（调度器）：决定下一个 batch 跑哪些请求。
• BlockManager（块管理器）：管理 KV cache 显存。调度时 Scheduler 先问它能否分配，它负责实际的分配与回收。
• ModelRunner（模型执行器）：拿到调度好的 batch，在 GPU 上跑模型。

1.3 核心数据结构

Config（nanovllm/config.py:6-25）

Config 是集中的静态配置，定义引擎与 tokenizer 的行为（模型路径、最大序列数、单 batch token 预算等）。与其把这些当松散参数到处传，nano-vLLM 把它们捆进一个对象，在各模块间共享同一份真相。

@dataclass(slots=True)class Config:    model: str    max_num_batched_tokens: int = 16384    max_num_seqs: int = 512    max_model_len: int = 4096    gpu_memory_utilization: float = 0.9    tensor_parallel_size: int = 1    enforce_eager: bool = False    hf_config: AutoConfig | None = None    eos: int = -1    kvcache_block_size: int = 256    num_kvcache_blocks: int = -1

几个值得知道的字段与约束：

• max_num_batched_tokens：单次前向可处理的最大 token 预算，chunked prefill 的切块依据。
• kvcache_block_size：每个 KV 块的 token 数，__post_init__ 断言必须是 256 的整数倍。KV cache 块参数正是 PagedAttention 相对早期引擎（如 ORCA）引入的核心。
• num_kvcache_blocks：总块数，启动时按可用显存测算（见 2.9 节）。
• hf_config：构造时即 AutoConfig.from_pretrained(model) 加载；max_model_len 会被钳制到 hf_config.max_position_embeddings 以内。
• eos：默认 -1，实际由 LLMEngine 从 tokenizer 填入（llm_engine.py:33）。

Sequence（nanovllm/engine/sequence.py:13-60）

Sequence 是引擎里最基本的对象，跟踪一个请求的全部状态：token id、状态、KV 块表等。

class SequenceStatus(Enum):    WAITING = auto()    RUNNING = auto()    FINISHED = auto()class Sequence:    block_size = 256    counter = count()def __init__(self, token_ids: list[int], sampling_params = SamplingParams()):self.seq_id = next(Sequence.counter)self.status = SequenceStatus.WAITINGself.token_ids = copy(token_ids)self.last_token = token_ids[-1]self.num_tokens = len(self.token_ids)self.num_prompt_tokens = len(token_ids)self.num_cached_tokens = 0self.num_scheduled_tokens = 0self.block_table = []self.temperature = sampling_params.temperatureself.max_tokens = sampling_params.max_tokensself.ignore_eos = sampling_params.ignore_eosself.is_prefill = True

关键属性：

• status：WAITING（未调度）→ RUNNING（进 batch）→ FINISHED（完成）。
• num_tokens：含输入在内的已处理 token 总数；num_prompt_tokens 仅输入。
• num_cached_tokens：已落 KV cache 的 token 数——既包括 prefix cache 复用的，也包括 chunked prefill 中已处理的进度游标。
• num_scheduled_tokens：本次 step 计划处理多少 token。这是 chunked prefill 的核心状态——一个长 prompt 可以被切成多段、跨多个 step 处理，每段的长度就记在这里。
• is_prefill：标记当前处于 prefill 还是 decode。
• block_table：逻辑块 → 物理块的映射表。
• block_size 是类属性，由 LLMEngine 初始化时注入：Sequence.block_size = config.kvcache_block_size（llm_engine.py:21）。

另外，Sequence 还定义了 num_blocks、last_block_num_tokens、block(i) 等便捷属性，以及一对为跨进程传输优化的 __getstate__/__setstate__（见 3.6 节）。

Scheduler 的两个队列（nanovllm/engine/scheduler.py:16-17）

Scheduler 维护 waiting 与 running 两个双端队列。所有请求先进 waiting；Scheduler 根据约束（KV cache 是否够、batch 预算是否够）把它们移入 running。

Context（nanovllm/utils/context.py:5-27）

@dataclass(slots=True)class Context:    is_prefill: bool = False    cu_seqlens_q: torch.Tensor | None = None    cu_seqlens_k: torch.Tensor | None = None    max_seqlen_q: int = 0    max_seqlen_k: int = 0    slot_mapping: torch.Tensor | None = None    context_lens: torch.Tensor | None = None    block_tables: torch.Tensor | None = None

Context 是给 GPU kernel 的 batch 布局信息，在前向前由 ModelRunner 准备好。它是 per batch 而非 per sequence 的：一个 batch 里各请求可能长度不同、KV 块散布在显存各处，Context 把这些局部信息打包成一次密集 kernel 启动所需的全部元数据：

• cu_seqlens_q / cu_seqlens_k：FlashAttention 变长布局的累积长度（cu = cumulative 前缀和）。
• slot_mapping：逻辑 token 位置 → 物理 KV slot。
• block_tables：每个序列的页表，供 attention 按块读历史 KV。

set_context 每次重建一个新的 Context 对象赋给全局 _CONTEXT（而非原地修改），避免跨 step 的状态残留。

1.4 KV cache 分页：从物理块到 slot 的八步

KV cache 本身很好懂：把已算过的 K/V 存下来，decode 时不必重算（想深入可看这两篇^[3][4]）。难的是管理——vLLM 的答案是像操作系统分页一样按块管理（理论部分可见 PagedAttention 笔记^[5]）。下面用一个具体例子走一遍完整机制：块大小 4，假设硬件最多容纳 100 个 token 的 KV，即 25 个物理块。

第 1 步：测算能造多少块。 可用于 KV cache 的显存 ≈ 总 VRAM − 权重 − 激活峰值 − 各类开销，除以单块字节数即块数（详细公式见 2.9 节）。假设算出来是 25 块。

第 2 步：预分配物理块。 引擎一次性在显存里建好 25 个空块：

physical blocks:[0] [1] [2] [3] ... [24]

第 3 步：把序列切成逻辑块。 新请求的 prompt 按块大小切分。长度 7 的 prompt [1, 2, 3, 4, 5, 6, 7] 切成两个逻辑块：

logical block 0: [1, 2, 3, 4]logical block 1: [5, 6, 7]

逻辑块只描述「需要几块」，还没有物理位置。

第 4 步：用 block table 把逻辑块映射到物理块。 调度时 BlockManager 从空闲池弹出物理块分配给该序列。若池子全空，逻辑块 0、1 分别拿到物理块 0、1：

seq.block_table = [0, 1]      # 逻辑块 i -> 物理块 block_table[i]

这里的关键认知是：KV cache 的归属权是序列级状态。BlockManager 管全局空闲/已用池，而每个序列自己携带它占用的物理块列表。由此物理块无需连续——序列看到的 KV 是逻辑连续的，物理上可散布显存各处。

第 5 步：把块级映射转成 token 级 slot mapping。 attention kernel 按 token 写 KV，所以跑模型前要把 block table 展开成 slot_mapping——每个 token 的精确落点：

slot = physical_block_id * block_size + offset_in_block

两个块 7 个 token 全部展开：

token 0 -> block 0, offset 0 -> slot 0*4+0 = 0token 1 -> block 0, offset 1 -> slot 0*4+1 = 1token 2 -> block 0, offset 2 -> slot 0*4+2 = 2token 3 -> block 0, offset 3 -> slot 0*4+3 = 3token 4 -> block 1, offset 0 -> slot 1*4+0 = 4token 5 -> block 1, offset 1 -> slot 1*4+1 = 5token 6 -> block 1, offset 2 -> slot 1*4+2 = 6slot_mapping = [0, 1, 2, 3, 4, 5, 6]

每个 slot 存一个 token 在每层的 K、V 向量，单 slot 单层约占 2 * num_kv_heads * head_dim * dtype_size 字节。

第 6 步：前向时把新算的 K/V 写进 slot。 QKV 投影产出新 K/V 后，attention 层按 slot_mapping 把它们写进持久 KV cache：token 0 的 K/V 写 slot 0，……，token 6 的写 slot 6。至此该序列 prompt 的 KV 落在物理块 0、1 里。

第 7 步：decode 时读旧写新。 decode 只输入最后一个 token，但 attention 需要完整上下文。比如 prefill 完的长度 7 序列，下一步 decode 的输入是 input_ids = [7]、positions = [6]，而 attention 要读 tokens 0–6 的 K/V——按 block table 从物理块 0（slot 0–3）和块 1（slot 4–6）读。设采样出 next_token = 8，序列变成 8 个 token，物理块 1 还有一个空位，新 token 写入 slot = 1*4+3 = 7：

读旧 K/V: slots 0, 1, 2, 3, 4, 5, 6写新 K/V: slot 7

第 8 步：序列结束后回收块。 命中 <eos> 或 max_tokens 后，该序列的块按引用计数归还空闲池（共享块要等所有引用者都释放，见 2.5 节）。

1.5 一次迭代的全貌

把组件和数据结构串起来，引擎的运行就两件事：

1. 启动时（一次）：Config 初始化并分发给各子模块；ModelRunner 加载权重、预热、测算并预分配 KV 块池。
2. 每次迭代（一次前向）：
1. 用户 prompt 被 tokenize、包成 Sequence，推入 Scheduler 的 waiting 队列；
2. Scheduler 按 FCFS^[6]（先来先服务）从 waiting 取序列，经 BlockManager 检查显存并分配块；长 prompt 可分块多次 prefill；
3. ModelRunner 准备 batch 布局，写入 Context；
4. ModelRunner 执行前向；
5. 重置 context，postprocess() 更新各序列状态、把新填满的块注册进 prefix cache 哈希表；
6. 命中 <eos> 或 max_tokens 的序列释放块、移出 running。

二、沿着一次推理走读代码

2.1 入口：example.py

example.py 展示了引擎的完整用法（注释为本文所加）：

def main():# 初始化 tokenizer 与推理引擎接口（LLM）    path = os.path.expanduser("~/huggingface/Qwen3-0.6B/")    tokenizer = AutoTokenizer.from_pretrained(path)    llm = LLM(path, enforce_eager=True, tensor_parallel_size=1)# 采样参数    sampling_params = SamplingParams(temperature=0.6, max_tokens=256)# 设置 prompt    prompts = ["introduce yourself","list all prime numbers within 100",    ]    prompts = [        tokenizer.apply_chat_template(            [{"role": "user", "content": prompt}],            tokenize=False,            add_generation_prompt=True,        )for prompt in prompts    ]# 生成    outputs = llm.generate(prompts, sampling_params)for prompt, output in zip(prompts, outputs):print(f"Prompt: {prompt!r}")print(f"Completion: {output['text']!r}")

用户面对的 LLM 类其实只是 LLMEngine 的空子类（nanovllm/llm.py：class LLM(LLMEngine): pass），全部实现都在 nanovllm/engine/llm_engine.py。

接下来的走读分三步：

• 2.2–2.3：LLMEngine 初始化引擎，跑 step() 主循环；
• 2.4–2.7：step() 内，Scheduler 在 BlockManager 配合下调度 batch；
• 2.8–2.12：ModelRunner 在 GPU 上跑模型，为每个序列产出下一个 token。

这部分容易迷路，所以全程用两个玩具序列做例子（token id 是示意值，不是真实 Qwen tokenizer 输出）：

 Write a travel plan for Seoul.    ->  [1, 2, 3, 4, 5, 6, 7] Write a travel plan for Tokyo.    ->  [1, 2, 3, 4, 8, 9, 10]

2.2 初始化：LLM 的六步

llm = LLM(model_path, **kwargs) 这一行内部做了六件事（llm_engine.py:11-41）：

1. 创建 Config。 从 kwargs 筛出 Config 的字段，构造配置对象，下传给所有子系统：

def __init__(self, model, **kwargs):    config_fields = {field.name for field in fields(Config)}    config_kwargs = {k: v for k, v in kwargs.items() if k in config_fields}    config = Config(model, **config_kwargs)

同时把块大小注入 Sequence 类属性：Sequence.block_size = config.kvcache_block_size。

2. 为张量并行 spawn worker 进程。 nano-vLLM 只实现了 TP（张量并行）。除 rank 0 外，额外 spawn tensor_parallel_size - 1 个 worker，每个 worker 跑一个 ModelRunner：

self.ps = []self.events = []ctx = mp.get_context("spawn")for i in range(1, config.tensor_parallel_size):    event = ctx.Event()    process = ctx.Process(target=ModelRunner, args=(config, i, event))    process.start()self.ps.append(process)self.events.append(event)

self.events 用于主进程与 worker 间的同步通信（协议见 2.8 节）。

3. 创建 rank 0 的 ModelRunner。 它在主进程内创建，持有 events 以便指挥 worker：

self.model_runner = ModelRunner(config, 0, self.events)

4. 创建 AutoTokenizer。 注意 tokenize 只在主进程做——worker 只碰张量（输入 token id，输出 logits）：

self.tokenizer = AutoTokenizer.from_pretrained(config.model, use_fast=True)

tokenizer 负责把 prompt 编码成 token id、把生成的 id 解码回文本。模型是 HuggingFace 格式时直接用 AutoTokenizer 即可（自写 tokenizer 的例子可参考 nanochat 的 Rust BPE tokenizer）。随后 config.eos = self.tokenizer.eos_token_id。

5. 创建 Scheduler。

self.scheduler = Scheduler(config)

6. 注册退出钩子。atexit.register(self.exit) 保证解释器关闭前清理现场：exit() 先通过 self.model_runner.call("exit") 跨进程广播退出指令（worker 的 loop() 收到 "exit" 后跳出循环），再 del self.model_runner，最后 join 各 worker 进程。

2.3 主循环：generate 与 step

llm.generate(prompts, sampling_params) 做四件事。

1. 添加请求。 把每个 prompt 编码成 token id、包成 Sequence 推入 waiting 队列。add_request 的 prompt 参数既可以是 str 也可以是 list[int]——预先 tokenize 过的可以直接传 id（llm_engine.py:43-47）：

if not isinstance(sampling_params, list):    sampling_params = [sampling_params] * len(prompts)for prompt, sp in zip(prompts, sampling_params):self.add_request(prompt, sp)

2. 反复跑 step 直到全部完成，顺带统计 prefill/decode 吞吐（带 tqdm 进度条）：

while not self.is_finished():    t = perf_counter()    output, num_tokens = self.step()if num_tokens > 0:        prefill_throughput = num_tokens / (perf_counter() - t)else:        decode_throughput = -num_tokens / (perf_counter() - t)

核心是这个 step()（llm_engine.py:49-55）：

def step(self):    seqs, is_prefill = self.scheduler.schedule()    token_ids = self.model_runner.call("run", seqs, is_prefill)self.scheduler.postprocess(seqs, token_ids, is_prefill)    outputs = [(seq.seq_id, seq.completion_token_ids) for seq in seqs if seq.is_finished]    num_tokens = sum(seq.num_scheduled_tokens for seq in seqs) if is_prefill else -len(seqs)return outputs, num_tokens

每个 step 三件事：schedule → run → postprocess——Scheduler 选序列，ModelRunner 跑前向，Scheduler 用采样结果更新状态。

3. 收集完成的序列。 每次 step() 返回本步完成的序列，按 seq_id 累积到 dict。

4. 排序、解码、返回。 全部完成后按 seq_id 排序恢复输入顺序，decode 回文本返回。引擎等 batch 内所有序列完成才返回——同样是出于简洁的取舍。

2.4 调度器：chunked prefill 下的 prefill 调度

Scheduler 持有两个队列和 BlockManager（scheduler.py:8-17）：

class Scheduler:def __init__(self, config: Config):self.max_num_seqs = config.max_num_seqsself.max_num_batched_tokens = config.max_num_batched_tokensself.eos = config.eosself.block_manager = BlockManager(config.num_kvcache_blocks, config.kvcache_block_size)self.waiting: deque[Sequence] = deque()self.running: deque[Sequence] = deque()

schedule() 先看 waiting（prefill），waiting 空了再看 running（decode）。prefill 调度的当前实现（scheduler.py:25-55）：

scheduled_seqs = []num_seqs = 0num_batched_tokens = 0while self.waiting and num_seqs < self.max_num_seqs:    seq = self.waiting[0]    remaining = self.max_num_batched_tokens - num_batched_tokensif remaining == 0:breakif not seq.block_table:# 新序列：检查显存并算 prefix cache 命中        num_cached_blocks = self.block_manager.can_allocate(seq)if num_cached_blocks == -1:break        num_tokens = seq.num_tokens - num_cached_blocks * self.block_sizeelse:# 继续未完成的分块 prefill        num_tokens = seq.num_tokens - seq.num_cached_tokensif remaining < num_tokens and scheduled_seqs:break# 只对队首序列允许切块if not seq.block_table:self.block_manager.allocate(seq, num_cached_blocks)    seq.num_scheduled_tokens = min(num_tokens, remaining)    num_batched_tokens += seq.num_scheduled_tokensif seq.num_cached_tokens + seq.num_scheduled_tokens == seq.num_tokens:        seq.status = SequenceStatus.RUNNING    # 整段 prompt 处理完才转 RUNNINGself.waiting.popleft()self.running.append(seq)    scheduled_seqs.append(seq)if scheduled_seqs:return scheduled_seqs, True

三个约束决定能调度多少：

• max_num_seqs：单 batch 最多多少序列；
• max_num_batched_tokens：单 batch 最多处理多少 token——太大吞吐高但 TTFT 变差；
• block_manager.can_allocate()：KV 块是否还够（不够就是显存耗尽）。

这段代码的关键是 seq.num_scheduled_tokens = min(num_tokens, remaining)：当 token 预算装不下整个 prompt 时，队首序列会被切成预算内的一段先跑，剩余部分留在 waiting 队首等后续 step 继续——这就是 chunked prefill。它把「一个长 prompt 独占一个 step」变成「长 prompt 分段、与其他请求共享 step」，直接改善 1.1 节说的 TTFT 与 GPU 利用率。只有累计处理量达到 seq.num_tokens 时，序列才真正进入 RUNNING。

一个具体例子。设 max_num_batched_tokens = 6，waiting 队首是长度 10 的 seq_x：

• 第 1 个 step：remaining = 6 < num_tokens = 10，但本轮还没调度任何序列（scheduled_seqs 为空），允许切块——num_scheduled_tokens = 6，本步只算前 6 个 token。postprocess 把游标推进到 num_cached_tokens = 6，因 6 < 10 不产出 token，序列仍留在 waiting 队首。
• 第 2 个 step：seq.block_table 非空，走续跑分支：num_tokens = 10 − 6 = 4，num_scheduled_tokens = 4。算完剩余 4 个 token，累计达 10 = 全长，转 RUNNING、移入 running，本步正常采样出第一个生成 token。

两个细节：其一，切块只切计算、不切显存——第一次调度时 allocate 已按全长把 KV 块全部分配好，分块只决定每步算多少；其二，只有每轮第一个被调度的序列允许切块——remaining < num_tokens and scheduled_seqs 时直接 break，后续序列要么整段进、要么不进，避免一个 batch 里堆多段碎块。

2.5 BlockManager：链式哈希、引用计数与懒惰失效

can_allocate：一次遍历同时回答两个问题

can_allocate（block_manager.py:58-73）回答「能不能分配」，顺便算出「prefix cache 能命中几块」：

def can_allocate(self, seq: Sequence) -> int:    h = -1    num_cached_blocks = 0    num_new_blocks = seq.num_blocksfor i in range(seq.num_blocks - 1):      # 注意：不含最后一块        token_ids = seq.block(i)        h = self.compute_hash(token_ids, h)        block_id = self.hash_to_block_id.get(h, -1)if block_id == -1 or self.blocks[block_id].token_ids != token_ids:break# 未命中（或哈希碰撞），后续块不再可能命中        num_cached_blocks += 1if block_id in self.used_block_ids:            num_new_blocks -= 1# 命中块正被占用，复用即可，不占新块if len(self.free_block_ids) < num_new_blocks:return -1return num_cached_blocks

细节：

• 不检查最后一块：最后一块通常没填满，内容还会变，不可复用，所以循环只到 num_blocks - 1。
• 哈希链：compute_hash(token_ids, prefix_hash) 用 xxhash（block_manager.py:35-41）——先把前一块的哈希值按 8 字节小端写入，再写入本块 token 字节。首块的 prefix 为 -1 不写。
• 碰撞防御：哈希命中后还要逐 token 比对 blocks[block_id].token_ids != token_ids，防止 64 位哈希撞车导致错误复用。
• 为什么只有在用的命中块才减 num_new_blocks：命中块正被其他序列持有（used_block_ids）时，复用它不消耗任何空闲名额，故扣减；命中块在空闲池里（已释放但哈希还在，见下文懒惰失效）时，把它捞出来仍要从 free_block_ids 拿走一个名额，故不扣减。最后 len(free_block_ids) < num_new_blocks 判断的就是「扣完免费复用后，还差几个真块」。

为什么必须链式哈希

如果不把前缀编进哈希，会出现「块内容相同但上下文不同」的误命中：

seq_c = [1, 2, 3, 4, 8, 9, 10, 7]    # "Plan a Seoul Trip" + "Include traditional markets."seq_d = [11, 12, 13, 14, 8, 9, 10, 7] # "Plan a Tokyo Trip" + "Include traditional markets."

两个序列的第二块内容完全相同（[8, 9, 10, 7]），但 KV cache 是带上下文的——同样的文字跟在「Seoul」和「Tokyo」后面，attention 算出的 K/V 完全不同。有一个用 Qwen3-0.6B tokenizer 做的最小实验^[7]验证了链式哈希的必要性：两个请求的首块相同，两种哈希结果一致；第三块文本同为 ' traditional markets.'，block-only 哈希相同（11326422787819615817），而 prefix-aware 哈希不同（9041253113613273947 vs 14881714884600895953）——只有前缀与内容都一致才允许命中。

allocate 与引用计数：prefix cache 的共享

allocate(seq, num_cached_blocks)（block_manager.py:75-92）做实际分配：前 num_cached_blocks 块走复用路径，其余从空闲池弹新块：

for i in range(num_cached_blocks):    ...    block_id = self.hash_to_block_id[h]    block = self.blocks[block_id]if block_id in self.used_block_ids:        block.ref_count += 1# 正被其他序列持有：引用 +1else:        block.ref_count = 1# 在空闲池里但哈希还在：捞起复用self.free_block_ids.remove(block_id)self.used_block_ids.add(block_id)    seq.block_table.append(block_id)for i in range(num_cached_blocks, seq.num_blocks):    seq.block_table.append(self._allocate_block())seq.num_cached_tokens = num_cached_blocks * self.block_size

共享机制用块大小 4 的例子画出来（该机制的可视化参考了 Inside vLLM 一文^[8]）：

块大小 = 4。seq_a 先跑（prompt 两块），seq_b 与 seq_a 共享首块。  ① seq_a = [1,2,3,4 | 5,6,7] 调度并跑完 prefill：     调度时空闲池弹出物理块 7、2 → seq_a.block_table = [7, 2]     前向结束后的 postprocess 才把首块哈希注册进表（见 2.7）     block[7]: hash=H([1,2,3,4]) 已注册, ref=1     block[2]: 未注册（最后一块不满，不入哈希表）  ② seq_b = [1,2,3,4 | 8,9,10] 到来，can_allocate 命中块 7     seq_b.block_table = [7, 9]     block[7].ref: 1 → 2          ← 共享前缀，KV 不重算、不复制  ③ seq_a 结束 deallocate：     block[2].ref = 0 → 归还空闲池     block[7].ref: 2 → 1          ← 仍被 seq_b 持有，不释放  ④ seq_b 也结束：block[7].ref = 0 → 归还空闲池，     但它的 hash 仍留在 hash_to_block_id 里。     下一个以 [1,2,3,4] 开头的序列依然能命中它——     直到该块被 _allocate_block 弹给别的内容时才清掉旧哈希。

第 ④ 步是容易忽略的精妙处：释放不等于失效（懒惰失效）。_deallocate_block 只把块还回空闲池，不动哈希表（block_manager.py:53-56）；_allocate_block 弹出空闲块时，若发现它带着旧哈希，才从 hash_to_block_id 删除（block_manager.py:47-48）。于是「暂时空闲但内容还有效」的块可以被后续请求零成本命中——对系统 prompt、few-shot 前缀这类高复用场景，这是白捡的命中率和显存。

2.6 decode 调度与抢占

waiting 排空后进入 decode 调度（scheduler.py:57-73）：

while self.running and num_seqs < self.max_num_seqs:    seq = self.running.popleft()while not self.block_manager.can_append(seq):if self.running:self.preempt(self.running.pop())else:self.preempt(seq)breakelse:        num_seqs += 1self.block_manager.may_append(seq)        scheduled_seqs.append(seq)assert scheduled_seqsself.running.extendleft(reversed(scheduled_seqs))return scheduled_seqs, False

decode 不再受 max_num_batched_tokens 约束——每序列每步只输入一个 token，token 预算没有风险，唯一的稀缺资源是 KV 块。

can_append 只有一行（block_manager.py:103-104）：

def can_append(self, seq: Sequence) -> bool:return len(self.free_block_ids) >= (len(seq) % self.block_size == 1)

len(seq) % block_size == 1 是布尔值：新 token 落在块首时（需要 1 个新块）为 True 即 1，否则为 False 即 0。may_append 同样只在这个时刻才真正 _allocate_block()——decode 不是每步都要新块，只有写满一块才追加。

抢占：当 can_append 失败（没有空闲块）且 running 队列还有别的序列，就驱逐队尾——FCFS 下最晚到达、优先级最低的序列：

def preempt(self, seq: Sequence):    seq.status = SequenceStatus.WAITINGself.block_manager.deallocate(seq)self.waiting.appendleft(seq)

注意用 appendleft 放到 waiting 队首：被驱逐的序列获得最高重调度优先级，保证公平。驱逐会 deallocate 释放它占用的全部块，换回显存给前面的序列续命。

2.7 postprocess：把新块写回 prefix cache

前向返回采样 token 后，postprocess 更新状态（scheduler.py:81-92）：

def postprocess(self, seqs: list[Sequence], token_ids: list[int], is_prefill: bool):for seq, token_id in zip(seqs, token_ids):self.block_manager.hash_blocks(seq)                    # ① 新填满的块注册进哈希表        seq.num_cached_tokens += seq.num_scheduled_tokens      # ② 推进已缓存游标        seq.num_scheduled_tokens = 0if is_prefill and seq.num_cached_tokens < seq.num_tokens:continue# ③ 分块 prefill 未收尾，不产生 token        seq.append_token(token_id)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)

三处都值得停顿：

1. hash_blocks 是 prefix caching 生效的另一半（block_manager.py:110-120）。can_allocate 只是「查」哈希表；真正把新填满的块「写」进哈希表的是这里——取上一个已哈希块的哈希作为前缀，对本步新填满的每块做链式哈希并注册。没有这一步，hash_to_block_id 永远是空的。
2. 游标推进：num_cached_tokens += num_scheduled_tokens，分块 prefill 的进度就是这样跨 step 累积的。
3. 中途不产出：分块 prefill 未跑完时直接 continue——中间段的 logits 不采样、不判 EOS，只有整段 prompt 处理完的那一步才 append token。这就是为什么 step() 要把 is_prefill 传进来。

2.8 ModelRunner：初始化

ModelRunner 是 step() 三部曲的第二环，负责把调度好的 batch 真正跑在 GPU 上。这是代码最密集的部分：分布式进程管理、模型初始化、KV 显存、输入准备、FlashAttention、CUDA Graph 都在这。

先记住拓扑：N 张卡就有 N 个 ModelRunner 实例，rank 0 在主进程，rank 1 到 N-1 是 worker。每个实例的初始化（model_runner.py:18-59）依次是：

1. NCCL 建组。 所有进程在 "tcp://localhost:2333" 汇合：

dist.init_process_group("nccl", "tcp://localhost:2333", world_size=self.world_size, rank=rank)

world_size = 进程数 = GPU 数。NCCL 是 NVIDIA 的多卡通信后端，后续 TP 的 all_reduce 就走它。

2. 设 device / dtype / 建模。

torch.cuda.set_device(rank)torch.set_default_dtype(hf_config.dtype)torch.set_default_device("cuda")self.model = Qwen3ForCausalLM(hf_config)

在默认 device 为 cuda 的状态下建模，权重直接落在 GPU 上。建完后又把默认 device/dtype 复位回 cpu 和原 dtype（model_runner.py:38-39）——避免主进程后续无关操作误占 GPU。

3. 加载权重。load_model(self.model, config.model) 把 safetensors 读进模型。融合权重（q/k/v → qkv_proj）与 TP 分片都在这一步完成，细节见 3.2 节。随后创建 Sampler。

4. 预热。 用「最坏情况」的假 batch 跑一次 prefill（model_runner.py:91-101）：

def warmup_model(self):    torch.cuda.empty_cache()    torch.cuda.reset_peak_memory_stats()    seq_len = min(self.config.max_num_batched_tokens, self.config.max_model_len)    num_seqs = min(self.config.max_num_batched_tokens // seq_len, self.config.max_num_seqs)    seqs = [Sequence([0] * seq_len) for _ in range(num_seqs)]for seq in seqs:        seq.num_scheduled_tokens = seq_lenself.run(seqs, True)    torch.cuda.empty_cache()

预热的两个目的：一是让 CUDA context、Triton kernel 编译等一次性开销在启动时付掉；二是量出激活显存的峰值——reset_peak_memory_stats 后跑一次最大负载，峰值统计就是下一步显存测算的输入。注意预热序列也设了 num_scheduled_tokens，与分块调度语义保持一致。

5. 测算并分配 KV cache。allocate_kv_cache() 决定整个块池的大小，是引擎侧最重要的一段算术，单独在 2.9 节展开。

6. CUDA Graph 捕获。enforce_eager=False 时执行 capture_cudagraph()，见 2.12 节。

7. 多卡时主从分叉。 至此各 rank 初始化相同，从这里开始分流（model_runner.py:51-59）：

if self.world_size > 1:if rank == 0:self.shm = SharedMemory(name="nanovllm", create=True, size=2**20)        dist.barrier()else:        dist.barrier()self.shm = SharedMemory(name="nanovllm")self.loop()

rank 0 创建 1 MB 共享内存后过 barrier；worker 过 barrier 后连上共享内存，进入 loop() 等指令。通信协议（model_runner.py:61-89）很朴素：

• 共享内存布局：前 4 字节小端序存数据长度 n，其后 n 字节是 pickle 负载；
• 主进程 write_shm()：pickle.dumps((method_name, *args)) 写入，然后 event.set() 通知所有 worker；
• worker loop()：event.wait() 阻塞，醒来后 read_shm() 反序列化出方法名和参数，执行同名方法；
• 主进程的 call() 先广播、再在本地执行同一方法，保证各 rank 步调一致。

每步 step() 里的 self.model_runner.call("run", seqs, is_prefill) 走的就是这个通道——seqs 会被 pickle 进共享内存（这里埋着一个为 decode 优化的细节，见 3.6 节）。

2.9 显存测算：能造多少个 KV 块

KV cache 块池多大，直接决定 batch 能做多大。通用公式（vLLM 文档^[10]）：

可用于 KV cache 的显存 ≈ 总 VRAM × gpu_memory_utilization − 权重 − 激活峰值 − 各类开销

一块显存的预算大致被五家分：模型权重、KV cache、激活显存（临时的，单步内生灭，经验上约为权重的 10–30%）、CUDA/PyTorch 系统开销、非 torch 分配（如 CUDA Graph 的显存池）。以 H100 80 GB 跑 Qwen3-32B（bf16）为例^[2]：

H100 80 GB，Qwen3-32B（bf16）：┌─────────────────────────────────────┬──────────────────────────┐│  模型权重 ≈ 64 GB                    │  余量 ≈ 16 GB            ││  32B 参数 × 2 字节                   │  KV cache + 激活 + 开销  │└─────────────────────────────────────┴──────────────────────────┘每 token 的 KV 占用（GQA：8 个 KV 头 × 128 维，64 层，bf16）：  2(K/V) × 64(层) × 8 × 128 × 2 字节 = 262,144 字节 ≈ 256 KB/token  → 2048 token 的上下文 ≈ 0.54 GB——2K 上下文下显存大头仍是权重，不是 KV

对照 nano-vLLM 默认的 Qwen3-0.6B（28 层、8 KV 头、head_dim 128、bf16）：每 token KV ≈ 2 × 28 × 8 × 128 × 2 = 114,688 字节 ≈ 112 KB；块大小 256 时单块约 28 MB。

代码实现（model_runner.py:103-121）：

def allocate_kv_cache(self):    free, total = torch.cuda.mem_get_info()    used = total - free    peak = torch.cuda.memory_stats()["allocated_bytes.all.peak"]    current = torch.cuda.memory_stats()["allocated_bytes.all.current"]    num_kv_heads = hf_config.num_key_value_heads // self.world_size    head_dim = getattr(hf_config, "head_dim", hf_config.hidden_size // hf_config.num_attention_heads)    block_bytes = 2 * hf_config.num_hidden_layers * self.block_size * num_kv_heads * head_dim * hf_config.dtype.itemsize    config.num_kvcache_blocks = int(total * config.gpu_memory_utilization - used - peak + current) // block_bytesassert config.num_kvcache_blocks > 0self.kv_cache = torch.empty(2, hf_config.num_hidden_layers, config.num_kvcache_blocks,self.block_size, num_kv_heads, head_dim)

三个显存读数来自不同视角，别混：

• used：CUDA 驱动视角的当前占用——权重、激活、其他进程、运行时开销全在里面；
• current：PyTorch 分配器视角的当前分配——主要是权重；
• peak：PyTorch 分配器见过的峰值——预热那次最坏前向留下的「权重 + 激活峰值」。

所以 − used − peak + current 的含义是：从目标额度（total × 利用率）里扣掉驱动视角的全部占用（used），再额外扣掉激活峰值——peak − current 恰好是「权重之外、前向高峰时还要多要的激活/工作区」。这一步依赖 2.8 节第 4 步的预热先跑出真实峰值。

两点 TP 相关的注意：num_kv_heads 要除以 world_size（每卡只持有自己那几片头的 KV）；GQA 下 KV 头数本来就少于 query 头数，这正是 GQA 省 KV 显存的地方。

最后把大张量切片挂到各层：

layer_id = 0for 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

Attention 初始化时 k_cache/v_cache 只是空占位张量（attention.py:57），这里才换成 KV 池的真实切片。

2.10 输入准备：prefill 与 decode

初始化完成后，看一次前向的入口（model_runner.py:160-168）：

def run(self, seqs: list[Sequence], is_prefill: bool) -> list[int]:    input_ids, positions = self.prepare_prefill(seqs) if is_prefill else self.prepare_decode(seqs)    temperatures = self.prepare_sample(seqs) if self.rank == 0 else None    logits = self.run_model(input_ids, positions, is_prefill)    token_ids = self.sampler(logits, temperatures).tolist() if self.rank == 0 else None    reset_context()return token_ids

只有 rank 0 采样（温度也只在 rank 0 准备）；worker 跑到 logits 为止，结果靠 TP 通信汇聚。

prepare_prefill 要把 Sequence 列表拍平成 FlashAttention 需要的变长布局。核心循环：

for seq in seqs:    start = seq.num_cached_tokens    seqlen_q = seq.num_scheduled_tokens    end = start + seqlen_q    seqlen_k = end    input_ids.extend(seq[start:end])    positions.extend(range(start, end))    cu_seqlens_q.append(cu_seqlens_q[-1] + seqlen_q)    cu_seqlens_k.append(cu_seqlens_k[-1] + seqlen_k)

四个局部量：start 是仍需计算的首个位置（有 prefix cache 或分块续跑时非 0）；seqlen_q 是本次算的 query 数（调度的那段）；end = start + seqlen_q；seqlen_k = end 是 attention 可见的完整上下文长度。prefix cache 命中时 seqlen_q < seqlen_k——query 只算未缓存后缀，key/value 覆盖全部历史。

input_ids/positions 描述「算什么」，cu_seqlens_q/cu_seqlens_k 描述「边界在哪」。剩余部分用 seq.block_table 把逻辑块翻译成 slot_mapping（新 K/V 写哪），并在有缓存可复用时准备 block_tables（旧 K/V 从哪读），最后 set_context(...) 打包。

走一个具体例子。设 seq_a = [1,2,3,4,5,6,7]（"Write a travel plan for Seoul."）已处理过，首块 [1,2,3,4] 已入缓存（尾块不满不入）。新的 prefill batch 有两个序列：

seq_c = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]   # 与 seq_a 共享首块seq_e = [21, 22, 23, 24, 25]             # 全新 prompt

seq_c: num_cached_tokens = 4, num_scheduled_tokens = 6       → start=4, seqlen_q=6, end=10, seqlen_k=10       → input_ids += [5,6,7,8,9,10], positions += [4,5,6,7,8,9]seq_e: num_cached_tokens = 0, num_scheduled_tokens = 5       → input_ids += [21,22,23,24,25], positions += [0,1,2,3,4]

拍平后的 batch：

input_ids    = [5, 6, 7, 8, 9, 10,   21, 22, 23, 24, 25]positions    = [4, 5, 6, 7, 8, 9,     0,  1,  2,  3,  4]cu_seqlens_q = [0, 6, 11]        # query 行边界cu_seqlens_k = [0, 10, 15]       # key/value 上下文边界

这里值得点破一处统一：prefix-cache 续跑与 chunked prefill 续跑是同一条代码路径。上例 seq_c 的 start = 4 来自 prefix cache 命中；换成 2.4 节那个「第一步只跑了 6 个 token」的长 prompt，第二步的 start 同样等于推进后的 num_cached_tokens，seqlen_k > seqlen_q 同样触发 block_tables 的准备，让 FlashAttention 读上一步写入的 KV。调度侧把两种情况统一成「已缓存游标 + 本次调度量」，算子侧统一成「query 只算后缀、K/V 覆盖全史」——prefix caching 和 chunked prefill 共享同一条前向路径，谁也不需要特判。

prepare_decode 简单得多：每序列恰贡献一个 token（model_runner.py:172-188）：

for seq in seqs:    input_ids.append(seq.last_token)    positions.append(len(seq) - 1)    context_lens.append(len(seq))    slot_mapping.append(seq.block_table[-1] * self.block_size + seq.last_block_num_tokens - 1)

每序列一行，无需 cu_seqlens。slot 从最后一个块的已用量推算：block_table[-1] * block_size + last_block_num_tokens - 1。注意 decode 总是准备 block_tables——读历史 KV 必须按页表索引，这与 prefix cache 无关。

回到玩具例子：seq_a 刚 prefill 完并采样出 token 8，则下一步 decode 的输入是 input_ids = [8]、positions = [7]，新 KV 写入 slot 7。

2.11 Attention：写 KV 的 Triton kernel 与 FlashAttention

前向本身是标准 Transformer 流水线：input_ids → VocabParallelEmbedding → [Attention → RMSNorm → MLP] × N → 最终 RMSNorm → ParallelLMHead → Sampler（Transformer 内部已有太多优秀资料^[9]，本文只盯引擎相关的部分；nano-vLLM 的 Qwen3 实现有几处非常规设计，见 3.1 节）。

Attention 层收到的输入形状：

q: [num_tokens, num_q_heads, head_dim]k: [num_tokens, num_kv_heads, head_dim]v: [num_tokens, num_kv_heads, head_dim]

第一件事是把本步新算的 K/V 写进持久池（attention.py:64-65）：

if k_cache.numel() and v_cache.numel():    store_kvcache(k, v, k_cache, v_cache, context.slot_mapping)

numel() 检查是为了跳过预热（那时 cache 还没分配）。store_kvcache 启动一个 Triton kernel，每个 token 一个 program，把该 token 的 K、V 向量拷到 slot_mapping 指定的物理 slot：

第一次读 Triton kernel（细节见 0.5）：这段代码不是普通 Python 函数——它在 GPU 上同时跑 N 份（N = token 数）。idx = tl.program_id(0) 问「我是第几份」，相当于手写并行循环时的循环变量；tl.arange(0, D) 生成 [0, 1, ..., D-1] 的下标向量；tl.load / tl.store 按「指针 + 下标」一次读写一批显存。整个 kernel 读作：第 idx 份 program 查到自己的槽位号 slot，把本 token 的 K、V 向量整体搬过去。

@triton.jitdef store_kvcache_kernel(key_ptr, key_stride, value_ptr, value_stride,                         k_cache_ptr, v_cache_ptr, slot_mapping_ptr, D: tl.constexpr):    idx = tl.program_id(0)    slot = tl.load(slot_mapping_ptr + idx)if slot == -1: return# CUDA Graph 填充的空位    key_offsets = idx * key_stride + tl.arange(0, D)    value_offsets = idx * value_stride + tl.arange(0, D)    key = tl.load(key_ptr + key_offsets)    value = tl.load(value_ptr + value_offsets)    cache_offsets = slot * D + tl.arange(0, D)    tl.store(k_cache_ptr + cache_offsets, key)    tl.store(v_cache_ptr + cache_offsets, value)

其中 D = num_kv_heads * head_dim 是单 token K/V 向量的扁平长度；key.stride(0) 是相邻 token 行之间的显存跨度。slot == -1 直接返回——那是 CUDA Graph 按固定 batch 重放时多出来的填充位（见 2.12）。

写完 KV 才做注意力，prefill 与 decode 走不同的 FlashAttention 入口：

if context.is_prefill:if context.block_tables is not None:    # 有 prefix cache 可复用        k, v = k_cache, v_cache    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)else:    # decode    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)

prefill 命中 prefix cache 时 k, v = k_cache, v_cache——直接把整块持久池当 K/V 传入，配合 block_table 让 FlashAttention 按页表读非连续物理块。这一行就是 PagedAttention 与 prefix cache 在算子层的接合点：引擎侧所有的块管理，最终都收敛成 attention kernel 手里的一张页表。

2.12 CUDA Graph：按 batch size 预捕获

run_model 里 prefill 永远走 eager，decode 可能走 CUDA Graph（model_runner.py:195-212）：

if is_prefill or self.enforce_eager or input_ids.size(0) > 512:return self.model.compute_logits(self.model(input_ids, positions))else:    ...    graph.replay()return self.model.compute_logits(graph_vars["outputs"][:bs])

两种模式的差别：

Eager：每次前向，CPU 逐个发射 kernel  CPU ──launch──launch──launch── ... ──launch──►  GPU  [emb]  [attn]  [linear]  ...  [lm_head]     ← kernel 之间是空泡CUDA Graph：先记录一次，之后整图重放  捕获: 在 torch.cuda.graph 里跑一次 → 录下完整 kernel 序列  重放: graph.replay()               → 一次调用，GPU 不间断跑完整段

decode 每步只产一个 token，计算量小，CPU 发射 kernel 的开销在占比上很可观；而 decode 的计算图形状稳定（每序列一个 token），正适合录下来重放。prefill 的 prompt 长度、总 token 数、attention 形状高度可变，没法高效捕获，所以 nano-vLLM 只对 decode 用图，且 batch 超过 512 也退回 eager。

捕获的实现（model_runner.py:222-257）比「录一次」精细：

• 按一组 batch size 各录一张图：graph_bs = [1, 2, 4, 8] + list(range(16, max_bs + 1, 16))，其中 max_bs = min(max_num_seqs, 512)。因为图的输入形状必须固定，只能为若干档 batch size 分别捕获。
• 共享显存池：首张图捕获后取 graph.pool() 作为 graph_pool，后续图都在这个池里录（torch.cuda.graph(graph, self.graph_pool)）——多张图复用同一份中间激活显存，否则十几张图各要一份工作区，显存吃不消。
• 重放时向上取整选图：next(x for x in self.graph_bs if x >= bs) 选不小于实际 batch 的图，把真实输入拷进预分配张量，多余空位的 slot_mapping 填 -1——正对应 2.11 节 kernel 里的 if slot == -1: return，填充位不写 KV。

三、模型与算子层

引擎骨架（调度、块管理、ModelRunner）之外，模型侧还有几处实现值得读：Qwen3 的两个非常规设计、张量并行 Linear 的切分与通信、RoPE 的预计算、采样器的 Gumbel-max 技巧，以及一个藏在跨进程传输里的瘦身优化。

3.1 Qwen3：QK-Norm 与融合残差

nano-vLLM 的模型实现是 Qwen3（nanovllm/models/qwen3.py）。相比教科书式 Transformer，它有三处值得注意。

(1) QK-Norm。 RoPE 之前对 q、k 各做一次 RMSNorm（qwen3.py:82-84）：

if not self.qkv_bias:    q = self.q_norm(q)    k = self.k_norm(k)q, k = self.rotary_emb(positions, q, k)

这是 Qwen3 相对 Qwen2 的主要注意力改动：归一化 query/key 的尺度，稳定注意力 logit。q_norm/k_norm 是作用在 head_dim 上的 RMSNorm（qwen3.py:68-70）。

(2) 残差与 RMSNorm 融合。Qwen3DecoderLayer.forward 的签名多带一个 residual（qwen3.py:146-159）：

def forward(self, positions, hidden_states, residual):if residual is None:        hidden_states, residual = self.input_layernorm(hidden_states), hidden_stateselse:        hidden_states, residual = self.input_layernorm(hidden_states, residual)    hidden_states = self.self_attn(positions, hidden_states)    hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)    hidden_states = self.mlp(hidden_states)return hidden_states, residual

残差不是「子层出口处 x = x + sublayer(x)」，而是作为显式状态在层间传递，在下一个 RMSNorm 里融合：add_rms_forward（layernorm.py:28-40）先 x.add_(residual) 再归一化，加法与归一化合进同一个 @torch.compile kernel，省一次显存往返。首层 residual is None，由 input_layernorm 初始化。

(3) 权重融合与词表共享。packed_modules_mapping（qwen3.py:187-193）声明了 HF 分存权重到推理融合权重的映射：

packed_modules_mapping = {"q_proj": ("qkv_proj", "q"),"k_proj": ("qkv_proj", "k"),"v_proj": ("qkv_proj", "v"),"gate_proj": ("gate_up_proj", 0),"up_proj": ("gate_up_proj", 1),}

q/k/v 融合成一个 qkv_proj（一次矩阵乘出 QKV），gate/up 融合成 gate_up_proj。另外 tie_word_embeddings 时 lm_head 直接复用 embed_tokens 的权重（qwen3.py:202-203），省一份词表大小的显存。

3.2 张量并行 Linear：列切、行切与 all_reduce 时机

TP 的核心在 nanovllm/layers/linear.py 的四种 Linear。理解它就理解了「为什么 attention 之后要 all_reduce」。

两种切分方向：

• ColumnParallelLinear（列并行）：按输出维切。权重 [out, in] 沿 out 切成 tp_size 份，每 rank 持 [out/tp, in]。各 rank 独立算自己那部分输出，无需通信——x·W = [x·W₀, x·W₁, ...]，各 rank 算的是不同列。
• RowParallelLinear（行并行）：按输入维切。每 rank 持 [out, in/tp]，各 rank 算 xᵢ·Wᵢ，最终输出是这些部分和——所以出口处要 dist.all_reduce(y)（linear.py:152-156）。

一个 attention/MLP 块的典型组合：

hidden → QKVParallelLinear（列并行，无通信）       → attention       → o_proj = RowParallelLinear（行并行，all_reduce）hidden → MergedColumnParallelLinear（列并行，gate+up 融合，无通信）       → SiluAndMul       → down_proj = RowParallelLinear（行并行，all_reduce）

「列并行进、行并行出」，每层只在两个出口各通信一次。这是 Megatron-LM 的经典模式：把通信压到每层两次，而非每个 Linear 都通信。

QKVParallelLinear 的精巧处（linear.py:96-128）：GQA 下 q、k、v 头数不同，融合后的输出大小是 (total_num_heads + 2 × total_num_kv_heads) × head_size。加载时按 shard_id（"q"/"k"/"v"）算各自在融合张量里的偏移，再按 rank 切自己的那片。这就是 packed_modules_mapping 里三者同映射到 qkv_proj 却带不同 shard_id 的原因。

权重如何从 safetensors 进来。load_model（utils/loader.py:12-29）遍历所有 safetensors 文件的每个权重名：

def load_model(model: nn.Module, path: str):    packed_modules_mapping = getattr(model, "packed_modules_mapping", {})for file in glob(os.path.join(path, "*.safetensors")):with safe_open(file, "pt", "cpu") as f:for weight_name in f.keys():for k in packed_modules_mapping:if k in weight_name:                        v, shard_id = packed_modules_mapping[k]                        param_name = weight_name.replace(k, v)                        param = model.get_parameter(param_name)                        weight_loader = getattr(param, "weight_loader")                        weight_loader(param, f.get_tensor(weight_name), shard_id)breakelse:                    param = model.get_parameter(weight_name)                    weight_loader = getattr(param, "weight_loader", default_weight_loader)                    weight_loader(param, f.get_tensor(weight_name))

注意内层是 for...else：权重名命中某个 packed key（如 q_proj）时，重映射到融合参数名（qkv_proj）、带着 shard_id 走该参数的 weight_loader 并 break；都不命中时走 else 分支的 default_weight_loader（直接 copy_）。每个融合参数的 weight_loader 内部再按 rank 取自己的 TP 分片（loaded_weight.narrow(tp_dim, start, shard_size) 或 chunk(tp_size)[tp_rank]）——所以**「融合 + 张量并行分片」在加载这一步一次完成**，同一份权重文件，N 个 rank 各取所需、互不重叠。

3.3 VocabParallelEmbedding 与 ParallelLMHead

词表相关的两个层也做了并行化（nanovllm/layers/embed_head.py）。

VocabParallelEmbedding（embed_head.py:9-42）按词表维切，每 rank 持 vocab/tp 行。前向时落在本 rank 区间外的 token id 被 mask 掉再查表，最后 all_reduce 汇总——词表大到单卡放不下时也能分布存。

ParallelLMHead（embed_head.py:45-66）继承它，前向多两步：prefill 时只取每序列最后一个 token 的 hidden（last_indices = cu_seqlens_q[1:] - 1）算 logits——中间 token 的投影纯属浪费，跳过；TP 下各 rank 的部分 logits 用 dist.gather 汇到 rank 0，与 2.10 节「只有 rank 0 采样」衔接。

3.4 RoPE：预计算 cos/sin 表

nanovllm/layers/rotary_embedding.py 的旋转位置编码在构造时预计算一张表：

inv_freq = 1.0 / (base ** (torch.arange(0, rotary_dim, 2, dtype=torch.float) / rotary_dim))t = torch.arange(max_position_embeddings, dtype=torch.float)freqs = torch.einsum("i,j -> ij", t, inv_freq)cache = torch.cat((freqs.cos(), freqs.sin()), dim=-1).unsqueeze_(1)self.register_buffer("cos_sin_cache", cache, persistent=False)

[max_position, 1, 2 × rotary_dim] 的 cos/sin 表，persistent=False 不进 state_dict（可由配置重算，省存储）。前向按 positions 索引出对应行，对 q、k 做旋转（最后一维对半切成 x1/x2，算 y1 = x1·cos − x2·sin; y2 = x2·cos + x1·sin 再拼回）。

两个优化点：apply_rotary_emb 带 @torch.compile（rotary_embedding.py:37），索引加旋转融合成一个 kernel；get_rope 带 @lru_cache(1)（rotary_embedding.py:51），同配置只建一个实例，N 层共享同一张 cos/sin 表。

3.5 Sampler：Gumbel-max 批量采样

采样器只有十几行（nanovllm/layers/sampler.py），但藏着一个优雅的批量技巧：

class Sampler(nn.Module):    @torch.compiledef forward(self, logits: torch.Tensor, temperatures: torch.Tensor):        logits = logits.float().div_(temperatures.unsqueeze(dim=1))        probs = torch.softmax(logits, dim=-1)        sample_tokens = probs.div_(torch.empty_like(probs).exponential_(1).clamp_min_(1e-10)).argmax(dim=-1)return sample_tokens

逐行看：

1. logits.float().div_(temperatures.unsqueeze(1))：升 fp32，按各序列自己的 temperature 缩放（unsqueeze 把 [bs] 广播到 [bs, 1]）——一个 batch 里每个序列可以用不同温度。
2. softmax 得概率分布。
3. 最后一行是 Gumbel-max 采样的变体。原理：对每个类别算 log pᵢ + gᵢ（gᵢ 为标准 Gumbel 噪声）再取 argmax，等价于按 p 采样；而指数分布噪声经 -log 变换恰是 Gumbel。所以 probs.div_(exponential_(1)) 等价于 log p + g，argmax 即采样结果。clamp_min_(1e-10) 防 log(0)。

这行为什么等价于按概率抽样（不想看推导可跳过）：直觉是「给每个词打分再加噪，谁高选谁」。普通加噪做不到「选中概率恰好是 pᵢ」，但 Gumbel 噪声有一个数学上恰好成立的性质：log pᵢ + gᵢ 取最大值的那个词，被选中的概率正好等于 pᵢ——这就是 Gumbel-max 定理。代码里没直接采 Gumbel，而是采了指数分布噪声（exponential_(1)），因为「指数分布取 -log」正是 Gumbel；而除法取对数后变减法，probs / 噪声 取 argmax 与 log p − log 噪声 取 argmax 等价——后者正是 log p + g 的形式。一行除法 + argmax，就完成了整个 batch 的并行抽样。

为什么不用 torch.multinomial：Gumbel-max 是纯张量运算，可融合、可 @torch.compile，天然支持 batch 内不同温度的并行采样，不需要逐序列 Python 循环。

一个推论：SamplingParams 断言 temperature > 1e-10（sampling_params.py:10-11），禁止把 temperature 设 0 做 greedy——除法会炸。要确定性输出得走别的路径。

3.6 Sequence 的 pickle 瘦身

最后是最容易忽略、却很精巧的一处（sequence.py:72-84）。

多卡时主进程要把 Sequence 经共享内存 pickle 给 worker。默认 pickle 会带上整个 token_ids 列表；而 decode 阶段序列可能已有上千 token，worker 却只需要 last_token 和 block_table。于是：

def __getstate__(self):    last_state = self.last_token if not self.is_prefill else self.token_idsreturn (self.num_tokens, self.num_prompt_tokens, self.num_cached_tokens,self.num_scheduled_tokens, self.block_table, last_state)def __setstate__(self, state):    ...if isinstance(last_state, list):self.token_ids = last_stateself.last_token = self.token_ids[-1]else:self.token_ids = []            # decode：不传历史 tokenself.last_token = last_state   # 只恢复一个标量

decode 时 last_state 只装一个 int，token_ids 置空——worker 收到的 Sequence 不带历史 token，共享内存占用大幅压缩。安全性来自 2.10 节的分工：decode 前向只需要 last_token（构造输入）与 block_table（读历史 KV），历史 token 本身用不到。prefill 时才传完整 token_ids（要算整段 prompt）。

四、结语

从 example.py 一路走到采样器，一个 LLM 推理引擎的本质就是这样一个循环：

1. Scheduler 选出下一批要跑的序列；
2. BlockManager 保证它们的 KV 块就绪（能复用的复用，不够的抢占）；
3. ModelRunner 把变长、散落的序列打包成一次密集的 kernel 启动；
4. 跑模型，采样；
5. 追加 token、写回 prefix cache、更新状态，进入下一步。

剩下的一切——分页、链式哈希、引用计数、chunked prefill、CUDA Graph——都是让这个循环在显存约束下转得更快的工程细节。而这 1200 行代码的价值在于：这些细节一个不少，却都读得完。

参考与延伸阅读

本文的写作受益于 junuxyz 的《Inside nano-vLLM》一文及其引用资料；下列每条附上内容说明，可按需取读。

1. nano-vLLM 仓库（本文所有代码引用的出处）：https://github.com/GeeeekExplorer/nano-vLLM
2. junuxyz，LLM 推理引擎在优化什么（TTFT/TPOT/goodput 的定义，prefill 计算受限 vs decode 显存受限，H100 + Qwen3-32B 的显存拆解，本文 1.1、2.9 节素材来源）：https://github.com/junuxyz/mlsys-notes/blob/main/notes/llm-inference-intro-p1.md
3. Sebastian Raschka，"Coding the KV Cache in LLMs"（从零手写 KV cache，含代码）：https://magazine.sebastianraschka.com/p/coding-the-kv-cache-in-llms
4. Hugging Face 的 KV cache 博客（KV cache 机制与显存开销）：https://huggingface.co/blog/kv-cache
5. junuxyz 的 PagedAttention 笔记（分页思想的理论部分）：https://github.com/junuxyz/mlsys-notes/blob/main/notes/pagedattention.md
6. FIFO / FCFS 背景（维基百科）：https://en.wikipedia.org/wiki/FIFO_(computing_and_electronics)
7. prefix 链式哈希的最小实验（验证「同内容不同前缀 → 不同哈希」，本文 2.5 节的实验数据）：https://gist.github.com/junuxyz/89501be9327da5e137515874d4c5b8e1
8. Aleksa Gordic，"vLLM"（prefix cache 的块共享与引用计数、懒惰失效的可视化，本文 2.5 节配图参考）：https://www.aleksagordic.com/blog/vllm
9. The Annotated Transformer（Transformer 逐行实现）：https://nlp.seas.harvard.edu/annotated-transformer/
10. vLLM 文档：gpu_memory_utilization（显存五构成与测算公式，本文 2.9 节素材来源）：https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/gpu_memory_utilization/