跳到主要内容
推理优化

3.2 配置、入口与一次 generate 的全流程

从 example.py 出发，拆解 Config / SamplingParams / LLM / LLMEngine，看懂 generate 主循环如何驱动 add_request、step 与 postprocess

nano-vllm LLMEngine 推理循环 源码解读

📑 目录

1. 从一段示例代码进入

读引擎最好的入口，是看它「怎么被用」。仓库里的 example.py 浓缩了全部对外 API：

from nanovllm import LLM, SamplingParams

# 1) 构造引擎：传入本地模型路径与若干全局开关
llm = LLM("/your/model/path", enforce_eager=True, tensor_parallel_size=1)

# 2) 定义采样策略
sampling_params = SamplingParams(temperature=0.6, max_tokens=256)

# 3) 批量生成
prompts = ["Hello, nano-vllm."]
outputs = llm.generate(prompts, sampling_params)

# 4) 取结果：每个元素是 {"text": ..., "token_ids": ...}
print(outputs[0]["text"])

四步就跑通了一次推理。注意两个细节：

  • 构造 LLM 时第一个参数是本地目录路径（不是 HuggingFace 仓库名），因为 nano-vllm 直接用 safetensors 加载本地权重，不负责下载。
  • generate 返回的是 dict 列表（含 "text" 和 "token_ids"），这是它与 vLLM 接口的一个小差异——vLLM 返回的是 RequestOutput 对象。

⚠️ 注意：示例里 apply_chat_template 通常在业务脚本里手动调用。和第 1 章讲的一样，直接把裸文本喂给 Chat 模型会导致输出质量下降；分词与模板是「引擎之外」的职责。

接下来我们顺着这四步往里看：SamplingParams 和 Config 是什么，LLM 构造时发生了什么，generate 内部又是怎样一个循环。

2. Config：全局配置的”宪法”

config.py 里的 Config 是一个 @dataclass(slots=True)，它把引擎运行需要的全局约束集中在一处。核心字段与默认值如下（示意，省略类型别名）：

@dataclass(slots=True)
class Config:
    model: str                              # 本地模型目录（必填）
    max_num_batched_tokens: int = 16384     # 单步最多处理的 token 总数（Token Budget）
    max_num_seqs: int = 512                 # 单步最多并发的序列数
    max_model_len: int = 4096               # 支持的最大上下文长度
    gpu_memory_utilization: float = 0.9     # 允许占用的显存比例
    tensor_parallel_size: int = 1           # 张量并行 GPU 数
    enforce_eager: bool = False             # True 则禁用 CUDA Graph
    hf_config = None                        # 延迟加载的 HuggingFace 配置
    eos: int = -1                           # EOS token id（初始化后由 tokenizer 填入）
    kvcache_block_size: int = 256           # 每个 KV 块容纳的 token 数
    num_kvcache_blocks: int = -1            # KV 块总数（运行时按显存推算）

构造后 __post_init__ 会做三条硬校验，并把上下文长度对齐到模型能力：

def __post_init__(self):
    assert os.path.isdir(self.model)                 # 必须是存在的本地目录
    assert self.kvcache_block_size % 256 == 0        # 块大小需 256 的整数倍
    assert 1 <= self.tensor_parallel_size <= 8       # TP 度在 1~8 之间
    self.hf_config = AutoConfig.from_pretrained(self.model)
    self.max_model_len = min(self.max_model_len,
                             self.hf_config.max_position_embeddings)

几个值得停下来想一想的设计：

  • 🔑 kvcache_block_size = 256：这是「每个 KV 块装多少个 token」。块越大，块表（Block Table）越短、管理开销越小，但内部碎片（最后一块没填满的浪费）越大；块越小则相反。256 是在碎片与管理开销之间取的折中，% 256 == 0 的断言也和底层内核的对齐要求配套。

⏰ 信息时效：这里的 256 是 nano-vllm 当前默认值，和 vLLM 常见的 block_size=16 并不相同——不要把两者的默认值混为一谈。

  • 🔑 max_num_batched_tokens 与 max_num_seqs：前者限制「一步最多算多少 token」（主要约束 prefill），后者限制「一步最多有多少条序列」（主要约束 decode）。它们共同构成调度器的 Token Budget，是 3.5 的主角。
  • 🔑 gpu_memory_utilization = 0.9：引擎不会按最大长度预留 KV，而是先加载权重、warmup 测峰值，再把「显存配额 − 权重 − 峰值开销」全部切成 KV 块（num_kvcache_blocks 因此初始为 -1，运行时才算出）。这正是高显存利用率的来源，3.6 会看到具体公式。

3. SamplingParams：采样超参

相比 vLLM 动辄几十个采样参数，nano-vllm 的 SamplingParams 极其克制，只有三个字段：

@dataclass(slots=True)
class SamplingParams:
    temperature: float = 1.0
    max_tokens: int = 64
    ignore_eos: bool = False

    def __post_init__(self):
        assert self.temperature > 1e-10, "greedy sampling is not permitted"
  • temperature：温度，控制采样随机性。
  • max_tokens：最多生成多少个 token，默认 64（示例里显式设成了 256）。它既是停止条件，也用于 benchmark 时固定生成量。
  • ignore_eos：是否忽略 EOS 继续生成，压测时常设为 True 保证每条序列都跑满 max_tokens，吞吐数据才可比。

📌 关键点：__post_init__ 直接断言温度必须大于 1e-10——nano-vllm 不支持贪心采样（temperature=0）。这是一个有意的简化：它的 Sampler 走的是统一的随机采样路径（见 3.3），没有为 temperature=0 单独写一条 argmax 分支，所以干脆在入口处挡掉。想要确定性输出，就把温度设到一个极小的正数。

4. LLM 与 LLMEngine：薄封装背后的初始化

LLM 这个类本身几乎没有逻辑——它是 LLMEngine 的薄封装，把 generate 等方法直接暴露给用户。真正干活的是 LLMEngine.__init__，它在构造时完成了四件大事：

class LLMEngine:
    def __init__(self, model, **kwargs):
        # 1) 从 kwargs 里挑出 Config 认识的字段，构造全局配置
        config_fields = {f.name for f in fields(Config)}
        config = Config(model, **{k: v for k, v in kwargs.items()
                                   if k in config_fields})
        # 把块大小同步给 Sequence（它用类变量记 block_size）
        Sequence.block_size = config.kvcache_block_size

        # 2) 张量并行：用 spawn 启动 rank>0 的 worker 子进程
        ctx = mp.get_context("spawn")
        self.ps, self.events = [], []
        for rank in range(1, config.tensor_parallel_size):
            event = ctx.Event()
            p = ctx.Process(target=ModelRunner, args=(config, rank, event))
            p.start()
            self.ps.append(p); self.events.append(event)
        # rank 0 的 ModelRunner 跑在主进程
        self.model_runner = ModelRunner(config, 0, self.events)

        # 3) 分词器只在主进程持有，并把 EOS 回填到 config
        self.tokenizer = AutoTokenizer.from_pretrained(config.model, use_fast=True)
        config.eos = self.tokenizer.eos_token_id

        # 4) 调度器 + 退出清理
        self.scheduler = Scheduler(config)
        atexit.register(self.exit)

这里有几个工程上的讲究：

  • 🔑 为什么用 spawn 而不是默认的 fork：CUDA 上下文不能被 fork 安全继承（子进程会拿到一个已初始化但不可用的 CUDA 状态）。spawn 启动的是干净的新进程，各自初始化自己的 GPU。
  • 🔑 主从进程靠 Event 做栅栏：rank>0 的 ModelRunner 在子进程里循环等待任务，主进程通过共享内存下发方法名与参数，用 Event 通知「可以读了」。张量并行的协同细节放在 3.6。
  • 分词器只在主进程：detokenize 是 CPU 任务，没必要在每个 TP rank 都跑一份。
  • atexit 清理：退出时给子进程发 "exit"，join 回收，避免僵尸进程。

💡 提示：Sequence.block_size = config.kvcache_block_size 这一行容易被忽略——Sequence 用一个类变量记录块大小，所以所有序列共享同一个值，不必每个实例都存一份。

5. 主循环：add_request → step → postprocess

终于到了引擎的心脏。generate 的骨架可以提炼成下面这样（去掉了 tqdm 进度条与吞吐统计）：

def generate(self, prompts, sampling_params):
    # 广播：单个 SamplingParams 复制给每条 prompt
    for prompt, sp in zip(prompts, sampling_params):
        self.add_request(prompt, sp)

    outputs = {}
    while not self.is_finished():
        finished, num_tokens = self.step()   # 跑一步，拿回本步完成的序列
        for seq_id, token_ids in finished:
            outputs[seq_id] = token_ids

    # 按 seq_id 排序后解码成文本
    outputs = [outputs[i] for i in sorted(outputs)]
    return [{"text": self.tokenizer.decode(ids), "token_ids": ids}
            for ids in outputs]

两个关键方法：

add_request：把一条 prompt 变成一个 Sequence 放进调度器。prompt 可以是字符串（内部先 encode）也可以是现成的 token id 列表：

def add_request(self, prompt, sampling_params):
    if isinstance(prompt, str):
        prompt = self.tokenizer.encode(prompt)
    self.scheduler.add(Sequence(prompt, sampling_params))

step：一次完整的「调度 → 前向 → 写回」：

def step(self):
    seqs, is_prefill = self.scheduler.schedule()          # 选 batch
    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

📌 关键点：注意 num_tokens 的符号编码——prefill 时返回正数（这一步处理的 token 总量），decode 时返回负的 batch 大小（-len(seqs)）。generate 就靠这个符号区分当前是 prefill 还是 decode 阶段，分别计算并显示 Prefill / Decode 两种吞吐（tok/s）。用一个数的正负同时编码「阶段」和「数量」，是个很省事的小技巧。

整个循环的分工一目了然：model_runner.call("run", ...) 负责「怎么算」，scheduler 负责「算谁、算完怎么办」。下面三节分别放大执行层（3.3）、KV 管理（3.4）和调度层（3.5），最后 3.6 回到 ModelRunner 把它们重新串起来。

📝 总结

  • example.py 展示了全部对外 API：LLM + SamplingParams + generate，返回含 "text" 的 dict 列表。
  • Config 是 @dataclass(slots=True)，集中管理显存比例、Token Budget、TP 度、KV 块大小，并在 __post_init__ 做硬校验、对齐上下文长度。
  • SamplingParams 只有三个字段，且断言温度 > 1e-10——不支持贪心采样。
  • LLMEngine.__init__ 用 spawn 启动 TP worker、只在主进程持分词器、注册 atexit 清理。
  • 主循环 generate → add_request → step → postprocess 中，step 用 num_tokens 的正负区分 prefill / decode 阶段。

🎯 自我检验清单

  • 能用 LLM 和 SamplingParams 写出一段最简推理代码，并取出 outputs[0]["text"]
  • 能说出 Config 的 max_num_batched_tokens、max_num_seqs、gpu_memory_utilization、kvcache_block_size 各自约束什么
  • 能解释 nano-vllm 为什么不支持贪心采样
  • 能说清 LLMEngine 为什么用 spawn 而非 fork 启动张量并行进程
  • 能复述 step() 的三步，并解释 num_tokens 符号编码的用途

📚 参考资料