附录张量形状参考

本附录集中列出现代 decoder-only 模型的数据流与张量形状，供阅读主线时查阅。完整原理从序列建模背景开始；这里采用的现代结构假设与资源分析章节一致。

1 从 token 到下一个 token

Tokenizer 将文本转换成 token ID。Embedding 用 ID 查表，将每个 ID 转成 $d$ 维向量。一个典型 decoder-only 模型的数据流是：

token IDs [B,S]
    ↓ Embedding [V,d]
hidden states [B,S,d]
    ↓ 重复 L 个 block
    │ RMSNorm → Attention → 残差相加
    │ RMSNorm → FFN 或 MoE → 残差相加
    ↓ final RMSNorm
    ↓ LM Head [d,V]
logits → 采样 → 下一个 token ID

残差相加要求分支输出回到 $d$ 维。LM Head 的 logits 是词表中每个候选 token 的未归一化分数。生成时一般只需要最后一个位置的 logits；训练时通常需要各预测位置的 logits。

矩阵乘法 $[n,a]\times[a,b]\to[n,b]$ 的权重有 $ab$ 个标量，主项计算量为 $2nab$ FLOPs。Embedding 查表只读取向量，不会对整张 $[V,d]$ 表做一次稠密矩阵乘法。

2 一层 Attention 的投影

暂时把 $B,S$ 合并成 token 数 $n$，输入 $X\in\mathbb{R}^{n\times d}$。

$$ Q=XW_Q,\qquad K=XW_K,\qquad V_{\rm val}=XW_V $$

权重	形状	输出形状	参数数
$W_Q$	$d\times d_q$	$n\times d_q$	$dd_q$
$W_K$	$d\times d_{kv}$	$n\times d_{kv}$	$dd_{kv}$
$W_V$	$d\times d_{kv}$	$n\times d_{kv}$	$dd_{kv}$
$W_O$	$d_q\times d$	$n\times d$	$d_qd$

Q 表达当前位置想找什么，K 表达各位置可供匹配的特征，Value 提供最终聚合的内容。它们都由可学习投影得到；这些含义是帮助理解的直觉，不是额外的数学约束。

工程实现常把 QKV 合成一个矩阵。这改变 kernel 与布局，不改变三个矩阵的参数总和。

3 Head 维度与 GQA

将 Q reshape 为 $[B,h_q,S,d_h]$，K、Value reshape 为 $[B,h_{kv},S,d_h]$。忽略具体转置布局，这些维度说明了每个轴的含义。

类型	KV head 数	一个 KV head 服务的 Q head 数
MHA	$h_{kv}=h_q$	1
GQA	$1<h_{kv}<h_q$	通常 $h_q/h_{kv}$
MQA	$h_{kv}=1$	$h_q$

GQA 的核心是组内 Q heads 共享 K 和 Value，但每个 Q head 仍然有自己的 Query 和注意力分布。GQA 原论文

若 $d=4096,h_q=32,d_h=128,h_{kv}=8$，则 $d_q=4096,d_{kv}=1024$。此时 Q 投影是 $4096\times4096$，K/V 投影各为 $4096\times1024$。不能把三者都写成 $d\times d$。

常见架构满足 $d=h_qd_h$，但这不是定义上的必然关系。如果题目显式给出 head dim，应优先使用给出的值；若不满足等式，要用 $d_q$ 的通用公式。标准分组实现一般还要求 $h_q$ 能被 $h_{kv}$ 整除。

4 Attention 的两次核心矩阵乘法

对一个 Query head，设 query 数为 $S_q$，可读取的 key 数为 $S_k$：

$$ A=\operatorname{softmax}\left(\frac{QK^\top}{\sqrt{d_h}}+M\right),\qquad O=AV_{\rm val} $$

形状依次为：

$$ [S_q,d_h][d_h,S_k]\to[S_q,S_k] $$

$$ [S_q,S_k][S_k,d_h]\to[S_q,d_h] $$

$M$ 是 mask。因果 mask 禁止看到未来 token；Softmax 沿 key 轴归一化，让一行权重的和为 1。$1/\sqrt{d_h}$ 用于控制点积分数的尺度。

每个 Q head 都需要这两次乘法，即使它与其他 Q head 共享 K/V。因此 GQA 降低 KV 存储与 K/V 投影成本，没有把注意力主项 FLOPs 按 $h_{kv}/h_q$ 缩小。

标准 Softmax、mask、reshape、残差加法都没有可训练参数，却都有运行成本。参数量不能覆盖全部计算。

5 FFN 与 SwiGLU

普通两矩阵 FFN 为：

$$ \operatorname{FFN}(X)=\phi(XW_1)W_2 $$

$W_1\in\mathbb{R}^{d\times m}, W_2\in\mathbb{R}^{m\times d}$，无 bias 时为 $2dm$ 个参数。FFN 对每个 token 独立执行，Attention 则在 token 之间交换信息。

SwiGLU 使用三组矩阵：

$$ \operatorname{SwiGLU}(X)=\left[\operatorname{SiLU}(XW_g)\odot(XW_u)\right]W_d $$

$W_g,W_u$ 均为 $d\times m$，$W_d$ 为 $m\times d$，所以有 $3dm$ 个参数。$\odot$ 是逐元素乘法；SiLU 为 $x\operatorname{sigmoid}(x)$。

如果普通 FFN 取 $m=4d$，参数是 $8d^2$；为了让 SwiGLU 参数规模接近它，可以取 $m\approx8d/3$，再按硬件友好的倍数取整。看到 FFN 就写 $8d^2$ 会在不同结构上算错。

6 Norm 与位置编码

在本文假设下，RMSNorm 为每个通道学习一个缩放值，所以一个 Norm 是 $d$ 个参数。传统带 scale 和 bias 的 LayerNorm 是 $2d$ 个参数。Norm 计算还包括归约、平方、开方等操作。

常规 RoPE 对 Q/K 进行位置相关旋转，不引入可训练位置 embedding 表；它仍有计算和缓存开销。若架构采用长度 $S_{\max}$ 的可学习绝对位置表，要另加 $S_{\max}d$ 参数。不要把“没有可训练参数”理解为“完全不占内存”。

7 Prefill 与 Decode 的区别

Prefill 一次处理已知 prompt 的所有位置。因果 mask 保证每个位置只利用合法前缀，但输入都是已知的，因此同一层中这些位置可以批量计算。

普通自回归 Decode 每次处理一个新输入 token，使用它预测后一个 token。由于后一个 token 要等采样结果出来，单条序列不能像 Prefill 一样把未来未知 token 全部并行计算。

对长度 $S$ 的 prompt，Prefill 的最后位置已经能预测第 1 个输出 token。随后输入这个输出 token，通过第 1 次 Decode 得到第 2 个输出 token。生成 $N$ 个输出 token 通常需要一次 Prefill 加 $N-1$ 次 Decode；最后一个输出若不再用于预测下一个 token，就不必写入 KV。这个边界关系在计算端到端延迟时很重要。

自测：给定 $B=2,S=128,d=1024,h_q=16,h_{kv}=4,d_h=64$，自行写出 Q/K/V 与单个 Q head 的分数矩阵形状。答案是 Q 为 $[2,16,128,64]$，K/V 为 $[2,4,128,64]$，单个序列的单个 Q head 分数矩阵为 $[128,128]$。

下一章：Dense 与 MoE 参数量。

附录 张量形状参考