1. 从标题到落地DSA计算过程到底在算什么第一次看到“DSA计算过程”这个标题很多人会下意识把它归到算法课里那种纯理论推导的范畴觉得无非是公式加伪代码。但真到工程里跑一遍就会发现DSADynamic Sparse Attention动态稀疏注意力的计算过程本质上是一套围绕稀疏模式动态选择和高效矩阵运算展开的工程流水线。它要解决的问题很具体当序列长度拉到几千甚至上万标准注意力那套全连接的 QK^T 计算量会以平方级膨胀显存和算力都扛不住。DSA 的思路是不把所有 token 对都算一遍而是先判断哪些位置值得算再只对这些位置做真正的注意力计算。这个标题背后牵扯到的关键词很能说明问题GEMM是稠密矩阵乘是算力的主要消耗点top-K是稀疏选择的核心操作决定保留哪些注意力连接gather是把选中的稀疏位置从大矩阵里捞出来的动作MLA则是多头潜在注意力这类结构它和 DSA 结合时会产生额外的计算约束。至于热搜里出现的“top-k 路由”其实和 DSA 里的 top-K 选择是同一类思想在不同场景下的叫法都是“先打分、再挑前 K 个、最后只算这 K 个”。所以这篇内容适合谁看如果你正在做长序列推理优化、想搞清楚稀疏注意力到底省在哪、或者准备自己实现一版 DSA 计算流程那这篇就是给你写的。我会把整个计算过程拆成可复现的步骤讲清楚每一步为什么这么设计、参数怎么定、坑在哪。下面直接从整体设计思路开始。2. 整体设计与思路拆解2.1 为什么 DSA 要拆成“选择 计算”两段标准注意力的计算可以概括为三步Q 乘 K 转置得到分数矩阵softmax 归一化再乘 V 得到输出。问题在于第一步QK^T 的复杂度是 O(n²·d)n 是序列长度d 是头维度。n 到 4096 时单头分数矩阵就是 1600 万个元素多头叠加后显存直接爆掉。DSA 的核心判断是大部分注意力权重其实接近于零真正起作用的只有少数几个位置。既然这样就没必要把整个分数矩阵都算出来再 softmax而是先估计哪些位置重要只对这些位置做完整计算。这就自然拆成了两段。第一段是选择阶段用一个轻量的打分机制给每个 query 找出最相关的 top-K 个 key 位置。第二段是计算阶段只对这 K 个位置做 QK^T、softmax 和加权求和。两段之间用 gather 操作把稀疏位置对应的 K、V 向量取出来组成一个小规模的稠密矩阵再交给 GEMM 去算。这样整体复杂度从 O(n²·d) 降到 O(n·K·d)K 通常取 32 到 256相比 n 小一到两个数量级。这个设计的精妙之处在于它把“判断重要性”和“精确计算”解耦了。选择阶段可以用低精度、低维度的近似打分计算阶段才用完整的头维度。两者用不同的计算资源选择阶段偏访存计算阶段偏算力正好能错开瓶颈。2.2 选择阶段为什么偏爱 top-K 而不是阈值过滤有人会问既然要稀疏为什么不设一个阈值分数低于阈值的直接丢掉原因有两个。第一阈值过滤的结果数量不可控不同 query 保留下来的 key 数量可能从 3 个到 300 个不等这会导致后续 GEMM 的 batch 维度参差不齐GPU 上很难做高效并行要么 padding 浪费算力要么动态 shape 拖慢 kernel。第二top-K 的数量是固定的K 一旦确定后续所有矩阵的形状都是确定的可以提前分配显存、编译固定 shape 的 kernel工程上友好太多。top-K 的另一个好处是它天然适配“路由”这个比喻。热搜里的“top-k 路由”说的就是这件事每个 query 像一个请求top-K 就是给它选出 K 个最该处理的专家这里是 key 位置。这种固定容量的选择机制在负载均衡和并行调度上都比变长选择好处理。2.3 gather 在整个流程里的位置和代价gather 是把选中的索引对应的 K、V 向量从原始大矩阵里取出来的操作。它看起来简单但代价不小。假设原始 K 矩阵形状是 [n, d]选中索引形状是 [n, K]gather 之后得到 [n, K, d]。这个操作的访存模式是随机的因为每个 query 选中的 K 个位置在内存里不连续。随机访存对 GPU 的 cache 很不友好实测下来 gather 阶段经常占到整个 DSA 计算时间的 20% 到 35%。所以工程上有个常见优化把 K、V 按块block组织选择阶段以块为单位做 top-K而不是以单个 token 为单位。这样 gather 的时候一次取一整块连续内存访存效率高很多。代价是稀疏粒度变粗可能多算一些不重要的位置但换来的访存收益通常更划算。块大小一般取 16 或 32和 GPU 的 cache line 对齐。2.4 MLA 结构给 DSA 带来的额外约束MLA多头潜在注意力的特点是它不直接存完整的 K、V而是存一个压缩后的潜在向量用的时候再投影出来。这个设计本来是为了省显存但和 DSA 结合时会产生一个矛盾DSA 的选择阶段需要拿到 K 来做打分而 MLA 的 K 是投影出来的如果每个 query 都去投影一遍所有 key那省显存的意义就没了。常见的解法是在选择阶段用压缩后的潜在向量做近似打分因为潜在向量保留了主要信息用它算出的相似度排序和用完整 K 算出的排序高度一致。只有到了计算阶段才对选中的 top-K 位置做完整投影。这样既保住了 MLA 的显存优势又让 DSA 的选择有依据。这个细节在实现时很容易被忽略如果选择阶段直接对完整 K 打分显存峰值会明显上升。3. 核心细节解析与实操要点3.1 打分函数的选型与计算选择阶段的核心是给每个 query-key 对一个分数然后按分数排序取 top-K。打分函数的选择直接影响稀疏质量。常见的有三种点积打分、L2 距离打分、以及学习出来的打分网络。点积打分最直接用 query 和 key 的内积作为相关性度量计算就是一次 GEMM形状是 [n, d] × [d, n] 得到 [n, n]。但注意这里不需要算出完整的 [n, n] 矩阵因为下一步就是 top-K可以用分块计算加在线 top-K 的方式每算完一块就更新当前的 top-K 候选避免存下整个分数矩阵。这个技巧能把选择阶段的显存从 O(n²) 降到 O(n·K)。L2 距离打分在某些场景下比点积更稳因为它对向量模长不敏感。但计算量稍大需要先算 query 和 key 的模长平方再做展开。实测在长序列上点积和 L2 的稀疏质量差距不大但点积的 GEMM 更容易被硬件加速所以默认推荐点积。学习出来的打分网络灵活性最高但引入了额外参数和训练成本一般只在有充足训练数据且对稀疏质量要求极高的场景下用。对大多数工程落地来说点积打分加块级 top-K 已经够用。3.2 top-K 的实现细节排序还是堆拿到分数后要取前 K 个。最朴素的做法是全排序复杂度 O(n log n)。但 n 很大时全排序没必要因为只要前 K 个。用大小为 K 的最小堆遍历一遍分数复杂度 O(n log K)当 K 远小于 n 时优势明显。K64、n8192 时堆方法比全排序快大约 3 倍。不过 GPU 上实现堆并不高效因为堆操作有大量分支和串行依赖。实际工程里更常用的是分块 top-K把分数分成若干块每块内用并行归约找出块内 top-K再把各块的 top-K 合并。这样每块内部是规整的并行操作块间合并的规模也小。块大小一般取 1024 或 2048和 shared memory 容量匹配。还有一个细节是稳定性。如果分数里有 NaN 或 Inf排序结果会乱掉。所以在打分之后、top-K 之前要加一步清洗把非有限值替换成一个很小的负数。这个检查看起来多余但在混合精度训练或推理里溢出是常事不加这步会偶发地选出完全错误的 top-K。3.3 gather 的访存优化前面提到 gather 是随机访存优化方向是提高局部性。除了块级选择还有两个技巧。第一个是索引预排序。选出的 top-K 索引在取之前先按索引值排个序这样 gather 的时候访存地址是递增的cache 命中率会高不少。排序本身有开销但 K 不大时比如 64排序成本远小于访存收益。实测预排序能让 gather 阶段提速 15% 到 25%。第二个是合并相同索引的读取。不同 query 可能选中相同的 key 位置如果能把它们合并成一次读取、多次使用就能减少重复访存。这需要先对索引做去重和计数实现稍复杂但在 K 较小、n 较大时收益明显。注意gather 之后的 K、V 矩阵形状是 [n, K, d]这个形状在后续 GEMM 里会被当成 batch 维度。如果 n 很大这个中间张量的显存占用是 n×K×d×2K 和 V 各一份n8192、K64、d128 时大约是 128MB单看不大但多头叠加后要乘头数规划显存时要把这部分算进去。3.4 GEMM 的形状与精度选择计算阶段的核心是两次 GEMM。第一次是 Q 和 gather 后的 K 做乘形状是 [n, d] × [n, K, d] 得到 [n, K] 的分数。第二次是 softmax 后的分数和 gather 后的 V 做乘形状是 [n, K] × [n, K, d] 得到 [n, d] 的输出。这两次 GEMM 的 batch 维度都是 n也就是每个 query 独立算自己的 K 个位置。这种 batched GEMM 在 GPU 上有成熟的 kernel但要注意 batch 数太大时kernel 启动开销会累积。n8192 时如果每个 query 单独启动一个 kernel开销很可观。实际实现里会把多个 query 合并成一个大的 batched GEMM让 GPU 的 SM 充分利用。精度方面选择阶段可以用 FP16 甚至 BF16因为只需要排序正确不需要数值精确。计算阶段的第一次 GEMM 也可以用低精度但 softmax 之后、第二次 GEMM 之前建议把分数转成 FP32 再算因为 softmax 对精度敏感低精度下容易出现权重分布失真。这个混合精度的策略能在几乎不损失质量的前提下把计算阶段的算力消耗降下来。4. 实操过程与核心环节实现4.1 完整流程的分步拆解把整个 DSA 计算过程按数据流走一遍大概是这样的输入 Q、K、V形状都是 [n, d]多头情况下是 [h, n, d]。选择阶段用 Q 和 K 的近似表示MLA 场景下是潜在向量算分数分块计算加在线 top-K得到每个 query 的 K 个索引形状 [n, K]。索引处理对索引做清洗去 NaN、预排序、可选去重。gather按索引从 K、V 里取出对应的向量得到 [n, K, d] 的稀疏 K 和稀疏 V。计算阶段第一次 GEMMQ [n, d] 和稀疏 K [n, K, d] 做 batched 乘得到分数 [n, K]。softmax对 [n, K] 做行归一化注意这里是在 K 个位置上归一化不是全序列。计算阶段第二次 GEMM归一化分数 [n, K] 和稀疏 V [n, K, d] 做 batched 乘得到输出 [n, d]。多头合并把各头的输出拼接或投影得到最终结果。这个流程里第 2 步和第 4 步是 DSA 特有的其余步骤和标准注意力类似只是作用范围从全序列缩小到了 K 个位置。4.2 参数选择K 取多大合适K 是 DSA 最关键的参数直接决定稀疏程度和计算量。K 太小稀疏过度注意力质量下降K 太大省不了多少算力。经验上K 的取值和序列长度 n 有关n 越大K 可以相对小一些因为长序列里真正相关的 key 占比更低。一个粗略的参考n1024 时 K 取 128 到 256n4096 时 K 取 64 到 128n8192 以上时 K 取 32 到 64。这个范围不是死的要根据任务调整。如果任务本身需要全局信息比如长文档摘要K 要偏大如果是局部依赖强的任务比如代码补全K 可以偏小。判断 K 是否合适有个实用方法对比 DSA 输出和全注意力输出的余弦相似度。在验证集上跑一批样本如果相似度稳定在 0.95 以上说明 K 够用低于 0.9 就要考虑加大 K 或改进选择策略。这个方法比看下游指标更直接因为下游指标受太多因素影响不好归因。4.3 块大小的选择与对齐块级 top-K 的块大小影响访存效率和稀疏粒度。块太小访存局部性提升有限块太大多算的位置太多浪费算力。块大小取 16 或 32 是比较稳的选择因为这两个值和 GPU 的 warp 大小32以及常见 cache line 大小对齐。如果硬件是其他架构块大小要相应调整。原则是让一次块读取正好填满一个 cache line 或一个 warp 的处理宽度。这个参数最好实测因为不同硬件的访存特性差异不小。测试方法是固定其他参数只改块大小看端到端延迟的变化选延迟最低的那个。4.4 一个可参考的实现骨架下面用伪代码把核心逻辑串一下语言用 Python 风格方便理解。实际部署时这些操作会用 CUDA 或 Triton 重写但逻辑是一样的。def dsa_attention(Q, K, V, K_top, block_size): # Q, K, V: [n, d] n, d Q.shape # 选择阶段分块打分 在线 top-K topk_indices online_topk(Q, K, K_top, block_size) # 索引清洗与预排序 topk_indices sanitize_indices(topk_indices) topk_indices sort_indices(topk_indices) # gather 稀疏 K、V sparse_K gather(K, topk_indices) # [n, K_top, d] sparse_V gather(V, topk_indices) # [n, K_top, d] # 计算阶段 scores batched_gemm(Q, sparse_K) # [n, K_top] scores scores.to(torch.float32) weights softmax(scores, dim-1) output batched_gemm(weights, sparse_V) # [n, d] return output这个骨架省略了多头处理和 MLA 的投影细节但主干逻辑完整。实际实现时online_topk和gather是最需要优化的两个函数前者决定选择质量后者决定访存效率。5. 常见问题与排查技巧实录5.1 输出质量突然下降怎么查DSA 输出质量下降最常见的原因是 top-K 选错了位置。排查顺序建议这样先看选择阶段的分数分布。如果分数几乎全相等说明打分函数失效了可能是 Q 或 K 的数值有问题比如被错误地归一化或缩放。正常情况下分数应该有明显的区分度top-K 的分数应该显著高于平均值。再看索引是否越界或重复。gather 时如果索引超出 K 的范围会取到错误的数据输出自然不对。索引重复则会导致某些位置被重复计算权重分布偏移。这两个问题在实现里加断言就能提前发现。最后看 softmax 的输入范围。如果分数里混入了未清洗的 Infsoftmax 会输出 NaN 或全零。前面提到的清洗步骤就是防这个的。5.2 速度没提升反而变慢理论上 DSA 应该比全注意力快但实测变慢的情况不少见。原因通常有三个。第一是 K 取太大。如果 K 接近 n稀疏带来的收益被 gather 和索引处理的开销吃掉了。这时候要重新评估 K 的取值或者检查是不是选择阶段本身太慢。第二是 gather 的访存效率太低。如果没做块级选择和索引预排序随机访存会把 cache 打爆gather 时间可能超过省下来的 GEMM 时间。用 profiler 看一下 gather 的耗时占比超过 30% 就要优化。第三是 kernel 启动开销。如果每个 query 单独启动 kerneln 大时启动次数太多开销累积。解决办法是合并 batch让一次 kernel 处理多个 query。5.3 显存峰值比预期高DSA 省的是计算量但中间张量不一定省。gather 后的稀疏 K、V 是 [n, K, d]如果 n 和 K 都不小这部分显存很可观。加上选择阶段的分数缓存、索引张量显存峰值可能比想象的高。降低显存的办法选择阶段用在线 top-K不存完整分数矩阵gather 后的稀疏 K、V 用完即释放不要一直持有多头之间可以串行处理避免所有头的中间张量同时存在。如果还紧张可以把 n 分块每次只处理一部分 query。5.4 常见问题速查表现象可能原因排查方法解决方向输出质量下降top-K 选错位置看分数分布是否有区分度检查打分函数、清洗 Inf速度变慢K 太大或 gather 低效profiler 看各阶段耗时调小 K、块级选择、索引预排序显存峰值高中间张量未释放显存快照看峰值来源在线 top-K、分块处理、串行多头结果不稳定索引越界或重复加断言检查索引范围清洗索引、去重softmax 输出异常分数含 Inf 或 NaN检查分数有限性打分后加清洗步骤提示这张表里的排查顺序是有讲究的先查质量再查速度因为质量不对的话速度优化没意义。质量里先查选择阶段因为 DSA 的问题八成出在选择上计算阶段和标准注意力一样出问题的概率低。5.5 几个容易踩的坑第一个坑是在选择阶段用了完整精度的 K。MLA 场景下这会直接吃掉显存优势而且没必要因为选择只需要排序正确。用潜在向量或低精度 K 打分质量几乎不受影响。第二个坑是忘了处理 padding。如果序列里有 padding token它们的分数应该被屏蔽掉否则可能被选进 top-K污染输出。屏蔽方法是在打分后把 padding 位置的分数设成负无穷再走 top-K。第三个坑是块大小和头维度不匹配。如果块大小是 32但头维度是 64gather 的时候一次取 32 个 token 的 64 维向量访存模式可能不是最优的。块大小最好和头维度一起考虑让一次读取的数据量是 cache line 的整数倍。第四个坑是top-K 的 K 在不同头之间共享。有些实现为了省事所有头用同一组 top-K 索引。但不同头的注意力模式可能不同共享索引会损失质量。如果算力允许每个头独立选 top-K 效果更好。折中方案是分组共享比如每 4 个头共享一组索引。6. 性能调优与扩展思路6.1 选择阶段的近似加速选择阶段虽然比计算阶段轻但在 n 很大时也不便宜。加速思路有两个方向。一是降维打分。把 Q 和 K 投影到更低维度比如从 128 降到 32再算分数计算量降到四分之一排序质量下降有限。投影矩阵可以固定为随机矩阵也可以用训练学出来的。随机投影在多数场景下够用且不需要额外训练。二是采样估计。不遍历所有 key而是随机采样一部分 key 算分数用采样结果估计整体排序。采样率取 10% 到 20% 时top-K 的召回率通常能到 90% 以上。这个方法在 n 极大时特别有效但要注意采样要覆盖所有位置不能只采前段。6.2 和 MLA 结合时的显存账MLA 省显存的核心是存潜在向量而不是完整 K、V。假设潜在维度是 d_c完整头维度是 dd_c 通常是 d 的四分之一到一半。DSA 和 MLA 结合时选择阶段用潜在向量显存占用是 n×d_c计算阶段对选中的 K 个位置做完整投影显存占用是 n×K×d。只要 K 远小于 n总显存就比全注意力的 n×n×d 小得多。算一笔账n8192、d128、d_c64、K64。全注意力存分数矩阵要 8192×8192×2 字节FP16约 128MB多头叠加后更大。DSAMLA 的选择阶段存潜在向量 8192×64×2 约 1MB计算阶段存稀疏 K、V 各 8192×64×128×2 约 128MB。看起来计算阶段不小但这是所有头共享的而且用完即释放。整体峰值显存比全注意力低一个数量级。6.3 后续可以扩展的方向DSA 的计算过程稳定之后有几个自然的扩展方向。一是动态 K不同 query 用不同的 K重要的 query 多选一些不重要的少选一些。这需要解决变长带来的并行效率问题但能进一步提升稀疏质量。二是跨层共享 top-K相邻层的注意力模式往往相似可以复用上一层的索引省掉重复的选择计算。三是和量化结合选择阶段用低比特量化进一步降访存计算阶段保持较高精度。这些扩展各有取舍落地时要根据实际瓶颈选。如果瓶颈在访存优先考虑量化和跨层共享如果瓶颈在算力优先考虑动态 K 和降维打分。我个人在实际操作中的体会是DSA 的调优八成精力花在选择阶段而选择阶段里八成花在 top-K 的实现和 gather 的访存上。把这两块做扎实整体性能自然就上来了。另外别一上来就追求极致的稀疏率先把 K 取大一点保证质量再逐步往下压这样排查问题会容易很多。