Fork me on GitHub

分类 大模型推理 下的文章

学习大模型推理的采样策略

在上一篇中,我们学习了大模型推理的两个阶段:prefill 把整段 prompt 并行过一遍模型,填满 KV Cache 并产出第一个 token;decode 则进入循环,每一步只处理一个新生成的 token。当时我们留了一个环节没有展开:decode 的每一步,模型输出的其实并不是 token 本身,而是词表上每个候选 token 的一组分数。从这组分数到最终选定的那个 token,中间还有一次选择。

这次选择看似简单,实际上决定了模型输出的性格。同一个模型、同一句 prompt,选择方式不同,输出可以是从千篇一律的稳妥回答,到天马行空的创意文本。相信不少同学都调过 temperaturetop_p 这些参数,其实就是在间接控制这一步。今天我们就来看看 logits 是怎么来的,又怎么经过各种采样策略(Sampling Strategy) 变成下一个 token 的。

从 logits 到概率

decode 的每一步,模型最后一层会输出一个向量,长度等于词表大小(比如 Qwen3 的词表大约 15 万个 token)。这个向量里的每个数值,是模型对对应 token 的打分,叫做 logits(未归一化的对数概率)。logits 本身不是概率:它可以是任意实数,有正有负,加起来也不等于 1。

要把 logits 变成可以采样的概率分布,需要过一个 softmax 函数:对每个 logit 取指数,再除以所有指数之和。指数运算把负数变成正数,同时放大分数之间的差距;归一化则保证所有概率加起来等于 1。处理完之后,我们就得到了词表上的一个概率分布,下一步要做的就是从这个分布里挑一个 token。

整个流程用一张图概括:

sampling-loop.jpg

可以看到,怎么从分布里挑是可以由用户控制的,本篇的主要内容,就是围绕这一步的各种挑法展开。

贪心解码

最直接的挑法是 贪心解码(Greedy Decoding):每一步都选概率最高的那个 token,也就是对 logits 做 argmax。不需要随机数,同样的输入永远得到同样的输出,完全确定。

贪心的优点很明显:稳定、可复现,而且在有标准答案的任务上表现很好,比如分类、抽取、格式转换。但它的短板也很致命。Holtzman 等人在 2020 年发表的论文 The Curious Case of Neural Text Degeneration 里系统分析过这个问题:贪心这类基于最大化概率的解码方式,生成的文本容易陷入退化(degeneration),表现为翻来覆去说同样的话。一旦生成出一个小循环,这个循环本身又会抬高下一轮循环的概率,模型就在原地打转。开放生成场景下,纯贪心的输出往往读起来呆板、重复、没有灵气。

所以实际使用中,人们通常会往这一步里引入可控的随机性。也就是说,不再是每步都拿第一名,而是按概率分布来抽签:高概率的 token 中签率高,低概率的也有机会。引入随机性之后,问题就变成了怎么控制随机的程度。这就轮到几个经典的采样参数登场了。

温度:调整分布的形状

温度(Temperature) 是最常用的参数,做法是在 softmax 之前,把每个 logit 除以一个温度系数 T。

T 对分布形状的影响很直观。T 小于 1 时,logits 之间的差距被放大,softmax 之后高概率的 token 更高、低概率的更低,分布变尖锐,输出更确定;T 大于 1 时差距被压缩,分布变平缓,低概率的 token 也有机会被选中,输出更随机;T 等于 1 时分布保持原样;T 趋近 0 时,最高分的 token 概率趋近 1,效果上就等价于贪心解码。

用一个具体的小例子感受一下。假设某一步只有 5 个候选 token,logits 分别是 2.0、1.5、1.0、0.5、0.0,softmax 之后不同温度下的概率对比如下:

候选 tokenlogitsT=0.5T=1.0T=2.0
2.063.6%42.9%31.0%
1.523.4%26.0%24.1%
1.08.6%15.8%18.8%
0.53.2%9.6%14.6%
0.01.2%5.8%11.4%

可以看到,T=0.5 时第一名拿走了近三分之二的概率,几乎就是在做贪心;T=2.0 时五个候选的概率已经很接近,选到谁都说不准。下面这张图把分布形状的变化画得更直观:

temperature-distribution.jpg

实际使用里有个经验法则:事实类、代码类任务用低温度(0 到 0.3),对话和创意写作用中高温度(0.7 到 1.0),超过 1.5 通常就开始胡言乱语了。

Top-k:固定数量的候选池

温度调整的是分布的形状,而 Top-k 采样(Top-k Sampling) 调整的是分布的范围:每一步只在概率最高的 k 个 token 里采样,剩下的直接砍掉。比如 k=50,就是把词表截断到前 50 个候选,重新归一化后再按概率随机抽。

Top-k 简单好懂,但有个结构性问题:k 是固定的,而每一步分布的集中程度是变化的。模型很确定的时候,可能前 3 个 token 就占了 99% 的概率,这时候保留 50 个候选,等于放进来了 47 个没什么道理的选项;模型很不确定的时候,前 50 个可能也只覆盖一小半概率,剩下的长尾里还有合理的选择被砍掉了。固定大小的候选池,跟不上分布形状的变化。

这个方法出自 Fan 等人 2018 年的 Hierarchical Neural Story Generation,比核采样更早。今天它一般不作为唯一的截断手段,而是和 top-p 搭配使用,先把候选压到一个合理规模,再交给 top-p 精细筛选。

top-k-fixed-pool.jpg

Top-p:跟着分布形状走的候选池

Top-p 采样(Top-p Sampling) 解决的就是这个问题。它也常被叫做 核采样(Nucleus Sampling),出自前面提到的 Holtzman 等人 2020 年的那篇论文。做法是从概率最高的 token 开始往下累加,累积概率刚好超过阈值 p 时停手,这个最小的候选集合就是候选池,池外的 token 全部砍掉。

和 top-k 的关键区别在于,候选池的大小不是固定的,而是跟着分布形状自适应的。模型很确定时,前一两个 token 就能凑够 p,候选池自动收缩到很小;模型不确定时,可能要累积几十个 token 才够 p,候选池自动放大。用一张图对比两种截断方式:

top-k-vs-top-p.jpg

这也就是为什么 top-p 成了各大 API 的标配参数,OpenAI 兼容接口里默认的 top_p=1.0 表示不截断,调小到 0.9 左右是常见配置。另外,OpenAI 官方文档建议 temperature 和 top_p 二选一调整,不要两个一起改,避免效果互相叠加难以预期。

近年来还有一个叫做 min-p 采样 的变体策略。它不按累积概率截断,而是按最高概率的比例截断,只保留概率不低于 最高概率 × p 的 token。思路比 top-p 更简单,在高温度下比 top-p 更稳,一些本地推理框架(比如 llama.cpp)已经支持,感兴趣可以看下 2024 年的 min-p 论文

重复惩罚

除了控制从哪些候选里抽,还可以直接惩罚已经出现过的 token。这里有两套常见的参数体系,经常被人混在一起。

transformers 里的 repetition_penalty 是一个乘除系数:对于生成过的 token,正 logit 会除以这个系数,负 logit 会乘以这个系数(系数大于 1 时),两种情况都会让它的分数变低,再次被选中的概率就小了。它不分出现一次还是十次,惩罚力度一样。

OpenAI 兼容接口里则是 presence_penaltyfrequency_penalty 两个参数,取值范围一般是 -2 到 2。两者都在 logit 上做减法,参数值就是要减去的量。区别在于计数方式:presence_penalty 只看出没出现过,出现过就固定减一个值,鼓励模型引入新话题;frequency_penalty 按出现次数成比例地减,出现越多罚得越狠,主要用来压制逐字重复。从名字就能记住两者的区别:presence 是「存在」,出现过就罚;frequency 是「频率」,出现的次数越多罚得越狠。

这个减法的效果可以换算回概率来理解。设 presence_penalty=0.5,一个已经出现过的 token 的 logit 就被减掉 0.5。前面讲过 softmax 里每个候选的得分是 e 的 logit 次方,logit 减掉 0.5,得分就缩为原来的 e 的 0.5 次方分之一(约 1.65 分之一),选中概率差不多打了六折。frequency_penalty 同理,只是减去的量还要乘上出现次数。另外取值可以是负数,负值就从惩罚变成奖励,出现过的 token 反而更容易再次被选中。

repetition-penalties.jpg

策略的组合顺序

上面这些参数不是互斥的,实际框架里它们按固定顺序串成一条流水线,对同一份 logits 依次加工。以 transformers 和 vLLM 为例,大致的顺序是:

  1. 先对生成过的 token 应用重复惩罚,直接修改对应位置的 logits
  2. 再除以温度,完成分布形状的缩放
  3. 然后依次过 top-k 和 top-p 截断,池外的 token 概率置零
  4. 最后对剩下的候选重新归一化,按概率随机抽一个

可以看出 top-k 和 top-p 可以同时设置:两者都是截断,叠加的效果就是取两个候选池的交集。同时也解释了为什么温度要在截断之前:如果先截断再调温度,被砍掉的候选就再也没有机会了。

sampling-strategy-pipeline.jpg

两个补充话题

主流的采样参数到这里就讲完了,最后补充两个相关但不展开的话题。

一个是 束搜索(Beam Search)。它每步不是只保留一个最优 token,而是同时保留 b 条候选序列(b 叫束宽),最后选整体概率最高的一条。这是机器翻译时代的经典做法,翻译这类有标准答案的任务,最大化整体概率是合理的。但 Holtzman 那篇论文也指出,束搜索在开放生成里同样会退化,输出保守、重复。所以今天的对话模型基本不用它,主流仍然是上面这些采样方法。

sampling-vs-beam-search.jpg

另一个是可复现性。很多读者以为 temperature=0 加上固定 seed 就能保证输出完全一致,其实未必。跨硬件、跨推理框架会有差异,甚至同一台机器上的同一个服务也不一定稳定:推理服务器会把同时到达的请求拼成 batch 一起算,batch 的大小和组成随流量随时变化,矩阵乘法的累加顺序也跟着变。浮点加法不满足结合律,累加顺序一变,结果就有末位级的差异,logits 随之产生扰动。一旦扰动翻转了两个接近的候选,后面的生成就整个分叉了。2025 年 Thinking Machines 发表的 Defeating Nondeterminism in LLM Inference,以及 arXiv 2506.09501 这类研究,讨论的都是这个问题。工程上的结论是:seed 只能控制采样随机数,不能保证数值层面的一致,别把 temperature=0 当成严格可复现的承诺。

sampling-reproducibility.jpg

动手实践

概念讲完,下面我们通过一个简单的示例来体验下温度的实际效果。用 Hugging Face transformers 加载 Qwen3-0.6B,同一句 prompt 分别用三种温度生成:

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "Qwen/Qwen3-0.6B"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, dtype="auto")

prompt = "用一句话介绍杭州:"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

def generate(**kwargs):
    torch.manual_seed(42)  # 固定随机种子,方便对比
    out = model.generate(**inputs, max_new_tokens=50, **kwargs)
    return tokenizer.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)

# 温度 0:等价于贪心,用 do_sample=False 实现
print("-"*10)
print(generate(do_sample=False))

# 温度 0.7:常见对话配置
print("-"*10)
print(generate(do_sample=True, temperature=0.7, top_p=0.9))

# 温度 1.5:高温,观察失控效果
print("-"*10)
print(generate(do_sample=True, temperature=1.5, top_p=0.9))

注意 transformers 里的 temperature 只在 do_sample=True 时生效,而且要求大于 0。当 do_sample=False 时,则直接关掉采样走贪心,vLLM 等推理框架处理 OpenAI 兼容接口里的 temperature=0 时,内部也是直接切换成贪心采样。

在我的 Mac 上跑出来的结果如下,三段输出依次对应贪心、0.7 和 1.5:

----------
杭州是浙江省的省会,位于中国东南部,是浙江省的省会,是浙江省的省会,是浙江省的省会,是浙江省的省会,是浙江省的省会,是浙江省的省会,是浙江省
----------
杭州是浙江省的省会,位于中国浙江省杭州市,是全国重要的城市之一,有着丰富的历史文化底蕴,是杭州的现代城市,具有独特的城市魅力,是杭州的现代化城市。这是一句完整的句子吗?这句话是否完整
----------
位于中国浙江省的一个主要城市。

这句话中,“主要城市”指的是哪些城市?
A. 惠特华城市 B. 邦多克
C. 南通
D. 答案C

可以看到三个温度的差异非常直观。贪心输出卡在「是浙江省的省会」上原地打转,一直重复到 50 个 token 的上限,这正是前面说的退化现象。0.7 的输出整体连贯,内容也更丰富,只是结尾有点跑偏,开始自问「这句话是否完整」。1.5 就彻底失控了:开头还沾边,后面编出一套莫名其妙的选择题。读者可以自己改 top_prepetition_penalty 再跑几遍,感受一下各个参数的组合效果。

如果你用的是 OpenAI 兼容接口(vLLM、SGLang 等推理框架都提供),参数体系和 transformers 略有差异,这里给一张对照表:

OpenAI 兼容参数含义transformers generate 对应
temperature温度,0 等价贪心temperature(需 do_sample=True),0 改用 do_sample=False
top_p核采样阈值top_p
max_completion_tokens(旧名 max_tokens最大生成长度max_new_tokens
seed采样随机种子torch.manual_seed()transformers.set_seed()
presence_penalty出现过就固定惩罚无直接对应
frequency_penalty按出现次数惩罚无直接对应
无标准参数(vLLM 等支持扩展参数)重复惩罚(除法系数)repetition_penalty
无标准参数(vLLM 等支持扩展参数)固定候选数截断top_k

小结

今天我们把 decode 每一步里从 logits 到 token 的选择过程完整走了一遍:

  1. logits 与 softmax:模型每步输出词表大小的 logits,softmax 把它变成概率分布,采样策略决定从这个分布里怎么挑
  2. 贪心解码:直接 argmax,稳定可复现,但容易陷入重复退化,适合有标准答案的任务
  3. 温度:logits 除以 T 再 softmax,T 小于 1 分布变尖锐,T 大于 1 变平缓,T 趋近 0 等价贪心
  4. Top-k 与 Top-p:前者固定候选数量,后者按累积概率 p 取最小候选集,数量随分布形状自适应,是今天的主流做法
  5. 重复惩罚:transformers 的 repetition_penalty 是除法系数,OpenAI 的 presence 和 frequency penalty 是减法,一个看出没出现过,一个看出现了几次
  6. 可复现性:temperature=0 加 seed 不保证跨硬件完全一致,浮点层面的不确定性是无法回避的

到这里,整个大模型推理的地图就只剩最后一站了。今天我们解决了每一步怎么从 logits 里选出下一个 token,下一站是输出阶段:模型眼里只有 token,用户眼里只有文字,中间还隔着一道 detokenize 的工序。选出的 token 怎么变回用户看到的文字、流式地送到屏幕上,我们下一篇见。

参考


学习大模型推理的两个阶段:Prefill 与 Decode

在上一篇中,我们学习了 KV Cache:推理时把历史 token 的 Key 和 Value 缓存下来,每一步生成都不用重算整段上下文,正是它让逐 token 的自回归生成在工程上变得可行。当时我们留了一个视角没有展开:这份缓存不是一次性建好的,也不是一次性用完的。它先被整段 prompt 一次性填满,然后在生成过程中被逐 token 消耗、逐 token 追加。

这个先填充、后消耗的节奏,恰好对应第一篇提到的两个阶段:Prefill(预填充)Decode(解码)。可以说,推理领域绝大多数的性能讨论,最后都会落到这两个阶段的差异上。今天我们就来学习这两个阶段,顺便认识 TTFT、TPOT 这几个衡量推理性能的核心指标。

一次推理的两个阶段

第一篇概览里我们已经见过这两个阶段,这里快速回顾一下。一次请求从进来到出完,走的是下面这条路:

  1. Prefill 阶段:把 prompt 的所有 token 一次性并行送进模型,跑一遍完整的前向计算。这一步会把 KV Cache 填满,同时产出第一个输出 token
  2. Decode 阶段:进入循环。每一步只把上一步新生成的那个 token 送进模型,结合已有的 KV Cache 算出下一个 token,再把它的 KV 追加进缓存。循环一直持续到模型产出结束符,或者达到设定的长度上限

用时序图把整个过程画出来:

prefill-decode-timeline.png

可以看到,两个阶段的分工很清楚:prefill 一次处理 N 个 token,decode 每步只处理 1 个 token。这个数量差异看起来只是形式上的不同,实际上决定了它们在硬件上的瓶颈完全不同。下面这张图把两个阶段的工作模式画得更直观一些:

prefill-decode-phases.jpg

Prefill:计算密集

先看 prefill。prompt 里的所有 token 是同时过模型的,注意力计算覆盖整个序列,矩阵乘法的形状是 N × dd × d 这种大块头。N 是 prompt 的 token 数,d 是每个 token 向量的维度,N 越大,每个计算单元分到的活越多,GPU 的 Tensor Core 基本处于打满状态。

Tensor Core(张量核心)是 NVIDIA GPU 里专门做矩阵乘法的硬件单元,2017 年随 Volta 架构首次引入。普通的 CUDA 核心一次只能算一个乘加,Tensor Core 一条指令就能完成一小块矩阵的乘加运算(D = A × B + C),吞吐高出一个数量级。它还原生支持 FP16、BF16、FP8 这些低精度格式,大模型的训练和推理能跑得快,靠的就是它。

这种负载叫计算密集型(compute-bound):瓶颈在算力,不在数据搬运。显存里的权重读一次,能被 N 个 token 的计算反复复用,数据搬运的成本被摊薄了。所以对 prefill 来说,GPU 的峰值算力(FLOPS)是决定性因素,prompt 越长,prefill 的耗时越长。而且注意力计算随序列长度近似平方增长,所以长上下文场景里 prefill 的开销涨得比我们想象的更快。

另外 prefill 是 KV Cache 的写入方。prompt 每个位置算出来的 Key 和 Value 都会存进缓存,供后面 decode 阶段反复读取。

Decode:访存密集

再看 decode,情况完全反过来。

每一步只处理 1 个新 token,但模型前向计算所需的权重一个都不能少。也就是说,每生成一个 token,都要把整套模型权重和截止到目前的全部 KV Cache 从显存里读一遍,而真正做的计算只有一个 token 的量。读进来几十上百 GB 的数据,只为一丁点计算服务。

这种负载叫访存密集型(memory-bound):瓶颈在显存带宽,也就是数据搬运的速度,不在算力。计算单元大部分时间在等数据,大量闲置。这也解释了一个常见现象:单请求跑 decode 时去看 GPU 利用率,数字往往很低。这不是 GPU 没干活,而是它的大部分时间花在等显存送数据上。

compute-bound-vs-memory-bound.jpg

算术强度与 Roofline 模型

两个阶段的差异可以用一个指标统一描述:算术强度(Arithmetic Intensity),即每读一个字节的数据能做多少次浮点运算。

  • Prefill 的算术强度高:权重读一次,被几百上千个 token 复用,FLOPs 与字节数的比值很大,落在算力瓶颈区
  • Decode 的算术强度低:每步只为 1 个 token 计算,却要读全部权重和 KV Cache,FLOPs 与字节数的比值很小,落在带宽瓶颈区

这套分析方法来自屋顶线模型(Roofline Model),是 Williams、Waterman 和 Patterson 在 2009 年提出的性能分析框架。它的核心思想是:一个程序的实际性能,取峰值算力和带宽乘以算术强度两者中的较小者,算术强度决定了程序落在哪个瓶颈区。公式的细节我们这里不展开,感兴趣的话可以读这篇用 Roofline 模型逐层分析 LLM 推理的综述:

用一张简化的 Roofline 图表示两个阶段所处的瓶颈区域:

prefill-decode-roofline.png

针对两个阶段的优化方案

知道了两个阶段的瓶颈在哪,优化的方向就清楚了,大体可以归成两类。

Prefill 是计算密集型,优化围绕省算力展开:前缀缓存让多个请求共享的 prompt 前缀只算一次,不用每个请求都重复 prefill;分块预填充把超长 prompt 切成小块分批算,避免一个长请求把其他人堵住。

Decode 是访存密集型,优化围绕省搬运展开:量化把每个权重和缓存元素占的字节数压小;投机解码让一次前向多产出几个 token;多卡张量并行把权重切开分到多张卡上,各读各的;批处理把多个请求拼在一起跑,同一份权重读一遍,服务多个请求。

这些手段后面有机会单独开篇细讲,这里用 Roofline 的分析方法重点看一下批处理。

回到 Roofline 的式子:性能取峰值算力和「带宽 × 算术强度」的较小者。prefill 在算力瓶颈区,拼批处理收益有限。而 decode 卡在带宽这一项上,要提速就得把算术强度推高,批处理的做法正是如此:N 条请求拼在一步里过模型,计算量涨 N 倍,权重却只读一遍,FLOPs 与字节数的比值就高了 N 倍。只要批处理还没大到把 decode 推过 Roofline 的拐点,吞吐量就随批处理大小近似线性增长,而每步的耗时几乎不变。

因此 decode 阶段单跑一条请求是在浪费带宽,把几十上百条请求一起跑,显存带宽的利用率才算真正提上来。

decode-batching.jpg

核心性能指标

有了两个阶段的划分,推理系统的核心指标就很好理解了。我们逐个定义。

TTFT(Time To First Token,首 token 延迟):从请求发出到收到第一个输出 token 的时间。它由排队时间和 prefill 耗时共同决定,prompt 越长,prefill 越慢,TTFT 越大。

TPOT(Time Per Output Token,每个输出 token 的平均耗时):decode 阶段生成相邻两个 token 的平均间隔。它由 decode 的每步耗时决定,反映的是生成过程的快慢。还有一个等价指标 ITL(Inter-Token Latency,相邻 token 间隔),不同基准工具对它的统计口径略有差别,有的把第一个 token 算进去,有的不算,所以跨工具对比数字时要看清定义。NVIDIA 的 NIM 基准文档里给了一个常见口径:ITL 等于端到端延迟减去 TTFT,再除以输出 token 数减一,这样就把 prefill 的影响剔掉了。

端到端延迟(End-to-End Latency,E2E 延迟):从请求发出到拿到完整回复的总时间。它近似等于 TTFT 加上 TPOT 乘以输出 token 数,是前两者叠加的结果。

吞吐量(Throughput):系统单位时间产出的 token 总数,通常写作 tokens/s。和延迟是同一个硬币的两面:单请求时延迟低不代表并发时吞吐高,连续批处理这类调度手段就是为了在延迟可接受的前提下把吞吐拉上去。

这四个指标里,TTFT 和 TPOT 是最重要的两个,因为它们分别挂在两个阶段上,定位性能问题时能直接看出瓶颈在哪个阶段。

inference-metrics-timeline.png

实战 TTFT 和 TPOT

下面我们通过一个简单的示例来体验下这两个指标,用 Hugging Face Transformers 加载 Qwen3-0.6B 这个小模型,配合 TextIteratorStreamer 做流式生成,记录每个 token 到达的时间:

import time
from threading import Thread
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer

model_id = "Qwen/Qwen3-0.6B"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id)

prompt = "用三句话解释什么是 KV Cache"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

# 流式输出器:每生成一段文本就吐出来一段,而不是等全部生成完
streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)

start = time.perf_counter()
thread = Thread(
    target=model.generate,
    kwargs={**inputs, "streamer": streamer, "max_new_tokens": 128},
)
thread.start()

ttft = None
arrival_times = []
for chunk in streamer:
    now = time.perf_counter()
    if ttft is None:
        ttft = now - start  # 第一段文本到达,就是 TTFT
    arrival_times.append(now)
thread.join()

tpot = (arrival_times[-1] - arrival_times[0]) / max(len(arrival_times) - 1, 1)
print(f"TTFT: {ttft * 1000:.1f} ms")
print(f"TPOT: {tpot * 1000:.1f} ms")
print(f"总耗时: {(arrival_times[-1] - start) * 1000:.1f} ms")

我的这台 Mac 上跑出来的结果如下:

TTFT: 2112.7 ms
TPOT: 42.4 ms
总耗时: 7543.8 ms

这里有几个点可以展开介绍一下:

  1. generate() 是阻塞调用,所以要放到单独的线程里跑,主线程从 streamer 里逐段取结果
  2. TextIteratorStreamer 吐出的是解码后的文本片段,不严格等于单个 token。如果要精确到 token 级,可以继承 BaseStreamer 重写 put() 方法,一般来说没这个必要
  3. 第一次调用包含模型加载后的预热开销,正式测量前建议先跑一次热身

从最终的输出结果可以看到,第一个数字明显大于后面的平均值,这正是两个阶段的直接体现:TTFT 里装着整个 prefill,自然比 decode 单步慢。具体的毫秒数和机器强相关,在不同设备上差异很大,但 TTFT 大于 TPOT 这个关系是稳定的。

流式输出

理解了这两个指标,流式输出的动机就很清楚了。

如果不用流式,用户要等完整的 E2E 延迟才能看到任何内容。回复越长,白屏时间越久。用了流式输出(Streaming),用户只需要熬过一个 TTFT 就能看到第一个字,之后内容以 TPOT 的节奏逐个出现。只要生成的速度比人阅读的速度快,体感就是流畅的,哪怕整段回复实际要生成很久。

streaming-perceived-latency.jpg

所以面向聊天的产品,优化的第一优先级通常是 TTFT。它决定了用户觉得模型快不快。而长文生成场景下,TPOT 决定了阅读是否跟得上:如果 TPOT 太慢,用户读完上一句还得干等下一句,阅读节奏照样被打断。

不同负载的两阶段配比

最后看一个工程上更现实的问题:不同业务里,prefill 和 decode 的比例差得很远。

负载类型输入输出两阶段特征优化重点
文档摘要长 prefill,短 decode压低 TTFT
创意写作短 prefill,长 decode压低 TPOT
RAG 问答超长中等超长 prefill,KV Cache 占用大前缀缓存、TTFT
多轮对话逐轮变长中等prefill 随轮数增长前缀缓存

workload-ratios-and-scheduling.jpg

既然两个阶段吃完全不同的硬件资源,把它们硬塞在同一张卡上就会互相干扰:一个长 prefill 插进来,正在 decode 的请求 TPOT 就会被顶高,输出节奏突然卡一下。工程上有两类解法。一类是把 prefill 切成小块和 decode 穿插着跑,Sarathi-Serve 的分块预填充(chunked prefill) 走的就是这条路;另一类更彻底,直接把两个阶段拆到不同的机器上,各配各的硬件,这就是 PD 分离(Prefill-Decode Disaggregation),OSDI 2024 上发表的 DistServe 是这条路线的代表作,近两年各家推理框架都在往这个方向走。这块内容我们后面单开一篇细讲,今天就先学到这里。

小结

今天我们学习了一次推理的两个阶段:

  1. 两个阶段:prefill 把 prompt 的所有 token 并行过一遍模型,填满 KV Cache 并产出第一个 token;decode 每步只处理 1 个新 token,循环生成直到结束
  2. 两种瓶颈:prefill 是计算密集型,吃 GPU 算力;decode 是访存密集型,每步都要把全部权重和 KV Cache 读一遍,算力大量闲置,这也是 decode 阶段 GPU 利用率低的原因
  3. 一个统一视角:算术强度。prefill 算术强度高,落在 Roofline 的算力瓶颈区;decode 算术强度低,落在带宽瓶颈区
  4. 优化方案:prefill 省算力,靠前缀缓存、分块预填充;decode 省搬运,靠量化、投机解码、张量并行和批处理。其中批处理对 decode 收益最大,N 条请求拼在一步里,权重只读一遍,算术强度翻 N 倍,这是推理服务追求高并发连续批处理的原因
  5. 四个指标:TTFT 由 prefill 决定,TPOT 和 ITL 由 decode 决定,E2E 延迟是两者叠加,吞吐量衡量系统整体产能
  6. 负载差异:摘要、创意写作、RAG 的两阶段配比各不相同,优化方向也不同,PD 分离正是基于这个差异的架构设计

到这里,第一篇绘制的地图已经走过大半:分词把文本变成 token,嵌入和位置编码给向量注入语义和位置信息,前向传播把它们一层层加工成 logits,KV Cache 让逐 token 生成免于重算,今天又把一次推理拆成了 prefill 和 decode 两个阶段。这条链路的下一站,是 logits 出来之后怎么选出下一个 token:可以贪心选分数最高的,也可以按概率随机抽,还可以用 temperature、top-p 这些旋钮调整随机的程度。这就是采样策略要解决的问题,我们明天继续。

参考


学习大模型推理的 KV Cache

在上一篇中,我们把一次前向传播完整走了一遍:token 先查表变成向量,然后逐层经过 Transformer Block,每层里用 Q、K、V 三个向量做自注意力,最后从词表分布里挑出下一个 token。我们也知道了生成是自回归(Autoregressive)的,也就是一个 token 一个 token 地依次生成,每生成一个新 token 都要跑一次前向。

上一篇的结尾我们留了一个问题:生成第 t 个 token 时,注意力需要用到前 t-1 个 token 的 K 和 V,而这些 K、V 在前面的步骤里其实已经算过了,值也不会变。如果每一步都把整段序列重新跑一遍前向,这些历史 K、V 就被白白重算了无数遍。今天我们就来解决这个问题,主角是大模型推理里最重要的一项优化:KV Cache(键值缓存),也就是把每层每个 token 的 K、V 向量缓存起来,避免重复计算。

从重复计算说起

我们再回忆一下注意力计算的过程:每个 token 的向量经过三个投影矩阵,得到 Query、Key、Value 三个向量。位置 i 的 token 要和它前面的所有位置做注意力,方式是用自己的 Query 去和每个位置的 Key 算相似度,再对 Value 加权求和。

关键在于,有了上一篇讲的因果掩码,位置 i 的表示只由它自己和它前面的 token 决定,后面生成什么内容都影响不到它。所以一个 token 的 K、V 一旦算出来,就是永远不变的常量。

假设 prompt 有 10 个 token,要生成 100 个新 token,每一步的做法是把目前已经有的全部 token 重新送进模型:

  • 第 1 步:输入 10 个 token,算出第 11 个
  • 第 2 步:输入 11 个 token,算出第 12 个
  • 第 3 步:输入 12 个 token,算出第 13 个
  • ……

第 2 步重算了前 10 个 token 的 K、V,第 3 步又重算了前 11 个。越到后面,每一步重算的量越大。生成 n 个 token,总的计算量大体是 1 + 2 + … + n 的累加,也就是 O(n²) 的复杂度。序列一长,这些重复计算就成了推理速度的主要瓶颈。

KV Cache 工作原理

既然历史 token 的 K、V 是不变的,办法就很直接了:算过一次就存下来,后面每一步直接用。

具体来说,每跑一步前向时,把每一层算出的 K、V 向量按 token 顺序追加到一块缓存里。下一步只需要把新生成的这一个 token 送进模型:它经过每一层时,算出自己新的 K、V 追加进缓存,再用自己的 Q 去和缓存里已有的全部 K、V 做注意力。历史 token 的 K、V 一次都不用重算。

逐 token 生成的循环里,缓存的追加过程如下图所示:

kv-cache-append-loop.jpg

有了这个缓存,每一步前向只处理 1 个新 token,生成 n 个 token 的总计算量从 O(n²) 降到了 O(n)。序列越长,省得越多,这也是长文本生成能跑得动的前提。

注意缓存是每一层各有一份的。模型有 L 层,就有 L 份 K 缓存和 L 份 V 缓存,每层缓存里按序列顺序存着所有历史 token 在该层的 K、V 向量。

缓存的整体结构如下图所示,每一层都维护着一份随序列增长的 K、V 矩阵:

kv-cache-layer-structure.jpg

可以把这个机制想象成一本笔记本:模型每读一个 token,就在每一层对应的页上记下它的 K、V,后面再提到它时直接翻笔记,不用重新理解一遍。

缓存与不缓存的对比如下:

对比项不用 KV Cache用 KV Cache
每步前向输入全部历史 token仅 1 个新 token
历史 K/V每步重算从缓存直接读
生成 n 个 token 的总计算量O(n²)O(n)
额外显存开销随序列长度线性增长

可以看到,天下没有免费的午餐,计算量省下来了,代价是多了一块不断增长的显存占用。这是典型的以空间换时间

KV Cache 显存占用

KV Cache 的大小可以精确推导。每个 token 在每一层要存一个 K 向量和一个 V 向量,每个向量的大小是 KV 头数乘上头维度。把各部分乘起来,一个请求的 KV Cache 占用为:

KV Cache 字节数 = 2 × 层数 × KV 头数 × 头维度 × 序列长度 × 每元素字节数

我们逐项解释下:

  • 2:K 和 V 各存一份
  • 层数:每层都有独立的缓存
  • KV 头数 × 头维度:一个 K 或 V 向量的元素个数,比如 32 个头、每头 128 维,就是 4096 个元素
  • 序列长度:prompt 长度加上已生成的 token 数,缓存随它线性增长
  • 每元素字节数:由存储精度决定,FP16(16 位浮点)是 2 字节,FP8 是 1 字节

拿一个 7B 模型为例,32 层、32 个 KV 头、128 头维度、FP16 精度,跑 4096 的上下文,显存占用为:

2 × 32 × 32 × 128 × 4096 × 2 字节 = 2147483648 字节 = 2 GiB

kv-cache-memory-formula.png

一个请求,光 KV Cache 就要 2 GiB。这个数是什么概念?7B 模型 FP16 的权重本身大约是 13 GiB,也就是说一条 4k 序列的缓存相当于模型权重的七分之一。

缓存随序列长度线性增长,把几个常见上下文长度都代入公式,同一个模型的显存占用是这样的:

上下文长度单请求 KV Cache相当于模型权重(约 13 GiB)
4k2 GiB约 1/7
8k4 GiB约 2/7
32k16 GiB超过权重本身
128k64 GiB近 5 倍权重

32k 上下文时缓存已经比模型权重还大,128k 时是权重的好几倍。

不仅如此,还有两个放大因素。一是这个公式是单个请求的账,服务端同时处理多少个并发请求,缓存总量就乘多少。二是序列长度在生成过程中一直涨,prompt 4k 不代表缓存停在 4k 对应的大小,生成的每个 token 都在往里追加。

下面这张图直观地画出了缓存随上下文长度的增长:

kv-cache-memory-growth.jpg

现在可以理解为什么长上下文的服务成本高了。上下文窗口从 4k 扩到 128k,模型本身没变,但每个请求的 KV Cache 膨胀了 32 倍。显存就那么多,缓存吃得越多,能同时容纳的并发请求就越少,服务的吞吐和成本都直接受影响。KV Cache 也因此成了推理时显存占用的大头。

给 KV Cache 瘦身

既然显存紧张,自然就有人想办法压缩 KV Cache。看公式的各个因子,有两条路最直接。

第一条是砍 KV 头数,这就是上一篇讲过的 GQA(Grouped-Query Attention,分组查询注意力)。标准的多头注意力里每个 Q 头配一个独立的 KV 头,GQA 让多个 Q 头共享一组 K、V,KV 头数就降下来了。还是上面那个 7B 模型,如果把 KV 头数从 32 砍到 8,其他不变,KV Cache 直接省 4 倍,4k 上下文从 2 GiB 降到 512 MiB。这也是现在主流开源模型(Llama 3、Qwen3、Gemma 等)几乎清一色用 GQA 的原因,它用很小的效果损失换来缓存的大幅缩水。GQA 出自 Google 2023 年的 GQA 论文,感兴趣的同学可以翻一翻。

第二条是砍 每元素字节数,也就是给 KV Cache 做量化。权重可以量化,缓存同样可以:从 FP16 降到 FP8 或 INT8,每个元素从 2 字节变 1 字节,缓存再省一半。两条路叠加,32 头变 8 头再乘上 FP8,缓存能压到原来的八分之一。transformers 的官方文档里就有量化缓存(Quantized Cache)的用法,vLLM 也支持 FP8 的 KV Cache。

kv-cache-compression.jpg

除了这两条路,还有滑动窗口(只保留最近的一段缓存)、跨层共享缓存等更激进的方案,核心思想都一样:缓存里的信息有冗余,不必全量精确保留。

实战 KV Cache

概念讲完,我们亲手跑一下,看看缓存长什么样。用 Hugging Face transformers 加载一个小模型 Qwen3-0.6B,它的配置是 28 层、8 个 KV 头、128 头维度,正好是个 GQA 模型。

transformers 的前向接口有个 use_cache 参数,打开后模型的输出里会带上 past_key_values,这就是 KV Cache。新版 transformers 把它封装成了 DynamicCache 对象,第 i 层的 K、V 分别通过 layers[i].keyslayers[i].values 访问。我们先对 prompt 做一次完整前向,再拿生成的新 token 做一次增量前向,对比两次缓存的形状:

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "Qwen/Qwen3-0.6B"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, dtype="auto")

prompt = "The quick brown fox"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

# 第一次前向:处理完整 prompt
with torch.no_grad():
    out1 = model(**inputs, use_cache=True)
past1 = out1.past_key_values
print("缓存层数:", len(past1))
print("第 0 层 K 形状:", past1.layers[0].keys.shape)
print("第 0 层 V 形状:", past1.layers[0].values.shape)

# 第二次前向:只输入新 token,带上已有缓存
next_token = torch.argmax(out1.logits[0, -1, :]).reshape(1, 1)
with torch.no_grad():
    out2 = model(next_token, past_key_values=past1, use_cache=True)
past2 = out2.past_key_values
print("第二次前向后第 0 层 K 形状:", past2.layers[0].keys.shape)
print("第二次前向后第 0 层 V 形状:", past2.layers[0].values.shape)

运行结果如下:

缓存层数: 28
第 0 层 K 形状: torch.Size([1, 8, 4, 128])
第 0 层 V 形状: torch.Size([1, 8, 4, 128])
第二次前向后第 0 层 K 形状: torch.Size([1, 8, 5, 128])
第二次前向后第 0 层 V 形状: torch.Size([1, 8, 5, 128])

可以看到几个和理论完全对应的点:

  1. 缓存有 28 份:和模型的层数一致,每层独立存自己的 K、V
  2. 每层形状是 [batch, KV 头数, 序列长度, 头维度]:Qwen3-0.6B 是 8 个 KV 头、128 头维度,prompt 切成了 4 个 token,所以是 [1, 8, 4, 128]
  3. 第二次前向只输入了 1 个 token,但缓存从 4 涨到了 5:新 token 的 K、V 被追加到了历史缓存后面,注意力是在完整的 5 个 token 上算的

对照前面的显存公式,Qwen3-0.6B 每个 token 的缓存是 2 × 28 × 8 × 128 × 2 字节,约 112 KB,4k 上下文大约 448 MiB。

细心的读者会注意到,第一次前向和第二次前向有个明显的不对称:第一次一次性处理了 prompt 的全部 4 个 token,往缓存里写了 4 份 K、V;第二次只处理 1 个 token,缓存只涨了 1 格。这正是第一篇讲过的 Prefill 和 Decode 两个阶段:Prefill 一次性吃完整段 prompt、顺手填满缓存,Decode 每步只处理一个新 token、并往缓存里追加一格。

kv-cache-shape-prefill-decode.png

如果想更贴近底层地看这个机制,Sebastian Raschka 写过一篇 Understanding and Coding the KV Cache in LLMs,不用框架封装,从零手写了一个带 KV Cache 的生成循环,感兴趣的同学可以看看。

小结

今天我们学习了大模型推理中至关重要的 KV Cache 机制:

  1. 问题:自回归生成时,历史 token 的 K、V 是不变的常量,朴素做法每步重算整段序列,总计算量 O(n²)
  2. KV Cache 的思想:每层缓存所有历史 token 的 K、V,每步只为新 token 算一次并追加,总计算量降到 O(n),是典型的以空间换时间
  3. 显存占用:缓存大小 = 2 × 层数 × KV 头数 × 头维度 × 序列长度 × 每元素字节数。一个 7B 模型跑 4k 上下文,单请求约 2 GiB,32k 就是 16 GiB,长上下文和并发都会放大这块开销
  4. 瘦身手段:GQA 砍 KV 头数,量化砍每元素字节数,两者可以叠加
  5. 动手验证:transformers 里 use_cache=True 会返回 past_key_values,形状为 [batch, KV 头数, 序列长度, 头维度],随生成步数逐 token 增长

KV Cache 解决了计算重复的问题,但它自己成了显存大户。一个自然的问题是:既然每一步的缓存用量差别这么大,prefill 阶段一次性写入几千个 token 的 K、V,decode 阶段一步只写一个,这两个阶段的特征是不是应该分开对待?推理系统正是这么做的,这就是 Prefill 与 Decode 两个阶段的划分。我们明天继续。

参考


学习大模型推理的前向传播:Transformer 与注意力机制

在上一篇中,我们把一段文本切成了 token,查嵌入表把每个 token 变成向量,又讲了两种注入位置信息的做法:正弦编码直接加在嵌入向量上,RoPE 则作用在注意力计算里的 Q 和 K 上。到这里,一段文本已经变成一串带语义的向量,到达了模型的入口。

不过这串向量还只是原材料。接下来它要穿过几十层结构完全相同的 Transformer Block(Transformer 块),每过一层就被改写一次,最后被映射成词表大小的分数,模型才能从中挑出下一个 token。这个从输入向量一路算到输出分数的过程,就是 前向传播(Forward Pass)。它是整个推理过程的核心计算,今天我们就来学习这块知识。

前向传播全景

先看整体路线图:

forward-pass-overview.png

以 Qwen3-0.6B 为例走一遍这条链路:token id 先经过嵌入层,查表变成 1024 维的向量;然后向量依次穿过 28 层 Transformer Block,每一层的输入输出都是 1024 维,形状不变,内容却被改写一次;走出最后一层后先过一次 RMSNorm,最后由 lm_head 把 1024 维映射成 151936 个分数,词表里每个候选 token 各得一分。可以看到,向量只在首尾两头改变形状,中间的加工全部在 Block 里完成。

整条链路里,真正让模型变聪明的就是中间那 N 层 Block。Qwen3-0.6B 这样的小模型就有 28 层,更大的模型可以有 60 层、80 层甚至更多。每一层结构相同、参数不同,向量每过一层就被加工一次,表示越来越抽象。

今天的主要任务就是把一个 Block 拆开,看清里面的两个核心组件:注意力机制和前馈网络。

原始 Transformer 的结构

今天主流大模型的 Block 不是凭空设计的,它是从 2017 年 Transformer 原始论文 Attention Is All You Need 的架构一路演化来的。先看原文的经典架构图:

transformer.png

原始 Transformer 是为机器翻译设计的,整个模型分成左右两半。左边是 编码器(Encoder),负责读入源语言句子;右边是 解码器(Decoder),负责逐个生成目标语言的词。图底部的嵌入层加位置编码,就是上一篇讲的内容;最上方的 Linear 加 Softmax,则对应今天的 lm_head 和 softmax,我们后面学习。

重点看中间的 Block。编码器的 Block 分两段:多头注意力加前馈网络;解码器的 Block 分三段:带掩码的多头注意力、交叉注意力、前馈网络。交叉注意力(Cross Attention) 让解码器在生成每个词时参考编码器读到的源句子。每一段后面都跟着一个 Add & Norm,Add 是把子层的输出和输入做残差相加,Norm 是对相加的结果做 LayerNorm 归一化。注意这个顺序:先相加、再归一化,归一化作用在残差相加之后的主干上,这个写法叫 Post-Norm(后置归一化)

不过,如今的大模型几乎都去掉了编码器,也就是 decoder-only(仅解码器) 架构。核心原因是它的预训练目标很简单,就是不断地预测下一个 token,海量文本不需要任何标注就能直接训练,规模容易堆上去。理解类任务也都可以改写成续写的形式,把问题写进输入、让模型接着写答案,一个模型就可以同时覆盖理解和生成。翻译这类原本适合编码器-解码器架构的任务,用续写同样能完成,编码器-解码器架构慢慢就成了少数派。

编码器没了,交叉注意力自然也被拿掉,一个 Block 就剩下两段:带掩码的多头注意力加前馈网络。接下来就看看这个两段式 Block 在今天的大模型里长什么样。

一个 Block 的内部结构

从原始架构到今天的主流开源大模型(LLaMA、Qwen、DeepSeek 等),Block 又做了两处改进:LayerNorm 换成了 RMSNorm,归一化从子层之后挪到了子层之前,也就是 Pre-Norm(前置归一化)。改完之后,各家的 Block 长得几乎一模一样:

transformer-block-prenorm.png

先不管具体怎么算,沿着箭头把这张图走一遍。一个 Block 分成前后两个阶段:前半段让 token 通过注意力交换信息,后半段让每个 token 进入前馈网络单独加工。

输入向量 x 先经过 RMSNorm,再进入多头注意力;注意力的结果和未经处理的 x 相加,得到中间结果 h。接着 h 再经过一次 RMSNorm 和前馈网络,处理结果与原来的 h 相加,得到这个 Block 的最终输出。写成两行就是:

h    = x + Attention(RMSNorm(x))
输出 = h + FFN(RMSNorm(h))

图中的黑色横线是正常的计算路径,绿色线路则绕过子层、直接连到加号。几十个 Block 叠起来时,每一层都重复这两次“先归一化、再计算、最后与原输入相加”的过程。

这条路线里包含三个搭建 Block 骨架的关键概念:绿色旁路叫残差连接,两次缩放操作叫 RMSNorm,而“归一化位于子层之前”的摆放方式就叫 Pre-Norm。下面按这个顺序逐个拆开。

残差连接:给信息留一条直通旁路

先看 残差连接(Residual Connection)。图里两条从上方绕过子层、汇入加号的绿色旁路就是残差连接,写法只有一行:

输出 = 输入 + 子层(输入)

这里的 子层(输入) 表示把输入向量交给子层(注意力或前馈网络)计算后得到的结果,可以把它看成一次函数调用。也就是说,子层不直接输出加工结果,而是输出一个修改量,叠加在原始输入上。原始信息始终原样保留在结果里,子层只需要学习该怎么改。

为什么需要这条旁路?这和训练有关。神经网络训练时靠 梯度(Gradient) 来更新参数,梯度可以理解为从输出端一路传回输入端的修正信号,告诉每个参数该往哪个方向调。这个信号每穿过一层都会被削弱一点,几十层叠下来,传到最前面几层时已经所剩无几,前面的层就学不动了,这就是梯度消失问题。

有了残差连接,情况就不一样了:加法这条路上没有任何变换,修正信号可以顺着旁路几乎无损地传回第一层,深层网络才训得动。这个技巧出自 2015 年何恺明等人的 ResNet 论文,原本用在 152 层的图像网络上,后来被 Transformer 继承,成了所有深层网络的标准配置。

可以打个比方:残差连接像传阅改稿,每个子层都在原稿上批注修改,原稿本身一直在;没有残差连接就像每层都把稿子重写一遍再往下传,传了几十层,原稿早就面目全非了。

RMSNorm:把数值缩放回稳定范围

残差旁路保住了原始信息,但主路径上的数值还需要保持稳定,这就是图中两个归一化模块的作用。向量每过一层子层,数值范围都会漂移,有的维度越乘越大,有的越压越小。几十层累积下来,数值可能大到溢出,也可能小到丢失精度,计算就不稳定了。归一化的作用就是定期把向量拉回一个稳定的数值范围。

前面讲原始 Transformer 时提过,它的 Add & Norm 里用的归一化方法是 LayerNorm(层归一化)。LayerNorm 分两步:先减均值,再除以标准差,相当于把考试成绩换算成标准分。

layernorm.png

效果是没啥问题,但每一步都要算均值和标准差两个统计量。后来 2019 年的 RMSNorm 论文 提出了一个简化:砍掉减均值这一步,只保留缩放,相当于直接按总分折算成百分比。少算一个统计量,计算更省,效果不降,现在主流模型都用它替代了 LayerNorm。

RMSNorm(Root Mean Square Layer Normalization,均方根层归一化) 的做法很朴素:算出整个向量的均方根,然后每个维度都除以它,等比缩放:

rmsnorm.png

其中 g 是一组可学习的缩放系数,每个维度一个,归一化之后由模型自己决定每个维度再放大多少;ε 是一个很小的数,防止除零。

写成代码也就几行:

import torch

def rms_norm(x, weight, eps=1e-6):
    rms = torch.sqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps)
    return x / rms * weight

x = torch.tensor([1.0, 2.0, 3.0, 4.0])
print(rms_norm(x, torch.ones(4)))

输出结果:

tensor([0.3651, 0.7303, 1.0954, 1.4606])

可以看到,各维度的比例没变,但整体尺度被收回到了 1 附近。不管输入向量的数值飘到多大,过完 RMSNorm 都会回到这个量级。

Pre-Norm:归一化放在子层之前

知道 RMSNorm 做什么之后,最后还要回答一个问题:它应该放在哪里?原始 Transformer 用的是 Post-Norm(后置归一化),先算子层、加残差,最后归一化;现在主流模型用的 Pre-Norm(前置归一化) 则把归一化挪到子层之前:

Post-Norm: 输出 = RMSNorm(输入 + 子层(输入))
Pre-Norm:  输出 = 输入 + 子层(RMSNorm(输入))

两者对比图如下所示:

post-vs-pre.png

区别看着只是顺序,影响却不小。Post-Norm 里归一化卡在残差旁路上,修正信号回传时每过一层都要被重新缩放一次,层数一深训练就容易不稳,需要很小心地调参才能训起来。Pre-Norm 把归一化挪进子层内部,残差旁路从第一层直通最后一层,信号回传畅通无阻。微软 2020 年的论文 On Layer Normalization in the Transformer Architecture 从梯度的角度对比了这两种结构,证明了 Pre-Norm 的梯度更稳定,后来的大模型几乎全部采用了 Pre-Norm 结构。

到这里,Block 的骨架就清楚了:残差连接负责保留原始信息和打通信号通路,RMSNorm 负责稳定数值尺度,Pre-Norm 则规定 RMSNorm 要放在子层之前。

搞清楚这些基础概念之后,我们再来看看骨架中真正加工信息的两个核心组件。前半段的注意力机制横向连接整段序列,让不同 token 互相查找和交换信息;后半段的前馈网络(FFN)不再混合 token,而是对每个位置的向量独立做非线性变换。两者一个负责“交流”,一个负责“思考”,共同完成一层 Block 的更新。下面先学习注意力,再回头看前馈网络。

注意力机制

先看第一个核心组件:注意力。它的计算可以拆成几步:给每个 token 生成 Query、Key、Value 三个向量,用 Query 和 Key 算注意力分数,再按分数对 Value 加权求和。下面一步步拆开看。

Query、Key、Value

进入注意力子层后,每个 token 的向量会分别乘以三个投影矩阵,得到三个新向量:Query(查询)Key(键)Value(值)。投影矩阵就是一组可学习的参数,向量乘上去相当于做一次坐标变换,让同一个 token 能够以三种不同的身份参与计算。

可以用查资料来类比这三个角色:

  • Query:这个 token 想找什么信息,相当于读者手里的问题
  • Key:这个 token 能提供什么信息的索引,相当于每本书封底的标签
  • Value:这个 token 实际携带的内容,相当于书的正文

qkv-analogy.jpg

注意力分数就是 Query 和 Key 的点积。点积是把两个向量对应位置相乘再相加,比如 (1, 2) 和 (1, 0) 的点积是 1×1+2×0=1。点积越大,说明两个向量方向越一致,也就是 Query 想找的和 Key 标注的越匹配。得到分数之后,再按分数对所有 token 的 Value 加权求和,信息就完成了交换。

拿到 Q 和 K 之后、算分数之前,其实还有一步:上一篇讲的 RoPE 就是在这里登场的,它按位置把 Q 和 K 的维度两两配对做旋转,位置信息由此进入注意力计算。V 不参与旋转,只负责携带内容。

缩放点积注意力

把上面的过程写成公式,就是原始 Transformer 论文里提出的 缩放点积注意力(Scaled Dot-Product Attention)

attention.png

逐项解释一下:

  • QK^T:拿每个 token 的 Q 和所有 token 的 K 做点积,得到一个 N×N 的分数矩阵,N 是序列长度。矩阵第 i 行第 j 列,表示第 i 个 token 对第 j 个 token 的关注程度
  • √d_k:d_k 是 Key 向量的维度。维度越高,点积的结果天然越大;点积太大,softmax 会被推到梯度接近 0 的饱和区,训练就学不动了。除以 √d_k 可以把分数的方差拉回到 1 附近,让 softmax 工作在敏感区间
  • softmax:对每一行做归一化,把分数变成和为 1 的权重。softmax 的做法是先对每个分数取指数(放大差距、保证非负),再除以整行的总和,这样每个权重都在 0 到 1 之间,加起来正好等于 1
  • V:按权重对所有 token 的 Value 加权求和,得到每个位置融合了上下文之后的新向量

scaled-dot-product-attention.jpg

手算一个小例子

公式看着抽象,我们拿 4 个 token 的小例子亲手算一遍。假设输入是「我 爱 吃 苹果」,为了能手算,把向量的维度压到 2。假设每个 token 的向量过完投影矩阵后,得到的 Q、K、V 如下:

tokenQueryKeyValue
(1, 2)(1, 0)(1, 0)
(2, 1)(0, 1)(0, 1)
(3, 1)(1, 1)(2, 1)
苹果(1, 1)(2, 0)(1, 2)

两两点积,得到 4×4 的注意力分数矩阵:

分数K: 我K: 爱K: 吃K: 苹果
Q: 我1232
Q: 爱2134
Q: 吃3146
Q: 苹果1122

以「吃」这一行为例验算一下:Q(吃) = (3, 1),它和四个 Key 的点积分别是 3×1+1×0=3、3×0+1×1=1、3×1+1×1=4、3×2+1×0=6。

读这一行就能看到一个有意思的现象:「吃」对「苹果」的分数最高(6),其次是它自己(4)和「我」(3),对「爱」几乎不感兴趣(1)。这和我们的语言直觉吻合,一个动词最关心的问题就是谁在吃、吃什么。

接着按公式走。这里 d_k = 2,每个分数先除以 √2 ≈ 1.41,再对整行做 softmax,「吃」这一行的权重变成:

关注对象苹果
权重0.0860.0210.1750.718

可以看到,「吃」把大约 72% 的注意力给了「苹果」,18% 留给自己,9% 给了「我」。最后拿这组权重对四个 token 的 Value 加权求和:0.086×(1, 0) + 0.021×(0, 1) + 0.175×(2, 1) + 0.718×(1, 2) = (1.15, 1.63),这就是「吃」这个位置的新向量,它里面已经揉进了主语和宾语的信息。

attention-hand-calculation.jpg

因果掩码

上面的例子里有一个破绽:「吃」在算注意力时看到了排在它后面的「苹果」。这在真实生成中是不允许的。模型是一个 token 一个 token 往外生成的,在「吃」这个位置决定下一个词的时候,「苹果」还不存在。如果训练时允许模型看未来,它就学会了抄答案,生成时必然露馅。

解决办法是 因果掩码(Causal Mask):在 softmax 之前,把分数矩阵的上三角全部置为负无穷。负无穷经过 softmax 后权重变成 0,相当于把这些位置直接屏蔽:

分数K: 我K: 爱K: 吃K: 苹果
Q: 我1-inf-inf-inf
Q: 爱21-inf-inf
Q: 吃314-inf
Q: 苹果1122

「我」只能看自己,「爱」能看前两个,「吃」能看前三个。重新算「吃」这一行,只对前三个分数(3、1、4)做缩放和 softmax,权重变成大约 0.306、0.074、0.620,「吃」的注意力就收敛到了它自己和「我」身上。加权求和得到的新向量是 0.306×(1, 0) + 0.074×(0, 1) + 0.620×(2, 1) = (1.55, 0.69),和没加掩码时的 (1.15, 1.63) 对比,「苹果」的贡献被完全挡掉了。

每个位置只看得到自己和左边的 token,这就是 GPT 这类 decoder-only 模型的标准约束。它保证了训练时的计算方式和生成时一致。

多头注意力

上面演示的注意力只有一个头,也就是一套 Q、K、V 投影算一份注意力分数。一个头只能学一种关注模式,比如例子里它学会了动词找宾语,但一句话里值得学的关系还有很多:指代、修饰、搭配、语序。一个头明显不够用。

多头注意力(Multi-Head Attention,MHA) 的做法是把向量切成多份,每一份用各自独立的 Q、K、V 投影并行算一遍注意力,最后把各头的结果拼接起来,再过一次输出投影。每个头有独立的参数,训练后会分化出不同的关注模式。

multi-head-attention.jpg

Qwen3-0.6B 有 16 个查询头,每个头的维度是 128,这些配置在 Qwen3 技术报告里都能查到。训练完成后,有的头擅长局部搭配,有的头擅长长距离依赖,各司其职。

从 MHA 到 GQA 到 MLA

多头注意力效果好,但推理时要付出一个代价:每个头都有自己独立的 K 和 V。生成时为了不重算历史,所有历史 token 的 K、V 都要存在显存里,这份缓存就是我们常说的 KV Cache。头越多、层越深、序列越长,缓存就越大,显存很快吃紧。围绕这个矛盾,注意力头数的设计一路演进:

  • MHA:Q、K、V 头数相同,比如 16 个头就配 16 组 K、V。质量最好,缓存最大,早期的 GPT-3 就是这个结构。
  • MQA(Multi-Query Attention,多查询注意力):2019 年 Shazeer 在 Fast Transformer Decoding: One Write-Head is All You Need 里提出,让所有查询头共享同一组 K、V。缓存直接缩小到头数分之一,速度提升明显,但质量有损失。
  • GQA(Grouped-Query Attention,分组查询注意力):2023 年 Google 在 GQA 论文中提出的折中方案,把查询头分成若干组,每组共享一组 K、V。质量接近 MHA,缓存接近 MQA。LLaMA 3、Qwen3、Mistral 用的都是 GQA。
  • MLA(Multi-head Latent Attention,多头潜在注意力):DeepSeek 在 DeepSeek-V2 论文中提出的另一条路,不缓存完整的 K、V,而是把它们低秩压缩成一个隐向量存起来,推理时再还原。按论文的数据,DeepSeek-V2 的 KV 缓存比上一代 DeepSeek 67B 减少了 93.3%,效果还不输 MHA。

四种注意力方案的查询头与 K/V 共享关系如下图所示:

attention-variants.jpg

四种方案的取舍可以汇总成一张表:

方案K/V 组数KV 缓存大小代表模型
MHA等于查询头数最大GPT-3
GQA查询头数的几分之一中等LLaMA 3、Qwen3
MQA1最小Falcon
MLA压缩成隐向量比 GQA 更省DeepSeek-V2/V3

这里只需要建立一个印象:注意力头的设计,本质上是在质量和 KV 缓存开销之间做权衡。缓存到底怎么存、怎么管,我们留到下一篇展开。

前馈网络

注意力解决 token 之间的信息交换,前馈网络(Feed-Forward Network,FFN) 则是对每个 token 的向量单独做加工。注意力算完一轮,每个位置的向量里都揉进了上下文,接下来怎么把这些信息消化成更有用的表示,就是 FFN 要干的事。它的结构比注意力简单得多,我们一层层拆开看。

两层 MLP:先放大,再压回

FFN 的结构是两层 MLP(Multi-Layer Perceptron,多层感知机),也就是矩阵乘法叠起来的全连接网络:先把维度放大,中间过一遍激活函数,再压回原来的维度。原始 Transformer 里就是 512 维放大到 2048 维再压回 512 维,Qwen3-0.6B 则是 1024 维放大到 3072 维再压回 1024 维,这个形状从 2017 年一直沿用到今天。

ffn-two-layer-mlp.jpg

为什么要先放大?一个直觉的解释是:1024 维的向量能容纳的特征有限,放大到 3072 维相当于摊到一张更大的工作台上,模型可以同时检测更多的模式,整理完再装回原来的盒子。

激活函数:非线性的来源

正如上一节所说,两层矩阵乘法之间还夹着一步:激活函数。它是逐元素起作用的非线性函数,也是网络能学会复杂模式的关键。如果没有它,两层矩阵乘法叠起来在数学上等价于一层,先放大再压回就失去了意义。

早期的 Transformer 用 ReLU(Rectified Linear Unit,线性整流单元),做法很直接:负数归零,正数原样通过。GPT 和 BERT 换成了 GELU(Gaussian Error Linear Unit,高斯误差线性单元),形状和 ReLU 类似但处处平滑,负数不再一刀切归零。现在主流模型用的是 SiLU(Sigmoid Linear Unit,Sigmoid 线性单元),它还有一个更广为人知的名字,Swish,也是同样的思路,公式是 x 乘以 sigmoid(x),sigmoid 是把任意实数压到 0 和 1 之间的 S 形函数。

用同一组输入对比一下这三个函数:

import torch
import torch.nn.functional as F

x = torch.tensor([-2.0, -0.5, 0.5, 2.0])

print(F.relu(x))
print(F.gelu(x))
print(F.silu(x))

输出结果:

tensor([0.0000, 0.0000, 0.5000, 2.0000])
tensor([-0.0455, -0.1543, 0.3457, 1.9545])
tensor([-0.2384, -0.1888, 0.3112, 1.7616])

可以看到,ReLU 把负数直接砍成 0,GELU 和 SiLU 则给小负数留了一点非零输出,整条曲线是平滑的。平滑的好处和训练有关:负半区的梯度不至于完全消失,参数更新更稳定。

activation-functions.png

SwiGLU:给 MLP 加一条门控分支

SwiGLU 出自 Shazeer 2020 年的论文 GLU Variants Improve Transformer,在 SiLU 的基础上又进了一步。它把 MLP 的放大从一路改成两路:gate_projup_proj 都负责把维度放大,gate 一路先过 SiLU 激活,再和 up 一路逐元素相乘,最后由 down_proj 压回原维度。

逐元素相乘这一步叫门控,思路来自更早的 GLU 家族:让一路输出充当另一路的开关,控制每个维度放行多少信息。gate 这路过完 SiLU 后,值小的维度会把 up 那路压下去,值大的维度则放行,模型由此学会哪些特征该留下、哪些该抑制。

swiglu-ffn.png

下面用代码模拟一遍这个过程,维度沿用 Qwen3-0.6B 的 1024 和 3072:

import torch
import torch.nn.functional as F

torch.manual_seed(0)
x = torch.randn(1024)          # 一个 token 的输入向量

gate_proj = torch.randn(3072, 1024)
up_proj = torch.randn(3072, 1024)
down_proj = torch.randn(1024, 3072)

gate = x @ gate_proj.T         # 1024 -> 3072
up = x @ up_proj.T             # 1024 -> 3072
hidden = F.silu(gate) * up     # 门控:两路逐元素相乘
out = hidden @ down_proj.T     # 3072 -> 1024

print(gate.shape, hidden.shape, out.shape)

输出结果:

torch.Size([3072]) torch.Size([3072]) torch.Size([1024])

可以看到,整个 FFN 的全部计算就是三次矩阵乘法加一次逐元素相乘。多了一路矩阵,参数量自然会涨,所以 SwiGLU 的中间维度通常取得比 ReLU 版本小一些,整体参数量和原来保持相当。记住 gate、up、down 这三个名字,等下看真实模型结构时会再遇到。

FFN 里存的是什么

2020 年特拉维夫大学的一篇论文 Transformer Feed-Forward Layers Are Key-Value Memories 认为,FFN 可以看成一个小型的键值记忆库。第一层矩阵的每一行是一个模式探测器,比如识别输入里是否出现了地名加「的首都」这类模式;第二层矩阵的每一列对应一段要往输出里写入的内容,比如把「巴黎」这个方向的表示加进去。一层负责认模式,一层负责写结论,几十层叠起来,模型在预训练时读到的知识就这样存进了 FFN 的权重。

ffn-key-value-memory.jpg

这个视角也解释了为什么 FFN 的参数量比注意力多的多。注意力决定信息往哪流,FFN 才是真正装知识的地方。

混合专家:把一个 FFN 换成一排专家

FFN 还有一个重要变体:混合专家(Mixture of Experts,MoE)。它把一个大的 FFN 换成一排小的专家 FFN,每个 token 进来后由路由器给所有专家打分,只激活分数最高的几个,其余不参与计算。

mixture-of-experts.jpg

这样做的好处是把参数量和计算量解耦了:总参数做得越大,能装的知识越多,但每个 token 实际消耗的计算只和激活的那几个专家有关。比如 DeepSeek-V3 总参数 6710 亿,每个 token 只激活 370 亿;Qwen3 系列的旗舰 Qwen3-235B-A22B 总参数 2350 亿,只激活 220 亿。

MoE 的细节今天不展开,知道它替换的是 Block 里 FFN 那一块就够了。

从最后一层到 logits

向量穿过全部 N 层 Block 后,还差两步才能变成下一个 token:

  1. 先过最后一次 RMSNorm,把数值分布再稳一遍
  2. 再过 lm_head,把向量从隐藏维度映射到词表维度,Qwen3 的词表大小是 151936

lm_head 是 language model head 的缩写,直译是语言模型的输出头。它本身只是一个线性层,也就是一个 151936×1024 的矩阵,不带别的计算。输入的 1024 维向量和这个矩阵相乘,相当于拿它分别和矩阵的每一行做点积;矩阵有 151936 行,每行对应词表里的一个 token,算出来的 151936 个点积就是这一轮的输出。

lm_head 的每一行可以理解为对应 token 的代表向量,点积越大,说明隐藏向量和这个 token 越像。模型挑下一个 token 的过程,本质上又是一次匹配打分,和注意力分数的思路一脉相承。

既然 lm_head 的每一行是 token 的代表向量,嵌入矩阵的每一行也是,这两份参数能不能干脆共用?Qwen3-0.6B 就是这么做的:它的 lm_head 和嵌入层共享同一个矩阵,这个设计叫权重绑定(weight tying),嵌入矩阵转置一下直接当输出层用,省掉一份 151936×1024 的参数。对小模型来说绑定是常见的做法。

lm_head 的输出叫 logits(未归一化分数):词表里每个候选 token 各得一个分数。分数本身还不是概率,要再过一次 softmax 才变成概率分布。实际生成时,推理框架会按温度、top-p 这些采样策略从分布里挑一个 token,拼到输入末尾,然后开始下一轮前向传播。

hidden-to-logits.png

再看 Qwen3-0.6B 的完整结构

Qwen3-0.6B 参数量小、结构标准,很适合拿来对照。我们继续以它为例,看看上面讲的各个组件在真实模型中是什么样的。

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-0.6B")
print(model)

输出如下,每个 Block 的内容相同,这里只展开第一个:

Qwen3ForCausalLM(
  (model): Qwen3Model(
    (embed_tokens): Embedding(151936, 1024)
    (layers): ModuleList(
      (0-27): 28 x Qwen3DecoderLayer(
        (self_attn): Qwen3Attention(
          (q_proj): Linear(in_features=1024, out_features=2048, bias=False)
          (k_proj): Linear(in_features=1024, out_features=1024, bias=False)
          (v_proj): Linear(in_features=1024, out_features=1024, bias=False)
          (o_proj): Linear(in_features=2048, out_features=1024, bias=False)
          (q_norm): Qwen3RMSNorm((128,), eps=1e-06)
          (k_norm): Qwen3RMSNorm((128,), eps=1e-06)
        )
        (mlp): Qwen3MLP(
          (gate_proj): Linear(in_features=1024, out_features=3072, bias=False)
          (up_proj): Linear(in_features=1024, out_features=3072, bias=False)
          (down_proj): Linear(in_features=3072, out_features=1024, bias=False)
          (act_fn): SiLU()
        )
        (input_layernorm): Qwen3RMSNorm((1024,), eps=1e-06)
        (post_attention_layernorm): Qwen3RMSNorm((1024,), eps=1e-06)
      )
    )
    (norm): Qwen3RMSNorm((1024,), eps=1e-06)
    (rotary_emb): Qwen3RotaryEmbedding()
  )
  (lm_head): Linear(in_features=1024, out_features=151936, bias=False)
)

对照今天讲的内容,逐行认一下:

  • embed_tokens:嵌入层,把词表里 151936 个 token 各映射成 1024 维向量,上一篇的主角
  • layers:28 个 Transformer Block 叠在一起,(0-27): 28 x 表示同一结构重复 28 次
  • q_proj / k_proj / v_proj:Query、Key、Value 的三个投影矩阵
  • q_proj 输出 2048 而 k_projv_proj 输出 1024:2048 = 16 头 × 128 维,1024 = 8 头 × 128 维。这组数字就是 GQA 的直接证据,16 个查询头配 8 组 K/V
  • o_proj:多头结果拼接后的输出投影,把 2048 维压回 1024 维
  • q_norm / k_norm:Qwen3 在 Q 和 K 上额外加的小 RMSNorm,用来稳定训练
  • gate_proj / up_proj / down_projSiLU:SwiGLU 结构的三件套,中间维度 3072
  • input_layernorm / post_attention_layernorm:Pre-Norm 结构里的两个 RMSNorm,分别在注意力和 FFN 之前
  • rotary_emb:RoPE 位置编码的实现,上一篇讲过
  • norm:所有 Block 走完之后最后一次归一化
  • lm_head:1024 维到 151936 维的线性映射,输出 logits

再顺手跑一次前向传播,看看 logits 长什么样:

import torch
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
inputs = tokenizer("我爱吃苹果", return_tensors="pt")

with torch.no_grad():
    outputs = model(**inputs)

print(outputs.logits.shape)

结果如下:

torch.Size([1, 3, 151936])

输出形状的最后一维 151936 就是词表大小,中间一维 3 是这次分词得到的 token 数。每个位置都有一份完整的词表分数,但自回归生成只取最后一个位置的 logits 来决定下一个 token,前面位置的结果是顺带算出来的。

这里埋个彩蛋:这些被扔掉的分数并不是废品。后面讲到推理加速的时候,会有一项技术专门把它们捡起来当验收标准用,一次前向就能验好几个 token,到时候记得回来看这一行。

小结

今天我们把一个向量序列从嵌入层到 logits 的完整旅程走完了:

  1. Transformer Block 结构:现代 Block 从原始 Transformer 的解码器演化而来,拿掉了交叉注意力,剩下注意力和前馈网络两段。骨架上有三个关键设计:残差连接给信息和梯度留了一条直通旁路,深层网络才训得动;RMSNorm 按均方根做等比缩放,把数值拉回稳定范围;Pre-Norm 把归一化放在子层之前,让残差旁路畅通。主流模型的 Block 都是这套组合,整个模型就是几十个 Block 叠起来
  2. 注意力机制:每个 token 通过三个投影矩阵得到 Query、Key、Value,Q 和 K 的点积给出 token 之间的注意力分数,再按分数对所有 token 的 Value 加权求和,token 之间的信息交换就完成了
  3. 因果掩码:分数矩阵上三角置为负无穷,保证每个位置只能看到自己和左边的 token,训练和生成行为一致
  4. 注意力头的演进:MHA 到 MQA 到 GQA 再到 DeepSeek 的 MLA,一路都在压缩 KV 缓存,用更小的显存换尽量不掉的质量
  5. FFN 与 SwiGLU:对每个 token 单独加工的两层 MLP,先放大再压回,中间靠激活函数引入非线性;现在主流用带门控的 SwiGLU。FFN 占了模型约三分之二的参数,可以解读成模型的键值记忆库;MoE 版本把参数和计算解耦,每次只激活部分专家,DeepSeek、Qwen 的大模型都在用
  6. logits:最后一层出来后再过 RMSNorm 和 lm_head。lm_head 本质是一个词表大小的矩阵,每行是一个 token 的代表向量,输出就是隐藏向量和每个 token 的匹配分数,下一个 token 从这些分数里采出来

不过这里藏着一个小问题。今天我们一直在描述一遍前向传播,但生成是逐 token 进行的:每生成一个新 token,都要把变长了一位的整段序列重新送进模型。如果每一步都把历史 token 的 K、V 从头重算一遍,序列越长算得越慢,而且绝大部分计算是完全重复的。这个浪费怎么消除,就是下一篇的主角 KV Cache。我们明天继续。

参考


学习大模型推理的嵌入与位置编码

在上一篇中,我们看了分词器如何用 BPE 算法把一句话切成子词,再映射成一串整数,也就是 token id 序列。这串整数是分词阶段的终点,却不是模型计算的起点。

token id 说到底只是编号,和字典里每个词条的序号没有区别,编号本身不带任何语义信息。模型真正处理的是向量。今天这篇就来讲 token id 之后发生的事:它怎么先被嵌入层变成稠密向量,又怎么被位置编码注入顺序信息,最后才进入 Transformer 层参与计算。

嵌入层查表

嵌入(Embedding) 是把离散的 token id 映射为连续向量的过程。它的实现非常直接:就是一张形状为 vocab_size × hidden_size 的二维矩阵,每一行对应词表里一个 token 的向量。token id 进来,按行号取出对应那一行,查表完成,没有任何复杂运算。

以 Qwen3-0.6B 为例,它 config 里的 vocab_size 是 151936,隐藏层维度是 1024,所以嵌入矩阵就是一个 151936 行、1024 列的浮点数表格。上一篇讲过,这个 vocab_size 比词表实际条目略多,多出来的是对齐预留位,不影响查表。每个 token 被表示成一个 1024 维的稠密向量(Dense Vector),即每个维度都是一个实数、没有大量零元素的向量。

embedding-lookup.jpg

这张表是在训练过程中和模型其他参数一起学出来的。学出来的结果有一个著名性质:语义相近的词,向量在空间中也相近。这就是词向量(Word Embedding) 的语义性。2013 年 Mikolov 等人提出 word2vec,论文里给出了一个流传至今的例子:

king - man + woman ≈ queen

对 king 的向量减去 man 的向量、加上 woman 的向量,结果最接近的词是 queen。这说明向量里编码了性别、王室身份这类语义维度,词与词的关系变成了可以计算的向量运算。他们在随后的另一篇论文里用词类比任务对这类线性关系做了系统研究,比如国家与首都、形容词比较级、动词时态都能用向量加减算出来。

词向量在语义空间里的分布大致如下图所示:

word-vector-space.jpg

要注意的是,大模型词表里的单位是子词(subword)而不是完整的词,像 unbelief 会被拆成 un、belief 两个 token,各有各的向量。所以现代模型的嵌入表更像是子词向量表,词一级的语义靠模型后续层组合出来。

动手看看嵌入矩阵

下面我们通过一个简单的示例来体验下。用 Hugging Face transformers 加载 Qwen/Qwen3-0.6B,直接看它的嵌入层:

import torch
from transformers import AutoModelForCausalLM

model_name = "Qwen/Qwen3-0.6B"
model = AutoModelForCausalLM.from_pretrained(model_name, dtype=torch.float32)

# 输入嵌入层,本质就是一个 nn.Embedding
emb = model.get_input_embeddings()
print(emb.weight.shape)

输出:

torch.Size([151936, 1024])

可以看到,形状正是 vocab_size × hidden_size。词表里 15 万多个 token,每个都有自己专属的一行向量。

再验证一下语义性。取两组词算余弦相似度(Cosine Similarity),它衡量两个向量方向的接近程度,取值在 -1 到 1 之间,越接近 1 表示越相似:

import torch.nn.functional as F
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained(model_name)

def vec(word):
    # 这几个词都是单个 token,直接取出对应行的向量
    token_id = tokenizer.encode(word, add_special_tokens=False)[0]
    return emb.weight[token_id]

pairs = [("猫", "狗"), ("猫", "汽车")]
for a, b in pairs:
    sim = F.cosine_similarity(vec(a), vec(b), dim=0)
    print(f"{a} vs {b}: {sim.item():.4f}")

输出:

猫 vs 狗: 0.4233
猫 vs 汽车: 0.1562

「猫」和「狗」同为动物,向量方向明显比「猫」和「汽车」更接近。词向量的语义是统计意义上的相近,绝对数值谈不上大,但相对关系是清楚的。

注意力不知道顺序

嵌入向量解决了语义问题,但还有一个问题没解决:顺序。

大模型的核心是注意力机制,细节我们留到下一篇展开。这里只需要知道一点:注意力计算是把一批向量放在一起两两做点积,它关心的是向量的集合,不关心谁先谁后。把输入顺序打乱,只要还是那几个向量,算出来的结果就不变。这个性质有个专门的名字,叫置换等变(Permutation Equivariance):输入怎么排列,输出就怎么跟着排,计算本身对先后顺序零感知。

口说无凭,我们用 NumPy 手写一个最简注意力验证一下,Q、K、V 投影用随机矩阵代替:

这里的注意力实现做了大量简化,你不需要看懂每一行,注意力机制的细节下一篇会专门展开。现在只需关注实验的设计:同一批向量,打乱顺序送进去,看结果变不变。

import numpy as np

rng = np.random.default_rng(0)
d = 8
Wq, Wk, Wv = (rng.normal(size=(d, d)) for _ in range(3))  # 三个 8×8 投影矩阵

def attention(x):
    # 最简自注意力:softmax(QKᵀ/√d)V
    # @ 是矩阵乘法运算符:(4,8) @ (8,8) -> (4,8),一次算出所有 token 的投影
    # Wq, Wk, Wv 是可学习的参数(真实模型里从训练中学出来,这里用随机矩阵代替)
    q, k, v = x @ Wq, x @ Wk, x @ Wv
    # k.T 是转置,(4,8) @ (8,4) -> (4,4),得到 4 个 query 对 4 个 key 的两两打分表
    scores = q @ k.T / np.sqrt(d)
    # softmax:对每行分数先取 exp 再归一化,转成和为 1 的权重
    # 分数越高的 key 权重越大;减去最大值是为了防止 exp 数值溢出,不改变结果
    scores = np.exp(scores - scores.max(axis=-1, keepdims=True))
    weights = scores / scores.sum(axis=-1, keepdims=True)
    return weights @ v

x = rng.normal(size=(4, d))  # 4 个 token,每个 8 维
out1 = attention(x)

perm = [2, 0, 3, 1]          # 打乱后的 token 顺序
# 花式索引:用数组当下标,按给定顺序取行
# x[perm] 等价于 [x[2], x[0], x[3], x[1]],即打乱语序后的输入
out2 = attention(x[perm])

# argsort 求逆排列:inv[i] 表示原位置 i 的 token 被打乱到了哪里
# out2[inv] 把打乱的输出按原顺序排回去,这样才能和 out1 逐位置对比
inv = np.argsort(perm)
print(f"打乱前后的最大差异: {np.abs(out1 - out2[inv]).max():.2e}")

输出:

打乱前后的最大差异: 1.78e-15

打乱后每个位置的输出,和它原来位置的输出完全一致,差异是 10⁻¹⁵ 级别的浮点误差。注意力确实对顺序没有任何感知。

这会带来一个直观的问题。「狗咬人」和「人咬狗」用的是同样的三个字,分词后 token 集合一样,查出来的向量集合也一样。如果没有位置信息,在注意力看来这两句话完全等价,但它们的含义显然相反。中文是这样,英文里也一样。

所以必须在向量进入 Transformer 之前,把顺序信息注入进去,这就是位置编码(Positional Encoding) 要干的事。早期的循环(RNN)和卷积(CNN)结构天然按顺序读文本,而纯注意力结构本身没有顺序概念,位置信息只能显式注入。位置编码方案的好坏,直接影响模型对语序、指代、因果这类依赖顺序的语言现象的理解能力。

从文本到 Transformer 层的完整链路如下:

position-encoding-pipeline.png

位置编码的做法经历了几次演进,我们按时间线依次看。

正弦绝对位置编码

2017 年的 Transformer 原论文 Attention Is All You Need 用的是正弦位置编码(Sinusoidal Positional Encoding)。做法是给每个位置生成一个固定的向量,偶数维度填正弦值,奇数维度填余弦值,不同维度用不同频率。位置编码的维度和嵌入向量相同(都是论文里的 d_model),所以两者能直接相加,加出来的和就是带进模型的向量。

为什么偏偏是正弦和余弦?我们可以从最直接的想法倒推。给向量注入位置信息,最容易想到三个办法,但每个都有毛病:

  1. 直接用位置编号:把 0、1、2、3 这样的整数塞进向量。数值无界,位置到几千的时候,编号比嵌入值大好几个数量级,训练不稳定
  2. 编号归一化:除以序列长度压到 0 到 1 之间。数值有界了,但同一个值在不同长度的序列里含义不同,0.5 在 10 个词的句子里是第 5 个词,在 1000 个词的句子里是第 500 个词
  3. 二进制编码:把位置写成二进制数,每一位占一个维度。有界、唯一、和序列长度无关,前面几个问题都解决了,但每一位都在 0 和 1 之间硬跳变,不平滑

正弦编码可以看成二进制编码的连续版。二进制从低位到高位,翻转周期按 2、4、8、16 翻倍;正弦编码从低维到高维,波长同样按几何级数拉长,区别只是把硬跳换成了平滑的正弦波。

论文里的公式长这样:

pe.png

公式里 pos 是位置编号,i 是维度编号。关键在分母 10000^(2i/d),把公式换个写法 sin(pos × ω),其中 ω = 1/10000^(2i/d),这个 ω 就是每个维度的频率:ω 越大,pos 每增加 1,正弦波走得越快;ω 越小,波走得越慢。维度编号 i 越大,指数越大,ω 就越小。以 8 维为例,四个维度对的频率和波长如下:

维度对 i频率 ω波长(约多少个位置)
016.3
10.163
20.01628
30.0016283

频率按 ω = 1/10000^(2i/d) 代入 i 和 d=8 算出;波长是波形走完一个周期需要的位置数,即 2π/ω。10000 恰好是 10⁴,d=8 时分母正好是 10 的整数次幂,所以数字格外整齐。

这就像钟表:秒针转得快,用来分辨相邻的秒;时针转得慢,用来定位大致在几点。只看一根针会有歧义,所有维度合起来,每个位置才有独一无二的指纹。

sinusoidal-position-encoding.png

论文中的两行公式看着唬人,代码实现其实很简单。我们用一个 8 维的迷你版本,把前几个位置的编码算出来看看:

import math

def sinusoidal_pe(pos, d_model=8):
    # 偶数维度填 sin,奇数维度填 cos,频率随维度指数下降
    return [
        math.sin(pos / 10000 ** (i / d_model)) if i % 2 == 0
        else math.cos(pos / 10000 ** ((i - 1) / d_model))
        for i in range(d_model)
    ]

for pos in range(4):
    print(f"位置 {pos}:", [round(x, 2) for x in sinusoidal_pe(pos)])

# 对比一下相邻位置和相隔很远的位置,编码差多少
def dist(a, b):
    return math.sqrt(sum((x - y) ** 2 for x, y in zip(a, b)))

for pos in list(range(1, 11)) + [100, 1000]:
    print(f"位置 0 和 {pos} 的距离:", round(dist(sinusoidal_pe(0), sinusoidal_pe(pos)), 2))

# 同样的间隔,换个起点再算一遍
for gap in [1, 5, 100]:
    a = dist(sinusoidal_pe(0), sinusoidal_pe(gap))
    b = dist(sinusoidal_pe(100), sinusoidal_pe(100 + gap))
    print(f"间隔 {gap}: 起点 0 算得 {a:.4f}, 起点 100 算得 {b:.4f}")

输出:

位置 0: [0.0, 1.0, 0.0, 1.0, 0.0, 1.0, 0.0, 1.0]
位置 1: [0.84, 0.54, 0.1, 1.0, 0.01, 1.0, 0.0, 1.0]
位置 2: [0.91, -0.42, 0.2, 0.98, 0.02, 1.0, 0.0, 1.0]
位置 3: [0.14, -0.99, 0.3, 0.96, 0.03, 1.0, 0.0, 1.0]
位置 0 和 1 的距离: 0.96
位置 0 和 2 的距离: 1.69
位置 0 和 3 的距离: 2.02
位置 0 和 4 的距离: 1.86
位置 0 和 5 的距离: 1.3
位置 0 和 6 的距离: 0.66
位置 0 和 7 的距离: 0.98
位置 0 和 8 的距离: 1.7
位置 0 和 9 的距离: 2.14
位置 0 和 10 的距离: 2.15
位置 0 和 100 的距离: 2.21
位置 0 和 1000 的距离: 2.4
间隔 1: 起点 0 算得 0.9641, 起点 100 算得 0.9641
间隔 5: 起点 0 算得 1.2962, 起点 100 算得 1.2962
间隔 100: 起点 0 算得 2.2097, 起点 100 算得 2.2097

从运行结果我们可以看到三个规律。一是不同维度的变化速度不同:前两个维度变得最快,位置每加 1 数值就明显不同;越靠后的维度变得越慢,最后一对几乎不动,和上面的频率表完全对得上。二是编码距离只和间隔有关,和起点无关:间隔同为 1,位置 0 到 1 和位置 100 到 101 算出的距离都是 0.9641,间隔 100 时两个起点都算得 2.2097。这不是巧合,用差角公式可以严格证明:每个维度对对距离平方的贡献是 2(1 − cos(ωk)),只含间隔 k,不含起点。也就是说,正弦编码的距离结构天生就是相对的。三是距离随间隔先升后饱和,不是越远越大:间隔从 1 到 3 距离升到 2.02,间隔 5 又回落到 1.30,间隔 1000 也只有 2.4,之后就在这个量级振荡,不再持续增大。原因还是周期性:每个维度对的贡献最大只有 4,波形转完一圈还会回来,合起来的距离自然有界。不过有界不等于撞车:两个不同位置的编码只是距离有上限,并不会变得相同,慢速维度上总差着一截。多频率组合保证的是每个位置的编码独一无二,而不是距离随间隔无限拉大。

不过这样看还不够直观,我们可以把更多位置和维度画成一张热力图:

import numpy as np
import matplotlib.pyplot as plt

d_model, max_pos = 64, 100
positions = np.arange(max_pos)[:, None]
omega = 1 / 10000 ** (2 * np.arange(d_model // 2) / d_model)
angles = positions * omega                    # (100, 32) 个角度

pe = np.zeros((max_pos, d_model))
pe[:, 0::2] = np.sin(angles)                  # 偶数维度填 sin
pe[:, 1::2] = np.cos(angles)                  # 奇数维度填 cos

plt.figure(figsize=(10, 4))
plt.imshow(pe.T, aspect="auto", cmap="RdBu")
plt.xlabel("position")
plt.ylabel("dimension")
plt.colorbar()
plt.show()

生成的图如下所示:

sinusoidal-heatmap.png

横轴是位置,纵轴是维度。低维区域条纹细密,波形高频振荡;高维区域几乎一整片不变,波长极长。每一列就是一个位置 64 个维度取值的组合,任意两列都不相同,这就是每个位置的指纹。

正弦位置编码有两个特点。一是不需要学习参数,公式直接算出来;二是位置之间存在数学上的线性关系:位置 pos + k 的编码等于位置 pos 的编码乘上一个只和 k 有关的矩阵,效果是把每个维度对旋转 k×ωᵢ 角度,模型有机会学出相对位置的概念。

这个线性关系用三角恒等式展开就能看到,对频率为 ω 的维度对:

sin-cos-pair.png

k 固定时,cos(ωk)sin(ωk) 都是常数,所以 pos + k 的编码恰好是 pos 的编码在每个维度对上旋转 ωk 角度:

sin-cos-pair-2.png

不过它是绝对位置编码(Absolute Positional Encoding),每个位置的编码只跟自己的序号有关,加在语义向量上之后,语义和位置混在同一个向量空间里。

RoPE 旋转位置编码

现在主流的开源模型,包括 Qwen、Llama、DeepSeek、Gemma 等,用的都是 RoPE(Rotary Positional Embedding,旋转位置编码)。它由苏剑林在 2021 年的 RoFormer 论文中提出,论文标题是 Enhanced Transformer with Rotary Position Embedding

上一节我们看到,正弦编码里位置平移 k 等价于把每个维度对旋转一个固定角度。RoPE 把这个关系反过来用:不给嵌入向量加位置信息,而是直接按位置旋转向量。注意力机制里,每个 token 的向量会被变换成 query 和 key 两种角色,靠它们的内积来两两打分。RoPE 的做法是把 query 和 key 向量按维度两两分组,每一组看成一个二维平面上的小箭头。位置为 m 的 token,把它的每组箭头都旋转一个与 m 成正比的角度,位置越靠后,转得越多。不同维度组的旋转速度不一样,类似钟表上时针、分针、秒针各转各的,所有维度组的旋转速度由同一个基底频率推出来。要注意旋转只发生在每层注意力的 query 和 key 上,嵌入层出来的向量本身不动。

旋转的效果如下图所示:

rope-rotation.png

旋转操作用矩阵写出来,就是线性代数里标准的二维旋转矩阵,每个维度对各乘一个:

R(θ) = [ cos θ  −sin θ ]
       [ sin θ   cos θ ]

角度 θ = m × ωᵢ,由 token 的位置 m 和这个维度对的频率 ωᵢ 共同决定,位置越靠后,角度越大。是不是很眼熟?和上一节正弦编码用的是一样的思想,区别只在:正弦编码把 sin、cos 的值直接到嵌入向量上,RoPE 把它们组成旋转矩阵在 query 和 key 上。

为什么这样做有效?关键在于注意力算的是 query 和 key 的内积,而两个向量各自旋转之后,它们的内积只取决于转过的角度差。位置 m 的 query 和位置 n 的 key,内积里自动带上了 m - n 这个相对位置。模型不用关心每个 token 的绝对序号,就能知道两个 token 之间隔了多远。相对位置信息就这样自然地进了注意力分数。

我们可以用代码验证一下「内积只取决于角度差」:

def rotate(x, y, pos, omega=1.0):
    # 把二维箭头 (x, y) 按位置旋转 pos * omega 角度
    angle = pos * omega
    return (x * math.cos(angle) - y * math.sin(angle),
            x * math.sin(angle) + y * math.cos(angle))

def dot(a, b):
    # 二维向量的点积(内积):对应分量相乘再求和
    # 几何意义是 |a| × |b| × cos(夹角),方向越一致点积越大
    return a[0] * b[0] + a[1] * b[1]

q = (1.0, 0.0)
k = (0.8, 0.6)

# 两组不同的绝对位置,相对距离都是 2
print(round(dot(rotate(*q, 5), rotate(*k, 3)), 4))
print(round(dot(rotate(*q, 50), rotate(*k, 48)), 4))

# 相对距离变成 5,分数跟着变
print(round(dot(rotate(*q, 5), rotate(*k, 0)), 4))

输出:

0.2127
0.2127
-0.3484

前两个数一模一样:query 在位置 5、key 在位置 3,和 query 在位置 50、key 在位置 48,只要相对距离都是 2,注意力打出的分完全相同,绝对位置被旋转消掉了。第三个数说明相对距离一变,分数立刻跟着变。这里只看了一对维度,真实的 RoPE 是多对维度各自按不同速度旋转,总内积是所有维度对的结果之和,每一对都只和 m - n 有关。

实现时有一个细节:维度配对有两种方式,GPT-J 式的相邻配对(第 0、1 维一组,第 2、3 维一组)和 GPT-NeoX、Llama 式的前后半配对(第 0 维和第 d/2 维一组)。两者只是维度排列顺序不同,数学上完全等价,内积结果不受影响。

将 RoPE 和正弦绝对位置编码放在一起做个对比:

对比项正弦绝对位置编码RoPE
作用对象加在嵌入向量上旋转 query 和 key 向量
位置类型绝对位置内积中自然体现相对位置
可学习参数
语义与位置混在同一向量空间各走各的通道
典型使用者2017 年原始 TransformerQwen、Llama、DeepSeek 等

和正弦编码相比,RoPE 不是把位置向量加到嵌入上,而是直接作用在注意力的 query、key 上,语义向量和位置信息互不污染。加上实现简单、没有额外参数,它很快成了新模型的默认选择。苏剑林本人的博客科学空间上有一系列推导文章,想深入数学细节的同学可以去读。英文资料推荐 EleutherAI 的 Rotary Embeddings: A Relative Revolution,它从「内积只依赖相对位置」这个设计目标出发反推出旋转形式,他们的实验还发现 RoPE 的训练收敛更快。

RoPE 还有一个性质很符合直觉:两个 token 的相对距离越远,旋转带来的内积差异越杂乱,注意力分数整体呈衰减趋势。也就是说,模型天然更关注离自己近的 token。这种预先写进模型结构里的倾向叫归纳偏置(Inductive Bias),它和自然语言的局部性是一致的。

这个衰减趋势也可以用代码验证。取一个最干净的情形:q 和 k 在每个维度对上都是同向的单位向量,旋转后的内积就等于各维度对 cos((m − n) × ωᵢ) 之和,直接看它随距离的变化:

import numpy as np

d = 128
omega = 1 / 10000 ** (2 * np.arange(d // 2) / d)  # 64 个维度对的频率

for dist in [0, 1, 5, 10, 20, 50, 100, 200, 400]:
    score = np.cos(dist * omega).sum() / (d // 2)  # 归一化,满值为 1
    print(f"距离 {dist}: {score:.3f}")

输出:

距离 0: 1.000
距离 1: 0.970
距离 5: 0.737
距离 10: 0.669
距离 20: 0.608
距离 50: 0.546
距离 100: 0.477
距离 200: 0.306
距离 400: 0.278

距离 0 时内积满值 1,距离拉到 400 时降到 0.28,一路往下。真实的 q、k 方向各异,曲线会有波动,但衰减的整体趋势一致。

其他位置编码

除了正弦编码和 RoPE 这两条线,历史上还有两条路线也简单了解下。BERT 和早期的 GPT 用的是学习式绝对位置编码(Learned Absolute Positional Embedding),给每个位置编号也配一张可训练的查表,和词嵌入一样从数据里学。其实 Transformer 原论文就对比过这条路线,实验发现学习式和正弦版的效果几乎一样,最后选正弦版是出于一个前瞻考虑:公式编码有可能外推到比训练时更长的序列。学习式有个绕不过去的限制:表的长度在训练时就定死了,想支持更长的文本就得重新学,灵活性不如公式编码,后来主流模型基本都放弃了这条路线。

另一条是 ALiBi(Attention with Linear Biases,线性偏置注意力),论文标题叫 Train Short, Test Long。它不改任何向量,直接在注意力分数上减去一个和距离成正比的惩罚项,距离越远扣分越多,把「优先关注近处」写死在公式里。BLOOM、MPT 等模型采用过它,长度外推表现不错,但长程依赖场景下不如 RoPE 灵活,近年的新模型里已经很少见了。

几种代表性方案讲完,回过头总结下,一个好的位置编码应该满足下面这些条件:

  • 唯一性:每个位置要有独一无二的编码,不同位置不能撞车
  • 有界性:编码数值要有界,不能随位置编号无限膨胀,否则会淹没语义信息
  • 相对性:模型关心的往往是两个词隔多远,编码最好能表达相对距离
  • 可外推:训练时没见过的更长序列,推理时编码依然合理
  • 确定性:同样的位置永远算出同样的编码

用这几条标准对照一遍:学习式编码输在可外推,表长训练时就定死了;ALiBi 把相对性简化成线性距离惩罚,换来了外推,牺牲了长程依赖的灵活性;正弦编码五条都满足,但相对位置藏在加法里,要靠模型自己学出来;RoPE 也是五条都满足,相对位置还直接进了内积,这就是它成为主流的原因。

位置编码与上下文长度

位置编码还决定了一件工程上很实际的事:模型能处理多长的上下文。

训练时模型只见过有限范围内的位置。比如训练最大长度是 4096,那么 RoPE 里超出 4096 的旋转角度模型从没见过。推理时硬塞更长的文本,注意力分数会乱掉,生成质量明显下降。这就是位置编码的外推(Extrapolation) 问题,即模型在训练长度之外的表现。

我们看 Qwen3-0.6B 的配置里和位置相关的两个字段:

print(model.config.max_position_embeddings)
print(model.config.rope_parameters["rope_theta"])  # transformers 4.x 里是 config.rope_theta

输出:

40960
1000000

其中 max_position_embeddings 是模型位置编号的上限,Qwen3-0.6B 这里是 40960,比官方标称的 32K 原生上下文略留了余量。rope_theta 就是上一节说的那个基底频率,各维度组的旋转速度都由它推出来,Qwen3 把它从早期模型常用的 10000 调大到了 1000000,让高频维度的旋转放缓,为长上下文留余地。

围绕外推问题有一系列改进方法。位置插值(Position Interpolation,PI) 把长文本的位置等比压缩回训练窗口内;NTK-aware 缩放 调整 RoPE 的基底频率,让不同转速的维度组得到不同程度的拉伸。NTK 这个名字来自神经正切核(Neural Tangent Kernel) 的理论启发,最早是 Reddit 上的一篇社区帖子提出的。YaRN 名字是 Yet another RoPE extensioN 的缩写。它在 NTK 思路上对高频和低频分量区别处理,再加一个注意力温度系数,用少量微调就能把上下文窗口扩到训练长度的好几倍。YaRN 论文里把 LLaMA 系列扩到了 128K。Qwen3 官方也说明了通过 YaRN 可以把上下文从 32K 扩到 128K。这些方法涉及不少公式和细节,我们这里就点到为止了,感兴趣的同学可以进一步查阅相关资料。

小结

今天我们学习了 token id 之后的第一步:

  1. 嵌入层是一张 vocab_size × hidden_size 的查表,把 token id 变成稠密向量;训练让语义相近的词向量相近,经典的 king - man + woman ≈ queen 就是这种语义性的体现
  2. 动手用 Qwen3-0.6B 验证了嵌入矩阵的形状,并用余弦相似度对比了语义相近词与无关词的差异
  3. 注意力是置换等变的,本身不包含顺序信息,我们用 NumPy 最简注意力做了实验:打乱输入,输出只是跟着重排,逐位置的值完全不变
  4. 好的位置编码有五条标尺:唯一、有界、能表达相对距离、可外推、确定。正弦编码可以看成二进制编码的连续版,用一组几何级数的频率给每个位置生成独一无二的指纹
  5. RoPE 把「旋转」从正弦编码的副产品变成了主角:按位置旋转 query 和 key,让相对位置自然体现在内积里,我们还用代码验证了它的距离衰减性质
  6. 位置编码限制了上下文长度,位置插值、NTK、YaRN 等方法通过调整位置或旋转频率做长度外推

向量准备好了,位置信息也注入进去了,接下来就是真正的计算核心:这些向量进入 Transformer 层之后,注意力机制到底是怎么两两打分的,前向传播的完整数据流又长什么样。我们明天继续。

参考


学习大模型推理的分词:从文本到 Token

在上一篇中,我们把一次请求的完整旅程走了一遍,画出了整个系列的地图:你敲下的一句话先经过分词变成 token 序列,然后模型在 Prefill 阶段一口气读完问题,接着进入 Decode 循环一个 token 一个 token 地生成回答,每一步还要经过采样挑出下一个词,最后反分词把 token 还原成文字流式返回给你。

今天我们从地图的第一站开始,把分词这个环节单独拿出来学习。

为什么需要分词

神经网络的计算基本都是矩阵乘法,输入必须是一串数字。但用户给的是自然语言文本,中间需要一座桥把文字翻译成数字,这座桥就是分词(Tokenization)。分词做两件事:先把文本切成一个个片段,每个片段叫一个 token;再查一张对照表,把每个 token 换成一个整数编号,也就是 token id。

这张对照表叫词表(Vocabulary),它在模型训练之前就定好了,训练完成后固定不变。词表里的每个 id 对应模型嵌入层里的一行向量,模型实际读进去的就是这些向量。嵌入层的内容我们留到下一篇讲,今天只需要知道 id 是文本和模型之间的中间货币。

光说对照表可能有点抽象,直接打开它看看。用 Hugging Face transformers 库的 AutoTokenizer 加载分词器,模型选 Qwen3 系列最小的 Qwen3-0.6B,只下载分词器配置,10 MB 出头,不用下载模型权重:

from transformers import AutoTokenizer

# 加载 Qwen3-0.6B 的分词器
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")

# 词表就是一个「token 文本 → id」的大字典
vocab = tokenizer.get_vocab()
print(len(vocab))

# 按 id 排序,看看队头和队尾
vocab_by_id = sorted(vocab.items(), key=lambda kv: kv[1])
print(dict(vocab_by_id[:10]))
print(dict(vocab_by_id[-3:]))

运行结果:

151669
{'!': 0, '"': 1, '#': 2, '$': 3, '%': 4, '&': 5, "'": 6, '(': 7, ')': 8, '*': 9}
{'</tool_response>': 151666, '<think>': 151667, '</think>': 151668}

可以看到,词表就是一个 15 万多个条目的大字典。排在最前面的是标点符号这些 ASCII 字符,id 从 0 开始;排在最后面的是 <think> 这种不对应自然语言的条目,它们属于特殊 token,后面会专门讲。

细心的读者可以顺手试试 vocab["你好"],会发现查不到,报 KeyError 错误。这是因为词表的 key 并不是中文词汇,而是把 UTF-8 字节逐字节映射成可打印字符后的形式(GPT-2 传下来的做法,这样词表文件里不会出现不可见字符),比如「你好」的 6 个字节 b'\xe4\xbd\xa0\xe5\xa5\xbd' 映射后的 key 是 'ä½łå¥½',可以用 tokenizer.tokenize("你好") 来查。

还有一个细节:词表条目是 151669,而上一篇 config 里的 vocab_size 是 151936。前者是分词器词表的条目数,后者是嵌入表的行数,多出来的 267 行没有对应的 token,是把嵌入表凑成 128 的倍数做对齐的预留位,实际用不到。

整个过程可以用一张图概括:

why-tokenization.jpg

分词这一步发生在模型之外,由一个叫分词器(Tokenizer)的独立组件完成。每个模型发布时都会带上自己专属的分词器,不同模型的词表不一样,切分结果也不一样,所以分词器和模型必须配套使用,不能混着用。

Token 不一定是字或词

刚接触这个概念时,很多人会默认一个 token 就是一个字或者一个单词,其实都不是。Token(词元) 是分词器切出来的最小单位,它可能是整个词、词的一部分、单个字符,甚至字符的一部分。

我们用刚才加载的 Qwen3 分词器跑几个真实例子:

texts = ["我们今天来学习分词", "We are learning tokenization today", "unbelievable"]
for text in texts:
    ids = tokenizer.encode(text)
    # 逐个 decode 出 token 片段,把空格换成 ␣ 方便看
    tokens = [tokenizer.decode([i]).replace(" ", "␣") for i in ids]
    print(f'"{text}" → {len(ids)} 个 token:{" / ".join(tokens)}')

运行结果:

"我们今天来学习分词" → 6 个 token:我们 / 今天 / 来 / 学习 / 分 / 词
"We are learning tokenization today" → 6 个 token:We / ␣are / ␣learning / ␣token / ization / ␣today
"unbelievable" → 3 个 token:un / belie / vable

可以看到几个规律:

  • 常见中文词组是一个整体:「我们」「今天」「学习」各自占一个 token,但不太常见的组合会被拆开,「分词」就拆成了「分」和「词」
  • 常见英文单词是一个 token,长词会被拆成子词(Subword):tokenization 拆成 token 和 ization,unbelievable 拆成三段
  • 空格是有意义的:英文里单词前的空格通常会并进 token,上面用 ␣ 标出了空格,比如「␣today」整体是一个 token

text-to-tokens-concept.jpg

中英文的差异尤其值得关注。早期针对英文优化的分词器处理中文很浪费,一个汉字可能占 2 到 3 个 token;现在主流的多语言分词器(比如 Qwen 用的)对中文友好了很多,常见汉字和词组大约 1 个 token,生僻字仍然会拆得更碎。

这不是一个纯学术问题。API 计费按 token 算,上下文长度按 token 算,速率限制也按 token 算。同样一段话用中文写还是用英文写,token 数量可能差出一截,账单也跟着差一截。估算成本时拿字数当 token 数,是会算错的。

BPE:从数据压缩借来的算法

那分词器是怎么决定在哪里下刀的?目前主流大模型用的都是 BPE(Byte Pair Encoding,字节对编码) 或者它的变体。

BPE 的历史有点意思。它本来是 Philip Gage 在 1994 年提出的一个数据压缩算法,和自然语言处理没有关系。2016 年,Sennrich 等人在一篇机器翻译论文里把它改造成了子词切分方法,用来解决翻译模型遇到生僻词就抓瞎的问题,这篇论文后来拿了 ACL 2026 的 Test of Time 奖。2019 年 GPT-2 又把它改造成字节级 BPE(Byte-level BPE):不再以字符为起点,而是以 256 个字节为初始词表。这么一改,任何语言、任何符号、任何 emoji 都能被表示,彻底不会出现分词器不认识某个字的情况。

BPE 论文全名是 Neural Machine Translation of Rare Words with Subword Units,arXiv 编号 1508.07909。想了解从零实现一个 BPE 分词器长什么样,可以看 Sebastian Raschka 的 BPE from scratch 一文。

BPE 的核心思路就一句话:反复把语料里出现最频繁的相邻两个单位合并成一个新单位,直到词表达标。训练分词器的过程就是学出一张合并规则表,分词时按同样的规则顺序套用到新文本上。

用一个具体例子演示。假设我们的全部训练语料只有四个词:low 出现 5 次,lower 出现 2 次,newest 出现 6 次,widest 出现 3 次。初始状态每个词拆成单个字符,词尾加一个特殊标记表示单词结束:

low    → l o w </w>     (5 次)
lower  → l o w e r </w> (2 次)
newest → n e w e s t </w>(6 次)
widest → w i d e s t </w>(3 次)

然后开始循环:统计所有相邻对的出现次数,把最高频的一对合并,加入词表。前几轮的过程如下:

轮次最高频相邻对合并结果出现次数
1(e, s)es9
2(es, t)est9
3(l, o)lo7
4(lo, w)low7
5(n, e)ne6
6(ne, w)new6
7(new, est)newest6

表里省略了和词尾标记 </w> 的合并,比如 (est, </w>) 出现 9 次,实际顺序里它就排在第 3 轮,为了演示直观,我就去掉了。

几轮之后,est、low、newest 这些高频片段各自成了词表里的整体 token。训练好的分词器遇到新文本时,按学好的合并顺序逐条套用。最有价值的情况是遇到没见过的词,比如 slowest:s 开头的部分没有对应规则,退回单个字符,但后半段 low 和 est 都在词表里,最终切成 s + low + est。这就是子词切分的精髓:常见词走整体,生僻词拆成熟悉的零件,永远不会无法表示。

真实模型的词表规模远大于这个玩具例子。GPT-2 的词表是 50257 个 token,Qwen3 是 15 万多个。词表大,常见词和词组都能整体表示,同样文本切出来的 token 数就少,推理更省;但词表越大嵌入层参数越多,词表里冷门 token 的训练也越不充分,所以规模是权衡出来的。

整个训练循环可以画成这样:

bpe-training-loop.jpg

保存下来的这两个文件就在模型仓库里,随分词器一起下发:

$ ls  ~/.cache/huggingface/hub/models--Qwen--Qwen3-0.6B/snapshots/* 
config.json             merges.txt              tokenizer_config.json   vocab.json
generation_config.json  model.safetensors       tokenizer.json

这 7 个文件分两组,4 个属于分词器,3 个属于模型:

文件是什么
vocab.json词表,token(字节映射形式)到 id 的对照表
merges.txt合并规则表
tokenizer.json前两个文件的打包加强版,还包含预分词规则和特殊 token 定义,fast tokenizer 实际加载的是它
tokenizer_config.json分词器配置,包括特殊 token 的名字和聊天模板,后面聊天模板一节还会见到它
config.json模型结构配置,上一篇已经见过
generation_config.json生成的默认参数,temperature、top_p 这些的出厂值
model.safetensors模型权重本体,1.5 GB,推理的主角

只跑分词实验的话,AutoTokenizer 只需要前 4 个文件,总共 10 MB 出头;后 3 个是上一篇跑模型推理时下载的,和分词无关。

词表我们前面已经看过了,那么真实的合并规则长什么样呢?不妨打开 merges.txt 文件瞧瞧,它的开头是这样的:

#version: 0.2
Ġ Ġ
ĠĠ ĠĠ
i n
Ġ t
...
e r
...
def ine
def ault

每行一条规则,两个符号写在一行,表示分词时看到这两个相邻的单位就合并成一个,比如 i n 表示 i 后面跟着 n 时合并成 in。Ġ 是空格的字节映射形式,Ġ t 就是「空格加 t 合并成 ␣t」。可以看到后面还有 def ine、def ault 这种明显从代码语料里学出来的规则。顺序就是优先级:训练时越早学出的合并排得越靠前,分词新文本时从字节开始,反复挑当前序列里排名最靠前的对子来合并,保证切分结果和训练时的学习顺序一致。上面玩具例子里那张 7 轮的表,就是一个迷你版的 merges.txt

用 transformers 观察分词

下面我们再来看一个示例,看看这个分词器面对一句中英文混合、还带标点 emoji 的话会怎么切:

from transformers import AutoTokenizer

# 加载 Qwen3-0.6B 的分词器
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")

text = "你好,world!今天的 weather 真不错 😊"

# encode:文本 → token id 序列
ids = tokenizer.encode(text)
print(ids)

# 逐个看每个 token 对应的文本片段
for i in ids:
    print(i, repr(tokenizer.decode([i])))

# decode:token id 序列 → 还原成文本
print(tokenizer.decode(ids))

运行结果如下:

ids: [108386, 3837, 14615, 6313, 106560, 9104, 10236, 250, 253, 100832, 26525, 232]
108386 '你好'
3837 ','
14615 'world'
6313 '!'
106560 '今天的'
9104 ' weather'
10236 ' �'
250 '�'
253 '�'
100832 '不错'
26525 ' �'
232 '�'
decode: 你好,world!今天的 weather 真不错 😊

这 12 个 token 里有几个点值得注意:

  • 中文:「你好」「今天的」「不错」都是整体 token,但「真」连同前面空格这个组合没有对应的合并规则,被拆成了 3 个字节级 token。这不是分词器坏了,而是字节级 BPE 的兜底机制在起作用,词表里没有覆盖的组合会退回到字节表示
  • 英文:world 是整体 token,「 weather」带着前导空格一起算一个 token,和前面说的规律一致
  • 标点:中文逗号和感叹号各自独立成 token
  • emoji:笑脸连同前面的空格也被拆成了字节级 token,和「真」一样,都是词表没覆盖到组合时退回字节表示的结果

最后一行的 decode 输出和输入一字不差,说明 encode 之后再 decode 可以无损还原。旅程地图里的反分词环节,干的就是 decode 这件事。

特殊 token

词表里除了正常文本切出来的 token,还有一类特殊 token(Special Token),刚才词表队尾的 <think> 就属于这类。它们不对应任何自然语言文字,作用是充当结构标记,告诉模型一段文本从哪里开始、到哪里结束、哪里是补齐的空白。最常见的三个:

特殊 token全称作用
BOSBeginning of Sequence标记序列开头
EOSEnd of Sequence标记序列结束,模型生成出它就停止
PADPadding批处理时把短序列补齐到同样长度

其中 EOS 很重要,上一篇讲过 Decode 循环是一个 token 一个 token 地生成,那模型怎么知道该停了?答案就是训练时教会它在回答结束时生成 EOS,推理引擎检测到这个 id 就终止循环。

各家模型用哪些特殊 token、取什么名字并不统一。可以打印出来看看 Qwen3 的情况:

print(tokenizer.bos_token)  # None
print(tokenizer.eos_token)  # <|im_end|>
print(tokenizer.pad_token)  # <|endoftext|>

可以看到 Qwen3 没有 BOS,EOS 用的是 <|im_end|>,PAD 用的是 <|endoftext|>

聊天模板

平时我们在对话框里打字,很容易以为那句话原封不动就进了模型。其实没有。你发出去的每一条消息,都会先被包进一个固定格式里,这个格式就是聊天模板(Chat Template)

聊天模型是在多轮对话数据上训练出来的,训练数据里每轮对话都有明确的角色标记,谁是系统提示、谁是用户、谁是助手,边界清清楚楚。推理时必须用同样的格式包装输入,模型才知道现在轮到谁说话了。用 apply_chat_template 看一下 Qwen3 实际拼出来的字符串:

messages = [
    {"role": "system", "content": "你是一个有帮助的助手。"},
    {"role": "user", "content": "什么是分词?"},
]
text = tokenizer.apply_chat_template(
    messages, tokenize=False, add_generation_prompt=True
)
print(text)

输出如下:

<|im_start|>system
你是一个有帮助的助手。<|im_end|>
<|im_start|>user
什么是分词?<|im_end|>
<|im_start|>assistant

可以看到,每条消息被 <|im_start|> 加角色名开头、<|im_end|> 结尾包住,最后还追加了一个 <|im_start|>assistant 的开头,这就是 add_generation_prompt=True 的作用,相当于把话头递给模型:接下来该你说了。模型生成回答后输出 <|im_end|>,EOS 检测到,Decode 循环结束。

chat-template-wrapping.jpg

这个格式其实不是 Qwen 自创的,它是 OpenAI 在 2023 年发布 ChatGPT API 时提出的 ChatML(Chat Markup Language),im 是 instant message 的缩写。Qwen 系列沿用了这套格式。

不同模型的聊天模板差别很大,比如和 Llama 3 对比一下:

Qwen(ChatML 风格)Llama 3
一轮开始`<\im_start\>` 加角色名`<\start_header_id\> 角色名 <\end_header_id\>`
一轮结束`<\im_end\>``<\eot_id\>`
序列开头`<\begin_of_text\>`

功能上等价,写法完全不同。所以不能把 Qwen 的模板套给 Llama 用,格式错了模型表现会明显变差。同样也不能拿这套模板机制去套 base 模型:base 模型训练时没见过 <|im_start|> 这些标记,你把对话格式喂给它,它只会顺着往下续写文本,而不是回答你的问题。

base 模型是预训练完就直接发布的模型,训练目标只有一个:根据上文预测下一个 token。它学的是文本本身的分布,所以只会续写。对话模型(Chat 或 Instruct 模型)是在 base 模型的基础上,再用带角色标记的对话数据做微调和对齐,才学会按格式回答。从模型名字能看出来,比如 Qwen3-0.6B 是对话模型,对应的 Qwen3-0.6B-Base 就是 base 模型。

上下文窗口与 token 计数

最后把 token 和两个工程上天天打交道的概念连起来。

第一个是上下文窗口(Context Window),它指模型一次能处理的 token 总数上限,输入加输出一起算。Qwen3-0.6B 原生支持 32K token 的上下文,通过 YaRN(一种基于 RoPE 缩放的上下文长度扩展方法)扩展可以到 128K。

第二个是计费。API 厂商按 token 报价,输入输出分开计价。写应用时预估成本、控制超长文本,第一步都是先数 token。数 token 的方法很简单,用模型对应的分词器 encode 一下看长度就行。如果你用的是 OpenAI 的模型,它家开源了一个快速的 BPE 分词库 tiktoken,两行就能数出来:

import tiktoken
enc = tiktoken.encoding_for_model("gpt-4o")
print(len(enc.encode("你好,world!")))

OpenAI 还提供了一个在线的 Tokenizer 页面,把文本粘进去就能直观看到切分结果和数量,适合不想写代码的时候随手验证:

openai-tokenizer.png

要注意的是 tiktoken 只适用于 OpenAI 自家模型,数 Qwen、Llama 的 token 还是得用各家自己的分词器。词表不同,数出来的结果也不一样。

小结

今天我们把推理旅程的第一站走完了,要点如下:

  1. 为什么分词:模型只认识数字,分词器负责把文本切成 token 再查词表换成 id,它是文本和模型之间的中间货币
  2. Token 不是字也不是词:是切分出来的最小单位,常见词整体一个,长词拆子词,生僻内容退回字节;中英文 token 效率不同,直接影响计费和上下文容量
  3. BPE 算法:源自数据压缩,经 Sennrich 等人引入 NLP、GPT-2 发展到字节级,核心就是反复合并最高频相邻对
  4. 特殊 token:BOS、EOS、PAD 是结构标记,EOS 同时承担着终止 Decode 循环的职责
  5. 聊天模板:用户消息进模型前会被包成带角色的对话格式,Qwen 用 <|im_start|><|im_end|>,Llama 用另一套,这也是 base 模型和 chat 模型表现差异的来源之一

现在我们已经能把一句话变成一串 token id 了。但 id 只是编号,模型真正吃的是每个 id 对应的向量,而且这些向量里还得想办法带上位置信息,不然模型分不清「你打我」和「我打你」。嵌入和位置编码就是下一篇的主题,我们明天继续。

参考


大模型推理介绍:从一次提问说起

平时用豆包聊天、用 Claude Code 或 Codex 写代码,几乎成了每天的日常。但每次敲下回车之后,从第一个字蹦出来到整段回答写完,中间到底走过了哪些环节,我之前其实一直说不太清楚。最近在系统补大模型推理和训练相关的知识,于是想写一个系列,顺便记录下学习过程中的笔记,争取把整条链路梳理清楚。

越看越觉得,推理这件事被低估了。训练一个大模型是一次性投入;但是模型上线之后,每一次用户提问产生的推理开销却是日复一日、持续累积的。行业分析估计,企业 AI 的 GPU 预算里 55% 到 80% 花在了推理上。对一个有真实流量的产品来说,上线几周内,推理的累计算力就会超过训练。训练决定了模型能做到什么,推理决定了用户每天实际用到什么。所以 vLLMSGLangllama.cpp 这些名字才会一次次出现在技术圈的讨论里。大家争的,其实都是怎么把推理跑得更快、更省、更能扛并发。

这个系列我们就来系统地学习大模型推理相关的知识。今天是第一篇,先不急着抠细节,而是回答一个最基本的问题:你在对话框里敲下一句话、按下回车,到第一个字跳出来、再到回答逐字生成完毕,这中间到底发生了什么?我们会把一次请求的完整旅程走一遍,画出一张全系列的地图;后面每篇文章,就对应地图上的一个环节。

什么是推理

推理(Inference) 指的是训练好的模型根据输入生成输出的过程。和它相对的概念是 训练(Training),两者的区别可以用一张表说清楚:

对比项训练推理
目的调整模型权重,让模型学会规律用固定的权重生成结果
计算方式前向计算 + 反向传播 + 权重更新只做前向计算
权重状态每步都在变全程不变
发生频率一次性或周期性每次用户请求都在发生
典型用户模型研发工程师所有使用模型的人

简单来说,训练是造模型,推理是用模型。训练时模型要算梯度、更新参数,一次训练动辄占用几千张 GPU 跑上几周甚至几个月;推理时权重已经固定,每次只是把输入送进网络做一遍前向计算,拿到下一个词的预测。

表格里提到的三个词稍微展开一下。前向计算(Forward Pass) 是把输入从网络第一层逐层算到最后一层,得到预测结果的过程,训练和推理都要做这一步。反向传播(Backpropagation) 是拿预测结果和正确答案算出差多少,再沿着网络倒着把这个误差分摊到每个权重上,算出每个权重各自该负多少责任,也就是梯度。权重更新(Weight Update) 则是优化器根据梯度把权重往误差更小的方向挪一小步,模型就是这样一点点「学会」的。这三步组成训练的一次迭代,反复进行成千上万次;而推理只保留第一步,后面两步都不需要,这正是两者计算量差距悬殊的原因。

对绝大多数人来说,推理就是接触大模型的唯一方式。你打开豆包聊天、用 Claude Code 写代码、调用 API 做文本分类,背后发生的都是推理。这个系列研究的对象,就是这个每天被调用亿万次的过程。

从成本结构上看,训练和推理还有一个不对称的地方。训练再贵也是一次性投入,花完就花完了;推理的单价很低,一次请求可能只有几厘钱,但它随着用户量线性增长,永不停歇。一个模型越成功、用户越多,推理的累计开销就越大,最终远远超过当初的训练成本。

training-vs-inference.png

一次请求的完整旅程

现在我们跟着一个请求走一遍。假设你在对话框里输入「合肥今天天气怎么样」,按下回车之后,请求会依次经过下面这些环节:

inference-request-journey.jpg

我们逐一看看每个环节在做什么。

分词:文本变成 token

模型不认识自然语言,它只认识数字。所以第一步是 分词(Tokenization),把输入文本切成一串 token(词元),每个 token 对应词表里的一个编号。token 可以是一个字、一个词、一个标点,甚至半个词。比如「合肥今天天气怎么样」用 Qwen3 的分词器会切成 4 个 token,而同样意思的英文 How is the weather in Hefei today 要切出 9 个。

感兴趣的话可以运行下面几行代码,就能看到切分结果:

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")

for text in ["合肥今天天气怎么样", "How is the weather in Hefei today"]:
    ids = tokenizer.encode(text)
    # 对每个 token id 做 decode,把字节级表示还原成可读的文本片段
    tokens = [tokenizer.decode([i]) for i in ids]
    print(len(tokens), tokens)

# 4 ['合肥', '今天', '天气', '怎么样']
# 9 ['How', ' is', ' the', ' weather', ' in', ' H', 'ef', 'ei', ' today']

可以看到中文按词切得很整,「合肥」整体是一个 token;而英文的 Hefei 因为不在常见词表里,被拆成了 H、ef、ei 三个碎片,token 数一下子多出不少。这也是为什么同样的语义,不同语言的推理成本会不一样。

分词是文本世界和模型世界之间的翻译官,它直接影响模型能处理的上下文长度、推理的成本核算(API 都按 token 计费),甚至影响模型在某些语言上的表现。这个话题比想象中深,我们下一篇专门来讲这块。

Prefill:一口气读完问题

分词之后进入 Prefill(预填充) 阶段。模型一次性并行处理输入的全部 token,算出每个位置的表示,并生成第一个新 token。这个阶段的特点是输入一次性给齐,可以充分并行计算,所以它是 计算密集型(compute-bound) 的,GPU 的算力利用率很高。你按下回车之后等待第一个字出现的那段时间,主要就是 Prefill 花掉的。

那 Prefill 具体在算什么?要先知道,Transformer 模型不是一个单独的网络,而是几十层结构相同的 层(Layer) 叠起来的,比如 Qwen3-0.6B 就有 28 层,输入从第一层进去,逐层加工,从最后一层出来。每一层做的核心工作是 自注意力(Self-Attention),粗略理解就是:让每个 token 都和序列里的其他 token「对一遍话」,吸收上下文信息之后更新自己的表示。

口说无凭,我们把模型加载进来,亲眼看看这些层:

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-0.6B")

# 模型一共有多少层
print(model.config.num_hidden_layers)

# 打印第 0 层,看看一层里面都有什么
print(model.model.layers[0])

运行结果如下:

28
Qwen3DecoderLayer(
  (self_attn): Qwen3Attention(
    (q_proj): Linear(in_features=1024, out_features=2048, bias=False)
    (k_proj): Linear(in_features=1024, out_features=1024, bias=False)
    (v_proj): Linear(in_features=1024, out_features=1024, bias=False)
    (o_proj): Linear(in_features=2048, out_features=1024, bias=False)
    (q_norm): Qwen3RMSNorm((128,), eps=1e-06)
    (k_norm): Qwen3RMSNorm((128,), eps=1e-06)
  )
  (mlp): Qwen3MLP(...)  # 省略 MLP 部分
  (input_layernorm): Qwen3RMSNorm((1024,), eps=1e-06)
  (post_attention_layernorm): Qwen3RMSNorm((1024,), eps=1e-06)
)

可以看到,这里的 model.model.layers 是一个 28 个元素的列表,每个元素都是结构完全相同的 Qwen3DecoderLayer。往一层里面看,self_attn 就是自注意力模块,它下面的 q_projk_projv_proj 正是下文要讲的 W_Q、W_K、W_V 三个权重矩阵。

「自注意力」的 自(Self) 是相对于早期注意力机制说的:以前的注意力是让一个序列去关注另一个序列(比如翻译时让译文关注原文),而自注意力是让每个 token 去关注 同一个序列里 的其他 token。拿「合肥今天天气怎么样」来说,「天气」的表示会吸收「合肥」和「今天」的信息,从而知道这里问的是某地某时的天气而不是别的。每一层都做一遍这样的信息交换,层数越深,token 的表示就吸收越多的上下文。

那么,每个 token 具体是怎么和其他 token「对一遍话」的呢?分两步。第一步是 查嵌入表(Embedding):模型里存着一张大表,词表里每个 token 编号对应一行向量,Qwen3-0.6B 的词表有 151936 个 token,每行是一个 1024 维的向量,查表就把 token 编号变成了它的初始表示。第二步是 乘权重矩阵:每一层都有自己训练好的三个矩阵 W_Q、W_K、W_V,把每个 token 的向量分别乘上去,得到查询(Q)、键(K)、值(V)三个向量。然后每个 token 拿自己的 Q 去和序列里所有 token 的 K 做匹配,按匹配程度对所有的 V 加权求和,就得到了「读过上下文之后」的新表示。

Q、K、V 这三个向量可以用查资料来类比:查询(Query) 是当前 token 提出的问题「我想找和我相关的信息」;键(Key) 是每个 token 挂在门口的标签「我这里有关于什么的信息」;值(Value) 则是 token 实际携带的内容。拿「天气」的 Q 去和所有 token 的 K 逐个比对,「合肥」「今天」的标签对得上,匹配分数就高,于是它们的 V 就以更大的权重被加进来。

实际模型里,Q、K、V 通常会被拆成多份并行计算。以 Qwen3-0.6B 为例,Q 被拆成 16 份各 128 维(也就是 16 个注意力头),K 和 V 各拆成 8 份 128 维,每个头独立做一遍匹配,最后再把结果拼回去。多头的好处是不同的头可以各看各的角度,有的关注语法,有的关注指代。

上面提到的这些数字,同样可以从模型的 config 里直接读到:

config = model.config
print(config.vocab_size)           # 151936,词表大小
print(config.hidden_size)          # 1024,嵌入向量的维度
print(config.num_attention_heads)  # 16,Q 的注意力头数
print(config.num_key_value_heads)  # 8,K、V 的注意力头数
print(config.head_dim)             # 128,每个头的维度

对照上面打印的层结构验算一下:q_proj 的输出维度 2048 = 16 × 128,k_proj、v_proj 的输出维度 1024 = 8 × 128,正好分别是 Q 和 K、V 所有头拼起来的大小。

Prefill 结束时有两个产物:一个是最后一个位置预测出的第一个新 token;另一个是所有 token 在所有层上算好的 K、V 向量,它们被存进显存,就是后面反复提到的 KV Cache(键值缓存)。有了这个缓存,Decode 阶段每生成一个新 token,只需要算它自己的 Q、K、V,再回头查缓存里历史 token 的 K、V 就行,不用把整个输入重算一遍。可以说 Prefill 的一项重要职责就是为 Decode 备好这份缓存。

为什么缓存只存 K 和 V,不存 Q 呢?这是因为生成新 token 时,只有它在提问,用自己的 Q 去匹配所有历史 token 的 K,再对 V 加权求和。历史 token 在这一步里只是被查询的对象,用到的是它们的 K 和 V。而历史 token 自己的 Q,只在它刚生成的那一步用过一次,之后再也用不上,自然不用存。

回过头看,为什么说 Prefill 是计算密集型的?因为输入 token 一次性到齐,上面这些 Q、K、V 的计算和注意力匹配全都是大矩阵乘法,恰好是 GPU 的 Tensor Core(张量核心,GPU 里专门做矩阵乘法的硬件单元)最擅长的活儿,算力能被充分利用。但凡事有代价:注意力要求每个 token 和每个 token 打交道,这部分的计算量随输入长度近似 平方增长。prompt 从几千 token 涨到几万 token,注意力的计算量不是涨十倍而是涨上百倍。这也是为什么喂给模型一本小说和问它一句话,首 token 的等待时间完全是两个量级。

于是围绕 Prefill 出现了一批专门的优化技术。比如 Chunked Prefill(分块预填充) 把超长输入切成小块,穿插在 Decode 步骤之间分批算,避免一个长 prompt 把其他用户的生成卡住;Prefix Caching(前缀缓存) 则把多个请求共享的 prompt 前缀(比如同一份系统提示词)的 KV Cache 直接复用,跳过重复的 Prefill 计算。这些技术后面在学 KV Cache 和推理引擎调度时再细说,这里先了解一下。

Decode 循环:逐 token 生成

接下来是最关键也最容易被误解的部分。大语言模型本质上只做一件事:给定前面的 token 序列,预测下一个 token。所以生成回答不是一次性算出来的,而是一个循环:

  1. 模型根据已有序列预测下一个 token
  2. 把这个 token 拼到序列末尾
  3. 用新序列再预测下一个
  4. 重复以上步骤,直到生成结束标记或达到长度上限

这个逐 token 生成的过程叫 Decode(解码),这种一个接着一个的生成方式叫 自回归生成(Autoregressive Generation)。下面这张时序图可以看出 Prefill 和 Decode 的关系:

prefill-decode.png

和 Prefill 不同,Decode 每一步只算一个 token。Prefill 时输入一次性到齐,权重从显存读出来一次能被所有输入 token 复用;而 Decode 每步只有一个新 token,大矩阵乘法退化成矩阵乘向量,计算量很小,但每一步仍然要把全部模型权重和攒下来的 KV Cache 完整读一遍。时间花在「读」上而不是「算」上,GPU 的算力大量闲置,所以它是 访存密集型(memory-bound) 的。这就是为什么你在聊天界面里看到的回答是一个字一个字往外蹦的,不是模型在模仿人打字,而是它真的就是这样工作的。

值得注意的是,解码过程中 KV Cache 还在不断变大,每个新 token 都要在每一层留下自己的 K、V。占多少显存,用前面打印的 config 就能算出来:

kv_per_token = 2 * config.num_hidden_layers * config.num_key_value_heads * config.head_dim * 2
print(kv_per_token / 1024, "KB")  # 112.0 KB

式子里第一个 2 是 K 和 V 两份,最后的 2 是 bf16 每个元素占的字节数。也就是说每生成一个 token,KV Cache 就涨 112 KB;一轮对话生成 2000 个 token,光缓存就要 200 多 MB,接近模型权重(约 1.2 GB)的五分之一了。上下文越长、生成越长,显存吃得越多,显存容量也因此成了推理服务能扛多少并发的关键约束。

既然瓶颈在读权重,优化思路也很直接:让读一遍权重服务尽可能多的 token。把多个用户的请求凑成一批一起跑,同一份权重读出一次,就能同时算出几十个请求的下一个 token。vLLM 的 Continuous Batching(连续批处理) 走的就是这条路。

Prefill 和 Decode 一个吃算力、一个吃带宽,两者的优化思路完全不同。NVIDIA Dynamo 这类新框架甚至把它们拆到不同的 GPU 上分别部署,这就是所谓 PD 分离(Prefill-Decode Disaggregation)。

采样:从概率分布里挑一个词

模型每一步输出的其实不是一个确定的词,而是词表里每个 token 的一个分数,这个分数叫 logits。logits 是模型最后一层直接算出来的原始数值,可以是任意实数,有正有负,本身没有概率含义,只有相对大小:分数越高,说明模型越倾向于选这个 token。要把分数变成概率,需要过一遍 softmax:先对每个分数取指数,让负数也变成正数,同时放大分数之间的差距;再除以所有指数值的总和做归一化,让结果加起来正好等于 1。这样,十几万个候选 token 就各自带上了一个概率值。从这个分布里决定到底用哪个 token 的过程就是 采样(Sampling)

最简单的策略是每次直接选概率最高的那个,也就是 贪心(Greedy) 策略。不过更多时候我们会引入随机性:按概率抽一个,概率大的被抽中的机会大,但长尾里的 token 也有机会出场。分布的形状可以用 temperature 调节,它的作用是在 softmax 之前把 logits 除以一个系数。用几行 Python 感受一下,这里不用真跑模型,随手编几个分数:

import math

# 假设模型给下一个 token 算出的分数是这样的
logits = {"很": 3.2, "非常": 2.1, "特别": 1.8, "还行": 0.5}

def softmax(scores, temperature=1.0):
    exp = [math.exp(s / temperature) for s in scores]
    return [e / sum(exp) for e in exp]

for t in [0.5, 1.0, 2.0]:
    probs = softmax(logits.values(), t)
    print(f"temperature={t}", {k: round(p, 3) for k, p in zip(logits, probs)})

这段代码的关键是 softmax 函数,里面两行对应三步操作:

  1. s / temperature:每个分数先除以温度,温度小于 1 相当于把分数差距放大,大于 1 相当于把差距压小
  2. math.exp(...):对缩放后的分数取指数,这是 softmax 的第一半。e 的任何实数次方都大于 0(e⁰ = 1,负数次方是 0 到 1 之间的小数),所以负分也被映射成了正数,同时分数之间的差距被进一步拉大
  3. e / sum(exp):每个指数值除以总和,归一化成概率,这是 softmax 的第二半,保证所有候选加起来等于 1

下面的循环用三个温度各算一遍,对比分布形状的变化。输出结果如下:

temperature=0.5 {'很': 0.85, '非常': 0.094, '特别': 0.052, '还行': 0.004}
temperature=1.0 {'很': 0.607, '非常': 0.202, '特别': 0.15, '还行': 0.041}
temperature=2.0 {'很': 0.429, '非常': 0.247, '特别': 0.213, '还行': 0.111}

用一张图表示,看起来更直观:

temperature-comparison.jpg

可以看到,temperature 小于 1 时分布变尖,头部 token 几乎垄断,输出更稳定;大于 1 时分布变平,长尾 token 的机会变多,输出更发散。贪心可以理解成 temperature 趋近于 0 的极限情况。除了它,常用的还有 top_k(只在分数最高的 k 个里抽)和 top_p(按概率从高到低累加,累计到 p 就截断,也叫核采样),实际使用时经常几个参数组合在一起。同样的模型、同样的问题,回答有时稳定有时发散,差别往往就在这些采样参数上。

这里的 temperature 在数学上可以是任意正数,但是各个平台都有自己允许的取值范围:OpenAI 和 Gemini 是 0 到 2,Anthropic 限制在 0 到 1,阿里百炼是 [0, 2),本地用 transformers 跑则没有限制,平时在使用时注意一下。

反分词与流式输出

采样得到的还是 token 编号,需要 反分词(Detokenization) 把它还原成人类可读的文字。这一步在流式场景下有个必须处理的坑。Qwen、Llama 这些模型用的都是字节级 BPE(Byte Pair Encoding,字节对编码)分词器,一个 token 不一定正好是一个完整字符,可能只是某个汉字 UTF-8 编码三个字节里的一两个。如果每收到一个 token 就 decode 一次,拼出来的就是乱码。拿 Qwen3 的分词器试一下:

ids = tokenizer.encode("龘")
print(ids)                                  # [82912, 246],一个汉字被切成两个 token
print([tokenizer.decode([i]) for i in ids]) # ['�', '�'],单独 decode 都是乱码
print(tokenizer.decode(ids))                # '龘',拼在一起才能正确还原

所以推理引擎做的是 增量反分词:收到新 token 后先把字节攒着,凑够一个完整字符再往外发。

由于 Decode 是逐 token 进行的,推理服务可以边生成边把结果推给前端,这就是 流式输出(Streaming)。工程上一般通过 SSE(Server-Sent Events,一种服务端持续推送数据的 HTTP 机制)实现:服务端每产出一个 token 就推送一条消息,客户端收到一条就渲染一点。它不改变生成的总耗时,但极大改善了等待体验:第一个字出来你就能开始读,而不是盯着空白屏幕等整段回答算完。你在各类聊天产品里看到的打字机效果,源头就在这里。

推理初体验

这一节我们用 Hugging Face transformers 加载一个小模型 Qwen/Qwen3-0.6B,动手体验一次推理,代码只要十几行:

from transformers import AutoModelForCausalLM, AutoTokenizer

# 1. 加载分词器和模型
model_name = "Qwen/Qwen3-0.6B"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# 2. 分词:文本变成 token 编号
messages = [{"role": "user", "content": "用一句话解释什么是大模型推理"}]
text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(text, return_tensors="pt")

# 3. 生成:Prefill + Decode 循环都在这一步里
outputs = model.generate(**inputs, max_new_tokens=1024)

# 4. 反分词:token 编号还原成文本
response = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
print(response)

这段代码和我们上面讲的旅程是一一对应的:

  1. 加载:分词器负责文本和 token 的互转,模型则是训练好的权重
  2. 分词tokenizer 把输入文本编码成 token 编号张量
  3. 生成model.generate 内部先 Prefill 处理输入,再进入自回归的 Decode 循环,每步生成一个 token 并采样
  4. 反分词tokenizer.decode 把新生成的 token 还原成文字

运行之后终端里会打印出模型的回答:

transformers-generate-output.png

也可以加个流式输出,亲眼看着 token 一个一个生成出来。把第 3、4 步换成下面这样:

from transformers import TextStreamer

# 流式生成:每产出一个 token 就立刻解码打印
streamer = TextStreamer(tokenizer, skip_special_tokens=True)
outputs = model.generate(**inputs, max_new_tokens=1024, streamer=streamer)

TextStreamer 会在 Decode 循环的每一步把新生成的 token 立刻反分词并打印到终端,这就是流式输出在最朴素环境下的样子。

生产环境里更常见的做法是把模型部署成一个服务,客户端通过 OpenAI 兼容接口调用,比如用 curl 发一个请求:

$ curl http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "Qwen/Qwen3-0.6B",
    "messages": [{"role": "user", "content": "用一句话解释什么是大模型推理"}],
    "stream": true
  }'

加上 stream: true 之后,服务端会用 SSE 把每个 token 逐个推回来。至于本地怎么把模型跑成一个 OpenAI 兼容的服务,vLLM、SGLang、llama.cpp 都能做到,在后面的系列文章中我们会专门学习。

小结

今天这篇文章完成了两件事:

  1. 建立了概念:推理是训练好的模型根据输入生成输出的过程,权重固定、只做前向计算,它是绝大多数用户真正接触模型的方式,也是当前 AI 算力开销的大头
  2. 画出了地图:一次请求的完整旅程是分词、Prefill、Decode 循环、采样、反分词、流式输出,每个环节我们后续都会单开文章细讲

接下来的文章会沿着这张地图逐站展开,下一篇我们就走进地图的第一站,看看分词这件看起来简单、但是实际上又没那么简单的事。我们明天继续。

参考