Transformer 推理的数学

训练时我们只关心一件事：每秒每张卡处理多少 token。推理多了一个维度——延迟。用户在等第一个字出来，在等后面的字一个个往外冒。好在推理可以拆成能算的东西：每一步要从显存里搬多少字节、做多少 FLOPs、什么时候受带宽限制、什么时候受算力限制。KV cache 压缩、张量并行、预填充/解码分离、continuous batching、PagedAttention、前缀缓存、投机解码和量化，解决的都是这本账里的某一项。

算例主要用 TPU v5e：它的规格简单、每美元算力高，数字也好算；需要时同时给出 H100 的对应数字。算术强度、硬件强度这两个概念来自 Roofline。

1. 自回归生成与 KV cache

语言模型做的事情是：给定前缀 \(x_{\lt t}\)，输出下一个 token 的分布 \(p(x_t \mid x_{\lt t})\)。生成就是循环：从分布里采一个 token，拼到序列末尾，再把整个序列送回模型。

最朴素的实现每一步都把完整前缀重新跑一遍前向。生成 \(n\) 个 token，第 \(t\) 步要处理 \(t\) 个 token：MLP 部分总共 \(\sum_t t = O(n^2)\) 的工作量，注意力部分每步是 \(O(t^2)\)，合起来 \(O(n^3)\)。这显然浪费——第 \(t\) 步重新算的前 \(t-1\) 个位置的结果和上一步完全一样，因为因果注意力里位置 \(j\) 只看得到 \(\le j\) 的位置，后面追加的 token 不会改变它。

真正被后续 token 需要的，只有每一层里每个历史位置的 key 和 value 向量：新 token 的 query 要和所有历史 key 做内积、再对所有历史 value 加权求和。把它们存下来，就是 KV cache。有了它，第 \(t\) 步只需要把一个新 token 跑过网络：线性层只算这一个位置，注意力用新 query 去读缓存。总代价降到 MLP \(O(n)\)、注意力 \(O(n^2)\)。

于是推理自然分成两个阶段：

• 预填充（prefill）：把整个 prompt 一次性跑一遍前向（和训练的前向一样，所有位置并行），把每层每个位置的 K、V 写进缓存，同时得到最后一个位置的 logits。
• 解码（decode / generation）：从 logits 采一个 token，把这一个 token 送进模型，读 KV cache 做注意力，把它自己的 K、V 追加进缓存，得到下一组 logits。重复直到遇到结束符或长度上限。

两个阶段的性质截然不同：预填充一次处理成百上千个 token，权重读一次能用很多次；解码每条序列每步只有一个 token，而且步与步之间有串行依赖，没法在时间维度上并行，只能靠把很多条请求拼成一个 batch来摊薄读权重的成本。

2. 推理在优化什么

先定几个量：

• TTFT（time to first token）：从请求到达到第一个输出 token 的时间，主要由排队和预填充决定。
• 每 token 延迟（也叫 TPOT / ITL）：解码阶段相邻两个 token 的间隔，约等于一次解码步（step）的时间。每用户的生成速度就是它的倒数。
• 吞吐：整个系统每秒生成多少 token；除以芯片数就是每芯片吞吐，直接决定每 token 的成本。

不同场景要的东西不一样：离线批量推理（跑评测、合成数据）只看总成本，单条延迟无所谓；聊天要在成本可控的前提下让 TTFT 短、生成速度高于人的阅读速度；端侧推理（笔记本上跑 llama.cpp）只服务一个用户，追求最低延迟。训练里"提高硬件利用率"总是好事，推理里它往往要拿延迟去换。

3. 线性层：解码为什么受带宽限制

回顾：roofline 与临界 batch

一次运算要花的时间至少是两样东西的较大者（假设二者能完全重叠）：算 FLOPs 的时间 \(T_\text{math}=\text{FLOPs}/C\)，以及把数据从 HBM 搬进计算单元的时间 \(T_\text{mem}=\text{Bytes}/W\)。其中 \(C\) 是芯片的峰值算力（FLOPs/s），\(W\) 是 HBM 带宽（Bytes/s）。算术强度 = FLOPs / Bytes；当它超过硬件强度 \(C/W\) 时，算子受算力限制（compute-bound），否则受带宽限制（memory-bound）。

看一个线性层：bf16 激活 \(X[B, D]\) 乘 bf16 权重 \(W[D, F]\)。\(B\) 是这一步里送进来的 token 数。

（每个乘加算 2 FLOPs；bf16 每个数 2 字节，分别是读 \(X\)、读权重、写输出。）推理时 \(B\) 通常最多几百，而 \(D, F\) 上万，所以分母里 \(2DF\)（读权重）占绝对主导：

TPU v5e 的 bf16 算力是 \(1.97\times10^{14}\) FLOPs/s，HBM 带宽 \(8.2\times10^{11}\) B/s，\(B_\text{crit}\approx 240\)。H100 SXM 的 bf16 稠密算力是 989 TFLOPs、HBM 带宽 3.35 TB/s，\(B_\text{crit}\approx 295\)。两者都是几百。

精度如何改变临界 batch

如果权重用 int8/fp8 存（每参数 1 字节），读权重的字节数减半，而计算仍用 bf16，那么条件变成 \(2BDF/DF=2B\ge C/W\)，即 \(B_\text{crit}\) 减半到 120。如果连计算也用 int8（v5e 的 int8 算力是 \(3.94\times10^{14}\)，翻倍），\(C\) 翻倍，\(B_\text{crit}\) 又回到 240。一般地：

直觉是：每个参数从 HBM 读进来要花 \(b_w/W\) 秒，在这段时间里芯片能做 \(C\cdot b_w/W\) 次 FLOPs；每个 token 对每个参数只做 2 次 FLOPs，所以要攒够 \(b_wC/(2W)\) 个 token 才能把读参数的时间"填满"。

4. 注意力：一个永远受带宽限制的算子

线性层的权重是所有序列共享的，batch 越大、每字节权重被复用的次数越多。注意力不一样：每条序列有自己的 KV cache。看一个头的标准多头注意力（用 FlashAttention 融合，不把注意力矩阵写回 HBM），query 长度 \(T\)，key/value 长度 \(S\)：

• 读 \(Q\)：\(\text{bf16}[B,T,D]\)；读 \(K, V\)：两个 \(\text{bf16}[B,S,D]\)；写输出 \(\text{bf16}[B,T,D]\)。
• \(QK^\top\) 做 \(2BSTD\) FLOPs，乘 \(V\) 再做 \(2BSTD\)。

预填充时 \(S=T\)，强度 \(=T/2\)，序列长到几百就受算力限制。解码时 \(T=1\)、\(S\gg1\)，强度 \(\approx 1\)——而且 \(B\) 和 \(D\) 都约掉了，增大 batch 完全无济于事，因为 batch 每多一条，要读的 KV 也多一份。

GQA 会把它稍微抬高。设 \(N\) 个 query 头共享 \(K\) 个 KV 头，组大小 \(G=N/K\)。同一个 KV 头被 \(G\) 个 query 头读，解码时强度变成

LLaMA-3 的 \(G=8\)，离 240 差得远。所以解码时注意力基本永远受带宽限制，它的耗时就是"读一遍 KV cache"的时间。这也说明了一个关键的不对称：batch 增大时，权重读取被摊薄，KV 读取不会。当一条序列的 KV cache 字节数和参数字节数可比时，增大 batch 带来的吞吐收益会迅速递减。

5. KV cache 有多大：MHA、MQA、GQA、MLA

基本公式

每层、每个 token 要存一个 key 和一个 value，各是 \(K\) 个头 × 每头 \(H\) 维。于是

2 是 K 和 V 两份，\(L\) 是层数，\(T\) 是 token 数。几种注意力变体的区别只在 \(K\)：

• MHA（多头注意力）：\(K=N\)，每个 query 头有自己的 KV 头。
• MQA（多查询注意力，Shazeer 2019）：\(K=1\)，所有 query 头共享一个 KV 头，缓存缩小 \(N\) 倍，但质量会掉一些。
• GQA（分组查询注意力，Ainslie et al. 2023）：\(1\lt K\lt N\)，每 \(G=N/K\) 个 query 头共享一组 KV。实践中质量接近 MHA、速度接近 MQA，已经是开源模型的默认选择。

几个具体数字（bf16）：LLaMA-2-13B 是 MHA，\(L=40, K=40, H=128\)，每 token \(2\cdot2\cdot40\cdot40\cdot128=819{,}200\) 字节 ≈ 0.82 MB，一条 8192 长的序列就是 6.7 GB，四条就超过了 26 GB 的参数。LLaMA-3-70B 用 GQA，\(L=80, K=8, H=128\)，每 token 327,680 字节 = 320 KiB；如果按 MHA（\(K=64\)）就是 8 倍。

MLA：存一个低秩隐向量

DeepSeek-V2 提出的多头隐注意力（Multi-head Latent Attention）换了一个思路：不减少头数，而是把 K、V 联合压缩成一个低维向量再缓存。记第 \(t\) 个位置的隐状态为 \(h_t\in\mathbb{R}^{d}\)，\(n_h\) 个头、每头 \(d_h\) 维：

\(W^{DKV}\) 是下投影，\(W^{UK}, W^{UV}\) 是把隐向量上投影回每个头的 key/value。推理时只缓存 \(c_t^{KV}\)。妙处在于根本不需要把 key、value 显式算出来：第 \(i\) 个头的注意力分数是

所以可以把 \(W^{UK}\) "吸收"进 query 一侧：先把每个头的 query 映射到 \(d_c\) 维隐空间，再直接和缓存里的 \(c_j\) 做内积。value 同理，\(W^{UV}\) 可以吸收进输出投影 \(W^O\)。

麻烦出在 RoPE 上。RoPE 在 query 和 key 上各乘一个依赖位置的旋转矩阵，分数变成 \(q_{t,i}^\top R_t^\top R_j W_i^{UK}c_j\)，中间夹着依赖 \(j-t\) 的 \(R_{j-t}\)，无法预先合并成一个固定矩阵，吸收就失效了。DeepSeek 的做法是解耦 RoPE：额外算一个很小的、所有头共享的带位置 key \(k_t^R=\text{RoPE}(W^{KR}h_t)\in\mathbb{R}^{d_h^R}\)，每个头的 query 也拆出一段带 RoPE 的 \(q_{t,i}^R\)，分数是两部分之和：

缓存里每层每 token 存 \(c^{KV}_t\)（\(d_c\) 维）和 \(k_t^R\)（\(d_h^R\) 维）。DeepSeek-V2 和 V3 的配置相同：\(d_c=512\)，\(d_h^R=64\)，\(n_h=128\)，\(d_h=128\)（对应 config.json 里的 kv_lora_rank、qk_rope_head_dim、num_attention_heads、v_head_dim）。于是每层每 token 是 \(512+64=576\) 个元素。换个角度看：\(d_c=4d_h\)、\(d_h^R=d_h/2\)，缓存量相当于只有 \(2.25\) 组的 GQA（\(2\cdot2.25\cdot128=576\)）。

以 DeepSeek-V3（61 层）为例，bf16 下每 token 的缓存是 \(61\times576\times2=70{,}272\) 字节 ≈ 69 KiB，比 LLaMA-3-70B 的 320 KiB 还小得多，而它有 128 个注意力头。如果同样的头配置用 MHA（key 是 128+64=192 维，value 128 维），每层每 token 要存 \(128\times320=40{,}960\) 个元素，是 MLA 的 71 倍。DeepSeek-V2 相对 DeepSeek 67B 把 KV cache 减少了 93.3%，最大生成吞吐提高到 5.76 倍。

MLA 还有一个常被忽略的好处：吸收之后，解码时 128 个头都在读同一份 576 维的缓存，相当于 \(G=n_h\) 的 MQA，注意力的算术强度随头数增长，而不是停在 GQA 的 \(G=8\)。代价是每个头的 query 要先做一次到 512 维的投影，注意力里的 FLOPs 变多。

其他减小 KV 的办法

• 局部（滑动窗口）注意力：一部分层只看最近 \(w\) 个 token，这些层的缓存上限是 \(w\)，与上下文长度无关。和少量全局层交错使用，长上下文时总缓存能大幅缩小。
• 跨层共享 KV：相邻几层共用一份 KV。缓存变小，batch 能开更大；但如果共享的缓存要被不同层各读一次，读取字节数不一定减少，单步时间未必变快。
• 量化 KV（见第 11 节）和分页/只读有效部分（见第 9 节）是运行时手段，与架构无关。

这些手段叠起来，可以把 KV cache 比标准 MHA 缩小一个数量级以上。下面的计算器按公式比较四种变体：

6. 延迟下界与吞吐上界

推理时显存里都有什么

训练要存参数、梯度、优化器状态和用于反向的激活 checkpoint；推理只存一份参数，没有梯度和优化器状态，激活也很小（不需要为反向保留；8k token 的预填充一份激活只有 \(8192\times5120\times2\) 字节 ≈ 80 MB 量级，长的预填充还可以切块做）。剩下的大头就是 KV cache，它的总量是 batch × 上下文 × 每 token 大小，上限只取决于你允许多长的上下文、多大的 batch。

一步解码要多久

把前两节合起来。设参数量 \(P\)、每参数字节 \(b_w\)，每 token 的 KV 字节 \(k\)，batch \(B\)、上下文 \(S\)，\(N\) 张芯片完全分片（每张卡读 \(1/N\) 的权重和缓存）。一步解码：

• 注意力：读全部 KV，永远受带宽限制：\(T_\text{KV}=\dfrac{B\cdot S\cdot k}{N W}\)。
• 线性层：读一遍权重，或者算 \(2BP\) FLOPs，取较大者：\(\max\!\Big(\dfrac{P b_w}{NW},\ \dfrac{2BP}{NC}\Big)\)。

推理性能的账，基本都在这个式子里。小 batch 时 max 取权重读取，得到理论最小步时：

这里忽略了注意力本身的 FLOPs（它只占小头）、非矩阵运算、通信和"读与算不能完美重叠"等因素，所以它是下界；实际系统一般在它的 1.5 倍以内就算不错。

举个例子：30B 稠密模型，int8 权重、bf16 计算，16 张 v5e，上下文 8192，每 token KV 100 kB。batch 4 时 \((4\times819\text{MB}+30\text{GB})/(16\times8.2\times10^{11})\approx2.5\) ms；batch 256 时 MLP 已经受算力限制，\(256\times819\text{MB}/(16\times8.2\times10^{11})+2\times256\times30\times10^9/(16\times1.97\times10^{14})\approx16+4.9\approx21\) ms。batch 扩大 64 倍，步时只涨 8 倍左右——这就是吞吐和延迟的交换。

什么时候解码才受算力限制

只看线性层：\(2BP/(NC)\ge Pb_w/(NW)\iff B\ge b_w C/(2W)\)，正是第 3 节的 \(B_\text{crit}\)，和芯片数无关——是每个模型副本的 batch，不是每张卡的。但注意 KV 那一项永远在：即使线性层进了算力区，步时里还有一块纯带宽的时间，所以整体 MFU 达不到 100%，上下文越长、KV 那块越大，MFU 天花板越低。

7. 解码计算器

计算器实现的是第 6 节的步时公式。选模型、硬件、精度、芯片数，拖 batch 和上下文，看每一步时间怎么分解、处在哪个区。左图是步时随 batch 变化（对数横轴），绿色底色是 MLP 受算力限制的区域，灰色是显存放不下的区域；右图是每芯片吞吐对每 token 延迟的帕累托曲线，沿曲线 batch 从 1 增大。

值得拖一拖的几个地方：

• 默认配置（LLaMA-3-70B、int8 权重和 KV、16 张 v5e、batch 32、8k 上下文）给出约 8.6 ms/步，和第 12 节手算的结果一致。把芯片数改成 8，步时翻倍，每芯片吞吐不变——加卡只买延迟，不买单卡效率。
• 把上下文拉到 32k 以上，KV 那条线很快压过权重线，而显存上限让 batch 根本到不了绿区。LLaMA-3-70B 在几乎所有实际配置下都受 KV cache 的带宽和容量限制。
• 把权重换成 bf16：权重读取时间翻倍，\(B_\text{crit}\) 从 120 回到 240；把计算换成 8-bit：\(B_\text{crit}\) 翻倍，但只在 batch 很大时才有意义。
• 换到 H100：带宽是 v5e 的约 4 倍、显存 5 倍，同样 batch 的步时大幅下降，但 \(B_\text{crit}\) 也更大（\(C/W\) 更高），想吃满算力需要更大的 batch。

8. 推理的分片：为什么只剩张量并行

预填充：和训练几乎一样

预填充是受算力限制的大矩阵乘，训练里能用的分片几乎都能用：先做张量（模型）并行，直到通信成为瓶颈（在 TPU 上，单个 mesh 轴的界约为 \(F/2200\) 路，一般是 4–8 路，推导见张量并行一篇）；再往上做序列并行（沿序列维切，注意力部分用 ring attention 之类的方式交换 KV），长上下文时它的额外通信相对较小。增加芯片换来更低的 TTFT，代价是通信开销、每卡吞吐略降。

解码：FSDP 不行，数据并行没用，没有序列可切

• FSDP 不可行：解码本来就卡在"把权重从 HBM 读进计算单元"上。FSDP 要在每一步把权重通过芯片间互联（ICI / NVLink）AllGather 过来，而互联带宽比 HBM 低一个数量级以上。正确的做法是让权重和 KV 待在原地不动（weight-stationary），只搬小得多的激活。训练完忘了关 FSDP 直接拿去推理，是一个常见且会带来数量级性能倒退的错误。
• 数据并行没意义：它复制参数，并不能让每份参数读得更快。与其做模型内的数据并行，不如直接起多个独立的服务副本。
• 没有序列维可切：每条序列每步只有 1 个 token。

剩下的就是各种形式的张量并行：MLP 权重按隐藏维 \(F\) 切（Megatron 风格），激活在各卡复制；注意力的 \(W_Q, W_O\) 按头切，KV 投影很小，KV 头数不够切时直接复制往往更便宜。

为了延迟，可以切得比训练更多

训练里张量并行的上限来自"通信时间不能超过计算时间"。解码时我们本来就不受算力限制，受的是权重读取。所以要比较的是：每张卡读自己那份权重的时间 vs 每层的激活通信时间。对一个 \([D,F]\) 的矩阵，切 \(Y\) 路：

令 \(T_\text{ICI}\gt T_\text{HBM}\)：\(\dfrac{2BD}{W_\text{ici}}\gt\dfrac{2DF}{Y\,W_\text{hbm}}\iff Y\gt\dfrac{F}{B}\cdot\dfrac{W_\text{ici}}{W_\text{hbm}}=\dfrac{F}{B\beta}\)，其中 \(\beta=W_\text{hbm}/W_\text{ici}\)，在 v5e/v6e 上约为 8。例如 \(F=16384\)、\(B=32\)，理论上可以切到 \(16384/(32\cdot8)=64\) 路而不明显损失吞吐。batch 越小，能切得越多——小 batch 的低延迟服务可以靠堆卡继续降延迟。

KV cache 也要分片

KV cache 是注意力时间的主要来源，绝不能复制。先按 KV 头切（Megatron 风格），但最多只能切 \(K\) 路（LLaMA-3 只有 8 个 KV 头，MQA 只有 1 个）；再多的卡就沿 batch 维切：缓存形状记为 \(\text{KV}[2, B_Z, S, K_Y, H]\)，完全分布。代价是每层注意力前后各一次 AllToAll：把按头切的 query 转成按 batch 切，算完注意力再转回来。这些通信作用在小小的激活上，换来的是不用搬巨大的、原地不动的 KV。上下文特别长、batch 又小时，还可以沿序列维切 KV：先 AllGather query，各卡算局部注意力，再像 FlashAttention 那样用 log-sum-exp 合并。

两个细节

• 延迟受限的集合通信：小 batch 时激活非常小，一次集合通信的耗时被每跳的固定开销（TPU 上约 1 μs 量级）主导，而不是带宽。例如在 v5e 上做 8 路张量并行时，缓冲区小于约 360 kB 就进入这个区，而 batch 16、\(D=8192\) 的 int8 激活只有 131 kB。
• 2D 权重驻留分片（来自 Efficiently Scaling Transformer Inference）：权重同时沿 \(D\) 和 \(F\) 切成接近方形的块，通信量随总芯片数 \(\sqrt{N}\) 下降。在 \(F=4D\) 时，芯片数超过约 72 它就优于 1D 张量并行。

MoE 顺带一提：每个 token 只激活 \(k\) 个 expert，但要读的是被任何 token 选中的 expert 权重。以 \(E=16, k=2\) 为例，参数多了 \(E\) 倍、FLOPs 只多 \(k\) 倍，所以受算力限制所需的 batch 要乘 \(E/k\)：v5e 上约 \(240\times8=1920\) 个 token。MoE 的解码比稠密模型更难吃满算力。

9. 推理引擎：分离、continuous batching、分页与前缀缓存

单次预填充和单步解码各自做快之后，真实服务还要把源源不断、长短不一的请求喂给这两个函数，在延迟-吞吐曲线上停在我们想要的位置。

从最朴素的方案说起

攒一批请求，一起做预填充，然后一起解码到全部结束，再处理下一批。问题很多：TTFT 被整批预填充拖慢；短的生成要等长的生成，batch 里越来越多空槽；预填充要 pad 到最长的 prompt，浪费算力；预填充和解码被迫用同一种分片（解码想要更多张量并行，预填充不需要）。

改进一步是交错（interleaved）：预填充一条一条做（单条就已经受算力限制），解码则把很多请求拼成 batch，一有空槽就插入新请求的预填充。这对离线批量生成很好，但在线服务有个问题：做某个用户的预填充时，所有正在解码的用户都要停下来等，生成速度变得一顿一顿的。

预填充 / 解码分离（disaggregation）

把两个阶段放到两组不同的芯片上：预填充服务器算出 KV cache，经网络发给解码服务器，后者把它插进自己的解码 batch。好处：

• 一个用户的请求不会被另一个用户的预填充阻塞；两组服务器的数量可以按负载各自扩缩。
• 两边可以用各自最优的分片和拓扑；预填充服务器除了正在处理的那条，不用常驻 KV，空出的显存可以拿来做前缀缓存。

代价是 KV cache 要过网络——又一个让 KV 变小的理由。DistServe、Splitwise 等系统在 GPU 上做了同样的事。

continuous batching

解决"短请求等长请求"的办法是以解码步为调度粒度（也叫 iteration-level scheduling，最早由 Orca 提出）。引擎编译两个函数：一个能处理变长 prompt、把结果插入 KV 缓冲区某个槽位的预填充；一个对当前所有活跃槽位各走一步的解码。调度器维护请求队列：每步之后，结束的序列立刻释放槽位，排队的请求立刻被预填充并插入。这样 batch 始终是满的，不用等最慢的那条。

和"交错"的问题相关，还有一种折中叫 chunked prefill（SARATHI）：把长 prompt 的预填充切成小块，每个解码步里顺带算一块，把预填充的算力"捎带"在受带宽限制的解码步上，同时避免长时间阻塞解码。它和分离是同一个问题的两种答案。

PagedAttention：像操作系统一样管理 KV

如果给每个槽位预留"最大上下文"长度的连续 KV 缓冲区，大部分请求根本用不到那么长，显存被大量浪费，batch 开不大；读取时如果也按最大长度读，还浪费带宽。第一步改进是"不规则读取"：注意力 kernel只读每条序列有效的那段。PagedAttention（vLLM）更进一步：把 KV 切成固定大小的块（比如每块 16 个 token），像虚拟内存的页一样按需分配，每条序列有一张块表把逻辑块映射到物理块。

这样每条序列只占用它实际需要的显存（浪费最多是最后一块没填满的部分），而且多条序列可以共享物理块（引用计数 + 写时复制），并行采样多个回答时共享 prompt 的 KV 不用复制。vLLM 在相同延迟下的吞吐比 FasterTransformer、Orca 等系统高 2–4 倍。代价是实现复杂度：kernel 要按块表做间接寻址。它是纯运行时的优化，与模型架构无关。

前缀缓存

预填充受算力限制、很贵，最好的优化是少做。两个 prompt 只要前缀相同，那段前缀的 KV 就完全相同（因果注意力，前面的位置看不到后面）。于是：

• 多轮对话：每一轮都是在上一轮后面追加。保存上一轮的 KV，新一轮只需预填充新增的 token。
• 系统提示 / few-shot 示例：大量请求共用同一段开头，算一次、反复用。

实现上，缓存的 KV 很自然地组织成一棵按 token 序列索引的前缀树（trie），共享前缀共享节点，按 LRU 淘汰。难点在容量：KV 动辄几 GB，要等到下一轮请求到来才有用。常见做法是用预填充服务器上空闲的 HBM，再加上主机内存（例如一台 8 卡 v5e 服务器有 128 GiB HBM，主机 DRAM 却有约 450 GiB）——主机内存对解码来说太慢，但对"读回一份缓存"足够。另外 KV 只存在处理过初始请求的那组卡上，所以需要亲和路由把同一会话的后续请求送回同一副本，这会和负载均衡冲突。KV 越小，能缓存的会话越多、读回越快——又是那个主题。

一个具体的实现是 Google 开源的 JetStream：一组"预填充引擎"和"解码引擎"（通常在不同的 TPU 切片上）由一个控制器编排；预填充线程、解码线程和负责把 KV 从预填充切片复制到解码切片的传输线程各司其职。引擎接口只有三个方法：prefill（输入 token，产出 KV cache）、insert（把一份 KV 插入解码 batch）、generate（对 batch 里每条序列各生成一个 token 并追加其 KV）。

10. 投机解码

解码受带宽限制，意味着在小 batch 时芯片的算力大量闲置：把 1 个 token 和把 5 个 token 送过模型，要读的权重是一样的，时间几乎相同。投机解码就利用这一点：用一个便宜的草稿模型先猜 \(k\) 个 token，再让大模型（目标模型）一次前向同时验证这 \(k\) 个位置。

贪心版本

1. 草稿模型贪心生成 \(k\) 个 token \(\tilde x_1,\dots,\tilde x_k\)。
2. 目标模型对"前缀 + 这 \(k\) 个 token"跑一次前向，得到 \(k+1\) 个位置的 logits。
3. 从左往右检查：第 \(i\) 个草稿 token 是不是目标模型在该位置的 argmax。找到第一个不一致的位置，前面都接受，这个位置换成目标模型的 argmax；如果全部一致，再用最后一个位置的 logits 免费多出一个 token。

输出和直接用目标模型贪心解码完全一样，只是一次验证可能前进多个 token。

采样版本：修正的拒绝采样

非贪心时要保证输出分布严格等于目标模型的分布 \(p\)。设草稿模型在该位置的分布是 \(q\)，草稿 token \(x\sim q\)。Leviathan et al. 和 Chen et al. 同时给出了做法：

• 以概率 \(\min\!\big(1, p(x)/q(x)\big)\) 接受 \(x\)；
• 否则拒绝，并从修正分布 \(p'(x)=\dfrac{\max(0,\,p(x)-q(x))}{\sum_y \max(0,\,p(y)-q(y))}\) 重新采一个。

证明只需一行。接受的总概率 \(\beta=\sum_x q(x)\min(1,p(x)/q(x))=\sum_x\min(p(x),q(x))\)，拒绝概率 \(1-\beta=\sum_x\big(p(x)-\min(p,q)\big)=\sum_x\max(0,p(x)-q(x))\)。于是输出 \(x\) 的概率是

所以 \(\alpha=\mathbb{E}[\beta]=\mathbb{E}\big[\sum_x\min(p,q)\big]=1-\mathbb{E}[\text{TV}(p,q)]\) 就是草稿和目标"有多像"的度量。贪心是它的特例。

每次验证期望前进几个 token

做一个简化假设：每个草稿 token 被接受的概率都是 \(\alpha\)，彼此独立。一次验证产生的 token 数 \(N\)：若第一个被拒的是第 \(i\) 个（\(1\le i\le k\)），产生 \(i\) 个（\(i-1\) 个接受 + 1 个修正）；若全部接受，产生 \(k+1\) 个。所以 \(P(N\ge n)=\alpha^{n-1}\)（前 \(n-1\) 个都被接受），\(n=1,\dots,k+1\)。用 \(\mathbb{E}[N]=\sum_{n\ge1}P(N\ge n)\)：

这是一个截断几何分布，\(k\to\infty\) 时趋于 \(1/(1-\alpha)\)。例如 \(\alpha=0.8\)、\(k=4\) 时 \(\mathbb{E}[N]=(1-0.8^5)/0.2\approx3.36\)。

时间上，每次验证要跑 \(k\) 次草稿模型和 1 次目标模型。记草稿一次前向与目标一次前向的耗时比为 \(c\)，并假设目标模型验证 \(k+1\) 个 token 的耗时约等于普通解码一步（受带宽限制时成立），则期望加速比为

它在 \(k\) 上有一个最优值：\(k\) 太小浪费了验证的并行度，太大时后面的草稿 token 几乎总被拒，草稿的成本白付。代价是总 FLOPs 增加——每次验证目标模型都算了 \(k+1\) 个位置，平均只留下 \(\mathbb{E}[N]\) 个。

草稿从哪里来：可以是同系列的小模型（最好蒸馏过），也可以是 n-gram 或从上下文里查找匹配；也可以把草稿器做进主模型里，例如在主模型后几层加几个额外的预测头（Medusa、EAGLE 一类），它们共享大部分参数，跑得快、分布也更接近。

11. 量化：权重与 KV

推理对参数和 KV 的精度通常没有训练那么敏感，而且量化可以在训练后直接做。按第 6 节的公式，它影响三件事：

1. 权重量化（int8 / fp8 / int4）：参数字节数 \(Pb_w\) 减半或减到四分之一。小 batch 时步时几乎就是权重读取时间，延迟直接按比例下降；参数占的显存变少，KV 可用的空间变多，batch 能开更大；同时 \(B_\text{crit}=b_wC/(2W)\) 下降，更小的 batch 就能进入算力区。PaLM 的推理测量（Pope et al.）显示：int8 权重带来更好的延迟-吞吐前沿，但并不提高最大吞吐——在算力区，速度由 FLOPs 决定，和权重存成什么无关。
2. 低精度计算（W8A8：权重和激活都是 int8 / fp8，用 8-bit 张量核心）：\(C\) 翻倍（v5e int8 394 TOPS，H100 fp8 约 1979 TFLOPs 稠密），提高的是算力区的速度，也就是最大吞吐和预填充速度；\(B_\text{crit}\) 随之回升。权重 int8 + bf16 计算是更常见的组合，因为无损地量化权重比做低精度算术容易得多。
3. KV 量化（int8 / fp8）：\(T_\text{KV}\) 减半，同样显存能放两倍的序列。长上下文时 KV 才是大头，这一项往往比权重量化更值。

实现上，量化就是 \(x\approx s\cdot\text{round}(x/s)\)，关键在缩放因子 \(s\) 的粒度：权重常按输出通道或按每 64/128 个元素一组取 \(s\)（int4 必须分组，否则离群值会吃掉精度）；KV 常按头、按 token 取 \(s\)。反量化在 kernel 里和矩阵乘融合，所以带宽节省是真实的。int4 这类激进量化一般需要 GPTQ、AWQ 之类利用校准数据的方法来控制误差。在上面的计算器里切换"权重 / KV / 矩阵计算"三个精度，就能看到这三种效果分别作用在哪条线上。

12. 实例：在 TPU v5e 上服务 LLaMA-3-70B

LLaMA-3-70B：\(L=80\)、\(D=8192\)、\(F=28672\)、\(N=64\)、\(K=8\)、\(H=128\)、\(V=128256\)。硬件选 TPU v5e，理由是每美元的 FLOPs 最多：按 2025 年 2 月的 Google Cloud 价格，H100 约 \(3.3\times10^{17}\)、v5p 约 \(3.9\times10^{17}\)、v5e 约 \(5.8\times10^{17}\) FLOPs/美元。代价是每张 v5e 只有 16 GB HBM，必须切得很碎。

KV 有多大

int8 下每 token \(2\cdot K\cdot H\cdot L=2\cdot8\cdot128\cdot80=163{,}840\) 字节，约 164 kB。一条 32k 的序列约 5.4 GB；batch 240 就是 1.3 TB，比 70 GB 的参数大得多，光装下这些就要约 \((70\text{GB}+1.3\text{TB})/16\text{GB}\approx85\) 张 v5e。

batch 32、上下文 8k、全 int8

KV 总量 \(164\text{kB}\times8192\times32\approx43\) GB，加上 70 GB 参数共约 113 GB。至少要 8 张卡（4×2，共 128 GB），非常紧，算上其他开销实际上得用 4×4。

步时：int8 权重 + bf16 计算，\(B_\text{crit}\approx120\)，batch 32 远在带宽区，步时就是把这 113 GB 读一遍。4×2 上 \(T\approx113\text{GB}/(8\times8.2\times10^{11})\approx17\) ms，吞吐 \(32/0.017\approx1860\) token/s，即每芯片约 230 token/s。还要检查张量并行有没有撞上通信界：用两个轴做模型并行，界约为 \(2F/2200\approx26\) 路，8 路没问题。换成 4×4：通信仍然不是瓶颈，步时减半到约 8.6 ms，每芯片吞吐不变。

最小拓扑与最大 batch

如果不在乎 batch，能装下参数就行。8k 上下文下：

三行的序列数都是 43，因为参数和 KV 按同样比例缩小。70B 模型能塞进 4 张 v5e，但 batch 只有 43，算力利用率很差。用满显存时，每步就是"把一整张卡的 HBM 读一遍"：\(16\text{GB}/8.2\times10^{11}\approx19\) ms。若中位生成长度 512，一次解码约 9.7 s，每芯片吞吐约 \(43/(0.019\times512\times N)\)：bf16 / int8 / int4 分别约 0.27 / 0.55 / 1.11 QPS/芯片。把拓扑各翻一倍，参数占比减半，KV 空间多出很多，batch 从 43 涨到约 138（例如 bf16 在 4×8 上：\((512-140)\text{GB}/2.68\text{GB}\)），而步时仍是"读一整张卡"的 19 ms，吞吐变成约 0.44 / 0.89 / 1.78。结论：受 KV 容量限制时，最小的拓扑不是最划算的拓扑。

分片能切多宽

bf16 在 4×8 上能不能纯张量并行？按受算力限制的界 \(Y\gt F\cdot M_Y/2200\)（\(M_Y\) 是用于模型并行的轴数），两轴约 26 路，所以 4×4 可以、4×8 不行。但解码受带宽限制时，按第 8 节的条件可以放宽到 \(Y\approx F/(8B)\)：batch 64 时约 56 路。具体算一下单个矩阵乘在 4×8 上的三个时间：\(T_\text{ICI}=2\cdot64\cdot8192/(9\times10^{10})\approx12\) μs，\(T_\text{HBM}=2\cdot8192\cdot28672/(32\cdot8.2\times10^{11})\approx18\) μs，\(T_\text{math}=2\cdot64\cdot8192\cdot28672/(32\cdot1.97\times10^{14})\approx5\) μs——仍然受 HBM 带宽限制，可以用 4×8 换更低的延迟。int8、int4 配置则更轻松。量化在这里多了一个好处：在撞上通信界之前能开更大的 batch。

预填充与两池配比

8192 token 的预填充在 16 张 v5e 上、40% MFU：\(2\cdot70\times10^9\cdot8192/(16\cdot1.97\times10^{14}\cdot0.4)\approx0.91\) s。比起十几毫秒的解码步，这相当可观。

中位 prompt 8192、中位生成 4096、batch 32 时，平均每步有 \(32/4096\) 条序列结束，每步被淘汰的 KV 约 \(B(P+G)/G=32\times12288/4096=96\) 个 token。

分离式部署时，预填充池的产出速率是 \(P_\text{srv}/0.91\) 条/秒，解码池的消耗速率是 \(B\cdot G_\text{srv}/(0.019\times512)\)。取 batch 32、生成长度 512，令两者相等得 \(P_\text{srv}\approx3G_\text{srv}\)：长 prompt、短回答的负载需要三倍于解码的预填充服务器。

延迟-吞吐前沿

把 int8 权重、16 路张量并行（4×4）的配置沿 batch 扫一遍，就得到它的前沿：每 token 延迟从小 batch 的约 5.5 ms 到大 batch 的约 20 ms，延迟翻倍左右就能换来每 token 成本约 100 倍的下降。2k 上下文时吞吐在 batch 120 附近饱和于约 1 token/ms/芯片；上下文更长时 batch 装不下，永远到不了饱和点。在所有超过 2048 的上下文长度下，KV 读取时间都超过 FLOPs 时间。上面的计算器可以画出这条前沿：选 LLaMA-3-70B、TPU v5e、16 卡、8-bit 权重，拖动上下文看右侧曲线。

13. 小结

• 推理 = 预填充（像训练，受算力限制，优化 MFU 和 TTFT）+ 解码（串行，受带宽限制，优化每步字节数）。
• 解码一步的下界：\(\big(B\cdot\text{KV}+\text{参数}\big)/\text{总带宽}\)；线性层要 batch 超过 \(B_\text{crit}=b_wC/(2W)\)（v5e bf16 为 240）才进入算力区，注意力永远不会。
• batch 是吞吐和延迟之间的旋钮；芯片数在带宽区买的是延迟，不是单卡效率；受 KV 容量限制时，加卡也能换来更大 batch 和更高单卡吞吐。
• KV cache 是推理特有的主角：GQA / MQA / MLA / 局部注意力缩小它，量化压缩它，PagedAttention 不浪费它，前缀缓存复用它，分离式部署要传输它。
• 解码分片只剩张量并行（加上 KV 按头再按 batch 切）；FSDP 和数据并行在解码时没有意义。
• 投机解码用闲置算力换延迟，每次验证期望前进 \((1-\alpha^{k+1})/(1-\alpha)\) 个 token，在小 batch 时划算。

参考与延伸阅读

文中的推导框架和 TPU 算例主要来自 How To Scale Your Model 的推理两章，想看更多习题和图，推荐直接读原书。

• Austin et al., How To Scale Your Model, Google DeepMind, 2025：Part 7. All About Transformer Inference（线性层、注意力、延迟与吞吐估计、KV 技巧、分片、推理引擎、投机采样）；Part 8. Serving LLaMA 3-70B on TPUs；Part 4 的 KV caching 一节。
• DeepSeek-AI, DeepSeek-V2（MLA，§2.1）；模型配置见 Hugging Face 上 DeepSeek-V3 与 DeepSeek-V2 的 config.json。
• Shazeer, Fast Transformer Decoding: One Write-Head is All You Need（MQA）；Ainslie et al., GQA。
• Pope et al., Efficiently Scaling Transformer Inference（2D 权重驻留分片、PaLM 推理前沿）。
• Kwon et al., Efficient Memory Management for LLM Serving with PagedAttention（vLLM）。
• Agrawal et al., SARATHI（chunked prefill）；Zhong et al., DistServe；Patel et al., Splitwise（预填充/解码分离）；Yu et al., Orca（iteration-level scheduling）。
• NVIDIA H100、H200 规格页（计算器里的 H100/H200 数字）。
• Leviathan et al., Fast Inference from Transformers via Speculative Decoding（期望 token 数与加速比公式）；Chen et al., Accelerating LLM Decoding with Speculative Sampling；Cai et al., Medusa；Li et al., EAGLE。
• JetStream（Google 的 TPU 推理引擎）。