一次推理：Prefill、Decode 与 KV Cache

一次回答有两个性能阶段

用户看到的是连续吐字，推理服务看到的是两个不同工作负载：

Prefill：并行处理完整提示词，把每层历史 token 的 K/V 写入缓存，得到第一个输出分布。
Decode：每轮只输入最新 token，复用历史 KV Cache，生成下一个，再循环。

Prefill 更像“快速读完资料”，吞吐受长序列矩阵计算影响；Decode 更像“一边写一边翻笔记”，每一步依赖上一步，延迟和内存带宽更关键。

INTERACTIVE FIGUREKV Cache：旧笔记只算一次

100%

提示词：Grok reads the

逐步生成 token，观察已缓存的 K/V 如何增长。若不缓存，每一步都要重算全部前缀；图中的工作量用 token 投影次数近似。

KV Cache 存的是什么

Attention 对历史 token 计算过 K/V 后，只要参数不变，后续生成无需重复投影。每层保留 key、value 和当前 step。Grok-1 的单份缓存元素数近似：

N_{KV}=2\times L\times B\times T\times H_{KV}\times d_h

$2$ 表示 K 与 V； $L=64$ 层； $B$ 是 batch； $T$ 是已占用序列长度； $H_{KV}=8$ ； $d_h=128$ 。若用 bfloat16，每元素 2 字节。单 batch、8,192 token 满窗口的理论数组约为：

2\times64\times1\times8192\times8\times128\times2\text{ bytes}\approx2\text{ GiB}

这只是 cache 张量，不含模型权重、激活、临时 buffer、分片和框架开销。GQA 把 KV heads 从 48 降到 8，使这部分约为独立 48 KV heads 时的六分之一。

Grok-1 的生成时序

InferenceRunner.initialize() 恢复参数、编译四类 JAX 函数：普通 forward、新缓存、prefill、sample step。run() 是一个 generator，接收请求、tokenize、分配 batch slot 并不断采样。

Request(prompt, temperature, nucleus_p, max_len)
  → SentencePiece encode
  → 选择 pad bucket / 写入 batch slot
  → hk_prefill_memory(prompt)
      → 64 层 forward，写入 KVMemory
      → sample_token 得到首 token
  → hk_sample_step(last_token, memory) × N
      → 只处理新增 token，更新 step
      → top-p / temperature / categorical sample
  → SentencePiece decode 已生成 token

这是基于真实源码的调用链摘要，不是可直接运行的替代实现。

Logits 如何变成一个 token

最后一层输出 131,072 个 logits。温度 $\tau$ 调整分布尖锐程度：

p_i=\operatorname{softmax}(z_i/\tau)

$z_i$ 是 token $i$ 的 logit； $\tau<1$ 使高分候选更集中， $\tau>1$ 更发散。Top-p（nucleus）先按概率排序，只保留累计概率达到阈值 $p$ 的最小候选集合，再随机采样。

Grok-1 top_p_filter() 排序 logits、对 softmax 概率累加并屏蔽尾部；sample_token() 再按 temperature 做 categorical sampling。temperature 为 0 时取 argmax。

随机不等于胡乱

低温适合代码和精确任务，但无法修正模型已把错误答案打成最高分；高温增加多样性，也增加不稳定。生产应用应把可验证约束放在 schema、工具和测试里，而不是只靠调低温度。

延迟与吞吐怎样取舍

长 prompt 提高 time-to-first-token，因为 Prefill 要处理更多 token。
长回答提高总时长，因为 Decode 串行执行很多步。
KV Cache 换取计算节省，却占用显存并限制并发。
Continuous batching 可把不同请求的 decode step 拼批，但会增加调度复杂度。
量化减小权重带宽与容量，可能带来数值误差和 kernel 约束。

本章检查点

KV Cache 不保存模型“想法”，而是每层历史 token 的 Attention K/V 投影。它让 Decode 不必重做旧 token 的 K/V，但上下文越长、并发越高，缓存内存越大。