跳到正文
LLM 推理
S01模型本身·137

生成循环

推理引擎就是一个把模型自己的输出再喂回去的 while 循环。这门课余下的全部内容,都是在优化这个循环。

  • 前向传播
  • 自回归
  • logits
  • EOS
  • prefill 与 decode

为什么重要

问题

在推理阶段,可以把语言模型看作一个函数:输入一串 token id,它为词表中的每个 token 返回一个分数。对现代模型而言,每行输出可能包含约 13 万个分数。模型不会把这些分数拼成句子,也不判断该在哪里停下;上一次调用发生过什么,它同样一无所知。

这些行为由推理引擎补上:调用模型、选择一个 token、把它追加到输入末尾,然后再次调用模型。 循环一直进行,直到模型生成序列结束标记,或请求达到预设上限。这就是自回归解码,本课程的每个引擎都从这个循环开始。

写出这个循环只需要几行代码,难点在于跑得快。直接实现会大量重复已经做过的计算,后续章节将逐个瓶颈拆解这些浪费。

核心思路

解法

引擎的核心逻辑可以浓缩成九行白板代码:

循环的本质python
def generate(model, tokenizer, prompt, max_tokens=64):
    ids = tokenizer.encode(prompt)

    for _ in range(max_tokens):
        logits = model(ids)          # (len(ids), vocab_size)
        next_logits = logits[-1]     # only the last position matters
        next_id = sample(next_logits)
        ids.append(next_id)

        if next_id == tokenizer.eos_id:
            break

    return tokenizer.decode(ids)

高亮的两处细节揭示了贯穿后续章节的两个约束。

模型为每个位置都返回 logits,而我们只留下一行,其余全部丢掉。 第一次调用时,多出来的行并不额外花钱:模型本来就要看完全部 N 个 prompt token,N 行输出出自同一次并行 matmul。之后每一次调用就都是纯粹的损失了,把早已处理过的 token 的表示重算一遍,只为拿到一行新的。S05 用 KV 缓存删掉这笔重复计算。

循环的长度依赖数据,事先无法得知。 请求跑完之前,你不知道它要迭代多少次。正因如此,LLM 请求的批处理才这么难;S10 里的调度器为此而生,连续批处理(S09)也才能比普通批处理快好几倍。

prefill 和 decode 是两种不同的负载

第一次前向一次处理 N 个 token。权重只加载一次,这笔开销被摊到大量算术上,此时 GPU 是算力瓶颈。之后每次前向只处理一个 token:模型里每一个权重都要流过一遍 ALU,换来的只是寥寥几次乘加,于是 GPU 变成显存带宽瓶颈,大部分时间在闲着。两个阶段对硬件的要求正好相反,本课程几乎所有技术都由此而来。

工作原理

工作原理

示意图自回归生成
一次迭代 = 一个 TOKENprompt"The cat"tokenizer文本 → id前向传播L × decoder blocklogits[vocab_size]采样器argmax / top-ptoken"sat"追加到上下文 —— 下一次前向就多看到一个 token若为 EOS解码 + 输出id → 文本3 次迭代后的上下文Thecatsatonthemat灰色 = prompt(一次性并行算完 —— 即 prefill)橙色 = 生成的(每次前向出一个 —— 即 decode)每一次前向都要重读整个上下文。这正是 S05 要消除的开销。
只有「前向传播」那个框是昂贵的,它左右两侧的一切都只花微秒级时间。决定 token/秒 的是绕回它的那条环路。

朴素版本为什么是二次方的

位置 i 的自注意力要看所有 ≤ i 的位置。因此对长度为 n 的上下文做一次前向,注意力操作数量大约是 。在没有缓存的情况下生成 T 个 token,意味着依次跑长度为 N、N+1、…、N+T 的前向,于是注意力总计算量为:

总代价text
no cache:   Σ n²   for n = N..N+T   →  O((N+T)³)
with cache: Σ n    for n = N..N+T   →  O((N+T)²)

prompt 有 2,000 个 token、随后生成 500 个 token 时,两种做法的算术量相差约 500 倍,接近三个数量级。生产引擎不会这么写。这份直接实现的价值在别处:它是一个简洁的参考版本,后续每项优化都可以拿它对拍。

用户真正能感觉到的三个数字

首 token 时间
TTFT
由 prefill 主导:prompt 长度 × 模型 FLOPs
token 间延迟
ITL
一次 decode:模型字节数 ÷ 显存带宽
吞吐
tok/s
所有并发请求的总和,不是单个请求的

吞吐和延迟互相拉扯,几乎全由 batch 大小说了算。batch 为 1 时 token 间延迟最好、吞吐最差,大 batch 则相反。S09S11 讲的就是如何同时守住两者。

动手观察

动手试试

模拟器跑的还是同一个自回归循环,只是把 Transformer 换成了词级 bigram 模型。每次前进一步,盯住两个变化:上下文每次增加一个 token;由于先前结果没有缓存,开销柱按二次方增长。

模拟器自回归解码
提示词
第 0 / 14 步
上下文(下一次前向时模型看到的内容)
thecat

第 1 次前向要把完整的 2 个 token 过一遍所有层。目前什么都没缓存。

下一个 token 的分布 · 14 个 logits
  • sat54%
  • ate14%
  • chased14%
  • slept14%
  • .0%

temperature 在 softmax 之前把 logits 除以温度值。为 0 时所有概率塌缩到一个 token;大于 1 时分布被抹平。

前向次数
1
输出 token
0
注意力计算量(无缓存)
4
有 KV 缓存
4
每次前向的开销

每根柱子是一次前向,高度 ∝ n²。最后一个 token 最贵,可它本该最便宜。

把 temperature 调到 0,循环变成确定性的贪心解码,相同的种子和 prompt 永远给出同一句话。调到 1.2 以上,分布被抹平,模型开始信马由缰。S04 讲的就是这一个旋钮。

亲手实现

实现

本章的可运行文件有三部分:一个微型模型、一个采样器,和这个循环。下面是循环本身,写成后续章节会继续扩展的形态。把它做成带显式状态的对象,后面的章节就能把缓存和调度器嵌进同一个结构里。

code/s01_generation_loop.py(节选)python
class Engine:
    """The smallest thing that deserves the name."""

    def __init__(self, model, tokenizer, max_context=512):
        self.model = model
        self.tok = tokenizer
        self.max_context = max_context

    def generate(self, prompt, max_tokens=64, temperature=1.0, seed=0):
        rng = np.random.default_rng(seed)
        ids = self.tok.encode(prompt)
        prompt_len = len(ids)
        stats = Stats(prompt_len=prompt_len)

        for step in range(max_tokens):
            # THE forward pass. Recomputes every position, every time.
            logits = self.model.forward(ids)      # (T, vocab)
            stats.record_pass(len(ids))

            next_id = sample(logits[-1], temperature, rng)
            ids.append(next_id)

            if next_id == self.tok.eos_id:
                stats.stopped = "eos"
                break
            if len(ids) >= self.max_context:
                stats.stopped = "context_full"
                break
        else:
            stats.stopped = "max_tokens"

        return self.tok.decode(ids[prompt_len:]), stats

这个形态里有三点值得留意,因为它们会在接下来十九章反复出现:

  • 停止是一种策略,不是一个事件。 EOS、token 预算、上下文上限、停止字符串、客户端断连,引擎每次迭代都要把它们全查一遍。真实引擎在这里会堆上十来个停止条件。
  • stats 对象有它的正经用途。 没有度量就没有优化,而每次前向的上下文长度是最值得记录的一个数字。后面每一章都会拿它做对比。
  • `logits[-1]` 就是浪费所在。 盯住这个下标。它在 S05 里会变成 logits[0],那时引擎已经快了一百倍。
在本地运行
在 bigram 模型上运行朴素循环,打印生成的文本,再打印一张开销对照表,把实际执行的注意力计算量与有缓存的引擎本该花的计算量并列。
$ python code/s01_generation_loop.py
预期输出: 一句简短的生成结果,随后是一张表,显示「无缓存 / 有缓存」的差距随每个 token 不断拉大。

只需要 NumPy — 查看环境准备.

生产实践

在生产环境中

在 vLLM 里,这个循环住在 EngineCore.step() 中,几乎认不出来。图里那五个框各自长成了一个子系统,但结构本身还在:

  • vLLM: EngineCore.step() 依次调用调度器、模型执行器、输出处理器。(在 V1 之前,同一个循环叫 LLMEngine.step(),多数早期文章用的是这个名字。)一次调用仍然是这个循环的一次迭代,只是同时服务很多请求。
  • llama.cpp / picoLM: 一个字面意义上的 C while 循环,调用 forward()sample(),再往 token 数组里追加。生产实现中没有比它更接近白板版本的了。
  • TensorRT-LLM: 循环本身被编译进了引擎图里,所以每一个 shape 都必须提前确定。

`for … else` 是在干正经事

Python 的 for…else 只在循环没有被 break 时执行,这恰好对应「撞到 max_tokens」的情况。写错了就是 finish reason 上典型的差一错误:OpenAI 风格 API 里的 finish_reason 字段,客户端要靠它判断是否继续。

练习

  1. 1
    加一个接受字符串列表的 stop 参数。停止字符串定义在文本上,你生成的却是 token,所以必须增量解码,还要处理跨越 token 边界的停止字符串。
  2. 2
    generate 改成生成器,边产出边 yield token。阻塞式 API 与流式 API 的全部区别就在这里,参见 S19
  3. 3
    给循环加上每次前向的墙钟计时。确认增长确实是上下文长度的二次方,再据此预测生成 4000 个 token 需要多久。

继续学习

接下来

这个循环调用 tokenizer.encodetokenizer.decode,仿佛它们既免费又显而易见。两者都不是。S02 会实现一个 byte-pair 编码器,说清楚流式解码为什么需要缓冲区,以及你数出来的 token 数为什么总和账单对不上。

自测

习题

先作答,再看解析。答错比答对更有价值,因为解析会指出你该回头重读哪一部分。

自测 1 题 / 共 6

在 decode 阶段(而非 prefill),GPU 为什么大部分时间是闲着的?

得分 0/6