动画图解 LLM 系统

Linear attention：把历史写进状态

不逐条翻 KV，能不能维护一个随时可查的“记事本”？

→ 逐步推进 · ← 回看 · 拖动时间轴观察数据流

动画图解 LLM 系统

先看普通 attention 在读什么

当前 Q 给每个 K 打分，softmax 归一化后，再把 V 加权相加。

一行 attention：y = Σ softmax(q·kᵢ / √d)ᵢ × vᵢ。每个历史位置都有一个权重。

新 Q 到来后，所有匹配分数都可能改变，所以通常要重新读取历史 K、V。

线性 attention 换一种相似度形式，让历史先合并成固定大小的状态，再由 Q 查询。

动画图解 LLM 系统

关键改写：让求和先发生

用 φ(q)ᵀφ(k) 作为相似度，并用权重总和归一化。

选非负特征映射 φ，例如 ELU(x) + 1。这样相似度非负，可以组成归一化加权平均。

分子：Σ [φ(q)ᵀφ(kᵢ)]vᵢᵀ = φ(q)ᵀ [Σ φ(kᵢ)vᵢᵀ]。把 Q 从求和里提出来。

分母同样提出来：Σ φ(q)ᵀφ(kᵢ) = φ(q)ᵀ [Σ φ(kᵢ)]。历史只需维护 S 和 z。

这是核化相似度下的等价重排。一般有限维 φ 不等于原始 softmax attention。

工程细节

先确认我们在重排哪一种 attention

记 Q̃ = φ(Q)、K̃ = φ(K)。每行 feature 维度是 r；V 每行有 dv 维。M 是只允许过去和当前的因果 mask。

A = (Q̃K̃ᵀ) ⊙ M，A ∈ ℝᵀ×ᵀ
Y = (AV) / rowsum(A)

逐行展开：分子 Σᵢ≤ₜ (q̃ₜᵀk̃ᵢ)vᵢᵀ，把 q̃ₜ 提到求和外，就得到 q̃ₜᵀSₜ。分母同理得到 q̃ₜᵀzₜ。

这里的除法按行广播。重排前后是同一个核化 attention；有限维 φ 通常并不精确等于 exp(qᵀk / √d)。

工程细节

value 有两维时，状态真的是矩阵

k̃₁ = [1, 0]，v₁ = [2, 1]；k̃₂ = [0, 1]，v₂ = [4, 3]；q̃ = [2, 1]。

S = k̃₁v₁ᵀ + k̃₂v₂ᵀ = [[2, 1], [4, 3]]
z = [1, 1]
q̃ᵀS = [8, 5]；q̃ᵀz = 3
y = [8/3, 5/3]

逐项算也是一样：两个相似度为 2、1，归一化权重为 2/3、1/3。输出是 2/3[2,1] + 1/3[4,3]。

动画图解 LLM 系统

每来一个 token，更新两个累加器

S 保存 key 与 value 的关联，z 保存用于归一化的 key 总量。

状态从 S₀ = 0、z₀ = 0 开始。假设 φ(k) 有 2 维，v 为 1 维，方便把每个数画出来。

第一个 token：φ(k₁) = [1, 0]，v₁ = 2。外积写入 S = [2, 0]，z = [1, 0]。

第二个 token：φ(k₂) = [0, 1]，v₂ = 4。累加后 S = [2, 4]，z = [1, 1]。

第三个 token：φ(k₃) = [1, 1]，v₃ = 3。累加后 S = [5, 7]，z = [2, 2]。

所有数字为手算示例；先更新到位置 t，再读出，就包含当前 token。

动画图解 LLM 系统

查询时，分子和分母一起读

继续用刚才的状态：S = [5, 7]，z = [2, 2]；当前 φ(q) = [2, 1]。

Q 不需要逐条拿出旧 token；直接对固定大小的 S、z 做两次乘法。

分子 = 2 × 5 + 1 × 7 = 17。分母 = 2 × 2 + 1 × 2 = 6。

输出 y = 17 / 6 ≈ 2.83。分母把“总贡献”变成“加权平均”，避免仅因历史变长而放大输出。

实际实现会处理分母接近 0 的数值稳定性，例如加 ε。

工程细节

没有 z，输出会随重复次数增长

让完全相同的 key / value 重复 m 次：每次相似度 a > 0，value 为 v。

分子 = m a v；分母 = m a
y = (m a v) / (m a) = v

如果只保留分子，结果变成 m a v。序列长一倍，输出也长一倍。基础归一化 linear attention 用 z 修正这个尺度。

一些现代线性模型采用其他归一化、门控或输出 norm，没有这个 z；它们是另一个明确定义的模型，不应混用公式。

工程细节

ε 是数值保护，也会改变很小分母的结果

y = numerator / (denominator + ε)

若分子 = 2×10⁻⁶、分母 = 10⁻⁶，理想比值为 2。ε = 10⁻⁶ 时结果变成 1；ε = 10⁻⁸ 时约 1.9802。

非负 φ 防止正负权重相消，但所有特征接近 0 时，分母仍可能很小。监测分母范围、状态精度与 NaN，比盲目增大 ε 更有效。

epsilon 与 max(denominator, epsilon) 是不同函数。实现、训练和参考计算必须使用同一种定义。

动画图解 LLM 系统

公式里的每个维度

St = St−1 + φ(kt) vtᵀ
zt = zt−1 + φ(kt)
ytᵀ = φ(qt)ᵀ St / [φ(qt)ᵀ zt]

状态的大小

φ(k) 有 dφ 维，v 有 dv 维。S 是 dφ × dv 的矩阵；z 是 dφ 维向量。

为什么叫“线性”

固定这些维度后，每个 token 更新与查询的工作量不随历史长度增长，总工作量随 token 数线性增长。

工程细节

固定大小的状态也可能不小

B 个请求、L 层、H 个状态 heads；每个 head 保存 S ∈ ℝʳ×ᵈᵛ 与 z ∈ ℝʳ。

Mstate = B × L × H × (r dv + r) × b

单请求、32 层、16 heads、r = dv = 128、FP32 状态：512 × (16384 + 128) × 4 = 33,816,576 bytes = 32.25 MiB。

这个数与历史长度无关。若另有卷积状态、门参数缓存或 full-attention 层，必须另外相加。

工程细节

短序列时，线性状态未必更省

按同一 head 比较：BF16 KV 的 K/V 各 128 维；FP32 线性状态 r = dv = 128。

KV bytes = 2 × T × 128 × 2 = 512T
state bytes = (128×128 + 128) × 4 = 66048
两者相等：T = 66048 / 512 = 129

T = 32 时，KV 只有 16 KiB，线性状态约 64.5 KiB；T = 8192 时，KV 是 4 MiB，固定状态的优势才明显。

这只是存储交点，不是速度交点。特征映射、投影、状态读写与 kernel 形状都会影响实际性能。

动画图解 LLM 系统

加一个门：让旧记忆逐渐淡出

把旧状态乘 α，再写入新信息；0 ≤ α ≤ 1 控制保留多少。

纯累加会把越来越多内容挤进同一状态。旧信息与新信息可能互相干扰。

先把旧 S 缩为 αS，再加入新的 φ(k)vᵀ。α 较小，遗忘更快；α 接近 1，保留更久。

在这个归一化示例里，z 也用相同门控更新。旧 token 的权重随之后各步门值的乘积衰减。

动画取 α = 0.5。现代门控线性模型还会采用向量门、不同归一化与更复杂的状态更新。

逐步算一遍

遗忘门怎样改变历史的权重？

一维特征都为 1，values 依次为 [2, 4, 8]；每步 α = 0.5，S 和 z 同时衰减。

t1：S = 2，z = 1，y = 2t2：S = 0.5×2 + 4 = 5，z = 1.5t3：S = 0.5×5 + 8 = 10.5，z = 1.750

所有数值为手算教学示例。

工程细节

门控递推为什么能做并行 scan？

将一次更新写作仿射变换 f(S) = aS + B，其中 B = k̃vᵀ。两个相邻更新可以合并成一个。

f₂(f₁(S)) = a₂(a₁S + B₁) + B₂
= (a₂a₁)S + (a₂B₁ + B₂)
(a₂, B₂) ∘ (a₁, B₁) = (a₂a₁, a₂B₁+B₂)

函数复合具有结合性，所以可以树形合并并求每个前缀。顺序不能交换：后来的门会衰减前面的写入。

z 使用同样的组合，只把 B 换成 k̃。并行 scan 降低依赖深度，但中间状态流量与矩阵维度仍决定效率。

动画图解 LLM 系统

Delta rule：先读旧答案，再写误差

把状态看成一张关联表。遇到同一个 key，先问“我现在记得什么？”

这里换成不带上页分母的 fast-weight 记忆。key k = [1, 0]，状态对它的旧预测是 2，新目标 value 是 5。

误差 e = v − Sᵀk = 5 − 2 = 3。只把需要修正的部分写回，而不是再完整叠上一份 5。

用 β = 0.5：S ← S + βkeᵀ。该方向的记忆从 2 变成 3.5；另一个正交 key 对应的 4 保持不变。

示例 key 为单位向量；一般 key 会相互相关，更新也会影响相关方向。DeltaNet 是另一类状态更新规则。

逐步算一遍

重复写同一个 key：累加和纠错差在哪？

一维 key = 1，目标 value = 5，学习率 β = 0.5，初始记忆 S = 2。每次先算误差 e = 5 − S。

第一次：e = 3；S ← 2 + 0.5×3第二次：e = 1.5；S ← 3.5 + 0.5×1.5第三次：e = 0.75；S ← 4.25 + 0.5×0.750

所有数值为手算教学示例。

工程细节

key 的范数决定了纠错力度

S′ = S + βk(v − kᵀS)
kᵀS′ = kᵀS + β||k||²(v − kᵀS)
新误差 = (1 − β||k||²) × 旧误差

若 ||k|| = 1，β = 0.5，误差减半。若 ||k|| = 2，仍取 β = 1，误差乘 −3：反向过冲且放大。

对一个固定 key 重复更新，收敛需 0 < β||k||² < 2。归一化 key 能让 β 的含义更稳定。

这里分析单 key 的误差递推；多个相关 key 的更新还会互相干扰，不能据此直接保证整个模型稳定。

动画图解 LLM 系统

训练时，状态更新也能按块并行

逐 token 递推适合 decode；prefill 和训练更适合把工作组织成大块。

最朴素的执行方式依次处理 t₁、t₂、t₃……，很难喂满 GPU。

对可结合的状态更新，可以用并行 scan 或 chunkwise 公式。每个 chunk 内做矩阵乘法，chunk 之间传递状态。

Delta rule 需要更专门的块算法；目标相同：保留递推语义，同时让 GPU 做更多规整的大块运算。

这是执行组织方式；具体块公式依赖累加、门控或 delta 更新规则。

工程细节

一个 chunk 的输出由两部分组成

先看不带门的累加形式。chunk 开始前已有 Sprev、zprev；chunk 内有 Q̃c、K̃c、Vc，M 是块内因果 mask。

分子 = Q̃c Sprev + [(Q̃cK̃cᵀ) ⊙ M] Vc
分母 = Q̃c zprev + rowsum[(Q̃cK̃cᵀ) ⊙ M]
Snext = Sprev + K̃cᵀVc
znext = zprev + Σrows K̃c

第一项从跨块状态读取更早历史；第二项处理块内 token。最后一次矩阵乘法把整块写入下一块的状态。

这给出了独立于逐 token 循环的参考计算路径。门控版本要加入门乘积；delta 版本不能直接套这条累加公式。

工程细节

最小递推里，每一行都有张量意义

S = zeros((feature_dim, value_dim), dtype=float64)
z = zeros(feature_dim, dtype=float64)
for q, k, v in zip(phi(Q), phi(K), V):
    S += outer(k, v)    # [r, dv]，包含当前 token
    z += k              # [r]
    denominator = q @ z
    y = (q @ S) / (denominator + epsilon)

先用 FP64、小尺寸确认公式；再替换 dtype 与 kernel，观察数值误差如何增长。只比较同一种特征映射与 epsilon。

NumPy 循环是教学参考，不能用它与优化后的 softmax GPU kernel 比吞吐。

动画图解 LLM 系统

它省掉了什么，又把难题放在哪？

Full softmax attention 基础 linear attention
Decode 的历史状态 O(n(dk + dv)) O(dφdv + dφ)
每步读历史的计算 O(n(dk + dv)) O(dφdv)
历史的形态 保留每个位置，逐项匹配 写进有限容量的状态
关键挑战 长历史的容量与带宽 精确检索、干扰与遗忘

混合架构可以让部分层维护线性状态，部分层保留 full attention。Kimi Linear 与 Qwen3-Next 可以接着看。

复杂度按一个 head、固定特征映射成本计算。FlashAttention 优化 softmax attention 的 IO，并未把它改成线性状态模型。

工程细节

长历史下，状态的数值范围会增长

非负特征累加让 z 持续变大；S 的正负项还可能相互抵消。低精度加法会逐渐丢失较小的新写入。

可复现的检查

固定 Q/K/V，用 FP64 递推作参考；扫描长度、输入尺度和 chunk 大小，记录最大误差与分母最小值。

工程上的处理

使用合适的累加精度；采用模型定义的归一化或遗忘机制。若重缩放 S / z，要同时调整后续写入尺度。

忽略 ε 时，S 和 z 同除常数保持当前比值；若保留 ε，它也要同比缩放。下一个 token 若仍原尺度写入，未来输出仍会改变。

工程细节

速度平了，质量却可能随历史变难

假想实验：decode attention 在 2K / 8K / 32K 都约 0.3 ms；但远距离精确检索准确率从 98% 降到 89%、61%。

现象优先调查对照实验
延迟基本不随 T 变固定状态确实生效测状态字节、kernel 读取量
近距离好、远距离差遗忘与状态干扰变证据位置，固定总长度
重复 key 容易覆盖写入规则与关联容量比较累加、门控、delta 结构
只有低精度长序列错数值累加误差FP64 / FP32 同输入参考

这些数字只用于教学诊断。换线性结构通常需要训练适配，不能把 softmax 层的权重原样替换后期待等价。

练习 1

不看逐项权重，算出二维输出

k̃₁=[1,0]，v₁=[2,1]；k̃₂=[0,1]，v₂=[4,3]；k̃₃=[1,1]，v₃=[3,2]。求 S、z，再用 q̃=[2,1] 读出 y。
展开参考解答

前两项 S=[[2,1],[4,3]]，第三个外积为 [[3,2],[3,2]]；所以 S=[[5,3],[7,5]]，z=[2,2]。

分子 [2,1]S=[17,11]，分母 [2,1]·[2,2]=6。y=[17/6,11/6]≈[2.8333,1.8333]。

先写出假设和单位，再代入数字。

练习 2

门控后的旧 token 还剩多少权重？

三个 token 的一维特征均为 1，values=[2,4,8]。α₂=0.5、α₃=0.25；S/z 从 0 开始。求最终输出。
展开参考解答

三个 token 的未归一化历史系数为 [α₂α₃, α₃, 1]=[0.125,0.25,1]。

分子=0.125×2+0.25×4+8=9.25；分母=1.375；输出=9.25/1.375=74/11≈6.7273。

先写出假设和单位，再代入数字。

练习 3

单位 key 的结论能直接套到 k=2 吗？

一维 S=1，k=2，目标 v=5，β=0.5。先算旧预测、误差、新 S、新预测。为什么这次没有把预测改成 5？
展开参考解答

旧预测 kS=2，误差 e=3；新 S=1+0.5×2×3=4，新预测 kS=8，发生过冲。

有效修正率 βk²=2，误差从 3 变为 −3。若 β=1/k²=0.25，则 S′=2.5、预测=5；key 范数必须进入分析。

先写出假设和单位，再代入数字。

动手实验

三种计算顺序，应该给出同一个输出

用 NumPy 在 CPU 上跑一个可检查的参考实现。小尺寸便于把中间张量打印出来。

python3 -m pip install numpy
python3 lab.py
python3 lab.py --seed 7
recurrent vs dense kernel: PASS
recurrent vs chunkwise: PASS
gated scan coefficients: PASS
delta unit-key convergence: PASS

改变 chunk 大小，验证边界不改变语义；把 φ 换成 ELU+1，再用错误的 softmax 参考比较，观察为什么两类 attention 并不相等。

下载实验代码 lab.py

这是正确性与成本的教学实验；耗时不代表 GPU kernel 性能。

参考资料

参考与延伸阅读

继续看：KV cache · Sparse attention · Kimi Linear