生成循环
推理引擎就是一个把模型自己的输出再喂回去的 while 循环。这门课余下的全部内容,都是在优化这个循环。
- 前向传播
- 自回归
- logits
- EOS
- prefill 与 decode
为什么重要
问题
在推理阶段,可以把语言模型看作一个函数:输入一串 token id,它为词表中的每个 token 返回一个分数。对现代模型而言,每行输出可能包含约 13 万个分数。模型不会把这些分数拼成句子,也不判断该在哪里停下;上一次调用发生过什么,它同样一无所知。
这些行为由推理引擎补上:调用模型、选择一个 token、把它追加到输入末尾,然后再次调用模型。 循环一直进行,直到模型生成序列结束标记,或请求达到预设上限。这就是自回归解码,本课程的每个引擎都从这个循环开始。
写出这个循环只需要几行代码,难点在于跑得快。直接实现会大量重复已经做过的计算,后续章节将逐个瓶颈拆解这些浪费。
核心思路
解法
引擎的核心逻辑可以浓缩成九行白板代码:
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 缓存删掉这笔重复计算。
prefill 和 decode 是两种不同的负载
工作原理
工作原理
朴素版本为什么是二次方的
位置 i 的自注意力要看所有 ≤ i 的位置。因此对长度为 n 的上下文做一次前向,注意力操作数量大约是 n²。在没有缓存的情况下生成 T 个 token,意味着依次跑长度为 N、N+1、…、N+T 的前向,于是注意力总计算量为:
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
- 所有并发请求的总和,不是单个请求的
动手观察
动手试试
模拟器跑的还是同一个自回归循环,只是把 Transformer 换成了词级 bigram 模型。每次前进一步,盯住两个变化:上下文每次增加一个 token;由于先前结果没有缓存,开销柱按二次方增长。
第 1 次前向要把完整的 2 个 token 过一遍所有层。目前什么都没缓存。
- 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 讲的就是这一个旋钮。
亲手实现
实现
本章的可运行文件有三部分:一个微型模型、一个采样器,和这个循环。下面是循环本身,写成后续章节会继续扩展的形态。把它做成带显式状态的对象,后面的章节就能把缓存和调度器嵌进同一个结构里。
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],那时引擎已经快了一百倍。
$ python code/s01_generation_loop.py只需要 NumPy — 查看环境准备.
生产实践
在生产环境中
在 vLLM 里,这个循环住在 EngineCore.step() 中,几乎认不出来。图里那五个框各自长成了一个子系统,但结构本身还在:
- vLLM:
EngineCore.step()依次调用调度器、模型执行器、输出处理器。(在 V1 之前,同一个循环叫LLMEngine.step(),多数早期文章用的是这个名字。)一次调用仍然是这个循环的一次迭代,只是同时服务很多请求。 - llama.cpp / picoLM: 一个字面意义上的 C
while循环,调用forward()和sample(),再往 token 数组里追加。生产实现中没有比它更接近白板版本的了。 - TensorRT-LLM: 循环本身被编译进了引擎图里,所以每一个 shape 都必须提前确定。
`for … else` 是在干正经事
for…else 只在循环没有被 break 时执行,这恰好对应「撞到 max_tokens」的情况。写错了就是 finish reason 上典型的差一错误:OpenAI 风格 API 里的 finish_reason 字段,客户端要靠它判断是否继续。练习
- 1加一个接受字符串列表的
stop参数。停止字符串定义在文本上,你生成的却是 token,所以必须增量解码,还要处理跨越 token 边界的停止字符串。 - 2
- 3给循环加上每次前向的墙钟计时。确认增长确实是上下文长度的二次方,再据此预测生成 4000 个 token 需要多久。
继续学习
接下来
这个循环调用 tokenizer.encode 和 tokenizer.decode,仿佛它们既免费又显而易见。两者都不是。S02 会实现一个 byte-pair 编码器,说清楚流式解码为什么需要缓冲区,以及你数出来的 token 数为什么总和账单对不上。
自测
习题
先作答,再看解析。答错比答对更有价值,因为解析会指出你该回头重读哪一部分。
在 decode 阶段(而非 prefill),GPU 为什么大部分时间是闲着的?