1. 从标题到落地DSA计算过程到底在算什么第一次看到DSA计算过程这个标题加上旁边挂着的GEMM、top-K、gather、MLA这几个词我脑子里第一反应是这大概率是在讲某种稀疏注意力Dynamic Sparse Attention的完整前向计算链路。后来跟几个做推理优化的朋友聊发现大家对这个词的理解其实挺分散的——有人以为是硬件里的数字信号加速有人以为是某种分布式存储架构还有人直接把它跟MLAMulti-head Latent Attention绑在一起当成一个整体模块。所以这篇我就按最主流的那条线来拆把DSA理解成一种动态稀疏注意力的计算流程它要解决的核心问题是——在长上下文场景下标准注意力那套O(n²)的复杂度实在扛不住得想办法只算该算的那部分。这个内容适合谁看如果你正在做推理引擎的算子优化、在折腾长文本场景的显存和延迟、或者单纯想搞明白稀疏注意力到底稀疏在哪一步那这篇应该能给你一些能直接抄的东西。我会把整个计算过程拆成几个阶段先讲清楚整体设计思路和为什么这么选型再逐个环节抠细节——GEMM怎么切、top-K怎么选、gather怎么搬数据、MLA在这里扮演什么角色最后给一份实操流程和踩坑记录。全程按我自己的理解来讲不保证跟某篇论文一字不差但保证逻辑能自洽、步骤能复现。先说结论性的判断DSA的计算过程本质上是一条打分 → 选路 → 搬运 → 重算的流水线。它不是一个单一算子而是一组算子的编排。理解这一点后面所有的细节就都能挂到这条主线上。2. 整体设计与思路拆解为什么要把注意力拆成四步2.1 标准注意力的瓶颈到底卡在哪要讲DSA得先把不稀疏的版本摆出来做对照。标准多头注意力在推理阶段decode的核心计算是当前token的query向量去和历史上所有token的key向量做点积得到一组注意力分数softmax归一化后加权求和value。问题在于历史长度n一旦上到几万甚至几十万这个和所有key做点积的步骤就是纯粹的线性扫描每生成一个token都要把整个KV缓存过一遍。这里有两个成本一个是计算成本点积的量级是n×d另一个是访存成本你得把整个KV缓存从显存里读出来。在长上下文场景下访存往往比计算更致命因为GPU的算力利用率会被内存带宽拖死。我实测过一个中等规模的模型上下文拉到32K之后decode阶段的瓶颈几乎全在KV读取上SM的利用率掉到30%以下。所以思路很自然能不能不读全部KV只读一小部分这就是稀疏注意力的出发点。而DSA里的DDynamic强调的是——选哪一部分不是固定的而是根据当前query动态决定的。2.2 四步流水线的分工与选型考量把只读一部分这个目标拆开就得到了四个必须回答的问题对应四个计算阶段阶段要回答的问题核心算子主要成本打分哪些历史token值得关注GEMMquery × key计算密集选路选出top-K个top-K排序/归约搬运把选中的KV取出来gather访存密集重算对选中的做完整注意力GEMM softmax计算密集这个拆法的好处是每一步职责单一方便分别优化。但代价也很明显你为了省掉全量注意力的成本额外引入了打分、选路、搬运三步的开销。如果这三步的开销加起来超过了省下来的部分那这套方案就是负优化。这是DSA落地时最容易翻车的地方后面会专门讲。那为什么还要用GEMM来打分因为打分本质上就是query和key的矩阵乘天然适合用GEMM这种高度优化的算子。哪怕只是粗筛用GEMM也比手写循环快一个数量级。这里有个关键取舍打分用的维度可以比完整注意力低。比如完整注意力用128维的head dim打分阶段可以只取其中32维做一个近似这样打分的计算量直接降到四分之一。这个技巧在MLA里体现得特别明显。2.3 MLA在这里的角色把KV压扁MLAMulti-head Latent Attention在这条链路里的位置很多人一开始会搞混。它不是DSA的替代品而是和DSA配合的一种KV压缩方案。标准MHA里每个head都有独立的K和VKV缓存大得吓人。MLA的做法是把KV投影到一个低维的latent空间缓存的时候只存这个latent用的时候再投影回来。这对DSA的意义在于gather阶段搬运的数据量直接变小了。假设原来每个token的KV是2×num_heads×head_dim用了MLA之后可能只剩一个latent_dim压缩比能到几倍甚至十几倍。搬运量一小gather的访存瓶颈就缓解了整个稀疏化的收益才真正体现出来。所以我的理解是DSA负责少读MLA负责读得小两者叠加才能把长上下文的成本压下来。单独用任何一个效果都会打折扣。这也是为什么这几个词总是绑在一起出现。3. 核心细节解析GEMM、top-K、gather各自的坑3.1 打分阶段的GEMM怎么切才不浪费打分阶段的GEMM是[1, d] × [d, n]这种形状decode阶段query只有一行输出是[1, n]的分数向量。这种瘦长的矩阵乘在GPU上其实不太好优化因为M维度只有1很难填满tensor core的阵列。我试过几种切法实测下来比较稳的是按n分块每块内部做完整的GEMM。比如把n切成每块4096每块算出一个[1, 4096]的分数然后立刻送去做局部的top-K候选筛选。这样做的好处是分块之间可以流水线并行而且每块的中间结果不用全部驻留显存压力小。注意分块大小不是越大越好。块太大单块的延迟高流水线填不满块太小kernel启动开销占比高。我一般从2048起步调根据实际GPU的SM数量往上试。另一个细节是打分用的精度。完整注意力通常用fp16或bf16但打分阶段其实可以用更低的精度甚至int8因为这一步只是粗筛不需要精确。我做过对比打分用int8量化后top-K选出来的集合和fp16版本的重合度能到95%以上但打分本身的速度快了将近一倍。这个trade-off在延迟敏感的场景下非常划算。3.2 top-K选多少、怎么选、什么时候选top-K这一步看着简单其实坑最多。第一个问题是K取多少。K太小选出来的token不够注意力质量下降生成结果开始胡言乱语K太大省下来的计算又不够抵消额外开销。我的经验是K的取值跟上下文长度和任务类型强相关。做检索类任务答案就在某几个位置K可以小到64甚至32做长文摘要这种需要全局信息的K可能得给到512以上。第二个问题是怎么选。朴素的top-K是全局排序但全局排序的复杂度是O(n log n)在n很大的时候本身就不可接受。实际用的多是分块top-K 归并每个块内部选出局部top-K再把所有块的候选归并出全局top-K。这样复杂度降到O(n m log m)m是块数乘以局部K。第三个问题是什么时候选。这里有个反直觉的点不是每个decode step都需要重新选。相邻token的query往往很相似选出来的KV集合高度重叠。所以可以每隔几步才重新做一次top-K中间的步骤复用上一次的选择。我实测过每4步重选一次质量几乎无损但top-K的开销直接降到四分之一。这个技巧在top-k 路由这个说法里其实就隐含了——路由不是每步都重算的。3.3 gather访存密集环节的优化空间gather是把top-K选中的KV从缓存里捞出来组成一个紧凑的[K, d]矩阵供后面的完整注意力使用。这一步是纯访存操作没有计算所以优化方向只有一个减少搬运的字节数提高访存效率。具体手段有几个。一是前面说的MLA压缩从源头减少每个token的字节数。二是让选中的KV在显存里尽量连续。如果KV缓存是按token顺序线性排布的那gather就是随机访问效率很低。有些实现会维护一个热KV的紧凑缓冲区把最近频繁被选中的token挪到一起gather的时候命中率更高。三是用向量化的load一次读128位或256位而不是逐个元素读。实操心得gather的耗时经常被低估。我在一个项目里做过profile发现gather占了整个DSA流程将近40%的时间比打分还高。后来把KV缓存改成按head分块、块内连续gather时间直接砍半。所以别光盯着GEMM优化访存这块的收益往往更大。3.4 重算阶段为什么不能直接用打分的结果有人会问既然打分阶段已经算出了分数为什么不直接softmax然后加权求和还要再算一遍完整注意力原因是打分用的维度和精度都不够。打分阶段可能只用了32维、int8精度得到的分数只是个近似。如果直接拿这个近似分数去做加权注意力的质量会明显下降。所以标准做法是用近似分数选出top-K的位置然后对这些位置用完整的维度、完整的精度重新算一遍注意力。这一步的GEMM形状是[1, d] × [d, K]因为K远小于n所以即使算得精确总成本也远低于全量注意力。这个先粗后精的两阶段设计是DSA能work的关键。粗筛负责把候选集从n降到K精算负责保证质量。两阶段的维度、精度、甚至数据类型都可以不一样给了很大的调优空间。4. 实操过程一条可复现的DSA计算链路4.1 环境与数据准备假设我们有一个已经训练好的模型KV缓存按标准方式存着。我们要实现的是decode阶段单步的DSA计算。输入是当前token的hidden state输出是注意力结果。为了讲清楚我用伪代码加注释的方式把每一步的参数和意图都标出来。先定义几个关键参数n是历史长度d是head dimd_score是打分用的维度通常小于dK是top-K的数量block_size是分块大小。这些参数没有万能值得根据实际场景调下面给的是一组我常用的起点。n 32768 # 历史长度 d 128 # 完整head dim d_score 32 # 打分维度取完整维度的1/4 K 256 # top-K数量 block_size 4096 # 分块大小4.2 第一步低维打分从KV缓存里取出key的低维部分前d_score维和query的低维部分做GEMM。注意这里只读key的前d_score维不是全部所以访存量也小。# query_low: [1, d_score] # key_low_cache: [n, d_score]只存了低维部分 scores gemm(query_low, key_low_cache.T) # 输出 [1, n]这一步的输出是一个长度为n的分数向量。如果n是32768d_score是32那这个GEMM的FLOPs是32768×32×2≈2M非常小。真正的开销在读取key_low_cache上32768×32×2字节fp16≈2MB这个量级在现代GPU上也就几十微秒。4.3 第二步分块top-K与归并把scores按block_size切块每块内部选局部top-K再归并。candidates [] for i in range(0, n, block_size): block scores[i:iblock_size] local_topk_idx topk(block, kK // 4) # 每块选K/4个候选 candidates.extend([i idx for idx in local_topk_idx]) # 归并所有候选选出全局top-K global_topk_idx topk(gather(scores, candidates), kK)这里每块选K/4是个经验值。块数大概是n/block_size88×64512个候选再从中选256个冗余度足够不会漏掉真正重要的token。如果块数更多每块的局部K可以相应调小。注意局部K不能太小否则真正的高分token可能在局部就被筛掉了。一般保证局部K至少是全局K的1/8到1/4。4.4 第三步gather选中KV拿到global_topk_idx之后把对应的完整KV取出来。# kv_cache: [n, 2, num_heads, d]完整精度 selected_kv gather(kv_cache, global_topk_idx) # [K, 2, num_heads, d]如果用了MLA这里的kv_cache存的是压缩后的latentgather之后再投影回完整的K和V。搬运量能小好几倍。4.5 第四步完整注意力重算对选中的K个位置做标准注意力。# query_full: [1, num_heads, d] # selected_k: [K, num_heads, d] # selected_v: [K, num_heads, d] attn_scores gemm(query_full, selected_k.transpose()) # [num_heads, 1, K] attn_probs softmax(attn_scores, dim-1) output gemm(attn_probs, selected_v) # [num_heads, 1, d]这一步的GEMM形状是[num_heads, 1, d] × [num_heads, d, K]K256时计算量很小。整个DSA流程的总耗时实测下来比全量注意力在n32K时能快3到5倍具体取决于K的取值和硬件。4.6 参数选择的计算过程K到底取多少可以用一个简单的成本模型来估。设全量注意力的成本是C_full n × d忽略常数DSA的成本是C_dsa n × d_score K × d gather_overhead。要让DSA划算需要C_dsa C_full即n × d_score K × d gather_overhead n × d代入n32768, d128, d_score32假设gather_overhead约等于K×d的0.5倍32768×32 K×128×1.5 32768×128 1048576 192K 4194304 192K 3145728 K 16384这个上界很宽松说明在n32K时K取到几千都还是划算的。但实际中K取太大注意力质量反而不一定好引入了太多噪声而且gather的随机访问效率会下降。所以K的实际取值更多是质量驱动的不是成本驱动的。我一般从128起步根据任务效果往上加。5. 常见问题与排查技巧实录5.1 生成结果突然变差怎么定位是不是DSA的锅这是最常见的问题。排查思路是逐阶段对比先把DSA关掉跑全量注意力记录输出再开DSA对比同一位置的输出差异。如果差异很大再逐个阶段查。现象可能原因排查方法输出完全乱码top-K选错了位置打印选中的idx看是否集中在某段输出轻微退化K太小或打分精度太低逐步增大K、提高打分精度长文后半段变差分块top-K漏掉了早期重要token增大局部K或改用全局top-K延迟没降反升gather开销超过预期profile各阶段耗时重点看gather我踩过最坑的一次是分块top-K的局部K设得太小导致一个关键token在它所在的块里排第9而局部K只取8直接被筛掉了。结果模型在需要引用开头信息的任务上全军覆没。后来把局部K调到16就正常了。这个bug很隐蔽因为大部分任务看起来都正常只有特定任务才暴露。5.2 gather效率低下的几种典型表现gather慢通常不是算子本身的问题而是数据布局的问题。几个典型症状和对策症状gather耗时随n线性增长但K没变。原因KV缓存是线性排布的选中的idx越分散随机访问越多。对策维护紧凑的热KV缓冲区。症状gather耗时波动很大。原因选中的idx分布不均匀有时集中有时分散。对策对idx做排序后再gather把随机访问变成顺序访问。排序的开销远小于随机访问的惩罚。症状gather后紧接着的GEMM变慢。原因gather出来的数据在显存里不连续GEMM的load效率低。对策gather时直接写到一块连续的临时缓冲区。实操心得对top-K的idx做一次排序看起来多了一步但实测能把gather加后续GEMM的总时间降20%以上。这个技巧在很多实现里都被忽略了。5.3 MLA压缩带来的精度损失怎么补MLA把KV压到低维必然有信息损失。补偿手段有几个一是提高latent维度压缩比别拉太满二是在重算阶段用完整的投影矩阵把latent还原得尽量准三是对压缩误差做校准比如在训练时加入重建损失。我的经验是压缩比控制在4到8倍之间比较稳再往上精度就开始明显掉了。而且不同层的敏感度不一样浅层可以压得狠一点深层要保守。这个得逐层试没有统一公式。5.4 常见问题速查表问题根因解决top-K选出的位置每步剧烈跳变query变化大或打分噪声大加平滑或降低重选频率显存不降反升候选缓冲区开太大复用缓冲区别每步新分配多卡并行时结果不一致gather的idx跨卡了确保KV缓存按卡分片idx本地化首token延迟特别高首次要建完整KV缓存预填充阶段就做好分块和索引6. 几个容易被忽略的工程细节6.1 重选频率与质量的平衡前面提过不必每步都重选top-K。但重选频率怎么定有个简单的判据看相邻两步query的余弦相似度。如果相似度高于0.95基本可以复用低于0.9就该重选了。实际实现里可以固定步数比如每4步也可以动态判断。动态判断更准但多了一次相似度计算得权衡。我一般用固定步数因为实现简单、延迟可预测。在延迟敏感的服务里可预测性比那一点点质量提升更重要。6.2 预填充阶段的处理DSA主要优化的是decode阶段但预填充prefill阶段也不能忽视。预填充时query是一整段不是单行所以打分阶段的GEMM形状是[seq_len, d] × [d, n]计算量大得多。这时候稀疏化的收益更明显但top-K的选择也更复杂——每个query位置都要选自己的top-K。一个简化做法是预填充阶段先用全量注意力把KV缓存建好同时记录下每个位置的重要token集合供后续decode复用。这样预填充的质量有保证decode又能享受稀疏化的好处。6.3 数值稳定性稀疏化之后softmax的输入分布会变。因为只对选中的K个位置做softmax如果这K个位置的分数普遍偏低softmax后的分布会很平输出接近均匀加权质量下降。解决办法是在打分阶段保留一个全局的max分数重算时用它做数值稳定而不是只用选中集合的max。这个细节很小但不做的话在某些输入上会出现明显的质量抖动。7. 我对这套流程的整体体会把DSA的计算过程拆完我最大的感受是它的难点不在单个算子而在算子之间的配合。GEMM、top-K、gather单独看都是成熟技术但拼在一起之后参数怎么传、数据怎么布局、精度怎么匹配每一步都有讲究。我见过不少实现单个算子都优化得很好但整体跑起来就是不如预期问题往往出在衔接处——比如gather出来的数据布局不适合后面的GEMM或者打分精度和重算精度不匹配导致质量抖动。另一个体会是稀疏化的收益高度依赖场景。在短上下文n4K下DSA的额外开销可能超过收益不如老老实实用全量注意力。只有在长上下文、且任务本身确实只依赖局部信息时DSA才真正划算。所以别一上来就无脑上稀疏先profile看瓶颈到底在哪。最后分享一个我常用的调试技巧把top-K选中的位置可视化出来。把每步选中的idx画成一条时间线上的点一眼就能看出选择是否合理——如果点集中在最近的一段说明模型在只看最近如果点均匀散布说明它在做全局检索。这个可视化帮我定位过好几次选择逻辑的bug比看数字直观得多。这套流程后续还能往几个方向扩展一是把top-K换成更灵活的路由机制比如按内容聚类而不是按分数排序二是把MLA的压缩和DSA的选择联合训练让压缩后的latent本身就编码了哪些位置重要的信息三是在多卡场景下做分布式的top-K每张卡先选本地的再跨卡归并。这几个方向我都在试有进展再单独写。
