第 10 章　局部与稀疏注意力：SWA 和 DSA

第 09 章讨论如何用更少的空间保存历史。本章换一个问题：一次 query 是否需要读取全部历史位置？SWA 按距离保留一个窗口，DSA 则由可训练的索引器打分，选择历史位置。这里的教学实现展示选择过程，并不意味着索引器已经学会了有效的检索策略。

读这两类实现时，要同时检查输出和执行过程：输出有没有使用被禁止的位置，程序又有没有真的少做计算？在稠密矩阵上加 mask，可以限制哪些位置参与输出，却不一定减少已经完成的计算。

10.1 SWA 的窗口到底包含谁

零基位置 $t$ 的因果滑动窗口宽度为 $W$，本章约定窗口包含当前 token。先说清这个约定，是为了避免实现时把窗口多算或少算一格。允许的位置满足

$$max(0,t-W+1)\le j\le t.$$

例如 $W=3,t=4$，可见位置是 2、3、4，数量为 3。源码 MultiHeadAttentionWithSWA 先形成差值 ${\mathrm{\Delta }}_{t,j}=t-j$，再以 (diff < 0) | (diff >= W) 标记禁止位置。第一项挡住未来位置，第二项挡住距离过远的历史。少了第二项，就只有因果注意力，没有滑动窗口；若写成 diff > W，又会多放进一个 token。见 固定 SWA 源码。

可切换原始尺寸，滚动查看细节

查看大图

尺寸读取中

适合窗口：完整显示图片，不放大超过原始尺寸。 按 Esc 关闭。

此图未提供文字说明

一个局部层读不到窗口以外的位置，但它读到的隐藏状态可能已经汇入更早的信息。这样叠加多层，信息仍能逐步传播。若每层窗口为 $W$，不考虑边界，经过 $L$ 层后，一个位置可能间接受到最多 $1+L(W-1)$ 个位置影响。这不等于拥有直接读取任意旧 token 的能力：信息必须穿过中间表示。如果需要保留远距离的直接连接，可以混合局部与全局层，让部分层读取整个历史。Longformer 和 Gemma 3 技术报告 展示了不同任务与架构中的局部/全局组合，不能把某一个比率当成通用最优值。

10.2 持久缓存可以只留 W，当前计算却未必只用 W

逐 token 解码时只需保留最后 $W$ 个 KV，持久缓存上界是 $2BWGds$ 每层。这个上界描述的是留给下一步的缓存；分块预填充时，当前这一步需要的数据可能更多。假设旧位置处理到 4，本次同时输入 5、6、7，$W=3$。位置 5 需要 3、4、5，位置 7 需要 5、6、7。当前计算的并集是 3—7，共 $W+T_q-1=5$ 个位置，如果一开始就裁成 5、6、7，前面的 query 就会失去本该能读取的历史。

源码先保存 old_cache 与 old_len，拼出 combined_k/v。注意力所需部分保留最多 W + num_tokens - 1；下一步的持久 cache 只留最后 W。前一个切片服务当前块中所有 query，后一个切片服务后续解码，不能互换。

裁剪还带来一个位置问题。丢弃开头若干项以后，用 k_start_pos_abs = ptr_current_pos - old_len + dropped 找回 key 的绝对起点。query 从 ptr_current_pos 开始。缓存即使只剩三个槽，也仍保存着原序列中的位置关系，不会把槽里的 token 误当成位置 0、1、2。

若有 $L_s$ 个局部层、$L_f$ 个全局层，缓存估算为

$$2BGds[L_{s}min(T,W)+L_{f}T].$$

memory_estimator_swa.py 解析比率、分配整数量的局部/全局层，再求和；不能把 $L_s/L_f$ 的小数比率直接当作实际层数。模型里则由 sliding_window_stride 决定哪些层使用局部注意力，所以还要核对模型实际分配的层数是否与估算一致。绘图脚本只是在不同长度上重复这项估算。训练时，教学实现仍构造 $T\times T$ 注意力分数，因而不具备专用滑窗内核的线性工作量。

10.3 DSA：先索引，再做注意力

固定窗口按位置筛选，窗口外的 token 即使与当前 query 有关，也无法被直接选中。DSA 尝试让选择依赖内容：索引器先为候选历史生成相关性分数，再选择 top-$k$。固定提交的 DSA 源码 使用输入隐藏状态，而完整 DeepSeek 模型将索引器与 MLA 组合。因此，分析这个文件时应按它实际缓存和投影的对象计账，不能直接套用完整模型的内存与投影成本。

设索引头数 $H_I$、索引宽度 $d_I$，当前 query 输入形状为 $[B,T_q,D]$，候选历史为 $[B,T_k,D]$。W_q_index 输出 $[B,T_q,H_I,d_I]$，W_k_index 输出共享的 $[B,T_k,d_I]$，W_weights 输出 $[B,T_q,H_I]$。索引分数为

$$I_{t,j}=\sum _{h=1}^{H_I}\frac{w_{t,h}}{\sqrt{H_I}}\mathrm{ReLU}\left(\frac{q_{t,h}^{I\mathrm{\top }}{k}_j^I}{\sqrt{d_I}}\right).$$

$w$ 是原始线性输出，不经过 softmax，可以为负。ReLU 只把各头的负点积清零；后面还要乘以可能为负的权重，再跨头求和。因此，最终分数未必非负，也不会由此得到“每行只剩 $k$ 项”的稀疏性。真正的离散选择来自 topk(...).indices。DeepSeek-V3.2 官方报告 是完整系统的来源，本文只对已列教学实现作代码级解释。

索引器必须在 top-$k$ 之前 加因果 mask，否则未来位置可能占满名额。选完以后，主注意力还要再次合并因果 mask。这并不多余：短序列开头可见位置少于 $k$，topk 为凑满固定输出宽度可能返回分数为负无穷的未来索引。最终只保留“被选中且不在未来”的位置，才能避免这些补足名额的索引参与 softmax。

主注意力根据被选位置构造 bool mask，把其余 logit 置为负无穷；随后缩放分数，执行 softmax，再对 V 加权。这里还需要检查是否出现整行都被挡住的情况。若所有分数都是负无穷，softmax 的归一化过程中会出现负无穷减负无穷，得到 NaN。把负无穷换成有限最小值，无法修复这个逻辑错误：整行数值相同时，softmax 可能给出均匀分布，反而让不允许的位置参与输出。有效 query 至少应有一个合法 key。

10.4 稀疏选择不等于稀疏执行，也不等于自动学会选择

教学 MultiHeadAttentionWithDSA 首先计算完整的 $Q{K}^{\mathrm{\top }}$，再应用稀疏 mask。它还缓存 cache_x，让索引器重投影历史输入。相较普通 MHA，这里还额外保存了历史隐藏状态。因此，这个文件虽然实现了稀疏选择，却没有据此减少 KV 存储。原理上稀疏 kernel 可只对选中键执行主注意力点积，约为 $O(Tkd)$；但索引器仍需要扫描候选位置、计算分数。要估算完整复杂度，必须把选择和主注意力的成本分别算进去。

即使选点逻辑正确，还有一个问题：索引器从哪里获得学习信号？topk 的索引是离散的，而该实现只把索引用于布尔 mask，主任务 loss 无法沿索引回传到索引器分数。如果没有额外的可微训练目标，仅把随机索引器接进 GPT、继续优化语言模型 loss，就不会沿这条路径学到预期的检索策略。因此，理解部署报告中的收益时，还要看索引器如何训练、稀疏计算如何执行；上游这个文件提供的是可以逐步检查的选点演示。

10.5 配套脚本与验证路径

SWA 的 TransformerBlock 根据层编号选择局部或全局模块，残差、LayerNorm、GELU 前馈没有改变。模型级位置 embedding、reset_kv_cache、缓存生成以及参数 CLI 与第 09 章同一主干。DSA 多出三个索引参数：头数、头宽和 top-$k$；它们影响选择器形状，不改变输出词表大小。两个目录的 README 给动机、命令与上游时间示例；本书未执行那些长序列性能实验。

tests.py 和 test_dsa.py 的核心检查包括完整输入与缓存预填充的一致性、跨 chunk 的窗口边界、选择不会越过因果边界、top-$k$ 足够大时退回密集行为。比较之前，两条路径都要关闭 dropout，并使用相同参数和绝对位置。只有这样，输出差异才反映计算路径的差异；不能要求独立随机初始化的两个模型逐项输出相同数值。

本书完整程序 examples/architectures/architecture_checks.py 中的 masks_and_sparse() 用 7 个 token、宽度 3、分块 [0:2],[2:5],[5:7] 验证 SWA。原稿记录，完整与分块输出误差为 0；修改未来 value 后，过去输出误差为 0。DSA 的稠密 mask 与逐 query 实际 gather 的结果最大误差为 $4.44\times {10}^{-16}$。程序还检查整行负无穷输入 softmax 后是否产生 NaN。原稿将环境、seed 与完整结果保存在 qa/architectures/algebra_checks.json。当前提供的文件包含程序，却缺少这份历史日志，所以上述精确误差只能作为原稿记录保留，不能当作已核验的运行结果。这里也没有速度测试。

10.6 补充练习与答案

题目。 ① 宽度 4、query 位置 9，可见哪些位置？② 分块长度 5 时，为什么当前注意力可能要用 8 个历史/当前位置？③ 如果先 top-$k$ 再删除未来项，为什么结果可能缺少本来合法的重要历史？④ 只返回 top-$k$ 索引的索引器，靠主注意力 loss 是否能获得常规梯度？

答案。 ① 6、7、8、9。② $W+T_q-1=4+5-1=8$，早期 query 和后期 query 的合法窗口并集更宽。③ 未来项已经挤占选择名额，删除它们不会自动把被挤掉的历史补回来。④ 不行，离散索引切断该路径；需单独设计训练目标或其他可微方案。

进一步检查真实上游 SWA、DSA 类的程序是 source_attention_checks.py。原稿报告，完整前向与三块缓存计算的最大误差均为 0；对 DSA 主输出平方和反传后，所有 indexer 参数的 grad 都为 None，与 10.4 所述训练边界一致。对应日志为 qa/architectures/source_attention_checks.json，当前文件中同样缺少这份历史记录。检查程序确实包含这些断言，但仅阅读源码还不能确认历史运行的精确结果。

参考资料与未验证边界

• 固定源码 ch04/06_swa、ch04/09_dsa，模型、估算、绘图和测试均纳入覆盖台账。
• Longformer：https://arxiv.org/abs/2004.05150。
• Gemma 3：https://arxiv.org/abs/2503.19786。
• DeepSeek-V3.2 官方报告：上文链接。
• 未验证：长上下文语言质量、稀疏 CUDA kernel、索引器训练、报告中的部署成本。