""" ============================================================================= 第 4 讲 · 代码 1：线性注意力 —— 一个"很傻但很重要"的结合律 ============================================================================= 对应讲义：lectures/lecture04_attention_alternatives_moe.md 第 2、3 节 ----------------------------------------------------------------------------- 【这个文件要解决什么问题？】 ----------------------------------------------------------------------------- 标准注意力 softmax(Q K^T) V 的代价与序列长度 n 的平方成正比： - 训练时要算 n × n 的分数矩阵； - 推理时 KV 缓存随 n 线性增长，每生成一个 token 都要读一遍。 如果去掉 softmax（或者换成可以拆开的形式），矩阵乘法的**结合律**就能帮大忙： (Q K^T) V = Q (K^T V) n×n 的中间矩阵 d×d 的中间矩阵 计算量从 n²d 降到 nd²。更妙的是，因果（causal）版本可以写成一个 **RNN**： S_t = S_{t-1} + k_t v_t^T （状态：一个 d_k × d_v 的矩阵，大小固定！） y_t = q_t^T S_t 这个文件： 1. 验证结合律，数一数 FLOPs； 2. 验证因果线性注意力的"并行形式"和"递推形式"完全等价（训练用并行、推理用递推）； 3. 实现实践中使用的"分块"算法，并测速：随序列变长，softmax 注意力 vs 线性注意力； 4. 加上"遗忘"：常数衰减（RetNet）和随输入变化的门控衰减（Mamba-2 的思路）。 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python 01_linear_attention.py （几秒钟） """ import math import time import torch import torch.nn.functional as F from torch.utils.flop_counter import FlopCounterMode def print_header(title): """打印一个带分隔线的标题。""" print() print("=" * 70) print(f" {title}") print("=" * 70) def count_flops(fn): """统计执行 fn() 的矩阵乘法 FLOPs（见第 2 讲 04 号文件）。""" with FlopCounterMode(display=False) as fc: fn() return fc.get_total_flops() # ============================================================================= # 第 1 部分：结合律 # ============================================================================= def demo_associativity(): """演示 1：(Q K^T) V = Q (K^T V)，但计算量天差地别。""" print_header("演示 1：结合律 —— 先算 K^T V") torch.manual_seed(0) n, d = 4096, 64 q, k, v = torch.randn(n, d), torch.randn(n, d), torch.randn(n, d) left = (q @ k.T) @ v right = q @ (k.T @ v) rel = ((left - right).norm() / left.norm()).item() f_left = count_flops(lambda: (q @ k.T) @ v) f_right = count_flops(lambda: q @ (k.T @ v)) print(f" n = {n}，d = {d}") print(f" (Q K^T) V ：中间矩阵 {n}×{n}，FLOPs {f_left:,}") print(f" Q (K^T V) ：中间矩阵 {d}×{d}，FLOPs {f_right:,}（少 {f_left / f_right:.0f} 倍 = n/d）") print(f" 两者结果的相对差异：{rel:.1e}（只是浮点误差）") print(""" 【边界】结合律只在"中间没有 softmax"时成立：softmax 作用在 n×n 矩阵的每一行上， 是非线性的，不能拆开。线性注意力的代价就是放弃 softmax（或用可拆分的"核函数"近似它）， 这会损失一部分表达能力 —— 后面几节的改进都是在想办法把能力补回来。""") # ============================================================================= # 第 2 部分：因果线性注意力的两种形式 # ============================================================================= def causal_linear_attention_parallel(q, k, v): """ 因果线性注意力的"并行形式"：Y = ((Q K^T) ⊙ M) V，M 是下三角的因果掩码（0/1）。 【形状】q, k : (n, d_k)；v : (n, d_v) → (n, d_v) 【复杂度】O(n² d)，和普通注意力一样；但可以一次性并行算完，适合训练。 """ n = q.shape[0] mask = torch.tril(torch.ones(n, n)) return ((q @ k.T) * mask) @ v def causal_linear_attention_recurrent(q, k, v): """ 因果线性注意力的"递推形式"（RNN）： S_t = S_{t-1} + k_t v_t^T y_t = q_t^T S_t 【直觉】S 是一个 d_k × d_v 的"记忆矩阵"：每来一个 token，就把"键 → 值"的关联 k v^T 累加进去； 读取时用 query 去"查询"这个记忆。 【复杂度】每一步 O(d_k d_v)，总共 O(n d²)；**状态大小固定**，与已经处理了多少 token 无关。 【边界】这是逐个 token 的 Python 循环，只为说明原理，真实实现会用分块算法（见下面）或专门的核函数。 """ n, d_k = k.shape d_v = v.shape[1] S = torch.zeros(d_k, d_v) ys = [] for t in range(n): S = S + torch.outer(k[t], v[t]) ys.append(q[t] @ S) return torch.stack(ys) def causal_linear_attention_chunked(q, k, v, chunk=64): """ 分块算法（实际训练中使用的方式）：块内用并行形式，块与块之间用递推传递状态。 【思路】把序列切成长度为 C 的块。对第 c 块： 块内部分：((Q_c K_c^T) ⊙ M) V_c —— 并行，O(C² d) 块间部分：Q_c S_{c-1} —— 之前所有块的信息都压缩在状态 S 里 更新状态：S_c = S_{c-1} + K_c^T V_c 总复杂度 O(n C d + n d²)，对 n 是**线性**的；而且块内的计算是大矩阵乘法，能充分利用 GPU。 这就是 CS336 课上说的"对偶性"：训练时用（分块的）并行形式，推理时用递推形式。 【边界】n 不必是 chunk 的整数倍，最后一块可以短一些。 """ n, d_k = k.shape S = torch.zeros(d_k, v.shape[1]) outs = [] for s in range(0, n, chunk): qc, kc, vc = q[s:s + chunk], k[s:s + chunk], v[s:s + chunk] intra = causal_linear_attention_parallel(qc, kc, vc) inter = qc @ S outs.append(intra + inter) S = S + kc.T @ vc return torch.cat(outs) def softmax_causal_attention(q, k, v): """标准的因果 softmax 注意力（参照组）。""" return F.scaled_dot_product_attention(q[None], k[None], v[None], is_causal=True)[0] def demo_two_forms(): """演示 2：并行形式 = 递推形式 = 分块形式。""" print_header("演示 2：因果线性注意力 —— 并行、递推、分块三种算法结果相同") torch.manual_seed(0) n, d = 300, 16 q, k, v = torch.randn(n, d), torch.randn(n, d), torch.randn(n, d) y_par = causal_linear_attention_parallel(q, k, v) y_rec = causal_linear_attention_recurrent(q, k, v) y_chk = causal_linear_attention_chunked(q, k, v, chunk=64) print(f" 并行 vs 递推：最大相对差异 {((y_par - y_rec).abs().max() / y_par.abs().max()).item():.1e}") print(f" 并行 vs 分块：最大相对差异 {((y_par - y_chk).abs().max() / y_par.abs().max()).item():.1e}") print("\n 推理时每一步要保存的东西（d_k = d_v = 128、BF16、单个头）：") for t in [1_000, 32_000, 1_000_000]: kv_cache = 2 * t * 128 * 2 state = 128 * 128 * 2 print(f" 已处理 {t:>9,} 个 token：softmax 注意力的 KV 缓存 {kv_cache / 1e6:>8.1f} MB；" f"线性注意力的状态 {state / 1e3:.0f} KB（不变）") def time_fn(fn, repeats=3): """简单计时：热身一次后取平均（秒）。""" fn() t0 = time.perf_counter() for _ in range(repeats): fn() return (time.perf_counter() - t0) / repeats def demo_scaling(): """演示 3：序列变长时的运行时间。""" print_header("演示 3：序列越长，差距越大") torch.manual_seed(0) d = 64 print(f" {'n':>7}{'softmax 注意力':>16}{'线性注意力（分块）':>20}") base = None for n in [1024, 2048, 4096, 8192, 16384]: q, k, v = torch.randn(n, d), torch.randn(n, d), torch.randn(n, d) t_soft = time_fn(lambda: softmax_causal_attention(q, k, v)) t_lin = time_fn(lambda: causal_linear_attention_chunked(q, k, v, chunk=256)) print(f" {n:>7}{t_soft * 1e3:>13.1f} ms{t_lin * 1e3:>17.1f} ms") print(""" n 每翻一倍：softmax 注意力的时间大约变成 4 倍（n²），线性注意力大约变成 2 倍（n）。 短序列时 softmax 注意力（高度优化的库函数）可能反而更快；序列越长，线性的优势越明显。 （具体数字取决于你的电脑。）""") # ============================================================================= # 第 4 部分：加上遗忘 —— 衰减与门控 # ============================================================================= def gated_linear_attention_recurrent(q, k, v, gamma): """ 带衰减的线性注意力（递推形式）： S_t = γ_t S_{t-1} + k_t v_t^T y_t = q_t^T S_t - γ_t 是常数 γ（0 < γ < 1）：RetNet； - γ_t 随输入变化（由当前 token 算出来）：Mamba-2 等"门控"模型的核心思路。 【直觉】γ 越小忘得越快。纯线性注意力（γ = 1）什么都不忘，记忆会被越来越多的信息"塞满"而变得模糊。 【参数】gamma : (n,) 张量，每个位置的衰减系数，取值在 (0, 1] """ n, d_k = k.shape S = torch.zeros(d_k, v.shape[1]) ys = [] for t in range(n): S = gamma[t] * S + torch.outer(k[t], v[t]) ys.append(q[t] @ S) return torch.stack(ys) def gated_linear_attention_parallel(q, k, v, gamma): """ 带衰减的线性注意力（并行形式）：Y = ((Q K^T) ⊙ D) V， D[t, s] = γ_{s+1} × γ_{s+2} × … × γ_t （s ≤ t），D[t, s] = 0（s > t） 【推导】展开递推：S_t = Σ_{s≤t} (Π_{r=s+1}^{t} γ_r) k_s v_s^T，所以 y_t = Σ_s D[t,s] (q_t·k_s) v_s。 【实现】用对数的累加和：log D[t,s] = cumsum(log γ)[t] - cumsum(log γ)[s]，避免连乘下溢。 """ n = q.shape[0] cum = torch.cumsum(torch.log(gamma), dim=0) log_d = cum[:, None] - cum[None, :] causal = torch.tril(torch.ones(n, n, dtype=torch.bool)) D = torch.where(causal, torch.exp(log_d), torch.zeros_like(log_d)) return ((q @ k.T) * D) @ v def demo_gating(): """演示 4：衰减让模型"忘记"遥远的过去。""" print_header("演示 4：加上遗忘 —— RetNet 的常数衰减与 Mamba-2 式的门控") torch.manual_seed(0) n, d = 200, 16 q, k, v = torch.randn(n, d), torch.randn(n, d), torch.randn(n, d) # 门控：γ_t 由当前 token 决定（这里用 k_t 的一个随机线性函数经过 sigmoid 模拟） w = torch.randn(d) gamma = torch.sigmoid(k @ w + 3.0) # 大多接近 1，偶尔明显小于 1 y_rec = gated_linear_attention_recurrent(q, k, v, gamma) y_par = gated_linear_attention_parallel(q, k, v, gamma) print(f" 门控线性注意力：并行 vs 递推，最大相对差异 {((y_rec - y_par).abs().max() / y_rec.abs().max()).item():.1e}") print("\n 一个 token 对 t 步之后的输出还有多大影响（常数衰减 γ）：") for g in [1.0, 0.99, 0.9]: row = " ".join(f"{g ** t:6.3f}" for t in [0, 10, 100, 1000]) print(f" γ = {g:<5}：t = 0/10/100/1000 时权重 {row}") print(""" - γ = 1（纯线性注意力）：永远不忘，但所有历史挤在同一个固定大小的状态里，信息互相干扰。 - 常数 γ < 1（RetNet）：越久远的信息权重越小，但"忘多少"对所有 token 都一样。 - 门控 γ_t（Mamba-2、Gated DeltaNet 等）：由内容决定"现在要不要忘"—— 遇到新段落可以大幅遗忘，遇到重要信息可以长期保留。CS336 课上的总结：门控是好东西（gating is good）。 并行形式依然存在（只是掩码 M 变成了衰减矩阵 D），所以训练仍然可以高效地并行。""") def main(): demo_associativity() demo_two_forms() demo_scaling() demo_gating() print("\n全部演示运行完毕。") if __name__ == "__main__": main()