LLM 推理全过程的维度变化与核心公式

把 Llama 3 风格的 dense decoder-only 模型从 embedding 到采样的整条路径一次写清楚，顺便把常见变体（MHA / MQA / GQA / MLA、RMSNorm / LayerNorm、SwiGLU / GeGLU、Flash Attention / Paged Attention）的数学位置都嵌进去 — 读完之后自己能把这张图默写出来。整篇围绕一个事实展开：这套结构二十年没怎么变过，所有”推理优化”都是在同一张骨架的某个位置做局部手术。

参数符号约定 — Llama 3 8B 为基准

全篇沿用同一套符号，举例用 Llama 3 8B：

所有 shape 标注都按 PyTorch 习惯写成 [B, ..., H]；权重矩阵按”输入维 × 输出维”约定写成 $W \in \mathbb{R}^{[\text{in}, \text{out}]}$。

核心公式速查 — Embedding · Norm · Attn · FFN · LM Head

Embedding

• $\mathbf{x}_i$ — 位置 $i$ 的 embedding 向量
• $E \in \mathbb{R}^{V \times H}$ — embedding 查找表（V 个 token，每个 H 维）
• $\text{token\_id}_i \in \{0, 1, \ldots, V-1\}$ — 第 $i$ 个输入 token 的整数 ID

许多实现把 $E$ 与 LM Head 的 $W_{\text{lm}}$ 共享（tied embedding），省显存也略微正则；Llama 系列默认不共享。

归一化：LayerNorm vs RMSNorm

标准 LayerNorm：

• $\mathbf{x} \in \mathbb{R}^{H}$ — 单 token 的隐状态向量
• $\mu, \sigma^{2} \in \mathbb{R}$ — 在 $H$ 维上算的均值和方差
• $\epsilon$ — 防除零小常数（通常 $10^{-5}$）
• $\boldsymbol{\gamma}, \boldsymbol{\beta} \in \mathbb{R}^{H}$ — 可学习的 scale / shift
• $\odot$ — 逐元素乘

RMSNorm（Llama / Mistral / Qwen 主流选择）：

• $\mathbf{x} \in \mathbb{R}^{H}$ — 单 token 的隐状态向量
• $\boldsymbol{\gamma} \in \mathbb{R}^{H}$ — 可学习的 per-channel scale（无 $\boldsymbol{\beta}$）
• $\epsilon$ — 防除零小常数
• 分母是 $\mathbf{x}$ 的均方根（root mean square），故得名 RMSNorm

RMSNorm 省了均值、也省了 $\boldsymbol{\beta}$，计算量和参数都约减半；实证发现对质量几乎无损。所有主流推理引擎都按 pre-norm 组织：norm 在残差分支内部，残差主干不经过 norm。

Q/K/V 投影 + 位置编码

• $X \in \mathbb{R}^{B \times S \times H}$ — 上一步 RMSNorm 的输出
• $W_Q \in \mathbb{R}^{H \times n_q d}$ — Q 投影权重
• $W_K, W_V \in \mathbb{R}^{H \times n_{kv} d}$ — K、V 投影权重（GQA 下 $n_{kv} < n_q$，所以更窄）
• $Q \in \mathbb{R}^{B \times S \times n_q d}$，$K, V \in \mathbb{R}^{B \times S \times n_{kv} d}$ — 投影输出，会再 reshape 出 head 维

RoPE（Rotary Positional Embedding）的核心想法是：与其把”绝对位置 $m$“加到 embedding 上（像原始 Transformer 的 sinusoidal），不如让位置以旋转的形式作用在 Q、K 上，使两个向量做内积时只剩下相对位置 $n - m$。

把每个头的 $d$ 维切成 $d/2$ 个二维子空间，第 $k$ 个子空间（$k = 0, 1, \ldots, d/2 - 1$）对应坐标 $(q_{2k}, q_{2k+1})$。位置 $m$ 上对这一对乘旋转角 $m\theta_k$ 的 2D 旋转矩阵：

• $m \in \{0, 1, \ldots, S-1\}$ — 当前 token 的绝对位置
• $k \in \{0, 1, \ldots, d/2 - 1\}$ — 二维子空间索引（在单头维度 $d$ 上每两维一组）
• $(q_{2k}, q_{2k+1})$ — 投影后 Q 向量的第 $k$ 个 2D pair
• $(q'^{(m)}_{2k}, q'^{(m)}_{2k+1})$ — 在位置 $m$ 旋转后的 pair
• $\theta_k$ — 第 $k$ 个子空间的基础角速度
• $\text{base}$ — 频率衰减基（Llama 默认 $10000$；YaRN 等会动态拉大）

展开成标量：

符号同上式 — 这只是把 2×2 矩阵展开成两条标量等式，方便对照实现。

这就是标准的逆时针旋转矩阵 $R(\phi) = \begin{pmatrix} \cos\phi & -\sin\phi \\ \sin\phi & \phantom{-}\cos\phi \end{pmatrix}$，作用在 2D 向量上等价于在复平面上乘 $e^{i\phi}$ — 也是为什么 Llama 官方 repo 直接把 $(q_{2k}, q_{2k+1})$ reshape 成 complex64 然后乘 $\cos(m\theta_k) + i\sin(m\theta_k)$（HuggingFace 则等价地拆成 $(q_k, q_{k+d/2})$ 的半旋转形式，二者只差一个坐标排列）。K 用同样的角度做同样的旋转，得到 $\mathbf{k}'^{(n)}$。

为什么旋转能编码相对位置？ 把整块 $d$ 维旋转记作 $\mathbf{R}_m$（分块对角，第 $k$ 个 $2\times 2$ 块用角度 $m\theta_k$）。它正交，且 $\mathbf{R}_m^{\top}\mathbf{R}_n = \mathbf{R}_{n-m}$，因此：

• $\mathbf{q}, \mathbf{k} \in \mathbb{R}^{d}$ — 单头的 Q、K 向量（未做位置旋转）
• $\mathbf{R}_m, \mathbf{R}_n \in \mathbb{R}^{d \times d}$ — 位置 $m$、$n$ 对应的分块对角旋转矩阵
• $\langle \cdot, \cdot \rangle$ — 标准内积
• 最后一步用了 $\mathbf{R}_m^{\top} \mathbf{R}_n = \mathbf{R}_{n-m}$：旋转矩阵正交且角度可加

attention 算 $QK^{\top}$ 时，每对 $(q_i, k_j)$ 的分数只依赖差 $j - i$ — 绝对位置在内积里被自动消掉，留下相对位置。这是 RoPE 的关键性质，也是它比加性位置编码更稳的根本原因。

频率谱设计：$\theta_k = \text{base}^{-2k/d}$（$\text{base}$ 一般取 $10000$；$d$ 是单头维度，不是模型 hidden dim）让 $d/2$ 个子空间分到从快到慢的角速度：

• $k = 0$：$\theta_0 = 1$，周期 $2\pi \approx 6.28$ token，承担”近邻”信号。
• $k = d/2 - 1$：$\theta \approx \text{base}^{-(d-2)/d} \approx 10^{-4}$，周期 $\approx 2\pi \cdot 10000 \approx 6.28\text{万}$ token，承担”远距离”信号。

几何级数分布让一个 head 内同时携带不同尺度的位置信号 — 与原始 Transformer 的 sinusoidal 同构，只是从加法搬到了乘法。

长上下文缩放：训练只见过 $m \le L_{\text{train}}$，低频子空间在 $L_{\text{train}}$ 内甚至跑不完一个周期；推理一旦 $m > L_{\text{train}}$，低频角度进入训练分布外，attention 立刻退化。三类常见解法都在改 $\theta_k$：

• Position Interpolation（Chen et al. 2023）：$m \to m/s$，等价 $\theta_k \to \theta_k/s$ 同比例压缩所有频率。简单但牺牲高频精度。
• NTK-aware scaling：只缩放低频、保留高频，等价于 $\text{base} \to \text{base} \cdot s^{d/(d-2)}$。
• YaRN（Peng et al. 2023）：分频段处理 — 高频（周期 $\ll L_{\text{train}}$，已学过完整周期）不变，低频（周期 $\gg L_{\text{train}}$，没学过完整周期）按 PI 缩放，中间段平滑过渡，再叠加一个 $1/\sqrt{t}$ 的温度修正抵消缩放带来的注意力熵漂移。Llama 3.1 / 3.2 从 8K → 128K 用的就是 YaRN。

RoPE 只作用于 Q、K，不作用于 V — V 是被加权求和的值本身，不需要位置信号。

Scaled Dot-Product Attention

• $Q, K, V$ — 投影并 reshape 出 head 维后的张量，最后两维为 $[S, d]$（多头维隐式 broadcast）
• $QK^{\top} \in \mathbb{R}^{S \times S}$ — 每对 $(q_i, k_j)$ 的相似度矩阵
• $\sqrt{d}$ — 缩放因子，防止 softmax 进入梯度近乎为零的区域
• $M \in \mathbb{R}^{S \times S}$ — causal mask，上三角为 $-\infty$（位置 $i$ 只能看 $\le i$ 的位置）
• softmax 沿 K 维归一化，输出 attention 权重，再与 V 加权求和回到 $[S, d]$

多头变体：MHA / MQA / GQA / MLA

四种变体的差别只在 K、V 这一支：Q 永远是 $n_q$ 个独立头，变的是有多少套独立 K、V，以及 K、V 是否被低秩压缩。下面用同一套记号一遍把原版 MHA 和三种变体的 head-level 公式写齐 — 省略 batch 和 layer 维，看位置 $i$ 上第 $h$ 个 head。

MHA — Multi-Head Attention（原版，Vaswani et al. 2017）

每个 Q head 都配一套独立的 K、V，$n_q$ 套全独立：

• $W_Q^{(h)}, W_K^{(h)}, W_V^{(h)} \in \mathbb{R}^{H \times d}$ — 第 $h$ 个 head 独立的 Q、K、V 投影（实现上拼成大矩阵一次算完，数学上等价）
• 每个 token 入 cache 的是 $n_q$ 套 $(\mathbf{k}^{(h)}, \mathbf{v}^{(h)})$ — 共 $2 n_q d$ 个数
• Llama 2 7B：$n_q = 32, d = 128$，单 token cache = $2 \cdot 32 \cdot 128 \cdot 2\text{B} = 16\text{ KB}$ / layer

MQA — Multi-Query Attention（Shazeer 2019）

所有 $n_q$ 个 Q head 共用同一套 K、V，只剩 1 套：

• $W_K, W_V \in \mathbb{R}^{H \times d}$ — 所有 head 共享同一份 K、V 投影
• cache 收缩到 $2d$，比 MHA 小 $n_q$ 倍；但表达能力受限，大模型直接套 MQA 容易掉点 — 工程上很少单独使用，更多被 GQA 替代

GQA — Grouped-Query Attention（Ainslie et al. 2023）

把 $n_q$ 个 Q head 切成 $n_{kv}$ 组，组内共享 K、V — 是 MHA 与 MQA 之间的连续插值：

• $W_K^{(g)}, W_V^{(g)} \in \mathbb{R}^{H \times d}$，$g = 1, \ldots, n_{kv}$ — 每组一份 K、V 投影
• $g(h)$ — 第 $h$ 个 Q head 所属的组索引
• $n_{kv} = n_q$ 退化为 MHA，$n_{kv} = 1$ 退化为 MQA
• kernel 里只实际算 $n_{kv}$ 套 K、V，分数计算时把 K、V 沿 group 维 broadcast 到 $n_q$，不真的复制张量
• Llama 3 70B：$n_q = 64, n_{kv} = 8, d = 128$，单 token cache = $2 \cdot 8 \cdot 128 \cdot 2\text{B} = 4\text{ KB}$ / layer — 比同样大小的 MHA 小 8 倍

MLA — Multi-Head Latent Attention（DeepSeek V2 2024）

GQA 只是按 head 数线性缩 cache；MLA 直接把 K、V 共同压成一个低秩潜向量 $\mathbf{c}^{KV}$，再为 RoPE 单独走一条共享的小分支。分成四步看：

(1) 内容分支 — K、V 共享一份 down-projection，入 cache 的只是潜向量 $\mathbf{c}^{KV}_i$：

(2) Q 端同样低秩 — 训练省显存，推理时 Q 不入 cache：

(3) RoPE 解耦分支 — K 侧只算一份共享的 $d_r$ 维向量，所有 head 共用：

(4) 拼接 + Attention — 内容部分和 RoPE 部分沿 head 维拼起来再算分数：

• $W^{DKV} \in \mathbb{R}^{H \times d_c}$ — KV 共用 down-projection（DeepSeek V3 取 $d_c = 512$）
• $W_K^{U,(h)}, W_V^{U,(h)} \in \mathbb{R}^{d_c \times d}$ — 第 $h$ 个 head 自己的 up-projection
• $W^{KR} \in \mathbb{R}^{H \times d_r}$ — 共享的 RoPE K 分支（DeepSeek V3 取 $d_r = 64$）
• $\mathbf{c}^{KV}_i \in \mathbb{R}^{d_c}, \mathbf{k}^{R}_i \in \mathbb{R}^{d_r}$ — MLA 实际缓存的就这两条 — 共 $d_c + d_r$ 个数
• DeepSeek V3：$n_q = 128, d = 128, d_c = 512, d_r = 64$，单 token cache = $(512 + 64) \cdot 2\text{B} = 1152$ 字节 / layer — 比同等规模 GQA 又小一个量级

为什么 RoPE 要单走一条？ 推理时有个”权重折叠”技巧 — 用矩阵结合律把 $W_K^{U,(h)}$ 折进 $W_Q^{U,(h)}$：

这样 K 内容部分根本不用真的重建，attention 直接在缓存的 $\mathbf{c}^{KV}$ 上算。但 RoPE 的旋转角依赖绝对位置 $m$，无法折叠进固定权重 — 一旦把 RoPE 加在重建后的 K 上，前面的折叠就失效。DeepSeek 的办法是把 RoPE 拆成独立的 $d_r$ 维浅分支，让”可折叠的内容部分”和”必须实时旋转的位置部分”互不干扰，KV-cache 压缩和相对位置信号才能同时保留。

对比总览：参数量 / 计算量 / KV cache

下面把每种变体的开销按步骤拆开。约定：单层 attention 子层，prefill 长度 $S$，忽略 norm / bias / softmax 等非矩阵乘项；linear 投影 $X W$（$X \in \mathbb{R}^{S \times m}, W \in \mathbb{R}^{m \times n}$）FLOPs 记 $2Smn$（每个 MAC 算 2 FLOPs），attention 的 $QK^{\top}$ 和 $AV$ 不扣 causal mask 的一半。MLA 额外用到 $d_q'$ = Q 端潜向量维度（DeepSeek V3 取 $1536$）。

参数量（单层）

Prefill FLOPs（单层，序列长度 $S$）

表里 “RoPE 分支” 在 MHA/MQA/GQA 指的是把旋转矩阵作用在 Q、K 上的逐元素 mul/add，量级 $\mathcal{O}(S n_q d)$，相对矩阵乘可忽略；在 MLA 里专指那条额外的 $W_Q^R$ / $W^{KR}$ 投影，属于真矩阵乘，不能忽略。

KV cache（单 token，单层，fp16）

把三张表横着读，能看出三件事：

1. 变体只动 K/V 一侧。Q 投影、$W_O$、$AV$ 在四种变体里完全相同 — 真正动手术的只有 K、V 这一支，所以同尺寸下 attention 的总参数和总 FLOPs 不会差到数量级。
2. MHA → MQA/GQA 是同时省 params、FLOPs、cache；MLA 是用 params/FLOPs 换 cache。GQA 通过减小 $n_{kv}$ 把 K、V 投影、cache 同步线性缩小；MLA 反过来，加了 DKV + 两段 UK/UV + 单独的 RoPE 分支，参数和 prefill FLOPs 跟同尺寸 GQA 同量级（不会显著降低），换来的是单 token cache 从 KB 量级压到 ~1 KB。
3. $S^2$ 项在 MHA/MQA/GQA 三者间没差。$QK^{\top}$ 和 $AV$ 都是 $2 n_q S^2 d$（K、V 共享只是 broadcast，不省二次项）。也就是说当 $S \gg H / d$ 时，三种变体的 prefill FLOPs 会趋同 — 它们的差距全部体现在 decode 阶段每生成一 token 要读多少 KV cache，而那是 memory I/O，不是 compute。

Output 投影 + 残差

• $\mathbf{x} \in \mathbb{R}^{H}$ — attention 子层的输入（pre-norm 之前的值）
• $Q', K', V'$ — Q、K 经过 RoPE 旋转后的版本；$V' = V$（V 不旋转，记号统一而已）
• $W_O \in \mathbb{R}^{n_q d \times H}$ — output projection，把多头 concat 后投回 $H$ 维
• $\mathbf{h} \in \mathbb{R}^{H}$ — attention 子层输出 + 残差

FFN 变体

经典双线性 FFN（GPT-2）

• $\mathbf{x} \in \mathbb{R}^{H}$ — FFN 输入
• $W_1 \in \mathbb{R}^{H \times I}$ — 升维投影
• $W_2 \in \mathbb{R}^{I \times H}$ — 降维投影
• $\phi$ — 逐元素的非线性激活（见下节）

激活函数

FFN 升维之后逐元素套的标量非线性。输入输出都是标量 $x$；它在每个 token、每个隐藏维度上独立应用。下面四个函数覆盖了 Transformer 时代用过的所有主流选择。

ReLU（Nair & Hinton 2010）

• $x > 0$ 时恒等映射，$x \le 0$ 时输出 0
• 计算最便宜，但负区梯度恒为 0 — “dead neuron” 问题
• 用于原始 Transformer 和早期 BERT 实现

Sigmoid

• 把 $\mathbb{R}$ 压到 $(0, 1)$，天然适合做 “门” 信号 — GLU 原版的 $\phi$ 就是它
• 单独当 FFN 激活已经基本不用 — 两端饱和导致梯度消失

GeLU（Gaussian Error Linear Unit，Hendrycks & Gimpel 2016）

实现里常用 OpenAI 给的 tanh 近似（数值上误差 $< 10^{-3}$，但省去 erf 调用）：

• $\Phi$ 是标准正态 CDF — 直观上 “按 $x$ 的尾部概率加权地放行 $x$”
• 处处可导、非单调（在 $x < 0$ 一侧有一段轻微的负向凹陷），比 ReLU 平滑
• 用于 GPT-2/3、BERT、ViT

SiLU / Swish（Ramachandran et al. 2017）

• 形状非常接近 GeLU（同样平滑、非单调、过原点），但表达式更简单 — 不用 $\text{erf}$ 也不用三次项
• 自门控（self-gated）：用自身的 sigmoid 控制信号通过率
• 用于 PaLM、Llama 全系（在 SwiGLU 里担任门控 $\phi$）

这四个里 ReLU 和 sigmoid 已经退出主流 FFN，活跃的是 GeLU（GPT 时代）和 SiLU（Llama / PaLM 之后）。两者在 bf16 下数值差别 < 1%，论文里挑哪个多是路径依赖；真正的工程拐点是把激活从 “套在投影上”（经典 FFN）换成 “套在门控上”（GLU 家族）。

GLU 家族（Llama、PaLM、Mistral 都在用）

• $\mathbf{x} \in \mathbb{R}^{H}$ — FFN 输入

• $W_{\text{gate}}, W_{\text{up}} \in \mathbb{R}^{H \times I}$ — 两个独立的升维投影

• $W_{\text{down}} \in \mathbb{R}^{I \times H}$ — 降维投影

• $\odot$ — 逐元素乘

• $\phi$ — 门控激活，按选择决定变体名：

  变体	$\phi$	代表模型
  GLU	$\sigma$	Dauphin et al. 2017 原版
  ReGLU	$\text{ReLU}$	—
  GeGLU	$\text{GeLU}$	T5 v1.1
  SwiGLU	$\text{SiLU}$	PaLM, Llama 1/2/3

SwiGLU 比经典 GeLU-FFN 多一个投影（三矩阵 vs 两矩阵），为了参数预算对齐，实现里一般把 $I$ 设成 $\tfrac{2}{3} \cdot 4H$（$4H$ 是 GPT-2 的惯例），Llama 3 8B 的 $I = 14336 \approx \tfrac{2}{3}\cdot 4\cdot 4096 \cdot 1.3$。

Mixture-of-Experts（MoE）

经典 dense FFN 里每个 token 都要过同一对 $W_{\text{up}} / W_{\text{down}}$ — 参数全用上，FLOPs 也全付。MoE（Shazeer et al. 2017）的核心是把 FFN 复制成 $N$ 份（“expert”），每个 token 只挑前 $k$ 份算，总参数线性放大，单 token 激活参数和 FLOPs 几乎不变 — 用容量换知识量，不换 compute。

形式上把 FFN 子层换成：

Router（gating network） 决定每个 token 走哪几个 expert：

• $\mathbf{x} \in \mathbb{R}^{H}$ — 当前 token 的隐状态（FFN 子层输入）
• $W_g \in \mathbb{R}^{H \times N}$ — router 投影，把 hidden state 映到 $N$ 个 expert 的 logit
• $\mathbf{s}(\mathbf{x}) \in \mathbb{R}^{N}$ — 所有 expert 的 router 概率
• $\mathcal{T}_k(\mathbf{x})$ — 被选中的 top-$k$ expert 索引集合
• $g_i(\mathbf{x})$ — combine 权重；Mixtral / DeepSeek 都对 top-$k$ 的 $s_i$ 再做一次归一化让它们和为 1
• $\text{FFN}_i$ — 第 $i$ 个 expert，结构通常就是 SwiGLU，仅权重不共享

MoE vs 传统 FFN

关键 trade-off：容量与 FLOPs 解耦。同样激活参数（即同 FLOPs 预算）下，MoE 可以塞 8–32 × 总参数 — 知识容量上去了。代价是显存（要装下所有 expert）、路由稳定性（避免热门 expert）、推理 batching（每 token 路径不同）。

Load Balance

最朴素的 router 训着训着会”塌缩”到少数热门 expert。GShard / Switch 用辅助损失（auxiliary loss）做软约束：

• $f_i$ — 当前 batch 内被路由到 expert $i$ 的 token 占比
• $\bar{s}_i$ — 当前 batch 内 expert $i$ 的平均 router 概率
• $\alpha$ — 辅助损失权重（Switch 用 $0.01$ 量级）
• 直觉：当 $f_i$ 和 $\bar{s}_i$ 都大时 = 这个 expert 既被选很多次又信心很高 → 惩罚 → 反向梯度把 router 的这条 logit 推回去

DeepSeek V3 进一步抛弃辅助损失，用无损失 load balance：给每个 expert 维护偏置 $b_i$，TopK 选择基于 $s_i + b_i$；被选少的 $b_i$ 慢慢调高，被选多的 $b_i$ 调低 — 只影响选择不污染梯度，避免 aux loss 对主任务造成性能下拉。

现代实现对照

几条规律：

1. 从粗到细。Mixtral 这代是”少而胖”（8 个接近 dense 大小的 expert）；DeepSeek 起把每个 expert 切小、数量推到上百，相当于”多而瘦” — 同样激活 FLOPs 下表达组合数指数级上升。
2. 共享 expert 成为标配。DeepSeek / Qwen-MoE / Llama 4 都给每层保留 1–2 个”所有 token 必过”的共享 expert，专门吃通用模式；剩下的稀疏 expert 才负责特化。
3. MoE 总和 KV cache 压缩配套。MoE 把总参数推到几百 B，long-context decode 同样压力大 — DeepSeek V3 同时上 MLA + 细粒度 MoE 不是巧合。
4. top-$k$ 走极端。Switch (k=1) → Mixtral (k=2) → DeepSeek V3 (k=8 但 expert 小) → Llama 4 (k=1 + 共享)。top-$k$ 小好做 batching 和容量约束，但要靠细粒度 / 共享 expert 把表达力补回来。

残差

• $\mathbf{h}$ — attention 子层输出（已含第一道残差）
• $\mathbf{x}_{\text{out}} \in \mathbb{R}^{H}$ — 整层 Transformer 输出，进入下一层
• RMSNorm 仍在残差分支内（pre-norm），主干 $\mathbf{h}$ 直接跨过

LM Head + 采样

• $\mathbf{x}_{\text{final}} \in \mathbb{R}^{H}$ — 最后一层 RMSNorm 后、最后一个位置的隐状态
• $W_{\text{lm}} \in \mathbb{R}^{H \times V}$ — 输出嵌入矩阵（可与 $E$ 共享，即 tied embedding）
• $\text{logits} \in \mathbb{R}^{V}$ — 词表上每个 token 的未归一化分数
• $T > 0$ — temperature，越大概率分布越平坦（$T \to \infty$ 趋近均匀，$T \to 0$ 趋近 argmax）
• $\mathbf{p} \in \mathbb{R}^{V}$ — 归一化后的概率分布

采样前通常叠几层 logits 变换：

重复惩罚（repetition / frequency / presence penalty）：

• $v$ — 词表中某个 token
• $\text{logits}_v$ — 该 token 的原始 logit
• $\alpha \ge 0$ — presence penalty（只要出现过就扣一次）
• $\beta \ge 0$ — frequency penalty（按出现次数线性扣）
• $\mathbb{1}[\cdot]$ — 指示函数（条件成立为 1，否则为 0）
• $\text{count}(v)$ — token $v$ 在已生成序列中的出现次数

Top-k：只保留最大 $k$ 个 logits，其它置 $-\infty$。

Top-p（nucleus）：按概率降序累加，保留累积概率 $\le p$ 的集合。

Min-p：保留 $p_v \ge p_{\min} \cdot p_{\max}$ 的集合，对低熵分布更友好。

Typical-p：基于与条件熵的偏差做截断，保留 $|-\log p_v - H(\mathbf{p})|$ 小的集合。

所有截断都作用在概率分布本身之前/之后，不改变公式的骨架。

Prefill 阶段维度流转 — S 个 token 一次性走一遍

输入：input_ids [B, S] = [2, 10]。下图是单层 Transformer 的前向，外面再套 32 层；shape 始终保持在 [B, S, H] = [2, 10, 4096]，残差边以橙色虚线表示。

Prefill 结束后，KV Cache 状态：每层已填入前 10 个位置。

Decode 阶段维度流转 — 第 t 步只走 1 个 token

前置状态：已有 $\text{cache\_len} = S + t - 1$ 个位置。

输入：input_ids [B, 1] = [2, 1]（上一步生成的 1 个 token）。

Prefill vs Decode 维度对照 — GEMM vs GEMV · 算力 vs 带宽

这张表是理解所有推理加速工作的起点：prefill 像训练的 forward，compute-bound；decode 是一串 GEMV，memory-bound，绝大部分时间在往 SM 里搬权重。两段的优化方向天差地别。

KV Cache 的形状与增长 — 单 token 几 KB · 长 context 几百 MB

每层一对 cache：

• $S_{\max}$ — 预分配的最大序列长度（一般 = 模型上下文上限或调度器允许的上限）
• 其它符号沿用顶部约定（$B, n_{kv}, d$）；4 个维度的顺序按 PyTorch [B, head, seq, head_dim] 习惯

每个 token、每层的 cache 大小（fp16）：

• 最左的 $2$ — K 和 V 两份
• $2\ \text{bytes}$ — fp16 每元素 2 字节（用 fp8 / int8 可减半到 1/4）
• 右侧代入 Llama 3 8B：$n_{kv} = 8$，$d = 128$

每个 token、全模型（32 层）：$4\ \text{KB} \times 32 = 128\ \text{KB} / \text{token}$。一个 4096-token 请求：$128\ \text{KB} \times 4096 \approx 512\ \text{MB}$。

几种工程优化：

• Paged Attention（vLLM）：把 cache 拆成固定大小 block（通常 16 token），用一张 block table 做虚地址到物理地址的映射，消除碎片。对应公式里没变化，只是张量布局和访问模式变了。
• Sliding Window Attention（Mistral）：只保留最近 $W$ 个 token 的 K、V。cache 上限从 $S_{\max}$ 降到 $W$，代价是信息截断，靠层间层叠传递远距离依赖。
• INT8 / FP8 KV Cache：把 fp16 cache 量化到 int8 甚至 fp8，per-channel 或 per-token 量化，误差可控，cache 占用减半到 1/4。代表工作 KIVI / KVQuant。
• KV 压缩 / 丢弃（H2O、StreamingLLM、SnapKV）：按注意力权重把不重要的位置踢出 cache，上下文超长时用。
• MLA：前面已经提过，是从模型结构层面改 cache 形状，不是后处理。

每步的计算量与内存开销 — Llama 3 8B fp16 · H100 拐点 ~330 FLOPs/byte

前面的维度图只说 shape，不说量级。推理优化 90% 的讨论都在算”这一步要花多少 FLOPs、搬多少字节”，所以这里直接把每步的开销拍成表。

以 Llama 3 8B、fp16、$B=1$ 为基准，prefill 取 $S=2048$，decode 取 $\text{cache\_len}=2048$（即生成到第 2048 个 token 时的某一步）。

参考硬件拐点：H100 SXM fp16 理论算力 ~989 TFLOPs，HBM 带宽 ~3 TB/s，roofline 拐点 $\text{AI}^{*} \approx 330\ \text{FLOPs/byte}$。高于它是 compute-bound，低于它是 memory-bound。

权重分布

全模型 fp16 权重约 16 GB，每次 forward 的”底价”就是把这 16 GB 从 HBM 里扫一遍。H100 @ 3 TB/s 下 $= 16/3000 \approx 5.3\ \text{ms}$，这就是单请求 decode 的物理下限。

每层每步的计算 / 内存读写

对同一层在 prefill（$N=S$）和 decode（$N=1$）下的各个子步做对照。“Weight HBM” 是要从显存搬的权重字节，“KV HBM” 是要读/写的 KV Cache 字节。中间 activation 默认被 kernel 融合，不单独算。

几个直接结论：

• FFN 是真正的主角。$W_{\text{gate}} + W_{\text{up}} + W_{\text{down}}$ 吃掉 ~75% 的 FLOPs 和 ~80% 的权重带宽。MoE、稀疏激活、FFN 量化全都冲着这块去。
• Attention 的 4 个投影（Q/K/V/O） 占 ~18%，真正的 $QK^{\top}$ 和 $\cdot V$ 只占 ~7% — prefill 里 attention 并不是瓶颈，投影才是。
• Decode 的 KV 读取 每层 8 MB，在 $\text{cache\_len}=2048$ 时只占权重读取的 ~2%。但当上下文拉到 64K、128K，它会翻几十倍、直接反超权重带宽，成为新瓶颈（这也是 Paged Attention、sliding window、KV 量化出现的原因）。

全模型一次 forward

把 32 层 + embedding + LM head 加起来：

Decode 的 1.05 FLOPs/byte 比 H100 拐点 330 低 两个半数量级 — 意味着理想情况下单请求 decode 的算力利用率只有 $1.05/330 \approx 0.3\%$。这就是 continuous batching 的数学依据：把 $B$ 拉到 32，同一批权重 read 被 32 个请求共享，arithmetic intensity 直接 ×32，decode 吞吐几乎线性增长，直到 attention 部分或算力先撞墙。

心算口诀

两条规则覆盖 90% 的推理性能估算：

1. FLOPs ≈ $2 P N$：$P$ 是参数量（~8B），$N$ 是这次 forward 处理的 token 总数。每个参数被每个 token 各用一次 MAC，一次 MAC 算 2 FLOPs。例如 prefill $S=2048$: $2 \times 8\text{B} \times 2048 \approx 33\ \text{TFLOPs}$，与分项加总的 31 TFLOPs 吻合。
2. 权重 HBM I/O ≈ $2 P$ bytes（fp16）：一次 forward 就是把模型扫一遍，约 16 GB。

Arithmetic intensity 本质是 $\frac{2 P N}{2 P} = N$ — forward 里一共参与的 token 数。Prefill 有 $S \cdot B$ 个 token，decode 只有 $B$ 个。这一个数字直接决定了 prefill / decode 瓶颈不同的根源。

计算复杂度总览 — 有 KV Cache vs 没 KV Cache 差三个数量级

Prefill（一次性处理 $S$ 个 token）：

• $L \cdot S \cdot H^{2}$ — 每层 4 个 $H \times H$ 投影 + FFN 三个 $H \times I$ 投影（$I \sim 4H$），各作用在 $S$ 个 token 上
• $L \cdot S^{2} \cdot H$ — attention 的 $QK^{\top}$ 与 $\cdot V$，含 $S \times S$ 的 score 矩阵
• 短序列下线性层主导；$S \gtrsim H$ 后 attention 二次项追上

Decode 每步（处理 1 个 token，历史 $\text{cache\_len}$）：

• $\text{cache\_len}$ — 当前已缓存的位置数（$= S + t - 1$）
• 单步只处理 1 个新 token，所以线性层把 $S$ 替换为 $1$；attention 仍要扫整段 cache，故随 cache 线性增长

生成 $T$ 个 token 总复杂度：

• $T$ — 生成 token 数（沿用顶部约定）
• $(S + T)$ — 平均 attention 跨距（prompt 长度 + 已生成段）
• 第一项是 decode 线性层对 $T$ 步累加；第二项是 attention 部分对 $t = 1, \ldots, T$ 的求和近似（精确是 $\sum_t (S + t)$）

对比无 KV Cache：$O(L \cdot (S+T)^{3})$，差异巨大。

把 FLOPs 和带宽一起看就更直观 — 这正是上一节那张表展示的：prefill 挑算力上限，decode 挑带宽上限；continuous batching 的意义就是把 $N$ 个请求的 decode 拼成一个大 GEMV，搬权重的成本被 $N$ 个请求摊薄，吞吐线性上涨直到算力或 attention 部分成为瓶颈。

工程优化怎么嵌进公式 — Flash Attn / Spec Decode / Continuous Batch

Flash Attention：数学上完全等价于标准 attention，公式一字不变。工程上把 softmax 和 matmul 融成一个 kernel，按 block 流式更新 softmax 的运行统计量（max、sum），避免把 $S \times S$ 的 attention 矩阵写回 HBM。复杂度不变，显存占用从 $O(S^{2})$ 降到 $O(S)$，速度提升主要来自 HBM 访问减少。FA-2 把切分粒度从 head 改到 query block；FA-3 在 H100/H200 上叠加 warpgroup MMA + producer-consumer 异步流水。

Flash Decoding：Flash Attention 在 decode 时 Q 只有 1 行，kernel 并行度不够。Flash Decoding 把 K、V 的 $\text{cache\_len}$ 维切成多块并行，再做一次 log-sum-exp 归约。公式还是同一个 softmax，只是拆成两段算。