Wafer发布AI性能工程系列内容

我们上周发布了全球最全面的 AI 性能工程仓库。现在我们会深入讲解每一项资源，这是 Wafer 的 AI 性能工程系列。链接在本串中。

第一部分：Transformer 推理机制。

给定一个提示和目前已经选出来的 token，因果语言模型会生成 logits：每个可能的下一个 token 对应一个分数。解码策略可以在选择 token 之前，通过温度、候选过滤或重复惩罚调整这些分数。

贪心解码取处理后分数最高的那个。采样则是从结果概率分布中抽样。将选中的 token 输回模型，预测下一个 token。

常规生成依赖这个选择，所以在选中输入 token 之前，无法计算下一个预测。你可以每次重新运行整个增长的前缀，这也会重复之前位置的计算。

给因果解码器追加一个 token，较早的位置仍然只能看到自己和它之前的位置。追加的 token 不贡献允许输入。在评估模式下，保持前缀、权重、掩码和位置计算固定，较早位置输入到第一层的内容保持不变，然后同样的输出会到达第二层。这个推理过程贯穿整个模型，保留了较早位置的状态。

在带 token 级归一化和 MLP 的常规因果解码器中，KV 缓存会保留每层计算出的键和值，方便你重用它们。这种等价性是数学上的：不同的浮点执行路径可能会引入数值差异，所以不要指望不同实现之间按位相等。

下一个 token 仍然需要它自己在每层的查询、键和值。它的查询会对缓存的键和当前键一起打分；注意力机制会对所有允许位置的分数做归一化，然后用这些权重组合对应的值。当前值也会参与计算。缓存存储了新查询需要从较早位置获取的信息。

旧的注意力输出反映了生成它的查询。新查询需要自己的分数和加权值总和，所以重用旧输出不能代替计算。不需要保留旧查询。保存 K/V 可以省去重新计算旧 token 状态的工作，只需要运行新 token 的投影、注意力和 MLP。

从已知的提示开始。预填充处理它并填充每层缓存，然后解码策略用最后位置的 logits 选择第一个输出。将输出输回模型：模型添加它的 K/V，然后生成第二个输出的 logits。重复这个过程直到生成停止。

你可以在模型处理完最终 token 之前就输出它，它的缓存条目可能还不存在。继续生成答案需要消费一次这个 token，这就是为什么输出 token 计数和已处理缓存长度需要分开记账。混淆二者会导致输入被跳过或被处理两次。