训练与推理：并行与自回归的不对称

4.1训练：teacher forcing + 全位置并行

训练时将真实 token 序列整体喂入，因果掩码（第 02 章）防止每个位置偷看未来，所有位置的 next-token loss 经一次矩阵乘同时计算完毕。

训练的输入是一段真实文本序列，例如 [BOS, the, cat, sat, on, the, mat]。模型不需要知道"正确答案是什么"——答案就是序列本身向右移一位：输入位置 0 预测位置 1，位置 1 预测位置 2，以此类推。这种做法叫 teacher forcing（教师强制）：每个位置的输入永远是真实 token，而不是模型上一步的预测结果，因此不存在误差累积。

能并行的关键在于：第 02 章介绍的因果掩码把注意力矩阵的上三角（未来位置）在 softmax 前置为 −∞，softmax 后权重归零——位置 t 只能看到位置 0…t 的信息，而这个约束对所有位置同时成立。于是一整个序列的 Q、K、V 可以同时投影、同时做矩阵乘，所有位置的注意力分数和 loss 通过单次前向传播同时得到，梯度也通过单次反向传播同时更新。

# tokens: (batch, seq_len)  — 整段真实序列
logits = model(tokens[:, :-1])      # 输入：去掉最后一个 token
targets = tokens[:, 1:]             # 标签：去掉第一个 token（右移一位）
loss = cross_entropy(logits, targets)  # 所有位置一次算完
loss.backward()                     # 单次反向传播，全部梯度

交叉熵 loss 的形式：对每个位置 t，取 logits[t] 经 softmax 后在真实 token 上的负对数概率，再对整个序列平均：

L = -1/T · Σₜ log P(xₜ₊₁ | x₁..xₜ)

  T       : 序列长度
  P(·|·)  : 模型 softmax 输出的概率
  xₜ      : 位置 t 处的真实 token

4.2推理：自回归循环

推理没有 ground truth 可供 teacher forcing，每步把刚生成的 token 追加到输入序列，重新前向传播，取最后一位的 logits 采样或 argmax，循环至 EOS 或达到长度上限。

推理的全过程：

1. 以 prompt（可以是 [BOS] 或用户的问题 token 序列）作为初始输入。
2. 前向传播，取最后一个位置的 logits，经 softmax 得到 next-token 概率分布。
3. 从分布中采样（temperature 采样、top-k、top-p 等）或取 argmax，得到 token t。
4. 把 token t 追加到输入末尾，回到步骤 2。
5. 直到生成 <EOS> 或序列长度达到上限，停止。

关键：步骤 3 生成的 token t 在步骤 2 之前根本不存在，所以步骤 2→3→4→2 的循环无法被展开成并行矩阵操作——这是数据依赖带来的串行性，不是工程局限。

4.3不对称是最大的认知陷阱

这套并行训练 / 串行推理的不对称，是初学者最高频的卡点：用训练式前向实现"推理"，输出永远是垃圾——因为少了把输出喂回输入的循环。

按上下文的 §5 所记：很多人实现 transformer 时，先写出 logits = model(tokens) 的训练前向，然后直接拿 logits[-1] 当"生成结果"，却发现无论怎么调参，输出都像随机噪声。症结在这里：

• 训练前向：一次喂入完整 token 序列，同时得到所有位置的 logits。这里的 logits[t] 预测 token t+1，是拿 t 之前的真实 token 算的。
• 推理前向：每次只取最后位置的 logits 采样一个新 token，把这个 token 追加进去，再跑整个序列——少了这个追加步骤，模型永远在对同一段 prompt 重复预测，不是在"生成"。

这也是 aha #3 的核心：两个阶段的前向传播在数学上调用的是同一套 attention + FFN 权重（第 03 章介绍的完整 block），差异纯粹在输入构造和调用方式上。理解了这一点，KV-cache 的必要性就立刻显现——既然推理每步都要重跑整个序列，所有过去 token 的 K、V 投影每次都要重算，这是巨大的浪费。

4.4KV-cache：别重算过去

推理时每一步新生成的 token 只需要计算自己的 Q，过去 token 的 K 和 V 已经算过且不会改变——KV-cache 把它们存下来，把每步复杂度从 O(n) 次矩阵乘降到 O(1) 次。

理解 KV-cache 需要先回忆 第 02 章的 scaled dot-product attention：

Q = X · Wq    (n, d_k)
K = X · Wk    (n, d_k)   ← 过去 token 的 K 每步都一样
V = X · Wv    (n, d_v)   ← 过去 token 的 V 每步都一样

scores = Q · Kᵀ / √d_k   (n, n)
out    = softmax(scores) · V   (n, d_v)

生成 token t+1 时，新输入只有 token t+1 本身。位置 0…t 的 X 没变，所以它们的 K 和 V 投影与上一步完全相同。KV-cache 做的事：在每层每头把已算过的 K、V 矩阵拼接存入缓存，新 token 只算自己的 Q（以及自身新的 K、V 追加入缓存），再和缓存里的 K、V 做 attention。

重要区分：缓存的是 K 和 V 的投影结果（形状 (past_len, d_k) 和 (past_len, d_v)），不是权重矩阵 Wk、Wv（权重在整个推理过程中始终不变，本来就无需重算）。

显存成本公式

KV-cache 的显存占用可以精确估算：

KV_cache_bytes = seq_len × layers × heads × head_dim × 2 × dtype_bytes

  seq_len    : 上下文窗口长度（已生成 + prompt 的 token 数）
  layers     : Transformer 层数
  heads      : 每层注意力头数（KV 头，GQA 下小于 Q 头数）
  head_dim   : 每头的维度 = d_model / heads
  × 2        : K 和 V 各一份
  dtype_bytes: fp16 = 2 bytes, bf16 = 2 bytes, int8 = 1 byte

以 7B 规模的典型配置（Llama-2-7B 同规格）具体计算：

seq_len    = 8 192
layers     = 32
heads      = 32        （MHA，KV 头 = Q 头）
head_dim   = 128       （d_model=4096 / 32头）
dtype      = fp16 = 2 bytes

KV = 8192 × 32 × 32 × 128 × 2 × 2
   = 8192 × 32 × 32 × 512
   ≈ 4 294 967 296 bytes
   ≈ 4 GB  ×  2  （K 和 V）
   ≈ 8 GB

结论：7B / 32层 / 32头 / d128 / 8K上下文 / fp16
     ≈ 8 GB KV-cache（不含模型权重本身的 ~14 GB）

这 8 GB 随序列长度线性增长：上下文翻倍到 16K，KV-cache 也翻倍到 16 GB。GQA（Grouped Query Attention，第 05 章）通过让多个 Q 头共享一组 KV 头，直接按比例压缩这个数字——32 头 GQA 分成 8 组，则 KV 头数从 32 降到 8，缓存降至约 2 GB。

4.5O(n²)：长上下文为什么贵

attention 的 O(n²) 根源在于 QKᵀ 产生 n×n 分数矩阵，这个矩阵必须落地（或用 FlashAttention 分块重算避免）；在长上下文下，显存瓶颈先于算力瓶颈到来。

回到 第 02 章的公式：scores = Q·Kᵀ / √d_k，其中 Q 形状 (n, d_k)，Kᵀ 形状 (d_k, n)，相乘结果是 (n, n)——n 个 query 对 n 个 key 各算一个相似度分数。

这个 n×n 矩阵的问题有两个维度：

• 算力：矩阵乘的浮点操作数是 O(n²·d_k)，随序列长度平方增长。
• 显存：n×n 矩阵要暂存在 GPU SRAM 或 HBM（高带宽存储）中以便后续 softmax 和乘 V。n=8192 时，每头分数矩阵有 8192×8192 ≈ 6700 万个 float16 元素 ≈ 128 MB；32 头 × 32 层 = 超过 100 GB，远超单卡 HBM 容量。

n = 8 192
每头 scores = n × n = 8192 × 8192 ≈ 67 108 864 个元素
fp16: 67 108 864 × 2 bytes ≈ 128 MB / 头

32 层 × 32 头 × 128 MB ≈ 128 GB（scores 矩阵总量）
→ 这是为什么朴素 attention 在长上下文下首先触达显存天花板
  而不是算力天花板

朴素实现将整个 n×n scores 矩阵写入 HBM，再读回来做 softmax，再写一次，再读回来乘 V——内存带宽是真正的瓶颈（HBM 带宽约 2 TB/s，但矩阵反复读写造成大量浪费）。FlashAttention（第 05 章）的核心思路正是把这个矩阵"分块留在 SRAM 里"，避免落地 HBM，从根本上规避这一瓶颈——且数学输出完全等价，不是近似。

4.6这些成本催生了下一章的优化

KV-cache 的线性显存增长和 O(n²) 的平方瓶颈，是第 05 章所有现代优化的直接动机——不理解成本的来源，就无法理解那些方案在解决什么。

本章揭示的两条成本线，在第 05 章各对应一簇优化方案：

第 05 章在这两条线上都会给出机制级解释：GQA 如何把 KV 头数从 32 压到 8；FlashAttention 如何用 SRAM 分块消除 HBM 落地；MLA 如何把 K/V 压成低秩潜向量再按需解压（DeepSeek-V2，2024-05）。把本章的成本公式和第 05 章的优化方案对照起来看，才能真正理解那些方案在省什么、代价是什么。