附录 B. 数学 ↔ 代码索引
| 数学概念 | 文件与函数 | 说明 |
|---|---|---|
| 内积 / 矩阵乘向量 | model/tensor.go MatVec | \(y=Wx\),W 为 \([out,in]\) |
| softmax | model/tensor.go Softmax | 减最大值;全 \(-\infty\) 返回全 0 |
| log-softmax | model/sample.go LogSoftmax32 / LogProbAt | 原始分布,不含温度/截断 |
| 交叉熵/困惑度 | 无运行时实现 | 训练侧概念;推理只算 logprobs |
| RMSNorm | model/norm.go RMSNorm | eps 在根号内,\(\gamma\) 直接乘 |
| per-head QK-Norm | model/attn.go headRMSNorm | 逐头归一化,且在 RoPE 之前 |
| SiLU / SwiGLU | model/tensor.go SiLU,model/mlp.go MLP | 门控逐元素乘 |
| 残差 / pre-norm | model/block.go Layer.ForwardWith | 两条残差,均加在归一化前的 h 上 |
| 缩放点积注意力 | model/attn.go Attention | 因果掩码用 \(-\infty\) |
| RoPE | model/rope.go ApplyRoPE | 劈半配对,只作用于 Q/K |
| GQA | model/config.go KVGroupSize;attn.go | Q 头 h 用 KV 头 \(h/G\) |
| KV cache | model/kv.go NaiveKV / CacheKV | 增量 append |
| 分页 KV | model/pagedkv.go BlockManager / PagedKV | 块表 + 引用计数 + COW |
| online softmax | model/pagedattn.go attentionPaged | 跨块合并,stage B |
| 采样 | model/sample.go Sample / SampleLogits | 温度→top-k→softmax→top-p→逆变换 |
| 批处理 / chunked prefill | cmd/server/scheduler.go,model/batchstep.go | token 预算、左 padding、绝对位置 |
结语:一条完整的公式链
把这一页浓缩成一次前向:
\[ \begin{aligned} x_0&=\text{embed}[t],\\ h&=x_0,\\ x&=\frac{h}{\operatorname{RMS}(h)}\gamma_1,\\ q&=R(p)\,W_qx,\quad k=R(p)\,W_kx,\quad v=W_vx,\\ \alpha&=\operatorname{softmax}\left(\frac{qK^\top}{\sqrt d}+M\right),\\ h&\leftarrow h+W_o\left(\alpha V\right),\\ h&\leftarrow h+W_d\Big(\operatorname{SiLU}(W_g\operatorname{RMS}(h)\gamma_2)\odot W_u\operatorname{RMS}(h)\gamma_2\Big),\\ \text{logits}&=W_{\text{lm}}\operatorname{RMS}(h)\gamma_f,\\ p&=\operatorname{softmax}(\text{logits}/T),\\ t'&\sim p\ \text{(top-k / top-p 截断后)}. \end{aligned} \]每一步为什么长这样,都能在上面对应的小节里找到推导。真正容易出错的通常不是公式本身,而是下标、顺序、归一化范围、绝对位置这四类约定;每个小节的“常见坑”都是仓库里实际踩过的。