Junxian Ma / Writings

TL;DR

• 稀疏程度高不一定更快。 Sparse Attention 少算了多少，只是问题的一半； 选中的数据能否被规则地搬运和复用，以及为了找到这些数据付出了多少代价， 同样决定端到端性能。
• NSA 在读取数据之前制造规则性。 它选择连续的 64-token blocks，并让 16 个 query heads 共享选择结果，使稀疏工作负载可以直接映射为规则 GPU tile。
• DSA 在读取数据之后恢复规则性。 它允许选择离散 token，但让 128 个主 heads 共享 token IDs，再由 FlashMLA 将随机 rows 搬入 SMEM、重组为规则 tile； 代价是 Discovery（Indexer + exact Top-K）仍需扫描全部候选。
• DSA 的瓶颈与地址随机性没有强关联，而在于找出这些 token 的过程。 在 本文测试的 H100 sparse-prefill 配置下，八组随机行序排列与有序访问的延迟 点估计相差不超过 0.47%；但在公开 DSA 链路中，Discovery（Indexer + Top-K） 在 64K 检查点已经比 FlashMLA 更慢，并且随着上下文继续增长而迅速扩大差距。
• 因此需要优化 Discovery，而不只是最后的 sparse kernel。 系统层面可以 融合 Indexer 和 Top-K，避免完整 FP32 scores 的写回与重复读取；算法层面则 需要避免每个 query 都扫描全部历史 KV。前者降低当前实现的常数开销，后者 才可能从根本上改变长上下文下的增长趋势。

1. Opening

Sparse Attention 的公式并不难。难的是：算法选出的 KV 不再构成规则矩阵， GPU 该怎样读取、计算和反向归约？本文用 NSA 与 DSA 回答这个问题。它们都在 减少 query–KV pairs，却选择了两条不同的 co-design 路线。

1. NSA 在 load 之前制造规则性。 连续 64-token blocks 与 16-head 共享 selection，使算法输出可以直接映射为规则 tile。
2. DSA 在 load 之后恢复规则性。 128 个主 heads 共享 token IDs，FlashMLA 将离散 rows 搬入 SMEM 后再计算；代价是候选发现（Indexer + Top-K）仍可能 二次增长。

两者共同指向本文的中心命题：

稀疏率只说明“少算了多少”。真正决定能否加速的，是选中的数据能否被 规则地搬运和复用，以及发现这些数据本身需要多少工作。

后文先统一 tensor 与 tile 的观察方式，再逐步拆解 NSA 和 DSA 的算法数据流、 kernel ownership 与访存路径。两组 H100 实验只在它们真正回答问题的位置出现： 一组检验 DSA sparse kernel 对 selected-row 地址模式有多敏感；另一组检验被 稀疏主核遮住的 Discovery 二次项会怎样随上下文增长。

本文不再重讲 dense Attention 或 FlashAttention，而从 block IDs 与 token IDs 怎样产生、又怎样被 kernel 消费开始。[PAPER]、[CODE]、[DERIVATION] 与 [MEASURED] 分别标记论文陈述、代码事实、本文推导与正式测量。

2. 阅读本文只需了解三个概念

后文会出现很多 shape，但观察方法只有三条。

先约定几个常用词：一个 program 在 Triton 中对应一次独立的 kernel 实例；在 CUDA 语境下可近似理解为一个 CTA（thread block）。SMEM 是 CTA 内部共享的片上内存，HBM/global memory 是容量更大但延迟更高的设备内存。

第一，logical tensor 不等于 HBM tensor。 以一条 Attention row 为例，

在语义上产生 $S,P:[1,T_K]$ 与 $O:[1,d_v]$。这里的 $S$ 和 $P$ 不一定写回 HBM。后文会明确区分：算法上存在的 logical shape、跨 kernel 保存的 materialized shape，以及一个 CTA/program 当前处理的 tile shape。

第二，每个输出都要有明确的 owner。 这里的 owner 是“最终负责写回这块 输出的 program”。假设一个 program 独占 $B_R$ 条 softmax rows，并分批读取 $B_{\mathrm{kv}}$ 条 KV：

其中 $\alpha$ 是 score scale。局部计算始终是两个规则矩阵乘：

Program 只需为自己拥有的 rows 保留 online-softmax 状态

每个 KV tile 到来后，先更新 row-wise maximum $m$，再把旧的 $z$ 与 $O_{\mathrm{acc}}$ 缩放到新基准，累加当前 tile 的概率和 value numerator。 因此 selected blocks 或 selected rows 可以分批进入，而无须物化完整概率矩阵。

第三，重点不是“跳过多少”，而是“不规则性停在哪里”。 对 每条路径，我们只追踪四件事：

1. selection 如何产生，是否物化中间 tensor；
2. 一个 program 独占什么 output；
3. 一次 KV load 被多少 rows 或 heads 复用；
4. 动态 index 在哪里重新变成规则 tile。

一次 FMA 按 2 FLOPs 计。除非特别说明，byte count 使用 BF16 useful payload； cache-line overfetch、page/TLB 与 allocator 影响会在相关位置单独讨论。

3. NSA：先让算法产生可执行的块级工作负载

Takeaway. NSA 的关键不是“选 16 个 block”，而是每个 ID 都指向连续的 64 个 tokens，并由同一 KV group 的 16 个 heads 共同使用。算法输出因此已经 接近一个规则的 $[16,64]$ 计算 tile。

NSA 将长程注意力拆成三种访问模式：

真正关键的推导链不是“三路相加”，而是

算法公式来自 NSA 论文。论文实验所用 kernel 没有完整公开，因此本文的 online Top-K、forward grid 与 backward 代码分析以 FLA 社区实现为例；凡是 实现特有的选择都会明确标出。[PAPER] [CODE/community] [GAP]

3.1 算法契约：三路稀疏各自提供什么

NSA 没有要求一条稀疏路径同时承担所有信息。Compression 保留远程概貌， selection 找回重要细节，sliding window 保证局部连续性。三路各自做 softmax， 最后再融合输出。

省略 batch 轴，使用 sequence-major layout。对长度为 $T$ 的序列，

NSA 面向 GQA。论文默认配置采用

第 $g$ 个 KV head 对应的 query-head group 为

三条分支的算法参数是：

其中 $L_t=t+1$ 是位置 $t$ 的因果前缀长度。它包含

个完整 compression windows，以及

个因果可见的 selection blocks。

三路不是同一份 KV 上的三个 mask。论文为它们提供独立的 K/V projection：

因此，即使三路访问了同一个 token，也不能简单求 union 后只算一次 Attention：它们使用不同表示、分别归一化，最后才组合输出。

Compression 回答“远处大致有什么”。 每个 KV head 将长度为 $\ell=32$ 的 overlapping window 压成一行 K/V，相邻窗口 stride 为 $d=16$。完整序列因此产生约 $C_T\approx T/d$ 行 compressed K/V。

位置 $t$、group $g$ 的局部 shapes 为

固定 stride $d$ 时，所有 query 的 compression Attention pair count 仍为

在相同 feature dimension 下，它相对 dense causal pair count 约减少 $d$ 倍， 但仍不是线性复杂度。这里比较的是算术规模，不是包含 compressor、routing 与 kernel efficiency 后的 wall-clock speedup。

Sliding window 回答“眼前发生了什么”。 局部分支只访问

提供固定、连续且易于 tiled attention 处理的局部上下文。

Selection 回答“远处哪些细节值得重新展开”。 它怎样产生 block IDs 是 下一小节的主线。三条分支分别完成 softmax，得到

再由独立 sigmoid gates 组合：

这些 gate 不要求和为 1。执行 DAG 上，sliding branch 可以与 compression/routing chain 并行，selection Attention 则必须等待 block IDs。 至于 compression probabilities 如何到达 routing，属于实现选择：可以物化、 与 Top-K 融合，或者像后文分析的 FLA 路径一样，根据 $Q$、compressed $K$ 与 LSE 分块重算。

3.2 Routing：从 Compression Probability 到 Top-16 Block IDs

Selection 没有另训练一套 router。它复用 compression Attention 已经得到的 概率，将“哪些压缩窗口重要”转换为“哪些 64-token blocks 值得展开”。整个过程 只有三步：空间重映射、16-head 求和、Top-16。

空间重映射。 令

为 query head $h$ 的 compression probability。论文 Eq. 9 将它变成 selection-block score：

越界的 compressed-window index 视为零。论文配置满足

直觉上，一个 64-token block 会与多个 overlapping compression windows 相交； 这个双重求和就是把所有相关窗口的概率累加到该 block。

GQA group reduction 与 Top-$n$。 同一 KV head 的 $G=16$ 个 query heads 必须共享 selection IDs，因此

论文 Sec. 4 的具体配置还规定：16 个 selection slots 中固定激活 1 个 initial block 和最近 2 个 local blocks，其余 slots 才由上述动态 ranking 填充。 [PAPER]

对成熟位置，

它是带有 slot 顺序的 index tensor，而不是数学集合。序列开头可能只有少于 16 个因果有效 blocks；kernel 可维持 16 个 selection slots，但必须用 $-1$ 或 block count 标记 padding。“固定 16 个 slots”不等于“始终存在 16 个有效 blocks”。

这三个固定 slots 也会影响 backward：initial block 会被大量 queries 共同选择， 而 local blocks 随 query 移动。到 3.4 节讨论 $dK/dV$ 时，这种 fan-in 差异 会直接变成负载不均。[DERIVATION]

完整 shape 链为

compression probabilities       [T,Hq,≈T/16]
  → Eq. 9 spatial remap          [T,Hq,≈T/64]
  → reduce G=16 query heads      [T,Hkv,≈T/64]
  → Top-16                       [T,Hkv,16]

但概率在哪里？ 标准 FlashAttention 不会把完整 probability matrix 写回 HBM，通常只返回 output 与 LSE。因此，“复用 compression probability”只说明 算法关系，还没有给出可执行的数据流。

在 $T=65536$ 上，完整窗口数实际为 $C_T=4095$。为说明固定 shape allocation 的上界，令 padded capacity

若按该 capacity 物化 per-head FP32 compressed scores，

一层需要

若在 Eq. 9 remap 之前先对每个 GQA group 求和，物化

仍需

若讨论的是 remap 之后的 group-reduced block scores，令 $M_T=\lceil T/\ell'\rceil=1024$，其 shape 应为

对应

这三个数字描述不同的数据流位置，不能混用。

当前 FLA 给出了一种可行机制：保存 compression LSE，根据 Q 与 compressed K 分块重算 probability，并用 bitonic merge 在线维护 Top-$n$，从而不物化上述 二次 score tensor。这条路径以重算换显存：Top-k kernel 必须再次用 Q 扫描 compressed K，因此 routing 仍包含一次 $\Theta(T^2/d)$ 的 QK scan 与候选 merge。它消除的是二次 score tensor 的 HBM materialization，不是二次 pair scoring 本身。

其配套路径还把 compression 简化为 non-overlap mean pooling、共享 K/V，并对 GQA group size 施加至少 16 且为 2 的幂等实现约束；这些都不是论文算法的通用 定义。因此，FLA 只能证明一种 infrastructure 路径可行，不能证明论文作者采用 了同样实现。[CODE/community] [GAP]

3.3 Selection Forward：逻辑 Gather，物理 Block Tile

现在 block IDs 已经产生。Forward 只剩一个问题：是否先把 16 个 blocks gather 成紧凑的 $[1024,d]$ tensor？答案是否定的。这个 tensor 只应存在于 数学视图中；kernel 应保存 IDs，并按 block 直接读取原始 KV。

从 16 个 IDs 到逻辑上的 1024 个 Tokens

Routing 输出离散 block IDs：

每个 $j_s$ 指向 selection K/V 中一段连续的 64-token block：

对成熟 query，如果 16 个 slots 都有效，数学上可以拼成

逻辑 attention shapes 为

这里的 1024 是 padded capacity，不是所有位置的有效 token 数。序列开头可能有无效 slots；包含当前位置的 causal block 也可能只有部分 token 可见。实现必须同时应用 ID mask、range mask 与 causal mask。

这个 $[1024,d]$ 只是逻辑 shape。原始 KV 中仍是 16 段离散 blocks；每个 block 内的 64 个 token rows 则可由一个 base 加规则 offsets 得到。NSA 改善的 不是“16 个 blocks 彼此连续”，而是把 1024 次 token-level 间接寻址收敛成 16 次 block-level 定位。

如果真的为每条 query/group 构造这两个 tensors，外部 gather 既要处理 16 个不规则 block bases，又要将它们组装为 compact selected-KV tensor 写回 HBM。 仅 K/V payload 就有

（BF16）。这些数据写入 HBM 后还要被 Attention 重新读取。因此 kernel 只保存 IDs，并在内部按 block 加载原始 K/V。

一个 Program 最终写回什么

先看论文默认配置。 此时 $d_v=B_V=128$，所以 $N_V=1$。一个 program 对应一个 query position 和一个 KV group，独占 16 个 query heads 的 完整 output：

这就是最重要的 ownership：同一个 program 既持有 16 条 softmax rows，也能 让它们共同复用每个 selected block。

再展开通用 grid。 当 value 维需要分成多个 tiles 时，完整 owner 还多一条 value-feature 轴：

它由 $(a,t,g,\nu)$ 唯一标识。这里 $a$ 是 batch，$t$ 是 query position，$g$ 是 KV group，$\nu$ 是 value-feature tile。

对 fixed-length 输入，FLA 使用

其中：

• $B$ 是 batch size；
• $T_Q,T_K$ 是每条序列在本次调用中的 query 与 KV token 数；
• $H_q,H_{kv}$ 是 query-head 与 KV-head 数；
• $G=H_q/H_{kv}$ 是一个 KV head 服务的 query-head 数；
• $d_k,d_v$ 是每个 head 的 Q/K feature 与 V/output feature 维度；
• $N_{\mathrm{slc}}$ 是 selection-ID 末维的 slot capacity，源码记作 S，本文配置为 16；
• $I^{\mathrm{slc}}$ 保存 block IDs；block_counts 若为 tensor，则记录 每个 query/group 的有效 slot 数，若为 scalar capacity，kernel 扫描全部 $N_{\mathrm{slc}}$ 个 slots；
• $B_V$ 是一个 program 沿 value feature 轴计算的 tile 宽度；
• $N_V=\lceil d_v/B_V\rceil$ 是覆盖完整 value 维所需的 tile 数。

FLA 还假设 $Q$ 对应每条 KV sequence 最后的 $T_Q$ 个 tokens。因此 fixed-length 路径中，query slot $t$ 在 KV sequence 内的绝对位置不是 简单的 $t$，而是

后面的 causal mask 与 selected-block 起点都相对于 $p_t$ 判断。

源码中的 launch grid 是一个三维 program-count tuple：

源码将三个 program IDs 命名为 i_t、i_v 与 i_bh。为避免 i_v 和 value tensor $V$ 混淆，本文把第二个 ID 改记为 $\nu$：

它们依次表示 query position、value-feature tile，以及被压平的 batch/KV-head 位置。第三个 ID 按

还原成 batch index $a$ 与 KV-head/group index $g$。这里特意用 $a$ 表示 batch，避免和后文的 KV block ID $b$ 混淆。第 $g$ 个 group 服务的 query heads 仍是

再定义 program $\nu$ 负责的 value-feature 区间

于是 program $(t,\nu,i_{bh})$，等价地说 $(a,t,g,\nu)$，是下面这个输出 slice 的唯一 writer：

物理 accumulator 按 $[G,B_V]$ 分配；最后一个不满 $B_V$ 的 tile 通过 feature mask 处理。由此也能看清 axis 顺序：去掉 batch flattening 后 grid 是 $(T_Q,N_V,H_{kv})$，而不是 $(T_Q,H_{kv},N_V)$。

如果 $N_V>1$，同一个 $(a,t,g)$ 会展开成 $N_V$ 个 programs；它们为各自 的 value-feature slice 重算相同的 QK/softmax，但写入的 $\mathcal D_\nu$ 两两不交，因此 output 没有写冲突。LSE 不沿 value 维切分， 源码只允许 $\nu=0$ 的 program 写回 $\mathrm{LSE}^{\mathrm{slc}}[a,t,\mathcal H_g]$。

回到论文配置，$d_v=128$，当前 FLA 也取 $B_V=128$，所以

此时每个 $(a,t,g)$ 恰好只对应一个 program，它独占 $[G,d_v]=[16,128]$ 的完整 selection output。该 program 随后加载并常驻 逻辑 query tile

一个 64-Token Block 怎样变成规则 Tile

随后循环 16 个 selection slots。本文用 $B_{\mathrm{kv}}$ 表示一次内循环 加载并计算的 KV-token row 数；在当前 FLA 路径中，它对应源码参数 BS，并 恰好等于 NSA 的 selection block length $\ell'$：

对一个有效 block ID $j_s$，program 计算 $\mathrm{base}=j_s\ell'$，并加载

这是算法层面的逻辑 tile。在当前 FLA 的 Hopper 分支中，

所以为了匹配 Triton dot 的 feature tile，实际构造

其中 feature 维 $192\rightarrow256$ 的尾部通过 mask/padding 处理。因此逻辑 contraction 是

而 QK 的物理 padded execution shape 是

这组 $192\rightarrow256$ 的执行 shape 是 当前 FLA Hopper 路径的 结论。非 Hopper 分支将 BK 上限设为 128；当 $d_k=192$ 时会得到两个 feature tiles，并触发源码的 NK == 1 断言，因此不能把这里的 $[16,256][256,64]$ 直接外推到该路径。

动态性只存在于每轮加载前的 block base；进入片上 tile 后，计算重新变成固定 shape。

这里特意不用 $B_K$ 表示 token tile。NSA 论文中的 token-tile 记号 $B_K$ 在当前 FLA 源码中对应 BS；FLA 的 BK 表示 padded feature width，例如

混用两套记号会把“64 个 tokens”误读成“256 个 feature elements”。

在线归一化。 Program 对 16 个 blocks 使用前文统一介绍的 exact online-softmax merge，只维护每个 query head 的 running max、normalizer 与 output accumulator，不写出 $P^{\mathrm{slc}}\in\mathbb{R}^{16\times1024}$。所有 selected blocks 完成后直接写回

因此 output 拥有唯一 writer，forward 不需要对 output 做 atomic reduction。

为什么 head sharing 是执行契约。 对一个 64-token block，QK 与 PV 的乘加量约为

若 K/V 为 BF16，本轮主要 K/V payload 为

只看这份被 16 个 heads 共享的 K/V payload，局部 arithmetic intensity 为

这不是整 kernel 的 roofline 数字：它没有计入 Q、output、indices、LSE、mask 和 launch 开销。它只说明，同一 K/V block 被越多 query heads 复用，一次不连续寻址带来的搬运就越容易被规则 matmul 摊薄。

对成熟 query，selection 主循环有 16 个 slots；序列开头仍运行相同控制结构，但无效 slots 被 mask。NSA 提供的是“固定容量、可 padding 的 workload”，而不是没有边界条件的绝对固定工作量。

可跳读的实现补充：variable-length 不改变 ownership。 公开路径将多条 不等长序列 pack 到 batch 轴为 1 的 tensor 中。第一条 grid 轴由 $T_Q$ 改成总 query 数 $T_Q^{\mathrm{total}}$，再通过 token_indices_q 找回每个 query 的序列号与局部位置；另外两条轴仍是 $(N_V,H_{kv})$。因此 varlen 只改变 $(a,t)$ 的定位方式，不改变“一个 program 唯一写一个 output slice”的 ownership。完整 packed shapes 为

[ Q:[1,T_Q^{\mathrm{total}},H_q,d_k],\quad K^{\mathrm{slc}}:[1,T_K^{\mathrm{total}},H_{kv},d_k],\quad V^{\mathrm{slc}}:[1,T_K^{\mathrm{total}},H_{kv},d_v]. ]

3.4 Forward 与 Backward 使用不同的稀疏图

Takeaway. Forward 的 query → selected blocks 关系只给 output 找到了 owner。若继续使用 query-centric ownership， $dK^{\mathrm{slc}}/dV^{\mathrm{slc}}$ 需要浮点 atomics 或 partial-gradient reduction；本文审计的 FLA 社区实现选择把它转置为 KV block → selecting queries，再让每个 KV-block program 完成自己的 many-to-one reduction。

Forward 的自然 owner 是 query，因为每个 program 能独立完成一个 $O^{\mathrm{slc}}_{t,\mathcal H_g,:}$。Backward 则有两种归约方向。

以下只讨论 selection branch 的 backward。完整 NSA 还包括 gates、compression branch、sliding branch，以及三个分支对共享 Q/input 的梯度合并。

重算 Probability，保留稀疏图

Selection backward 不需要保存完整 probability matrix。当前 FLA 首先由 forward output 与上游梯度计算每个 softmax row 的

随后根据 $Q,K^{\mathrm{slc}},I^{\mathrm{slc}}$ 与 $\mathrm{LSE}^{\mathrm{slc}}$ 分块重算 score/probability。对 batch/sequence $a$、query $t$、group $g$ 与 selected block $b$， 有

这里 $V^{\mathrm{slc}}_{a,g,b}$ 表示 block $b$ 内 $B_{\mathrm{kv}}$ 行 value vectors，shape 为 $[B_{\mathrm{kv}},d_v]$。相应的局部量 shapes 为

若不使用 activation checkpoint，FLA selection autograd state 至少包括

以及 variable-length 场景所需的 block counts 和 sequence metadata。当前 FLA 的 autograd 路径实际保存 q, k, v, o, lse，并在 context 中保留 block indices 等 metadata。准确结论是“不保存 $P^{\mathrm{slc}}$”，而不是 “forward 只需保存 $O^{\mathrm{slc}},\mathrm{LSE}^{\mathrm{slc}},I^{\mathrm{slc}}$”。

$dQ$ 继续由 Query 拥有

对 query $t$，只有它自己选择的 blocks 参与

每个 block 的局部 shape 为

同一 program 可以循环全部 slots。逻辑上的最终结果为

而当前 Triton kernel 取

所以物理 accumulator b_dq 的 shape 是 $[16,256]$；末尾 64 个 padded feature lanes 由 mask 处理，最终只写回前 192 维。因此 dQ 仍适合 query-centric ownership。

源码中 fixed-length 的 dQ launch grid 为

Backward 取

使 softmax backward 中的 $\Delta$ 始终对应完整 value 维。因此当前路径中 program $(t,0,i_{bh})$，等价地说 $(a,t,g)$，独占一个 $[G,d_k]$ 的 dQ slice。variable-length 路径只把第一条 grid 轴换成 packed token 位置，ownership 不变。

$dK/dV$ 为什么不能继续由 Query 拥有

一个 KV block $b$ 可能被许多 queries 选中。令

它的梯度是跨 queries 的 many-to-one reduction：

若沿用 query owner，不同 programs 会同时更新相同的 $dK^{\mathrm{slc}}_b,dV^{\mathrm{slc}}_b$，只能使用大量 floating-point atomics，或写 partial buffers 再归约。更自然的办法是先转置稀疏关系：

forward adjacency:   query → selected KV blocks
backward adjacency:  KV block → all selecting queries

然后让一个 program 拥有一个 $(a,g,b)$ KV block。

转置稀疏图，让 KV Block 成为唯一 Owner

这里先看 fixed-length 路径。令完整序列的 block 数为

将一个 KV block $(a,g,b)$ 压平为 row

CSR 用两个一维 arrays 保存“每个 block 被哪些 queries 选中”：

其中 csr_indices 分配到 padded edge capacity；只有 csr_offsets[-1] 之前的 prefix 有效。源码存入的是压平后的 absolute query position

因此 row $r=r(a,g,b)$ 满足

当前 FLA 通过 counting-scatter 构造它：

遍历有效 selection edges
  → atomic count 每个 (batch,group,block) 的 fan-in
  → prefix sum 得到 csr_offsets
  → 再遍历 edges，用 atomic cursor scatter query IDs

atomic 只作用于紧凑整数 metadata；昂贵的 $dK^{\mathrm{slc}}/dV^{\mathrm{slc}}$ 浮点归约则获得唯一 owner。variable-length 路径改用 packed block/query IDs，但“一个 CSR row 对应一个 KV-block owner”的 关系不变。这里真正重要的不是 CSR 格式本身，而是它完成了从 query-centric 到 KV-centric 的稀疏图转置。

KV-centric dKV tile。 fixed-length 的 dKV launch grid 为

由于当前 backward 路径中 $N_V=1$，program $(0,r(a,g,b))$ 是 KV block $(a,g,b)$ 的唯一 writer。variable-length 路径将第二条 grid 轴替换为“packed KV blocks 数 $\times H_{kv}$”，但仍然 由一个 program 独占一个 KV block。该 program 持有

和同 shape 的 gradient accumulators。它从对应 CSR row 每次读取 $B_{\mathrm{qry}}$ 个 query positions，并展开每个 query 的 $G$ 个 heads：

批内 query 与 output-gradient shapes 为

局部重算与归约重新形成规则 matmul：

例如 $B_{\mathrm{kv}}=64,B_{\mathrm{qry}}=4,G=16$ 时，$R=64$。逻辑上的 score-recompute contraction 为

当前 Triton kernel 的 $\texttt{BK}=256$，所以实际 tl.dot tile 是

其中末尾 64 个 feature lanes 由 mask 补零；b_dk 也相应使用 $[64,256]$ 的物理 accumulator，再只写回前 192 维。

CSR 的作用不只是压缩 metadata。它把“多个 queries 向同一 KV gradient 地址写入”的冲突，改写为“一个 KV owner 顺序读取自己的 queries”，再在 owner 内部恢复规则 Tensor Core tile。

写冲突消失了，Fan-in 长尾还在

定义

当 block_counts=16 是固定标量 capacity 时， $dQ^{\mathrm{slc}}$ program 扫描固定的 16 个 padded slots；FLA 也接受逐 query 的 tensor block_counts，此时运行时 $N_S$ 可以缩短。dKV program 的循环长度则由 $NQ_{a,g,b}$ 决定。fixed initial/sink block，以及由动态 ranking 形成的热门 blocks，可能拥有远高于平均值的 fan-in，形成长尾 CTA；移动的 local slots 本身只产生有界的局部 fan-in。拆分热门 row 可以改善负载均衡，却会重新引入 partial reduction 或 atomic。Backward 因而把问题从“跳过多少 FLOPs”转成了 “如何调度 many-to-one reduction”。

本文没有继续把 fan-in 长尾做成性能实验。一方面，论文作者使用的 backward kernel 没有公开；另一方面，社区实现中改变 routing topology 时，fan-in skew、 query gather locality 与 CSR row 顺序会一起变化，很难把 latency 差异归因给 单一变量。因此，FLA backward 将不规则性从浮点写冲突转移成了 CSR row length、query-gather locality 与 CTA load balance。本文没有对这些因素报告 性能测量；这里只把它作为源码可确认的执行后果，而不声称它解释了论文作者 kernel 的 backward latency。[DERIVATION] [CODE/community] [GAP]

3.5 NSA 的规则性来自算法契约，不来自 Kernel 魔法

在论文默认配置下，NSA selection 能映射为高效 kernel，是因为算法主动提供 三项执行契约：

Routing 输出的每个动态 ID 已经对应一段连续的 64-token block。因此，kernel 在 load 之前就知道：不规则性只存在于 block base，block 内 64 个 token 可以 通过规则的 affine offsets 展开。Program 加载规则的 $[64,d]$ K/V tile， 再与共享该组 IDs 的 16 个 query heads 形成 $[16,64]$ score tile。本文分析 的 FLA community backward 则进一步转置稀疏图，让 KV block 成为 $dK^{\mathrm{slc}}/dV^{\mathrm{slc}}$ 的唯一 owner。

这也给出三个不能越过的边界。

第一，NSA 不是“任意 block sparsity 都会自动加速”的证据。真正产生硬件效率的是连续 block、head sharing、固定 slot capacity 与明确 ownership 同时成立。

第二，selection 与 sliding 的 token budget 在成熟位置近似固定，但 compression 仍访问约 $L_t/d$ 行 compressed K/V。固定 $d$ 时，全层 compression 仍为 $\Theta(T^2/d)$。

第三，FLA 展示了 online Top-k、group-centric forward 与 CSR backward 的一套完整可行机制；它不能证明论文作者的未公开实现采用相同 routing materialization、tile 参数或 inverse-adjacency construction。

论文报告，在其 $G=16,\ell'=64,n=16$ 配置与 A100 测试环境中，作者的 Triton NSA 相比 Triton FlashAttention-2 在 64K 达到 9.0× forward 与 6.0× backward speedup。这个结果可以归因于论文报告的完整系统，不能直接归因于本文从社区代码拆出的某一个 kernel 选择。[PAPER] [GAP]

换句话说，NSA 最值得复用的不是某一条稀疏公式，而是一种共同设计方法：

先让算法输出 GPU 能消费的稀疏单位，再分别为 forward output 与 backward gradient 选择唯一 owner。

4. 从 Block 到 Token：DSA 为什么能走得更细

NSA 与 DSA 不是简单的“旧方案”和“新方案”。NSA 选择 block，是为了让训练 和推理都能直接使用连续 tile；DSA 选择 token，是因为它在 MLA 上找到了另一条 足够强的复用轴。这里的变化不是取消硬件约束，而是改变算法满足硬件约束的方式。

4.1 Infra：连续 Block 不再是规则性的唯一来源

先看 NSA 面对的问题。任意 token selection 会产生大量离散 KV row；在训练中， 这些 rows 还要参与 $dK/dV$ 的反向归约。NSA 因而把选择单位固定为连续 64-token block。代价是 selection 只能表达“这一段值得看”，不能只取 block 中真正相关的几个 token。换句话说，NSA 用选择粒度换取了 load 前的规则性。 [PAPER]

DSA 的出发点不同。DeepSeek-V3.2 已经建立在 MLA 上，每个 token 只保存一条 共享的 latent KV entry。DSA 又让同一个 query 的 128 个主 heads 共用一组 token IDs。于是，一条离散 KV row 虽然难以读取，却能在读入后同时服务很多 heads；FlashMLA 可以把随机寻址留给 producer，再让 Tensor Core consumer 处理规则的 SMEM tile。此时，连续 block 不再是获得硬件效率的唯一办法。 [PAPER] [CODE]

从 infra 角度看，设计重心因此发生了三步变化：

这给 DSA 更细的内容选择能力，但不是免费的升级。NSA 的 routing 在压缩后的 token 轴上进行；DSA 的 Indexer 保留完整 token 分辨率，随后还要执行 exact Top-K。它把主 Attention 限制为 $O(L\kappa)$，却没有消除 Discovery 的 $O(L^2)$。所以这次变化更准确的描述是：用更强的跨 head 复用和专用 gather kernel，换取 token-level selection；再把系统瓶颈从 sparse main Attention 推向 Indexer 与 Top-K。 [DERIVATION]

4.2 Training：把新 Router 接入已有 MLA Checkpoint

DeepSeek 没有从头训练 DSA，而是在已扩展到 128K context 的 DeepSeek-V3.1-Terminus 上继续训练：

1. Dense warm-up： 保留 dense Attention、冻结主模型，用 dense Attention 分布产生的 KL loss 训练 Indexer。
2. Sparse training： 启用 Top-2048，主模型用 language-modeling loss 适应稀疏模式；Indexer 继续使用独立 KL loss。它的输入从主模型计算图中 detach，使两路梯度彼此隔离。

这套流程的意义是：保留已有 MLA 参数化，只接入新的 router 和 sparse kernel， 而不必从头训练整套 Attention。[PAPER]

证据边界。 DeepSeek 的报告并没有给出一项“NSA 对 DSA”的直接消融， 也没有逐条声明放弃 NSA 的原因。上面确认无疑的是两篇论文各自的 algorithm contract、DSA 对 MLA/MQA sharing 的选择，以及它的 continued-training 流程； “设计重心为什么能够从 block 转向 token”则是基于这些公开事实作出的系统分析。

下面进入 DSA 的实际数据流，依次追踪 Indexer score、Top-K token IDs、MLA 执行表示与 FlashMLA kernel。

5. DSA：把任意 Token Selection 重新变成规则矩阵

Takeaway. DSA 允许任意 token selection，但让 128 个主 heads 共享同一 组 IDs。FlashMLA 只让 producer 处理离散地址；进入 SMEM 后，consumer 看到的 仍是规则矩阵。

DSA 的完整前向链只有三个阶段：

下面逐阶段展开，并保留理解 shape 所需的记号。

这里的 $D_V^{\mathrm{MHA}}=128$ 不能和 FlashMLA 内部的 latent value 维度 $D_C=512$ 混写。主 Attention 也有两套代数等价、执行 shape 不同的表示：

5.1 Indexer：低维，但仍扫描全部历史

第一步：为每个 Query–Token Pair 打分

Indexer 的任务只有一个：为 query $t$ 和每个历史 token $s$ 计算便宜的 routing score，再选出 Top-$\kappa$。论文公式是

其中

对全部 query 与历史，最终 routing-score matrix 为

后面的 projection、RoPE、Hadamard、FP8 和 DeepGEMM fusion 都在回答同一件 事：怎样以更低代价计算这个 $[T_Q,T_K]$，而不把更大的 $[T_Q,64,T_K]$ 中间量写进 HBM。

DSA 一层里同时存在两种 score：

以及

$I$ 是 Lightning Indexer 的 routing score，只负责产生 token address；$A^{\mathrm{slc}}$ 是主 MLA 在 selected tokens 上重新计算的 attention logit，只有它会进入主 Attention 的 softmax。

Indexer 决定“去哪里看”；主 MLA 再决定“看见以后怎样加权”。

第二步：生成低维 Query、共享 Key 与 Head Weight

归一化 hidden states

先生成主 MLA 与 Indexer 共用的低秩 query latent：

Indexer 不再从 7168-D hidden 独立做一次大投影，而是从 $C^Q$ 得到：

reshape 后：

发布版公开 DeepGEMM 接口以物化的 FP8 Q 为输入；生产系统是否把这段 projection、旋转和量化进一步融合，公开代码没有给出答案。

历史 token 的 Indexer key 则只有一条共享向量：

持久 cache 的逻辑 shape 为：

这意味着 64 个 Indexer query heads 共享同一个历史 key，而不是为每个 head 维护一份 K cache。

Indexer 还从当前 hidden state 直接生成 query-dependent head weights：

其中 $u^I_{t,j}$ 是 query $t$ 对第 $j$ 个 Indexer head 的基础权重； 这条分支不经过 softmax。把官方 reference 的两个固定 normalization 吸收后， 定义

需要注意，只有 Indexer Q 复用 $C^Q$；Indexer K 与这里的 weights 都直接 由 $X$ 生成。

第三步：RoPE、Hadamard 与 FP8

每条 128-D Indexer Q/K 被拆成：

前 64 维应用 non-interleaved RoPE，与主 MLA 使用的 interleaved RoPE 不同。这不是排版差异：若两边使用不同的 pair layout，对应位置的 dot product 就会改变。

旋转后再应用归一化 Hadamard 变换 $\mathcal H_{128}$：

所以量化前有：

Hadamard 不改变精确实数内积；它的作用是扩散 outlier，使后续 E4M3 量化更稳定。“内积不变”不能外推成“量化后逐 bit 不变”。

关于两个 $H$ 的记号。 $H_I=64$ 是一个标量，表示 Indexer head 数；$\mathcal H_{128}\in\mathbb R^{128\times128}$ 才是 Hadamard 变换 矩阵。它的元素为 $\pm1/\sqrt{128}$，是固定、正交且不参与学习的线性 变换。实现无需在内存中物化这张 $128\times128$ 矩阵，而是用 fast Hadamard transform 通过分层加减在 $O(128\log128)$ 时间内完成。Q/K 使用 同一个 $\mathcal H_{128}$，所以精确算术下内积不变；它位于 partial RoPE 之后、FP8 量化之前，量化后的内积则只近似保持。

令 $q^{(8)},k^{(8)}$ 表示 E4M3 数据，$s^Q_{t,j},s^K_s$ 表示相应 scale。官方 reference 将 query scale 与两个 normalization 一起折入 query-dependent weight：

于是 scorer 的实现语义可写成：

这里的 $q^{(8)},k^{(8)}$ 明确指经过 partial RoPE、Hadamard 与 FP8 量化后的向量，不再和变换前的 $q^I,k^I$ 混用。

Indexer K cache 的逻辑 payload 是：

与完整主 KV row 相比，这是全历史扫描带宽能够显著下降的基础。

第四步：在片上融合 64 个 Heads，只写最终 Score

按数学定义，固定 query $t$ 时先有：

全局 logical shape 是：

再经过 ReLU、weight 和 head reduction：

如果把 $[T_Q,64,T_K]$ 写入 HBM，低维 Indexer 会先制造一个比最终 score 大 64 倍的中间量。发布版 DeepGEMM scorer 因而把矩阵乘、ReLU、query-dependent weighting 与 64-head reduction 融合起来，只写最终的 FP32 $I$。

它的一种核心 tile 按代码中的 GEMM 方向可以概括为：

最后一步对 64 个 Indexer heads 做 ReLU、query-dependent weighting 与 reduction，再将 candidate-major accumulator fragments 直接映射到逻辑 $[2,256]$ 的输出地址。这里没有单独的 transpose 阶段。中间的 64-head 结果 $Z$ 停留在片上；HBM 只看到每个 query 对 256 个 candidates 的最终 score tile。

但这个发布版 scorer 不包含 Top-K。公开 DeepGEMM pipeline 仍然为

分配并物化完整的矩形 carrier。主 scorer 写入每行的合法 candidate 区域； 启用 clean_logits 时，invalid/future 区域由独立 cleanup kernel 填充。 随后 exact Top-K 仍需重新扫描这张完整矩形结果。

因此 DeepGEMM 已经完成的是 scorer 内部融合：FP8 matrix multiply、ReLU、 query-dependent weighting 与 64-head reduction 在同一 kernel 中完成， $[T_Q,64,T_K]$ 的 per-head scores 不写入 HBM。但这不是 Indexer–TopK fusion。发布版公开 pipeline 仍然物化完整的 $I:[T_Q,T_K]_{\mathrm{FP32}}$，后续 exact Top-K 仍由独立实现重新扫描。

第五步：Top-K 将 Score 变成 Token 地址

对 query 位置 $t$，令 $p_t$ 表示最后一个 causal 可见位置。数学上：

$J_{t,r}$ 是 token address，不是 routing score。同一个 query 的 128 个主 heads 共享 $J_t$，不同 query 的集合则可以完全不同。

公开接口还需要两步 shape/dtype 适配。

第一，官方 inference reference 的 torch.topk 返回 INT64 indices，并以一次调用的 min(index_topk, end_pos) 作为固定末维；causal mask 负责让未来位置保持无效。

第二，在本文的 DSA 配置下，FlashMLA sparse-prefill 接收固定宽度为 $\kappa=2048$ 的：

若某行只有 $\kappa_t<\kappa$ 个合法 tokens，其余 slots 必须改写为 invalid sentinel，例如 -1 或不小于 $T_K$ 的地址。不能重复某个合法 token 来填满， 否则 softmax 会把该 token 计算多次。

从 reference 的 INT64 Top-K 到 kernel 的固定宽度 INT32 tensor，因而是一项真实的数据准备工作，不是可以在 shape 图里省略的类型标注。

Top-K IDs 是离散控制流，language-modeling loss 不能直接穿过“某个 token 是否 入选”这一决策。4.2 节介绍的两阶段训练用独立 KL loss 解决 Indexer 学习问题； 它不改变推理时 exact Top-K 的执行成本。

至此，Indexer 阶段结束。它交给下一阶段的不是 Attention 概率，而是 $J^{\mathrm{kernel}}:[T_Q,1,2048]_{\mathrm{INT32}}$ token addresses。

5.2 MLA Bridge：为什么主 Kernel 是 576-D MQA

Indexer 只产生地址。真正的 Attention 权重仍由主 MLA 计算。本节只做一次代数 重排：把每个 head 独有的 K/V up-projection 移到 Attention 两侧，使 selected KV row 能被所有 heads 共享。结果是 FlashMLA 看到 576-D shared key 和 512-D latent value，而不是 128 份展开后的 MHA K/V。

原始 MHA 表示会重新展开 KV Cache

共享 query latent 经过主 MLA 的 query up-projection：

每个 head 拆成：

$q^R$ 使用主 MLA 的 interleaved RoPE。另一侧，hidden states 先投影出 512-D KV latent 与 64-D RoPE key：