Attention Block 推理路径：从输入 Shape 到 Vocab Logits（单请求版，含数字例子）

本文只讨论**推理（inference）**路径，不涉及训练反向传播、参数更新、dropout。目标是把 input -> attention block -> logits over vocab 串成一条连续 shape 链路，并在每个步骤给出具体数字例子。全文默认单请求，不引入 batch 维。

1. 记号与统一数字配置

通用符号：

$S$ ：本次输入 token 数（prefill 常见 $S > 1$ ，decode 常见 $S = 1$ ）
$T_{hist}$ ：历史 KV cache 长度
$T$ ：当前可见总长度， $T = T_{hist} + S$
$d_{model}$ ：隐藏维度
$H$ ：attention 头数（标准 MHA 下，Q/K/V 头数相同）
$d_{head}$ ：每头维度，通常 $d_{model} = H \cdot d_{head}$
$V$ ：词表大小

本文数字例子（decode 一步）统一使用：

$S = 1$ （单请求）
$T_{hist} = 128$
$T = 129$
$d_{model} = 4096$
$H = 32$
$d_{head} = 128$ （ $32 \cdot 128 = 4096$ ）
$V = 32000$

输入 token id：

这是一个长度为 $S$ 的 token id 序列。
数字化 shape： $𝐈 : 1$ 。

2. 从 Token 到首层输入

token embedding + 位置编码（RoPE/ALiBi 等）可写成：

𝐗_{0} = E m b e d (𝐈) + P o s E n c

嵌入输出是二维隐藏表示，shape 为 $S \times d_{model}$ 。
数字化 shape： $𝐗_{0} : 1 \times 4096$ 。这一步可以理解为：把离散 token id 变成连续向量，并注入位置信息，供后续注意力计算。

3. 单个 Attention Block（Pre-Norm）

以第 $ℓ$ 层为例，输入：

$𝐗_{ℓ}$ 的 shape 为 $S \times d_{model}$ 。
数字化 shape： $𝐗_{ℓ} : 1 \times 4096$ 。

3.1 LayerNorm + QKV 线性映射

{\hat{𝐗}}_{ℓ} = L N (𝐗_{ℓ})

LN 不改 shape： $1 \times 4096$ 。作用是先把数值尺度稳定下来，让 Q/K/V 投影更平稳。

更展开地看，对单个 token 向量 $𝐱 \in ℝ^{d_{model}}$ ，LayerNorm 在特征维上做：

μ = \frac{1}{d_{model}} \sum_{i = 1}^{d_{model}} x_{i}, σ^{2} = \frac{1}{d_{model}} \sum_{i = 1}^{d_{model}} (x_{i} - μ)^{2}

L N (x_{i}) = γ_{i} \frac{x_{i} - μ}{\sqrt{σ^{2} + ϵ}} + β_{i}

其中 $γ, β \in ℝ^{d_{model}}$ 是可学习参数， $ϵ$ 是数值稳定项。
直觉上，它先把激活“标准化”，再用 $γ, β$ 恢复可学习的缩放与平移能力。

𝐐 = {\hat{𝐗}}_{ℓ} 𝐖_{Q}, 𝐊 = {\hat{𝐗}}_{ℓ} 𝐖_{K}, 𝐕 = {\hat{𝐗}}_{ℓ} 𝐖_{V}

其中：

$𝐖_{Q}, 𝐖_{K}, 𝐖_{V} \in ℝ^{d_{model} \times (H d_{head})}$

先看 projection 后最后一维：

$𝐐_{raw} : S \times (H d_{head})$
$𝐊_{raw} : S \times (H d_{head})$
$𝐕_{raw} : S \times (H d_{head})$

数字化：

$𝐐_{raw} : 1 \times 4096$
$𝐊_{raw} : 1 \times 4096$
$𝐕_{raw} : 1 \times 4096$

reshape 成多头：

$𝐐, 𝐊_{new}, 𝐕_{new} : H \times S \times d_{head}$

数字化：

$𝐐 : 32 \times 1 \times 128$
$𝐊_{new} : 32 \times 1 \times 128$
$𝐕_{new} : 32 \times 1 \times 128$ 直觉上，这是把同一个 token 的表示切成 32 个子空间并行建模，每个头各自关注不同关系。

拆分过程可写成（以 $𝐐_{raw}$ 为例）：

𝐐_{raw} \in ℝ^{S \times (H d_{head})} \overset{reshape}{\to} 𝐐 \in ℝ^{H \times S \times d_{head}}

在本例里， $H d_{head} = 32 \times 128 = 4096$ ，所以本质只是把最后一维 4096 重排为 32 x 128。

3.2 与 KV Cache 拼接（推理关键）

历史缓存：

$𝐊_{cache} : H \times T_{hist} \times d_{head}$
$𝐕_{cache} : H \times T_{hist} \times d_{head}$

数字化：

$𝐊_{cache} : 32 \times 128 \times 128$
$𝐕_{cache} : 32 \times 128 \times 128$

拼接后：

𝐊_{all} = C o n c a t (𝐊_{cache}, 𝐊_{new}) \in ℝ^{H \times T \times d_{head}}

𝐕_{all} = C o n c a t (𝐕_{cache}, 𝐕_{new}) \in ℝ^{H \times T \times d_{head}}

数字化：

$𝐊_{all} : 32 \times 129 \times 128$
$𝐕_{all} : 32 \times 129 \times 128$ 因为是 decode 场景（ $S = 1$ ），每一步只新增一个位置，历史部分复用 cache，不重复算旧 token 的 K/V。

标准 MHA 下，Q/K/V 头数一致，均为 $H$ 。

3.3 注意力分数到权重

𝐒 = \frac{𝐐 𝐊_{all}^{⊤}}{\sqrt{d_{head}}} + 𝐌

其中 $𝐌$ 是因果 mask。它保证当前位置只能看见历史与当前，不能看见未来 token。

$𝐒$ 的 shape 为 $H \times S \times T$ 。
数字化： $𝐒 : 32 \times 1 \times 129$ 。

softmax 后：

𝐏 = s o f t m a x (𝐒, last dim)

数字化： $𝐏 : 32 \times 1 \times 129$ 。也就是每个头都会产出一条长度 129 的注意力分布。

3.4 权重加权 Value 得到上下文

𝐂 = 𝐏 𝐕_{all}

$𝐂$ 的 shape 为 $H \times S \times d_{head}$ 。
数字化： $𝐂 : 32 \times 1 \times 128$ 。 $𝐂$ 可看作“按注意力权重汇总后的上下文摘要”。

合并多头后， $𝐂_{merge}$ 的 shape 为 $S \times (H d_{head})$ 。
数字化： $𝐂_{merge} : 1 \times 4096$ 。

合并过程与拆分相反：

𝐂 \in ℝ^{H \times S \times d_{head}} \overset{transpose + reshape}{\to} 𝐂_{merge} \in ℝ^{S \times (H d_{head})}

实现上通常先把维度从 $(H, S, d_{head})$ 调整到 $(S, H, d_{head})$ ，再把后两维展平。
因为元素总数不变（本例每个位置仍是 $32 \times 128 = 4096$ ），所以这一步是无信息损失的重排。

输出投影：

𝐎_{attn} = 𝐂_{merge} 𝐖_{O}, 𝐖_{O} \in ℝ^{(H d_{head}) \times d_{model}}

数字化： $𝐎_{attn} : 1 \times 4096$ 。

残差连接：

𝐘_{ℓ} = 𝐗_{ℓ} + 𝐎_{attn}

数字化： $𝐘_{ℓ} : 1 \times 4096$ 。残差连接让模型在“保留原信息”和“引入新信息”之间平衡，也让深层堆叠更稳定。

3.5 MLP 子层（推理前向）

{\hat{𝐘}}_{ℓ} = L N (𝐘_{ℓ})

𝐔 = ϕ ({\hat{𝐘}}_{ℓ} 𝐖_{up}) ⊙ ({\hat{𝐘}}_{ℓ} 𝐖_{gate})

𝐎_{mlp} = 𝐔 𝐖_{down}

𝐗_{ℓ + 1} = 𝐘_{ℓ} + 𝐎_{mlp}

数字化（关注输入输出主 shape）：

$L N (𝐘_{ℓ}) : 1 \times 4096$
$𝐎_{mlp} : 1 \times 4096$
$𝐗_{ℓ + 1} : 1 \times 4096$ 可以把这部分理解成逐位置的非线性特征变换：attention 负责“找信息”，MLP 负责“加工信息”。

4. 多层堆叠到最终 Logits

经过 $L$ 层后：

$𝐗_{L}$ 的 shape 为 $S \times d_{model}$ 。
数字化： $𝐗_{L} : 1 \times 4096$ 。

最终归一化：

𝐇 = {L N}_{final} (𝐗_{L})

数字化： $𝐇 : 1 \times 4096$ 。

映射到词表：

𝐙 = 𝐇 𝐖_{lm}^{⊤} + 𝐛_{lm}

其中：

$𝐖_{lm} \in ℝ^{V \times d_{model}}$
$𝐙 \in ℝ^{S \times V}$

数字化： $𝐙 : 1 \times 32000$ 。 $𝐙$ 中每个值都是对应词表 token 的未归一化分数（logit）。

推理采样通常取最后位置：

$𝐙_{last} = 𝐙 [:, - 1, :]$
数字化： $𝐙_{last} : 32000$

5. 一条完整的 Shape 链（通用 + 数字）

通用链路可读作：
$𝐈 (S)$ → $𝐗_{0} (S, d_{model})$ → $L$ 层 Block → $𝐇 (S, d_{model})$ → $𝐙 (S, V)$ → $𝐙_{last} (V)$ → next token。

数字例子对应：
$(1)$ → $(1, 4096)$ → $L$ 层 block → $(1, 4096)$ → $(1, 32000)$ → $(32000)$ 。

6. Prefill 与 Decode 的区别（推理视角）

Prefill： $S$ 较大，一次性处理多个 token 并构建初始 cache。
Decode：通常 $S = 1$ ，每步只算新 token 的 Q/K/V，再与历史 cache 拼接；（单请求省略 batch 维）代价主要在长度 $T$ 的注意力读取。

两者公式一致，主要区别是 $S$ 、 $T$ 和 cache 状态。