4. 线性层、参数量与计算量
4.1 一次矩阵乘向量到底算了什么
线性层 \(y=Wx\)，\(W\in\mathbb{R}^{m\times n}\)，\(x\in\mathbb{R}^n\)。按 0.2 的定义，每个输出分量是一次长度为 \(n\) 的内积：\(n\) 次乘法、\(n-1\) 次加法，约 \(2n\) 次浮点运算（FLOP）。\(m\) 个输出就是约 \(2mn\) 次。
所以“参数量”和“每 token 计算量”可以分开算：参数量是 \(mn\)，计算量是 \(2mn\) 每 token。矩阵乘矩阵 \(XW\)（\(B\) 个 token 一起）则是 \(2Bmn\)，权重只读一次，被 \(B\) 个 token 复用，这是 batching 提速的根源（第 13 节）。
4.2 Qwen3-0.6B 的账
本项目默认模型是 Qwen3-0.6B：\(h=1024\)、\(L=28\) 层、\(H=16\) 个 Q 头、\(H_{kv}=8\) 个 KV 头、\(d=128\)、中间层 \(3072\)、词表 \(V=151936\)。每层参数：
| 张量 | 形状 | 参数 |
|---|---|---|
| q_proj | \(2048\times1024\) | 2,097,152 |
| k_proj | \(1024\times1024\) | 1,048,576 |
| v_proj | \(1024\times1024\) | 1,048,576 |
| o_proj | \(1024\times2048\) | 2,097,152 |
| gate_proj | \(3072\times1024\) | 3,145,728 |
| up_proj | \(3072\times1024\) | 3,145,728 |
| down_proj | \(1024\times3072\) | 3,145,728 |
| 4 个 norm 的 \(\gamma\) | 两个 \(1024\)，两个 \(128\) | 2,304 |
| 每层合计 | 15,730,944 |
28 层：\(15{,}730{,}944\times28=440{,}466{,}432\)。embedding 是 \(151936\times1024=155{,}582{,}464\)，且 tie_word_embeddings=true，lm_head 与它共享，不重复计。总计
这就是“0.6B”的来历。用 float32 存：\(596{,}048{,}896\times4=2{,}384{,}195{,}584\) 字节 \(\approx2.22\) GiB，和 TODO 里“常驻 2.2 GB”一致；bf16 存储则约 1.19 GB（safetensors 文件 1.4 GB 左右，还含少量元数据和对齐）。
4.3 每 token 的 FLOP 估算
按 4.1 的 \(2mn\)：
- Q/K/V/O 投影：\(2(1024\cdot2048+1024\cdot1024+1024\cdot1024+2048\cdot1024)=12{,}582{,}912\)；
- MLP 三个矩阵：\(3\times2\times1024\times3072=18{,}874{,}368\)；
- 每层合计约 \(31.46\) MFLOP；28 层约 \(0.881\) GFLOP；
- lm_head：\(2\times151936\times1024=311{,}164{,}928\approx0.311\) GFLOP。
所以 decode 一个 token 约 \(1.19\) GFLOP（不含注意力随历史长度的部分）。作为对照，注意力在长度 \(T\) 时每层约 \(4HT^2d=8192T^2\) FLOP：\(T=1024\) 时约 8.6 GFLOP/层，远大于线性部分，这就是长上下文 prefill 昂贵的原因。
4.4 embedding 是“one-hot 矩阵乘”
查表 E[id] 可以看成：构造 one-hot 向量 \(e_{id}\)（只有第 id 位是 1），再算 \(E^\top e_{id}\)。矩阵乘的结果正好取出第 id 行。这样想的好处是：反向传播时梯度是“把梯度写到对应行”，而不是真的乘一个巨大的单位矩阵。