2.4.1 · NSA: 压缩、选择与滑窗的训练原生稀疏 attention

Native Sparse Attention(NSA)从预训练开始就采用稀疏attention结构, 无法当作推理插件直接追加到稠密checkpoint上. 它把历史信息分到压缩、选择和滑窗三条并行分支。 压缩分支负责「全局摘要」, 选择分支恢复「远距细节」, 滑窗分支维持「局部依赖」. 三个分支分别softmax, 再由可学习gate合并输出. 论文的目标同时包含算法与硬件. Prefill和训练通常受算力限制, 需要减少实际QK与AV运算并支持反向传播; Decode通常受KV读取带宽限制, 需要让同一GQA组的query heads共享连续KV块. NSA的block设计让候选单位、GQA共享和kernel tile从训练开始对齐, 与token级top-k的物理执行不同.

1. 三分支计算图

1.1. 统一输出与独立归一化

对位置 $t$ 的query $q_t$, NSA构造三组KV, 分别记为 compression、selection与window. 输出为

其中gate由输入特征经过MLP与sigmoid得到. 每个分支拥有自己的K/V投影并单独计算softmax. 这点决定它不能等价改写为把三组positions拼起来做一次attention: 三个分母、三套表示和gate都会改变输出. 独立分支也是训练约束. 局部模式更容易学习, 若所有候选共用一套表示与softmax, recent tokens可能形成捷径, 压缩和远距选择得不到足够梯度. 专门的window分支吸收局部模式, 让另外两条分支分别学习全局摘要与远距细节. 代价是需要维护三套投影状态, 逻辑KV条目数不能只按最终候选求和.

1.2. 压缩分支

压缩分支用长度 $l$、步长 $d$ 的滑动块组织历史. 每一块的keys经过带块内位置编码的可学习MLP $\varphi$ 映射成一个compressed key, values同理:

论文通常取 $d<l$, 让相邻压缩块重叠, 减少边界处信息被切断. 实验配置为 $l=32,d=16$, 因而长历史大约每16个token产生一个压缩entry. 压缩attention读取全部entries, 保留对全历史的覆盖, 复杂度仍随长度线性增长, 但系数显著低于读取全部原token. 压缩不是平均池化. $\varphi$可学习, 且有块内位置信息, 目标是为后续query保存可检索语义. 一个entry无法无损恢复32个token的所有细节, 所以该分支承担粗粒度全局感知, 而不是替代精细检索.

选择与滑窗分支。 选择分支把原始KV划为长度 $l'$ 的连续块, 保留importance最高的 $n$ 个块. 论文实验取 $l'=64,n=16$, 即最多读取1024个原始token; 其中包含固定激活的首块与两个局部块. 连续block使加载和矩阵乘法能使用规则tile, 代价是关键token会带入同块邻居. window分支读取最近 $w$ 个token, 实验中 $w=512$. 它不依赖top-k, 为短程语法、邻近复制和最新状态提供确定路径. 由于window输出独立归一化并门控, 即使selection也选中局部块, 两条分支仍有不同语义, 不能为去重而随意删掉其中一份.

2. 压缩分数怎样驱动选择

2.1. 复用已经计算的注意力分数

NSA没有为selection增加一套独立轻量indexer. 压缩分支本来就计算

这些分数同时提供粗粒度saliency. 当compression block与selection block使用相同划分时, 可以直接用压缩分数排序. 划分不同时, 论文根据两类块在序列上的空间覆盖关系聚合多个压缩分数, 得到selection block importance. 这样选择开销被并入一条必须存在的全局分支. 分数复用建立了明确耦合: 压缩表示既要贡献全局输出, 又要为原始token块排序. 如果压缩器保留了适合摘要的信息却无法区分细节块, selection recall会下降; 如果只为排序优化, compression输出可能变差. 原生训练让两类目标通过最终语言建模损失和分支梯度共同适配.

2.2. 不同block划分的映射

设压缩块长度为 $l$, 步长为 $d$, selection块长度为 $l'$, 且 $d$整除二者. 一个selection块会与多个重叠compression blocks相交. 论文把相应 $p_t^{cmp}$ 项求和得到 $p_t^{slc}[j]$. 这一步复用已有粗分数并按空间关系归约, 无需重新对原始keys做QK. 以 $l=32,d=16,l'=64$ 为例, 一个64-token选择块跨越多个起点间隔16的压缩窗口. 边界附近的压缩窗口还可能同时覆盖相邻选择块, 因而importance不是简单每四项分一组. 实现必须严格复现公式的索引偏移、序列开头处理和因果边界; 用直觉分桶会在块边界产生系统偏差. top-$n$之后, selected blocks按整数位置gather原始selection K/V. 候选个数大致固定, 但历史较短或固定首块、局部块与动态top块重叠时, 实际unique blocks会减少. kernel需要规定去重与顺序, 训练反向也要把梯度scatter回正确原位置.

GQA组内共享选择。 GQA中多个query heads共享一组KV heads. 若每个query head各选不同块, 实际KV读取是所有选择的并集, 算术稀疏却不一定减少带宽. NSA把同一GQA组内各query head的block importance相加, 形成共享排序, 然后所有heads读取同一批KV blocks. 这种聚合牺牲部分head专属性. 某个head独有的重要块可能被其他heads的总分压过, 但共享候选避免重复HBM读取, 还能把整组queries同时放入片上存储. 质量与物理效率的折中发生在选择规则本身, 不能等到kernel阶段再弥补. 验收时应同时测per-head oracle recall与group-union读取量. 独立head top-k可作为更宽松候选预算下的诊断参照, 共享top-k是实际路径. 若只比较FLOPs, 会遗漏GQA下候选并集导致的读放大; 若只比较字节, 又看不到少数专用head的质量损失.

3. 训练、Prefill与Decode一致性

3.1. 原生训练意味着什么

NSA论文用27B总参数、3B激活参数的MoE骨干进行实验, 30层, 4个GQA组与64个attention heads. NSA与Full Attention基线都在8k长度数据上预训练约270B tokens, 再用32k数据继续训练与监督微调. 论文报告训练损失稳定, 并在通用、长上下文与推理评测上比较. 这一组实验展示了三分支稀疏结构在该MoE骨干上的端到端训练结果. 三分支拥有独立KV、压缩MLP、gate与共享block选择, 参数和优化路径从训练开始就不同. 对已有checkpoint做迁移需要新增参数初始化、蒸馏和继续训练, 论文结果不能直接作为迁移保证. 原生还指稀疏存在于前向与反向. 若前向只计算selected blocks, 反向也应只对对应QK/AV路径计算梯度, 才能节省训练计算. top-n索引本身是离散的, 未入选blocks不会通过selection分支获得梯度, 但compression分支仍对全局压缩entries传播梯度, window分支则稳定局部学习.

3.2. Prefill为何能够稀疏

许多后处理方法需要先用稠密Prefill生成attention map或索引, 只在Decode稀疏. NSA的importance来自压缩分支自身, 而压缩分支本就是稀疏架构的一部分. Prefill可对每个query生成压缩分数、归约block importance并执行selected attention, 不需要先跑一遍完整原token attention. Prefill中query数多, attention通常compute-bound. 因此重点是减少实际矩阵乘法并保持Tensor Core利用率. 连续selection blocks允许按块加载KV; 压缩和window可沿用规则attention kernel. 但top-n、索引映射和不规则query-to-block关系仍有开销, 理论非零元素比例不会原样变成加速比. 因果训练还要求每个query只选可见blocks. selection block若跨过当前位置, 必须掩掉未来token; 压缩窗口若包含未来位置也不能供当前query使用. 短序列单元测试应覆盖块首、块中和块尾, 对比显式朴素实现的输出与梯度.

Decode为何强调共享读取。 Decode每步query很少, 主要成本是读取历史KV. NSA只读取全部compressed KV、固定window KV与少量selected original blocks. 当上下文足够长时, 读取量近似 $t/d+w+nl'$, 相比全历史 $t$ 降低; 压缩分支仍线性增长, 因而不是常数内存访问. 同一GQA组的heads共享selected blocks, kernel一次加载连续KV块到片上存储, 再与整组queries计算. 论文kernel把query/output循环放到grid调度, 内层顺序读取选中的连续块. 这种group-centric布局同时提高算术强度并避免每个head重复读取KV. Decode状态包含三套K/V以及压缩器尾部. 新token到来时, window追加并淘汰最旧项; selection原始KV需保留供未来块选择; compression每隔步长产生或更新entry. 因此NSA减少每步读取, 不必然减少所有KV存储. 计算cache字节时应分别列出三套投影、重叠压缩entry和量化元数据.

4. Kernel与成本核算

4.1. 为什么按block而不是token

随机token gather会形成大量短事务, 难以填满矩阵乘单元. 64-token选择块把每个index转成一段连续地址, kernel块 $B_k$只需满足 $B_k\mid l'$, 就能分块加载到SRAM. 多算同块低价值token换取更少事务和规则矩阵形状, 是NSA明确的硬件折中. block size太大时, 候选预算浪费在邻居上; 太小时, index数量、排序与事务数增加. 最佳值依赖head维度、dtype、page布局与设备. 论文在A100与特定配置上的结果不能自动外推到其他加速器, 复现应重新profile而不是只保持相同理论稀疏率.

4.2. 三条分支的实际开销

设历史长度为 $t$, 压缩步长为 $d$, window为 $w$, 选择 $n$ 个长度 $l'$ 的块. 单query参与attention的逻辑KV数量近似

代入论文实验值, 长上下文时为 $t/16+512+1024$. 64k上下文约为5632个逻辑位置, 显著少于65536. 但三条分支分别softmax并有独立投影, 还存在压缩MLP、importance归约、top-n、索引和gate, 所以不能把 $65536/5632$ 直接称为端到端加速倍数. 训练要分别核算forward与backward, Prefill记录QK、AV和选择工作区, Decode记录HBM实际读取字节. selection blocks若分散在许多pages, 物理读取可能大于有效载荷; compression blocks重叠会增加构造计算; window与selected local blocks重复并不会消除双分支计算.

可复现的性能实验。 性能实验应固定模型shape、dtype、batch、GQA组数和序列长度, 比较Full Attention与NSA完整三分支. 只测selection kernel不能代表总结构, 只测端到端模型又会被MoE通信与MLP掩盖. 两类结果都要提供, 并报告warm-up、编译配置和峰值显存. 序列长度从短到长扫描能够找到交叉点. 短序列上top-n与三次launch可能比稠密kernel更贵; 长序列上压缩比例与固定选择预算才占优势. Decode还应扫描batch size, 因为更多并发query会提高算术强度, 改变带宽瓶颈. 训练验证不能省略梯度. 对小shape使用双精度朴素实现, 固定top-n indices后比较Q、K、V、压缩MLP和gate梯度; 再单独测试选择边界的离散变化. 最终用长序列观察反向workspace和是否出现数值溢出.

5. 质量边界与相邻路线

5.1. 三类失败与相邻路线

压缩失败发生在重要细节没有进入compressed representation, 会同时削弱全局输出与block排序. 选择失败发生在importance映射或group聚合漏掉关键原始块. window失败发生在局部跨度超过 $w$ 或局部分支形成过强捷径. 三者可能在最终logit上叠加, 需要分支级消融定位. 可以记录gate分布、compressed attention质量、selected blocks相对于稠密oracle的覆盖率以及window距离直方图. gate低不必然代表分支无用, 因为少数关键query可能高度依赖它; 平均recall高也可能漏掉唯一证据. 测试集应包含多证据检索、长代码引用、长推理与局部复制. 固定激活首块和局部块为极端模式提供保底, 但也占用 $n$ 的预算. 当上下文包含大量分散证据时, 16个blocks可能不足; 增大 $n$ 提升召回, 同时线性增加selected attention与读取. 这是真实质量—成本旋钮, 不是仅调kernel就能消除的限制.

与DSA、CSA/HCA的区别。 DSA使用独立轻量indexer对历史token打分, 再从主MLA KV中读取token级top-k. 它的代理目标、索引K和主attention可分别训练与量化. NSA不另建selector, 而是复用compression attention分数, 候选单位是连续原始token blocks, 并保留独立window分支. CSA同样先压缩token轴再稀疏选择. V4公开报告中的CSA在压缩entries上执行主attention, 与NSA选择分支回读原始连续blocks不同. HCA对更强压缩后的全部entries做attention, 没有动态top-k漏选, 误差主要来自压缩. NSA则让compression全覆盖与selected原始细节并行存在. 比较这些结构时, 应直接核对压缩对象、实际读取的K/V和分支归一化方式.

与HySparse及跨层共享的区别。 HySparse的oracle来自周期性full attention layer产生的真实block scores, 后续多个sparse layers跨层复用indices与该full层KV. NSA在每个层内由compression分支产生importance, 没有依赖上一full层的跨层oracle. 前者用full刷新换真实saliency, 后者用原生压缩分支保持所有层可训练稀疏. HySparse保留独立SWA分支, 与NSA专门的window分支有表面相似性; 但其global sparse分支共享上游full KV, NSA三分支拥有独立K/V并在本层训练. HySparse2进一步加入self/cross decoder的KV Bridging与Prefill早退, NSA论文没有这一跨decoder机制. 选型应先问是否拥有原生预训练机会. 若有, NSA可以让训练、Prefill和Decode使用同一稀疏结构; 若只有既有checkpoint, 后处理或继续训练路线更现实. 再问硬件是否能从连续block与GQA共享中获益, 以及业务是否接受固定block预算的召回边界.

工程验收清单。 工程验收先核对compression窗口数量、重叠关系、块内位置编码和因果边界. importance从compression映射到selection时, 序列开头、未满末块与固定激活块都要单独检查. GQA路径还要核对组内聚合、top-n稳定排序、去重与连续KV地址. 分支输出测试要确认三次softmax、三套K/V和gate顺序. 将某一gate置零应只移除对应分支; 将selected budget覆盖全部历史时, selection分支应退化为相应投影上的dense attention; window扩展到全历史时也应与该分支的dense参考对齐. 同一张验收表应列出质量、训练吞吐、Prefill TTFT、Decode时延、逻辑FLOPs、实际读字节和峰值显存. NSA让粗粒度全局覆盖、细粒度连续块和局部窗口共同进入原生训练, 并按同一组block约束执行. 缺少任一分支或只在Decode启用选择, 都不能直接沿用论文结论.

一个短序列的完整手算。 设历史长度为128, compression block长度32、stride 16, 则完整窗口起点为1、17、33、49、65、81、97, 共7个compressed entries. selection block长度64时只有前后两个原始块. 对位于序列末尾的query, compression分支对7个entries做attention; selection分支根据7个压缩权重向两个64-token块归约; window若为32则读取最近32个token. 假设归约后前块分数0.35、后块0.65, 动态预算只允许一个block, selection选择后64个token. 这个结果与直接对128个原始token取top-64不同: 后块内即使只有一个关键位置, 其余63个邻居仍会参与selection softmax. 同时window覆盖的最近32个token也位于后块, 但它们通过独立window K/V和独立softmax再次贡献. gate决定两种表示的组合, 不能对positions去重后只算一次.

再考虑因果位置70. 起点49的32-token compression窗口跨到token 80, 其中未来部分对该query不可见; 起点65窗口也未完整可见. 实现可只使用已经完整的窗口, 或对尾部做带mask的部分压缩, 但必须与训练定义一致. selection的第二个64-token块也只能访问65到70. 这说明「按块读取」仍需token级因果mask, 不能因block被选中就暴露完整物理块.

论文实际配置还固定首块与两个局部块. 当历史短到这些固定块覆盖全部序列时, 动态top-n会出现大量重叠, unique blocks少于16. 性能测试若始终按16乘64估算读取会高估短序列成本; kernel若不去重则会重复计算并改变独立selection分支内部的softmax质量. 两种行为都应由参考实现明确规定.

梯度与选择稳定性。 在固定indices条件下, selection attention对选中K/V与query完全可微, compression和window分支也有普通attention梯度. 不连续点发生在两个block importance交换排序的位置. 常规自动微分把top-n结果视为当前前向的离散路由, 未选block不会通过selection路径获得梯度; 它仍可能通过compression路径影响损失和未来importance. 因此训练诊断不只看总loss. 应记录候选集合在相邻checkpoint之间的Jaccard相似度、top-n边界分差、各分支gate和压缩分数熵. 如果边界分差长期接近零, 很小数值扰动就会频繁换块, 造成梯度噪声与不可复现排序. 稳定排序规则可确保相同输入得到相同indices, 却不能替代模型学习更清晰的importance.

混合精度会放大边界问题. compression score在低精度下归约到selection blocks, 多个GQA heads再求和, 累积顺序会影响接近的分数. 验收应以高精度归约作为参考, 分别比较importance误差、indices一致率与最终输出误差. 若只比较logits, 不容易判断是候选变化还是attention数值变化. 反向kernel还要处理同一原始block被batch内多个queries选中. gradient scatter可能产生原子写冲突或额外归约workspace. 训练速度报告应包含这一真实反向路径, 不能用固定随机indices且无梯度的前向microbenchmark代替.

Cache布局与生命周期。 三分支独立K/V意味着每个新token通常产生selection K/V与window K/V, compression K/V则按stride生成. window只需保留最近 $w$ 项; selection历史要保留, 因为未来任一query都可能选中远距block; compressed历史也要保留用于全局attention与importance. 忽略量化时, 长序列主要状态近似为

这里 $T$ 是历史长度, $n$ 仍表示选择块预算, $b_e$ 为每个元素的字节数; K与V维度相同才可使用系数2. 三个分支的head维度可能不同, 不能直接合并成一个统一系数. 公式还未包含压缩尾部原始状态、索引、位置编码与scale. prefix cache复用需要三条分支来自同一checkpoint和同一blocking配置. 改变 $l,d,l',w$ 后, 旧compressed entries与selection block编号都不再兼容. cache key应编码这些结构参数以及K/V投影版本. 请求从共享前缀继续时, 未满compression窗口的原始尾部也必须恢复, 否则下一个entry会遗漏前缀token. continuous batching中不同请求的top-n blocks不同. selection kernel可把每个请求与GQA组作为调度单元, 但过多短任务会增加grid开销; 将它们合并又可能引入padding. 工程上应记录每次launch的有效queries、selected block数和加载字节, 用真实分布选择batch策略.

speculative decoding回滚时, window尾部、selection原始KV和compression部分窗口要同步截断. 已完成且完全位于确认前缀内的compressed entries可保留; 包含草稿token的entry必须重建. 这种生命周期规则是实现推论, 论文没有给出生产状态机, 因而需要逐token参考测试确认. 分布式部署还要保持GQA组与KV分片对齐. 如果同组query heads分散到多个tensor parallel rank, 先聚合importance再做全局top-n会引入通信; 若每个rank局部选择, 最终KV读取并集又可能扩大. 更合适的切分通常让共享一个KV head的queries共置, 但还需兼顾投影与输出通信. pipeline parallel不会改变层内三分支语义, 却要求每个stage保存本层完整的压缩尾部和原始selection cache. 这些布局成本应计入端到端结果, 不能只引用单卡kernel速度.

上线前可用故障注入验证状态管理: 改错一个selection block编号、丢弃一个未满compression尾部、交换两个GQA组的indices, 确认校验与canary输出能发现. 静默读错KV往往仍产生有限数值, 比显式崩溃更难定位. 元数据中保留请求、层、分支、位置范围与结构版本, 能把错误缩小到具体数据路径.

5.2. 从张量形状走到一次完整执行

三条分支分别保存什么。 设batch为 $B$, 序列长度为 $N$, query head数为 $H_q$, KV head数为 $H_{kv}$, 每个head维度为 $D$. 输入 $X\in\mathbb{R}^{B\times N\times d_{model}}$ 先经过三套K/V投影. selection分支得到 $K^{slc},V^{slc}\in\mathbb{R}^{B\times H_{kv}\times N\times D}$; window分支的投影shape相同, 推理时只保留最近 $w$ 个位置; compression分支先生成逐token的K/V, 再沿序列轴把长度 $l$、步长 $d$ 的窗口压成 $M$ 个entries:

当 $N=65536,l=32,d=16$ 时, $M=4095$. 这个数字比粗略写成 $N/d=4096$ 少1, 差别来自第一个完整窗口需要先容纳32个token. 流式实现还会保存不足32个token的尾部; 尾部达到窗口长度后生成entry, 随后每进入16个新token再生成一个. 因果训练若允许部分窗口, 计算定义会改变, 训练和推理必须采用同一种规则. GQA的映射可写为 $g(h)=\lfloor hH_{kv}/H_q\rfloor$, 第 $h$ 个query head只读取第 $g(h)$ 个KV head. 若 $H_q=64,H_{kv}=4$, 每16个query heads共享一个KV head. NSA在这16个heads之间归约block importance, 选出共同的block IDs. kernel据此加载一次K/V, 让16个query heads复用; 若每个head各自top-$n$, 最坏情况下会出现 $16n$ 个不同block, GQA原有的带宽收益会被候选并集抵消.

图1. 参考NSA论文v2 Figure 2与第3节展开的三路数据流. 三个query框是同一 $q_t$ 的视觉副本; compression attention产生概率 $p^{cmp}$ 与输出 $o^{cmp}$, 概率经窗口映射和GQA组内求和产生整数块索引, 索引控制gather读取原始selection K/V. 三个分支独立归一化, 再分别乘上当前输入产生的gate并求和. 三个gate框对应左下公式中的三个标量, 不表示独立缓存. 实线为数据流, 虚线为整数索引控制流; 三种缓存的维护范围见图底部.

放大查看NSA数据流

• compression attention覆盖全部压缩entries, 其softmax权重同时参与最终输出和block importance归约.
• top-$n$只决定selection分支读取哪些原始连续块, 不裁剪compression与window分支.
• 反向传播沿三个输出分支回到各自投影; 离散索引在当前step固定, selection梯度只写回被选block.

三条路径的投影独立, 所以cache也不能只保留一份原始K/V再临时切片. selection cache保留全部历史原始块, window cache只需循环保存最近 $w$ 项, compression cache保存 $M$ 个压缩entries和一个尚未封口的原始尾部. 若三个分支head维度都为128、KV head数为4、bf16每元素2字节, 64K序列中selection K/V约占 $2\times4\times65536\times128\times2=128$ MiB; window K/V约1 MiB; compression K/V约8 MiB. 这组数只计算单层裸张量, 未含页表、scale、对齐填充与临时workspace.

6. 路由训练、系统成本与实验判读

6.1. Selector 梯度与长上下文成本

selector的梯度到底走到哪里。 令compression logits为 $a_{t,i}$, 压缩注意力权重为 $p_{t,i}=\operatorname{softmax}(a_t)_i$. 第 $j$ 个selection block的分数由与它相交的compression窗口归约:

其中 $R_{j,i}$ 表示第 $i$ 个压缩窗口对第 $j$ 个原始块的贡献. 若窗口跨越两个selection blocks, 对应的一列可能向两个block贡献分数, 具体系数由论文的空间覆盖映射决定. GQA还会在top-$n$之前对组内query heads的 $s_{t,j}$ 求和. $I_t$ 是整数索引集合, gather出的原始K/V进入selection attention. 固定 $I_t$ 后, selection分支与普通attention一样可微. 选中block的K/V得到QK与AV两条梯度, query也从selection输出得到梯度; 未选block在这条路径上的梯度为0. TopN在排序交换点不连续, 常规反向不会对“某块差一点入选”产生梯度. NSA没有用straight-through estimator把离散索引伪装成连续变量.

compression分支解决了训练信号中断的一部分. $p_{t,i}$直接参与compression输出, 因此语言建模损失会更新compression Q/K/V和压缩MLP; 同一组权重又决定 $s_{t,j}$, 参数变化会在后续step改变候选集合. 这属于通过共享表示产生的间接适配, 不是对top-$n$求导. window分支始终可微, 在训练早期selection尚不稳定时仍能传递局部梯度. gate梯度也需要单独观察. 设三个分支输出为 $u^{cmp},u^{slc},u^{win}$, gate logits为 $z^c$, 且论文使用逐分支sigmoid $g^c=\sigma(z^c)$. 对某个分支有

三个gate不要求和为1, 因而可以同时较高或同时较低. 若某个 $z^c$过早进入sigmoid饱和区, $g^c(1-g^c)$接近0, 对应分支即使偶尔有用也难以恢复. 训练日志应按层、head组和token位置记录gate分布, 还要结合遮蔽分支后的loss变化; 单看平均gate会漏掉少量但关键的远距query.

64K上下文的计算与读取手算。 取论文常用配置 $N=65536,l=32,d=16,l'=64,n=16,w=512$, 暂时忽略固定块与动态块重叠. 对序列末尾的一个query, compression读取4095个entries, selection读取 $16\times64=1024$ 个原始位置, window读取512个位置, 合计5631个逻辑KV位置. 与Full Attention的65536个位置相比, QK和AV主体约为11.64分之一. 这个比例只描述attention主体. 若 $H_q=64,D=128$, 每个逻辑位置在所有query heads上的QK与AV合计约需 $4H_qD$ 次浮点运算, 三分支主体约为

Full Attention对应约 $2.15\times10^9$ FLOPs/query. NSA还要计算压缩MLP、importance归约、组内reduce、top-$n$、三次online softmax与gate. Decode下这些固定工作占比更明显; Prefill含大量queries, block attention主体更容易把排序与launch开销摊薄. 带宽要按KV heads计算. 仍取 $H_{kv}=4,D=128$, bf16 K/V每个逻辑位置需要 $2\times4\times128\times2=2048$ 字节. 若三条分支具有相同head维度, 单步裸载荷约为 $5631\times2048=11.0$ MiB, Full Attention约128 MiB. 实际selection读取取决于页布局: 16个连续64-token blocks若各自跨页, 会额外读取对齐区域和页表; window通常连续, compression entries也规则排列. profiler应同时给出有效字节与HBM事务字节.

短上下文会出现另一种结果. 当 $N=2048$ 时, compression只有127个entries, 固定的selection预算最多1024个token, window又读512个token. 三条分支合计接近1663个位置, 再加排序与三套softmax, 相对Full Attention的优势很小. 如果固定首块、局部块和动态块大量重合, 去重会降低selection实际读取; 若实现保留重复索引, 相同block会在selection softmax中重复出现并改变输出, 已经不只是性能问题.

训练阶段还要乘上反向成本与保存状态. online softmax通常保存每行的log-sum-exp, selection保存block IDs, 压缩器保存反向所需激活. 朴素实现若为每个query保存完整4095维compression概率, workspace会随 $N^2/d$ 增长; fused kernel可重算部分统计量或分块保存, 用更多计算换显存. 所以训练报告需要列出峰值显存和backward时间, 单独的forward稀疏率无法代表可训练性.

与DSA、MoBA和KDA放到同一组坐标。 这些方法都减少长序列attention成本, 但它们削减的对象并不相同. 比较时可固定五个问题: 选择信号从哪里来、候选粒度是什么、训练时是否执行同一路由、局部信息由谁保证、部署时主要保存什么状态.

NSA与DSA都包含hard选择, 训练信号来源不同. DSA的indexer拥有自己的表示和监督目标, 可以单独测索引召回; NSA让compression权重兼任全局attention和路由信号, 少了一次独立历史扫描, 也把压缩质量与选择质量耦合在一起. NSA以连续block换取规则访存, DSA的token级候选纯度更高, kernel需要处理更细碎的地址. MoBA与NSA都选择连续block, MoBA把block routing本身作为主attention的稀疏拓扑, 当前block提供确定的因果局部连接. NSA保留完整的compression和window输出, selected blocks只是三条并行路径之一. 相同的top-$k$ block数无法代表相同成本: NSA还读取compressed entries和window, MoBA则受block代表计算、路由负载与当前block处理影响.

KDA属于线性attention的recurrent路线. 它把历史压入随时间更新的状态, 每个token都参与状态更新, 没有“候选集合”“top-k召回”或“gather哪些KV”这几个步骤. 因而不能用selector recall评价KDA, 也不能把KDA的固定状态大小写成一种hard sparse attention. 对比指标应换成状态大小、每token更新成本、长程信息保留和并行训练方式. 混合模型可以让KDA承担全局状态, 再配局部softmax attention; 这与NSA用compression、selection和window分工的出发点相近, 计算图却完全不同.

消融怎样定位三条路径。 最直接的消融是分别关闭一个分支, 同时保持其余参数、训练token和推理预算不变. 关闭compression不仅去掉一个全局输出, 还会让selection失去importance来源, 因而不能把结果解释成单纯的“全局摘要贡献”. 若要测compression输出本身, 可以保留它产生的indices并把 $g^{cmp}$置0; 若要测路由作用, 则保留compression输出, 用随机块或固定块替换top-$n$. 两组实验回答不同问题.

selection消融可以逐步改变 $n$ 与 $l'$. 固定总selected tokens $nl'$ 时, 小块提高定位精度并增加索引数量, 大块提升连续加载效率并带入更多邻居. 固定block数只会同时改变信息预算和计算量, 难以区分质量变化来自哪一项. GQA共享消融应比较每head top-$n$、组内共享top-$n$和组内候选并集, 同时报per-head oracle recall与实际读取字节. window消融要按任务距离切片. 把 $w$ 从512降到256会影响局部复制、代码缩进、短程语法和最近对话状态; 长检索下降也可能来自证据附近的局部组合被破坏. 将window扩大到全历史可验证该分支投影的dense参考, 却不能代表完整NSA退化成Full Attention, 因为compression与selection仍经独立softmax和gate参与输出.

压缩器的消融应控制 $l/d$ 与entry数量. 把重叠窗口改成不重叠窗口会同时减少压缩构造计算并改变边界信息; 用平均池化替换可学习MLP则直接测试压缩表示能力. 除最终困惑度外, 还应记录compressed attention对远距证据的覆盖、selection block recall、top-$n$边界margin与分支遮蔽后的loss增量. 四类数结合起来, 才能区分“压缩内容丢了”“排序选错了”和“选对后gate没有使用”.

一次可复现的长上下文实验至少给出 $B,N,H_q,H_{kv},D,l,d,l',n,w$, dtype、块页大小和是否去重. 训练结果还需给出预训练token数、长度日程、Full Attention基线是否同数据同算力, 以及forward/backward分别采用何种kernel. 系统结果按训练、Prefill和Decode分开, 再报告端到端模型吞吐; 这样才能看出收益来自少算QK/AV、少读KV, 还是其他模块占比变化.

还可以用合成样本把错误定位到具体路径. 第一组样本把唯一证据放在selection block边界两侧, 逐个移动query位置, 观察importance映射是否随重叠窗口平滑变化. 第二组让同一GQA组的16个heads分别偏好不同block, 比较独立top-$n$、共享top-$n$与候选并集; 这组样本能直接暴露组内求和是否压掉少数head. 第三组把证据放在window之外但compression可见的位置, 再把compression gate置0, 检查模型是否确实依赖远距分支.

数值正确性应分两层验证. 小shape下用显式dense mask构造三条参考attention, 固定indices后逐项比较输出、log-sum-exp以及Q/K/V梯度. 随后放开top-$n$, 比较compression logits、block scores和最终indices. 前一层失败说明attention kernel或scatter有错; 后一层失败通常来自边界映射、GQA归约精度或排序规则. 两类误差混在最终logits中时, 一个错误候选就足以掩盖后面的数值问题.

性能曲线也要保留形状变化. 序列从2K扫到128K, batch从1扫到服务常用并发, 分别记录compression、selection、window、top-$n$和gate耗时. 若总时延在某个长度后不再近似线性, 通常需要检查selection pages是否过于分散、临时索引是否触发额外拷贝, 或并发请求的block数差异是否造成SM尾部空转. 这些数据最终会反过来约束 $l,d,l',n,w$ 的选择: 算法预算与硬件布局必须在同一组配置上验证.

线上监控可以沿用相同分解. 每层记录实际unique blocks、压缩entry数、三分支gate分位数和HBM读取量, 再按请求长度分桶. 候选数没有变化但读取量突然升高, 多半是物理页分散或前缀复用失效; 读取量稳定而长检索下降, 则优先检查importance分布与边界margin. 这样定位时可以直接落到路由、布局或模型质量中的一层.

6.2. 三分支的信息分工与容量边界

压缩分支是一条可寻址的全局记忆。

重叠窗口改变了 token 的统计权重.

当压缩块长度 $l=32$、步长 $d=16$ 时，序列内部token通常出现在两个compressed entries中，靠近首尾的token出现次数较少。重叠能缓解块边界，让跨越16倍数位置的短语至少在某个窗口中完整出现；它也让中部信息通过两条压缩路径影响attention，形成隐含的覆盖权重差异。

若压缩器近似线性，一个内部token对所有entries的总贡献可能约为边界token的两倍。学习型MLP和位置编码可以校正，前提是训练长度和padding分布覆盖这些边界。测试同一证据放在序列开头、中部、尾部与不同block offset，能检查重叠是否产生位置偏置。

重叠还增加构造成本。长度n大约产生 $n/d$ 个entries，每个entry读取l个原始状态，朴素压缩工作约为 $nl/d$。在32/16配置下，输入读取约为2n。若用卷积式滑窗复用中间结果可降低常数，仍要保存块内位置关系。压缩attention省下的QK不能与构造成本混在一起。

一份表示承担输出与路由.

Compressed key先与query打分，得到compression branch的attention质量；同一批分数再聚合成selection block importance。于是它有两个目标：为压缩value提供正确读取权重，同时给原始KV块排序。一个只保留主题的key可能适合全局摘要，却分不清两个包含相似主题、只有细节不同的blocks；一个极度强调词法尖峰的key适合检索，作为全局语义memory又可能不稳定。

可把损失梯度分成直接路径与间接路径。直接路径来自compression output对最终loss的影响，连续可微；间接路径经过top-n indices，硬选择对未选项通常没有普通梯度。Selection分支仍会推动当前入选块中的原始KV学习，却难以直接告诉compressed key“另一个块本应入选”。原生训练的优势在于长期共同适应，不表示离散路由问题自动消失。

训练中扩大候选、使用软教师或周期性稠密oracle，可以给边界外blocks更多信号；这些属于可能的训练策略，不能在没有论文证据时当成NSA默认实现。分析重点是认清双重职责，再用compressed output质量与selection recall分别验收。

压缩容量约束全局粗粒度路径.

Selection只读固定数量原始blocks，未选历史仍可通过compression branch影响输出。压缩memory提供了一条全局粗粒度路径：即使细节块漏选，模型仍可能知道相关主题、文档或粗关系。它与window共同让selection无需承担全部信息。

固定维度entry不可能无损保存32个token。若一块含多个实体、数字和顺序关系，压缩器要决定哪些特征值得未来query使用。重叠让信息有两次编码机会，却不增加每个query最终读取的独立原始细节。多证据任务仍可能需要selection补回。

压缩率由stride而非block长度单独决定。减小d增加entries和全局读取，降低信息间隔；增大l但保持d会让每个entry看到更宽上下文，构造更贵、局部信息竞争更强。$l$ 与 $d$ 是表示范围和采样密度两个旋钮，不能只报告“16倍压缩”。

6.3. Selection分数的空间映射

重叠覆盖形成一张线性映射。 将query t的compressed probabilities写成向量 $p_t\in\mathbb R^{N_c}$，selection block importance可写成

矩阵A由compression窗口与selection blocks的空间重叠关系决定。若一条compression窗口覆盖两个selection blocks，权重可分配给两边；若直接向所有相交块累加，同一概率会被重复计算。具体A必须按论文定义实现。

线性形式便于测试。对很短序列显式构造A，检查每个compression entry的质量流向哪些selection blocks；再与kernel归约比较。序列开头、尾部未满块和causal截断会让A的行列数变化，最容易出现off-by-one。

若A只含0/1，block得分与被多少compression windows覆盖相关。内部selection block通常接收更多完整窗口，边界block较少。归一化每列覆盖数可消除这种先验，却改变论文算法。复现应先严格一致，再通过受控实验判断边界偏置。

Top-n 之前还要处理固定块。 首块与局部blocks固定激活，为attention sink、开头指令和最近上下文提供保底。若它们计入总n，动态槽位减少；若在动态top-n之外追加，总token预算增加。短序列中固定集合大量重叠，去重时机也会改变最终数量。

合理执行顺序要由定义确定：先把固定blocks mask为已选，再从其余blocks取剩余名额；或先全局top-n，再做并集。前者保持总block数，后者可能超过n。两种结果在长序列平均差异不大，边界样本和性能口径会不同。

固定首块还可能形成捷径。模型知道它永远可见，倾向把全局摘要写入BOS附近；这能提高稳定性，也让所谓动态检索的一部分能力来自sink中继。屏蔽首块或打乱其内容可以测量依赖程度，但部署结构若固定保留它，不必追求完全消除。

离散集合的容量冲突。 同一query可能需要多个语义独立blocks。固定n限制可访问地址数，即使ranking完全正确也无法覆盖超过n个分散证据。Block长度64还意味着每条证据带入邻居，真正用于独立地址的容量至多n。 定义达到教师概率质量阈值 $\tau$ 所需最小block数

若 $n_{0.95}$ 随证据数或上下文长度增长，固定16会形成结构上限。平均语言建模中attention可能尖锐，长文归纳与多证据验证更容易触及上限。 动态n可依据importance margin或累计质量调整，训练与kernel会面对变长形状。实际系统可只允许8、16、32几个档位。模型若只用固定16训练，推理时扩大集合会改变softmax分母，质量也未必严格单调，动态预算应进入训练分布。

6.4. GQA共享选择与候选预算

求和分数优化平均覆盖。 同组H个query heads对block j的importance为 $r_{hj}$，NSA用聚合分数 $R_j=\sum_h r_{hj}$ 排序。选定集合S最大化

这恰好优化组内总质量，却没有保证每个head的最低覆盖。一个分布尖锐但只在少数样本活跃的head，平均贡献可能输给多个普通heads共享的blocks。 因此要同时报告平均和最差head覆盖。若最差head承担远程复制或代码符号，少数失败会在特定任务集中出现。增加n、按head子组选择或给稀有head加权都能缓解，读取并集与元数据同步增加。

独立Top-k并集扩大了候选预算. 让每个head独立选n个blocks，再取并集，提供候选规模最多增至Hn的诊断参照。并集大小介于n与Hn之间，实际值反映head需求重合。若并集接近n，共享几乎免费；若接近Hn，固定共享集合必然牺牲专属性。

可画“并集大小—概率质量”曲线：先为每head选择较小预算，再逐步合并。它比单一Jaccard更直接回答多读多少blocks能恢复多少质量。Kernel也能利用组间重复，只加载唯一blocks。

独立集合不一定是部署候选，因为随机读取和排序开销可能很高。它主要作为诊断：共享质量下降来自容量冲突，还是compressed importance本身没有找到教师blocks。

共享读取怎样转化成带宽收益。 同一GQA组共享KV，本就希望一次加载供多个query heads使用。若heads选不同blocks，KV cache物理仍只有一份，但kernel需要加载并集，再为各head应用不同mask；片上复用下降，控制流更复杂。共享S让所有heads使用规则矩阵tile。

理论FLOPs按H个heads乘n blocks计算，不因共享集合直接减少；主要收益来自K/V只从HBM加载一次、地址元数据更少、矩阵形状一致。Prefill中query rows多，算力占比可能更高；Decode中每步query少，读取复用尤其重要。

性能实验应比较三条路径：完全共享S；每head独立S但合并唯一KV；每head独立加载。第一条展示NSA实际设计，第二条分离集合约束与理想kernel，第三条是简单基线。只比较第一和第三会把算法与实现收益混在一起。

三次 Softmax 的信息分工。

6.5. 独立归一化与门控输出

假设远程selection logits整体比local logits低2，但其中含关键证据。统一softmax会让大量local token占据分母，远程概率很小；独立softmax让selection分支内部仍产生归一化输出，再由gate决定整条通路权重。模型可以在分支级保留远程信息。

Compression同理。它的entries是聚合表示，logit尺度未必与原始KV相同。独立归一化避免三种表示强行校准到同一数值范围。代价是每个分支即使整体不相关，也会产生和为1的输出，只能靠gate抑制。

Gate若采用sigmoid，各分支权重互不排斥。三个gate都高时输出范数增大，后续输出投影与归一化可以适应；若用softmax gate，三者总和固定，会形成竞争。NSA的具体gate是计算图的一部分，复现不能替换为看似更自然的混合而仍声称等价。

分支重复不是重复计算同一语义.

Recent token可能同时出现在window、selected block与覆盖它的compressed entry中。三次出现分别使用原始window表示、selection表示和压缩摘要，softmax上下文也不同。重复地址不代表输出可去重。

Window负责当前层局部模式，selection关注少数原始远程块，compression提供全局粗景。模型可能在同一token上抽取不同特征。强行把候选并集只算一次，会删除分支专用K/V和独立分母，得到另一种架构。

系统层可复用底层输入读取或融合投影，前提是不改变数学结果。比如同一hidden state同时生成window与selection K/V时可合并GEMM，输出仍要按分支存储。优化应通过数值对照确认。

Gate 的可解释边界.

Gate大小能显示模型在某个query上偏向哪条通路，但不能直接等同于因果贡献。某分支输出范数很小，即使gate高，实际影响仍低；两个分支输出方向相反还会抵消。更可靠指标是门控后向量范数、屏蔽分支造成的logit变化与任务结果。

按距离和任务统计仍有价值。局部复制时window贡献应突出，远程精确查询selection增强，全局归纳compression更重要。若所有任务都被一条分支垄断，其他分支可能训练不足，或数据不需要其能力。

分支dropout可迫使模型避免单一路径依赖，但会改变训练算法。没有论文支持时，应作为后续实验而非NSA事实。默认结构能否自然分工，需要从公开消融和复现指标判断。