【infra之路】LLM 预测一个 Token 的完整流程：从文本输入到概率输出
前言
当你向 ChatGPT 输入"今天天气真"，模型回答"好"。这个看似简单的"接一个字"，背后经历了分词、向量化、几十层 Transformer 的矩阵运算、概率采样等一系列精密步骤。而且模型不是在"想好一整句话再输出"——它是一个 token 一个 token 地生成，每次生成一个 token 都要把整个模型跑一遍。
这篇文章用一个具体的例子，跟踪一个 token 从输入到输出的完整旅程。以 LLaMA-7B（32 层 Transformer）为例，拆解每一步发生了什么、数据怎么变换、计算量有多大。
一、总览：一次前向传播的六个阶段
不管输入多长的 prompt，模型预测"下一个 token"的流程都可以拆成六步：
① 分词（Tokenization） 文本 → 整数 ID 序列 ② 嵌入（Embedding） 整数 ID → 稠密向量 ③ 位置编码（Position） 给每个向量加上位置信息 ④ Transformer 层 × N 自注意力 + 前馈网络逐层处理 ⑤ 输出投影（LM Head） 最后一层的隐藏状态 → 词汇表维度的 logits ⑥ 采样（Sampling） logits → 概率分布 → 选出下一个 token下面用一个具体的例子走一遍。假设输入是“今天天气真”，模型需要预测下一个 token。
二、第一步：分词（Tokenization）
LLM 不认识中文字符，它只认识整数。分词器（Tokenizer）把文本切分成一系列 token ID。
以 LLaMA 使用的 SentencePiece（BPE）为例：
输入文本: "今天天气真" 分词结果: [2521, 1083, 29576, 17582]注意几点：
- 一个汉字不一定对应一个 token。"今天"可能被分成"今"和"天"两个 token，也可能被合并成一个 token（取决于 BPE 训练的合并规则）。高频词更容易被合并，生僻字更容易被拆成多个子词。
- LLaMA 的词汇表大小是32,000（
vocab_size=32000），意味着每个 token ID 是 0~31999 之间的一个整数。 - 分词是确定性的——同一段文本永远被分成相同的 token 序列。
三、第二步：嵌入（Embedding Lookup）
拿到 token ID 后，需要从嵌入矩阵中"查表"，把每个整数变成一个稠密向量。
嵌入矩阵 E: 形状 [32000, 4096]（vocab_size × hidden_dim） token ID 2521 → 取 E 的第 2521 行 → 一个 4096 维的浮点向量 token ID 1083 → 取 E 的第 1083 行 → 一个 4096 维的浮点向量 ...对于 4 个 token 的输入，嵌入层的输出是一个[4, 4096]的矩阵——4 个 token，每个 4096 维。
嵌入矩阵本身是模型训练出来的参数，占了模型参数的很大一部分。对于 LLaMA-7B，嵌入矩阵有 32000 × 4096 × 2 bytes ≈262 MB。
四、第三步：位置编码
嵌入向量本身不包含"这个词在句子的什么位置"的信息。Transformer 需要额外注入位置信息，否则"今天天气真"和"天今真气天"会被认为是一样的。
主流 LLM 使用RoPE（Rotary Position Embedding），它的做法是：在每一层计算注意力时，对 Query 和 Key 向量做旋转，旋转角度取决于 token 的位置。
位置 0（"今"）: Q₀ 和 K₀ 旋转 0° 位置 1（"天"）: Q₁ 和 K₁ 旋转 θ° 位置 2（"气"）: Q₂ 和 K₂ 旋转 2θ° 位置 3（"真"）: Q₃ 和 K₃ 旋转 3θ°RoPE 的优势在于：两个 token 之间的注意力得分只取决于它们的相对位置（位置差），而不是绝对位置。这让模型可以泛化到比训练时更长的序列。
五、第四步：Transformer 层（核心）
这是 LLM 计算量最大的部分。LLaMA-7B 有32 层 Transformer，每一层的结构相同，数据从第 1 层流向第 32 层。
图1：Decoder-only Transformer 的单层结构。输入向量经过自注意力机制和全连接网络，每步都有残差连接和层归一化。
5.1 自注意力机制（Self-Attention）
这是 Transformer 最核心的计算。它的目的是让每个 token "看看"其他所有 token，决定该关注谁、关注多少。
计算 Q、K、V
输入矩阵 X（形状 [4, 4096]）分别乘以三个权重矩阵，得到 Query、Key、Value：
Q = X × W_Q 形状: [4, 4096] （4 个 token，每个产生一个 4096 维的 Query） K = X × W_K 形状: [4, 4096] V = X × W_V 形状: [4, 4096]实际中 4096 维会被拆成32 个注意力头（每个头 128 维），多头并行计算。
计算注意力分数
Attention(Q, K, V) = softmax(Q × K^T / √d_k) × V拆开看：
Q × K^T：每个 token 的 Query 和所有 token 的 Key 做点积，得到一个 [4, 4] 的注意力分数矩阵。分数越高，表示两个 token 之间的"关联度"越强。/ √d_k：除以 √128 ≈ 11.3 做缩放，防止点积值太大导致 softmax 梯度消失。- Causal Mask：把"未来位置"的注意力分数设为 -∞。位置 0 只能看位置 0，位置 1 能看 0 和 1，位置 2 能看 0、1、2……这是 Decoder 的关键约束——模型不能偷看后面的内容。
softmax：把每一行的分数归一化成概率，和为 1。× V：用注意力权重对 Value 做加权求和，得到每个 token 的"融合了上下文信息的输出"。
Causal Mask 矩阵（4 个 token）: tok0 tok1 tok2 tok3 tok0 [ ✓ -∞ -∞ -∞ ] ← "今"只能看自己 tok1 [ ✓ ✓ -∞ -∞ ] ← "天"能看"今"和"天" tok2 [ ✓ ✓ ✓ -∞ ] ← "气"能看前三个 tok3 [ ✓ ✓ ✓ ✓ ] ← "真"能看所有输出投影
多头注意力的结果拼接后，再乘以一个输出权重矩阵 W_O，得到自注意力的最终输出。加上残差连接：
X_attn = LayerNorm(X + MultiHeadAttention(X)) 形状: [4, 4096]5.2 前馈网络（Feed-Forward Network）
自注意力处理的是 token 之间的交互，前馈网络处理的是每个 token 自身的特征变换。
LLaMA 使用SwiGLU激活函数（比 GELU 效果更好），计算过程：
FFN(x) = (x × W₁ ⊙ Swish(x × W₃)) × W₂其中 W₁ 和 W₃ 把维度从 4096 扩展到 11008（约 2.7 倍），W₂ 再压缩回 4096。⊙表示逐元素相乘，Swish(x) = x × σ(x)是激活函数。
同样加上残差连接和层归一化：
X_ffn = LayerNorm(X_attn + FFN(X_attn)) 形状: [4, 4096]5.3 32 层堆叠
上面这个过程（自注意力 + 前馈网络）重复32 次。每一层的输出作为下一层的输入，信息逐层传递和提炼。最后一层输出的 [4, 4096] 矩阵，就是模型对整个输入序列的"最终理解"。
每一层都有独立的权重参数（W_Q、W_K、W_V、W_O、W₁、W₂、W₃ 以及 LayerNorm 的参数），LLaMA-7B 的大部分参数都花在了这 32 层的权重矩阵上。
六、第五步：输出投影（LM Head）
经过 32 层 Transformer 后，取最后一个 token（“真”）的隐藏状态（一个 4096 维向量），通过一个线性层投影到词汇表维度：
最后一个 token 的隐藏状态 h: 形状 [4096] LM Head 权重矩阵 W_lm: 形状 [32000, 4096] logits = h × W_lm^T 形状 [32000]得到的logits是一个 32000 维的向量，每个值对应词汇表中一个 token 的"原始得分"。比如：
logits[6521] = 8.7 （"好"的得分） logits[29892] = 5.3 （"热"的得分） logits[1523] = 3.1 （"冷"的得分） logits[892] = 0.2 （"狗"的得分） ...有些模型（如 LLaMA）使用权重共享（Weight Tying）：LM Head 的权重矩阵和嵌入矩阵是同一个矩阵 E。这样可以节省 262 MB 参数。
七、第六步：采样（Sampling）
logits 只是原始分数，需要转成概率才能选 token。
Softmax 归一化
prob_i = exp(logit_i / T) / Σ exp(logit_j / T)其中 T 是温度参数（Temperature），控制概率分布的"锐利度"。
温度参数的作用
Temperature = 0.1（低温）: "好" 98.2%, "热" 1.3%, "冷" 0.3%, ... ← 几乎确定 Temperature = 1.0（常温）: "好" 72.5%, "热" 15.1%, "冷" 5.2%, ... ← 有一定随机性 Temperature = 2.0（高温）: "好" 42.3%, "热" 28.7%, "冷" 15.1%, ... ← 很随机温度越低，模型越"保守"，倾向于选概率最高的 token；温度越高，分布越平坦，低概率 token 也有机会被选中。
常见采样策略
Greedy（贪心）：直接选概率最高的 token。确定性输出，但容易生成重复和无聊的文本。
Top-K：只保留概率最高的 K 个 token，从中采样。比如 K=50，就把 32000 个候选缩减到 50 个。
Top-P（Nucleus）：保留累积概率达到 P 的最小 token 集合。比如 P=0.9，就取概率从高到低累加到 90% 的那些 token。好处是候选数量动态变化——当模型很确定时候选少，不确定时候选多。
Top-P=0.9 的例子: "好" 72.5% → 累积 72.5%，继续 "热" 15.1% → 累积 87.6%，继续 "冷" 5.2% → 累积 92.8%，超过 90%，停止 候选集合: {"好", "热", "冷"}，重新归一化后从中采样Repetition Penalty（重复惩罚）：对已经出现过的 token 降低概率，防止模型陷入重复循环。
选出 Token
采样完成后，得到下一个 token 的 ID。把它追加到序列末尾：
原始序列: [2521, 1083, 29576, 17582] → "今天天气真" 新 token: 6521 → "好" 更新序列: [2521, 1083, 29576, 17582, 6521] → "今天天气真好"然后把这个新序列再次送入模型，预测下一个 token。如此循环，直到模型生成EOS（End of Sequence）token 或达到最大长度限制。这就是自回归生成（Autoregressive Generation）。
八、Prefill 与 Decode：推理的两个阶段
在实际推理中，LLM 的工作分为两个截然不同的阶段：
图2：LLM 推理的两个阶段。Prefill 阶段并行处理所有输入 token 并初始化 KV Cache；Decode 阶段逐 token 生成，每步只处理一个 token 但需要读取完整 KV Cache。
Prefill（预填充）阶段
用户提交 prompt 后，模型一次性并行处理所有输入 token。这个阶段的特点是：
- 所有 token 同时通过 Embedding、Transformer 层，GPU 的计算单元被充分利用（compute-bound）
- 计算出所有 token 的 K 和 V 向量，写入KV Cache
- 产出第一个生成的 token
- 性能指标：TTFT（Time to First Token），即用户从提交到看到第一个字的延迟
Decode（解码）阶段
Prefill 产出第一个 token 后，进入逐 token 生成阶段。每生成一个 token：
- 只有一个新 token通过 Transformer 层（batch_size=1）
- GPU 的计算单元利用率很低，瓶颈变成了从显存读取权重和 KV Cache 的速度（memory-bound）
- 每生成一个 token 的时间基本固定（不像 Prefill 那样和 prompt 长度成正比）
- 性能指标：TPS（Tokens Per Second），即每秒生成多少个 token
这就是为什么 ChatGPT 有时候"想了很久才蹦出第一个字"（Prefill 处理长 prompt），但一旦开始输出就"一个字一个字很稳定"（Decode 阶段）。
九、KV Cache：避免重复计算的关键
每次 Decode 生成新 token 时，如果从头重算所有 token 的 K 和 V，计算量会随着序列长度平方增长。KV Cache通过缓存已计算的 K、V 向量来避免这个问题。
图3：KV Cache 的工作原理。Prefill 阶段计算并缓存所有输入 token 的 K/V；Decode 阶段只需计算新 token 的 K/V 并追加到缓存中，注意力计算时用新 Q 和缓存中的所有 K/V 做点积。
工作原理
第 t 步生成新 token 时: 1. 新 token 经过线性投影得到 Q_new, K_new, V_new 2. K_new, V_new 追加到 KV Cache 3. 注意力计算: Attention(Q_new, [K_cached, K_new], [V_cached, V_new]) 4. 只有 Q_new 是本次新算的，K 和 V 大部分来自缓存计算量对比
| 无 KV Cache | 有 KV Cache | |
|---|---|---|
| 每步向量投影计算 | O(t × d²) | O(d²) |
| 每步注意力计算 | O(t² × d) | O(t × d) |
| 生成 N 个 token 总计算量 | O(N³) | O(N²) |
对于生成长度为 1000 token 的文本，KV Cache 带来的加速约为1000 倍。代价是显存占用——每个 token 在每一层都要缓存 K 和 V。
KV Cache 的显存占用
LLaMA-7B 的 KV Cache 显存公式：
mem = 2 × num_layers × num_heads × seq_len × head_dim × dtype_bytes × batch_size = 2 × 32 × 32 × seq_len × 128 × 2 × 1 ≈ 0.5 MB × seq_len一个 2000 token 的序列，KV Cache 需要约1 GB显存。这也是为什么 LLM 有上下文长度限制——不是模型"记不住"，而是 KV Cache 的显存撑不下。
十、计算量估算：预测一个 Token 到底做了多少运算
以 LLaMA-7B 为例，估算 Decode 阶段生成一个 token 的计算量：
| 计算环节 | FLOPs（近似） |
|---|---|
| 每层自注意力（QKV 投影 + 注意力 + 输出投影） | ~2 × 4096² × 3 ≈ 1 亿 |
| 每层 FFN（SwiGLU 两次线性变换） | ~2 × 4096 × 11008 × 2 ≈ 1.8 亿 |
| 32 层合计 | ~32 × 2.8 亿 ≈90 亿 FLOPs |
| LM Head（4096 → 32000） | ~2.6 亿 |
| 单 token 总计 | ~93 亿 FLOPs |
每生成一个 token，模型要做约93 亿次浮点运算。一张 A100 的峰值算力是 312 TFLOPs，理论上每秒能生成约 33,000 个 token——但实际受限于显存带宽（Decode 阶段是 memory-bound），A100 上 LLaMA-7B 的实际生成速度大约是100-200 tokens/s。
十一、总结
LLM 预测一个 token 的过程，本质上是一次完整的前向传播：文本经分词变成整数，经嵌入变成向量，经过几十层 Transformer 的注意力和前馈网络逐层提炼，最后通过 LM Head 投影到词汇表空间得到一个 32000 维的 logit 向量，经 Softmax 转成概率后采样选出下一个 token。
一个实用的理解框架：Embedding 把离散符号变成连续空间中的点，Transformer 层在这个空间中做上下文感知的特征变换，LM Head 把最终特征映射回离散的词汇表空间。整个模型做的事情就是在连续空间中找"最合适的下一个词"。Decode 阶段通过 KV Cache 缓存历史计算避免重复，使得逐 token 生成在工程上可行。
参考资料
- 大模型推理机制解析：预填充与生成阶段
- LLM 模型推理全流程解析：从输入到输出的技术实现
- KV Cache 深度解析：从原理到显存优化
- 从输入到输出：大语言模型一次完整推理简单解析
- 大语言模型 Next Token Prediction 与 Transformer 架构