架构与成本：从教学模型走向现代实现

先分类：改模型函数，还是重排同一计算

前十章已经形成一台可以训练、恢复、评价和生成的完整解码模型。本章把常见扩展放回它的计算图，而不是罗列缩写。目标是判断一项修改究竟改变了表示、参数共享、激活函数、存储，还是运算调度。先修为矩阵乘法、范数、注意力、缓存，以及基础正弦余弦恒等式。涉及三角旋转时会写出所用矩阵和代数关系；不把硬件优化的实验结果当成数学定理。

本章的扩展是对照研究，默认教学项目仍保持学习式绝对位置、普通多头、LayerNorm和GELU。把RoPE或GQA直接塞进旧检查点，通常会改变模型给出的分数，需要相应模型定义与训练或转换协议；它们不是无损换皮。相反，分块计算同一softmax注意力可以在实数算术下保持函数不变，但浮点误差、反向传播和性能仍各有验收。这一区分能防止把“更现代”误解成“随时可替换且必然更好”。

把参数、激活、缓存和运算分开记账

沿用 $B$ 为批量、$T$ 为当前长度、$d$ 为宽度、$f$ 为前馈宽度、$N$ 为层数、$V$ 为词表大小、$C$ 为最大位置数。主项目参数总数已在第五章得到 $P=2Vd+Cd+N(4d^2+2df+f+5d)+2d$。这个数与本次输入长度 $T$ 无关；但注意力概率张量的形状为 $[B,h,T,T]$，会随长度平方增长。一个是持久权重，一个是运行期间的中间数据，不能相加成一个没有单位的“模型大小”。

训练时还要保存梯度和优化器状态。若权重、梯度和AdamW的两个矩都以float32保存，只计这四类张量，每个参数需要四加四加八，共十六字节。它还没有计入激活、暂存、分配器和数据批次。混合精度可能额外保存主权重，框架也可能用不同状态精度；因此十六字节是所述实现假设下的账本，不能无条件套给每个训练系统，更不能把推理只读权重内存与训练内存直接比较。

为估计矩阵乘运算，约定一次乘法和一次加法分别计一个FLOP。每块Q/K/V/O四个投影合计 $8BTd^2$；前馈两个投影合计 $4BTdf$；完整稠密分数与值汇总合计 $4BT^2d$。这个估计明确不含softmax、GELU、归一化、偏置和访存。主项目先做完整矩阵乘再遮罩，所以遮掉未来不意味着已经省掉上三角的乘法。因果专用内核可以改变常数，不能由一个mask张量自动推得。

举一个量级对照：宽度和前馈宽度固定时，长度加倍，投影项约加倍，注意力矩阵乘项约四倍；若只把批量加倍，两项都约加倍。实际时间却可能受显存带宽、内核启动、并行度和缓存影响，不必按这些倍率变化。账本用于提出瓶颈假设，计时用于检验；二者不一致时应检查假设，而不是把公式直接写成设备上的毫秒数。

RoPE：把相对位置写进查询与键的点积

先看二维的一对坐标。定义旋转矩阵 $R(\alpha)$，把查询在位置 $m$ 旋转 $m\omega$，键在位置 $n$ 旋转 $n\omega$；$\omega$ 是为这一坐标对固定的角频率。旋转矩阵满足转置等于逆矩阵，因此两个旋转后的点积，只依赖位置差 $n-m$ 所对应的相对旋转。这里“相对”来自明确的代数消去，而不是给位置贴了一个名字。

把这个乘法展开，可以直接检查符号。$R(a)^{\mathsf T}R(b)$ 的左上项为 $\cos a\cos b+\sin a\sin b=\cos(b-a)$，右上项为 $-\cos a\sin b+\sin a\cos b=-\sin(b-a)$；另外两项分别为正的 $\sin(b-a)$ 和 $\cos(b-a)$。于是整矩阵正是 $R(b-a)$。再把 $a=m\omega,b=n\omega$ 代入，得到上式。换用别的旋转方向约定时，位置差的符号也可能改变，不能脱离矩阵定义记一个正负号。

每头宽度为偶数时，可以把坐标两两配对，第 $i$ 对使用频率 $\omega_i=10000^{-2i/d_h}$，其中 $i=0,\ldots,d_h/2-1$。这个频率表是模型设计约定，不是从注意力公式唯一推出来的定律。通常旋转Q和K，V不作同样旋转。与主项目直接相加的位置向量相比，它改变了进入点积的表示方式；即使两版参数数接近，也不是相同函数。

若查询和键都位于同一个位置，它们受到同一旋转，点积应与未旋转时相同，这是第一条检查；把两个位置同时加一个相同整数，位置差不变，点积也应保持，这是第二条独立检查。较长位置上仍能计算正弦余弦，只说明公式有定义，不证明训练长度之外的预测质量保持。外推还涉及学习到的频率使用、数值精度和数据分布，不能把“能算出来”写成“无限上下文”。

import math
def rotate(pair, angle):
    x,y = pair
    return [math.cos(angle)*x-math.sin(angle)*y, math.sin(angle)*x+math.cos(angle)*y]
def dot(a,b): return sum(x*y for x,y in zip(a,b))
q,k,omega = [1.0,2.0],[3.0,-1.0],.17
m,n = 2,7
observed = dot(rotate(q,m*omega),rotate(k,n*omega))
relative = dot(q,rotate(k,(n-m)*omega))
shifted = dot(rotate(q,(m+11)*omega),rotate(k,(n+11)*omega))
assert abs(observed-relative)<1e-12 and abs(observed-shifted)<1e-12
assert abs(dot(rotate(q,.8),rotate(k,.8))-dot(q,k))<1e-12
print({'relative_dot':observed,'translation_difference':abs(observed-shifted)})

GQA：哪些头共享K和V

普通多头有 $h_q$ 个查询头，且 $d=h_qd_h$，同时有同样多的K/V头；分组查询注意力改用较少的 $h_{kv}$ 个K/V头，并令每组查询共享一个K/V头。设 $h_q$ 能被 $h_{kv}$ 整除，每组查询数为 $g=h_q/h_{kv}$，第 $i$ 个查询头使用第 $\lfloor i/g\rfloor$ 个K/V头。查询仍可不同，头之间并非所有计算都被合并。

例如八个查询头、两个K/V头，每组四个查询，共享映射为零到三用第零组，四到七用第一组；缓存是八个K/V头时的四分之一。若在实现里先把两组K/V真正复制成八组再保存，张量接口可能“看起来匹配”，存储节省却已经丢掉。逻辑广播、临时展开和持久缓存的实际分配需要分别查看。

投影参数要另算。Q仍从宽度 $d$ 映射到 $h_qd_h=d$，O也保持 $d$ 到 $d$，这两项共有 $2d^2$ 个权重；K与V各从 $d$ 映射到 $h_{kv}d_h$，两项合计 $2dh_{kv}d_h$。所以无偏置的注意力投影总数为 $2d^2+2dh_{kv}d_h$。取 $d=128,h_q=8,h_{kv}=2,d_h=16$，Q/O共三万二千七百六十八，K/V共八千一百九十二，总四万零九百六十；原普通多头是六万五千五百三十六，参数比例为八分之五，不能混用缓存的四分之一比例。

共享K/V也改变了参数化：原来每个头可独立投影出的键值，现在被组内约束。即使把旧权重求平均作为初值，也没有保证在所有输入上保持原输出。主论文研究了从已有检查点继续训练的转换方法；本课只推导共享映射和资源账本，不伪称已经完成那种训练。若扩展项目，应先写K/V形状和映射测试，再单独做质量与预算对照。

RMSNorm与门控前馈：少一个统计量，多一条乘法通路

RMSNorm把均方根作为尺度：$r=\sqrt{d^{-1}\sum_i x_i^2+\varepsilon}$，输出通常为 $\gamma_i x_i/r$。它没有减去均值，也通常没有主项目LayerNorm中的偏移参数。取 $x=(1,3)$，忽略很小稳定项时，均方根为 $\sqrt5$，标准化约为 $(0.447,1.342)$；LayerNorm则约为 $(-1,1)$。具体数值说明它们不是同一函数。

给整个向量加十后，LayerNorm的中心化坐标保持，RMSNorm的分母和各分子都会变化，一般不保持原输出。对正比例缩放，在忽略或相应处理稳定项时，二者都具有某种尺度不变性；固定的稳定常数使这只是有条件的近似。减少均值计算可能有工程收益，但不能由少做一个归约直接宣称任意设备上必然更快，速度仍需测量。

门控前馈引入两条宽化投影再逐坐标相乘。以无偏置的SwiGLU为例，$\operatorname{SiLU}(u)=u/(1+e^{-u})$，$F(x)=[\operatorname{SiLU}(xW_g)\odot(xW_u)]W_d$；前两矩阵形状为 $[d,f]$，最后为 $[f,d]$，$\odot$ 表示同位置相乘。于是门控分支会根据输入改变另一条通路的幅度，而不是单纯把GELU换一个函数名。

这里的三分之八只是在同权重预算下解出的宽度比例。实际实现需要整数宽度，并可能为硬件对齐取整；加偏置、改变输出维度或共享矩阵时，账本都要重算。比较GELU前馈与门控前馈，应同时报告参数和运算，不能把额外增加的一半权重当作免费机制改进。主项目没有暗中启用这些替换，避免理论、代码和检查点分别使用三种不同架构。

在线softmax：不保存整行概率也能得到相同汇总

对一个查询，设合法分数为 $s_j$、值向量为 $v_j$。最终输出是 $\sum_j e^{s_j}v_j/\sum_j e^{s_j}$。为了避免溢出，可以减去同一个最大值。若按块读取分数，不预先知道全行最大值，就维护已读最大值 $m$、缩放后的分母 $a$、缩放后的未归一化向量分子 $o$。注意 $o$ 不是已经除过分母的注意力输出。

读取一块新分数后，先令 $m^\prime$ 为旧最大值与新块最大值的较大者。旧分子、分母原本按 $e^{-m}$ 缩放，现在要改按 $e^{-m^\prime}$ 缩放，所以都乘 $\alpha=e^{m-m^\prime}$；新块每项直接乘 $e^{s_j-m^\prime}$ 加入。这样所有项重新拥有同一个尺度，处理完所有块才计算 $o/a$。这一步没有丢掉小概率项，数学目标仍是原来的稠密注意力。

用循环不变量把理由闭合。假定已读集合为 $I$，保存状态满足 $a=\sum_{i\in I}e^{s_i-m}$、$o=\sum_{i\in I}e^{s_i-m}v_i$。乘 $\alpha$ 后，每个旧项都变成 $e^{s_i-m^\prime}$；新块中的项也使用相同的 $m^\prime$，相加后恰好覆盖 $I\cup\mathcal B$，所以不变量保留。首个合法块直接由有限和初始化；有限个互不重叠块处理完后覆盖全部合法键，再除以分母即回到定义的注意力。

这还给出两个可反驳实现的检查：如果某块被重复读取，部分键会被算两次，权重体系已经改变；如果遮罩在块最大值计算之后才随意处理，全遮罩块可能引入未定义值。块顺序可以调整，但覆盖集合、每项出现次数和合法性不能改变。浮点加法不满足精确结合律，所以数学不变量仍要与数值误差记录配合，不能仅因推导成立就跳过实机前后对照。

import math
scores, values = [1000.,1001.,999.,1002.],[2.,-1.,3.,4.]
maximum = max(scores)
weights = [math.exp(s-maximum) for s in scores]
reference = sum(w*v for w,v in zip(weights,values))/sum(weights)
m,den,num = -math.inf,0.0,0.0
for begin,end in [(0,1),(1,3),(3,4)]:
    new_m = max(m,max(scores[begin:end]))
    rescale = 0.0 if m == -math.inf else math.exp(m-new_m)
    local = [math.exp(s-new_m) for s in scores[begin:end]]
    den = rescale*den+sum(local)
    num = rescale*num+sum(w*v for w,v in zip(local,values[begin:end]))
    m = new_m
observed = num/den
assert abs(observed-reference)<1e-12
print({'dense':reference,'blocked':observed})

FlashAttention利用分块与在线归一化，减少高带宽显存与片上存储之间的读写，并避免显式保存完整注意力矩阵。这不等于把每对查询键的计算全部删掉，也不意味着注意力变成线性算术复杂度。上面的Python验证只展示归一化机制，没有实现GPU内核或测量加速。若在项目中改用融合内核，应另做前向、反向、mask、不同长度和精度的对照，速度与数值误差一起报告。

扩展实验的边界与交付

本章交付四张分开的表：参数、训练状态、KV缓存和矩阵乘估算；再运行旋转点积与分块softmax两个独立实验。选一项改变模型函数的扩展，写清新参数和训练方案；选一项保持目标函数的调度优化，写清数值容差与计时方法。主项目的默认训练结果继续使用原架构，不能把此处的机制演示当成已经训练并验证的新版语言模型。第十二章将把所有已完成部分收拢成可复跑的最终交付。