从Stem稀疏注意力看 大模型推理的架构取舍

ARCHITECTURE INSIGHT

从Stem稀疏注意力看 大模型推理的架构取舍

为什么大多数稀疏注意力停在论文里，而Stem跑通了全栈？答案不在算法本身，在架构思维。

25%

算力占比

3.7×

首字延迟降低

0.54%

精度损失

稀疏注意力有一个反常识的真相：最难的不是"算出哪些token可以跳过"，而是"让跳过的部分在GPU上真的变快"。

前者是算法问题，发论文足够了；后者是工程问题，决定了你的方案能不能跑在生产环境里。过去两年，稀疏注意力领域的论文井喷——思路各异，但大多数共享一个尴尬的处境：论文里跑出30%的加速比，部署到真实GPU上，纹丝不动。

这不是个例，是常态。原因很简单：稀疏化打乱了GPU最擅长的连续内存访问模式，跳块判断本身引入了额外开销，计算粒度又填不满Tensor Core。算法层面省下的算力，被工程层面的开销吃干抹净。

所以当腾讯混元团队说Stem用25%的计算量逼近稠密注意力精度、128K上下文首字延迟降低3.7倍的时候，一个架构师的第一反应不应该是"算法真厉害"，而应该是——"他们怎么把理论加速比兑现成实测性能的？"

这个问题的答案，才是Stem真正值得聊的东西。

跑不通的墙

要理解为什么大多数稀疏注意力方案跑不通，得先理解GPU是怎么干活的。

GPU的内存层次像一座金字塔：顶层是寄存器，极快但极小；中间是SRAM（片上共享内存），快但有限；底层是HBM（显存），大但慢。这座金字塔对连续、规则的访问模式高度优化——你顺序读取一整块数据，它能预判你的意图，提前把数据搬到高速缓存里，计算单元几乎不用等。

稠密注意力天然就是连续的：一个注意力头挨个处理KV Cache里的所有token，内存访问整整齐齐，Tensor Core满载运行。但稀疏注意力打破了这个规则——你跳过了某些块，KV Cache的访问就变成了从分散的内存位置gather数据。缓存命中率骤降，内存带宽利用率低下，实际延迟远高于理论预测。

第一堵墙

内存访问模式不规则

第二堵墙更隐蔽。GPU的Tensor Core以固定粒度执行矩阵运算——16×16等矩阵块。稀疏化之后，实际需要计算的块可能填不满一个Tensor Core计算单元，计算单元利用率低，kernel launch开销占比反而上升。更麻烦的是，"跳块"判断本身不是免费的——你得先判断哪些块该跳过，这个判断也是计算，也有开销。

第二堵墙

计算粒度不匹配

两堵墙叠加，结果就是：算法层面省下30%的计算量，工程层面的额外开销吃掉25%，净收益5%——还要考虑精度损失，可能连5%都不值。

行业里有个更极端的案例。MInference的思路是动态选择稀疏模式（Vertical-Slash等），在特定设置下精度确实不错。但它的模式估计本身就有显著开销，在16K/32K短序列上，延迟反而比稠密注意力更高。换句话说，你花了力气做稀疏化，结果比什么都不做还慢。它的精度优势还伴随着另一个代价：预算需求高达55%到81%——这已经不是"稀疏"了，更像是"半稠密"。

另一个普遍但容易被忽视的问题是：大多数稀疏注意力论文报告的是BF16精度下的学术基线数据。但工业部署几乎都用量化模型（W8A8-FP8），量化后的精度表现和BF16不是一回事。论文里漂亮的数字，到了生产环境可能直接崩塌。

所以问题不是"稀疏注意力能不能省算力"——理论上当然能。问题是：理论稀疏比能不能在真实GPU上变成端到端实测加速？对大多数方案来说，答案是不能。

两堵墙的根源是同一个：稀疏注意力方案在设计时没有把GPU的执行模型当约束。Stem的解法不是绕墙，而是从算法设计阶段就让方案天然适配GPU——先判断谁重要（TPD和OAM），再让"跳过不重要的人"在GPU上真正变快（HPC算子）。接下来的"道"讲前者，"术"讲后者。

用架构师的眼睛看稀疏注意力

Stem的名字来自"树干"（Stem）的隐喻。在因果注意力架构中，初始位置的token如同一棵树的主干——第一个Value向量参与了每一个输出的计算，而最后一个Value向量只参与最后一个输出。如果剪掉第一个token，误差会逐层递归累积，传播到所有位置；剪掉最后一个，只影响局部。这不是直觉，是因果约束下的数学必然。

理解了这一点，再看Stem的两个核心设计——TPD和OAM——就会发现它们不是什么新鲜的算法技巧，而是架构师再熟悉不过的决策模式。

TPD：优先级排序

现有稀疏注意力方法有一个共同的隐含假设：序列中每个位置分配相同的Top-k预算。均匀分配，一刀切。

这在架构设计里有个对应的错误：所有请求等量对待。任何一个做过高并发系统的人都知道这是反模式——关键路径上的请求和边缘请求怎么能享受同等资源？于是有了优先级队列，有了SLA分级，有了热数据缓存。

TPD做的事情一模一样。它不再均匀分配预算，而是从初始位置的最大预算k_start线性衰减到末尾的k_end。初始token获得最多资源，因为它们的剪枝误差会逐层放大；末尾token逐步缩减，因为误差只影响局部。关键点：TPD不增加总计算量，而是在相同预算下重新分配资源。

最优衰减比μ=0.7——这意味着末尾位置保留初始位置70%的预算。这跟架构师熟悉的80/20法则是同一件事：20%的token承载了80%的信息流权重，资源倾斜给它们才是效率的来源。

消融实验 · Qwen3-8B

均匀分配 + 传统评分　29.41% + TPD　　　　　　　　 31.43%↑2.02 + TPD + OAM　　　　　 31.64%↑0.22

不花更多钱，只是把钱花在刀刃上。

OAM：价值评估

传统方法用注意力分数（QK^T）来挑token——谁的分数高，谁留下。但分数高不等于贡献大。一个token可能和Query高度相关（高分），但Value向量幅值接近零，对输出几乎无贡献；反过来，分数中等但Value幅值很大的token，才是真正的信息富矿。

这又是架构师熟悉的场景：请求调度不能只看"谁在等"（排队时间），还得看"等了值不值"（请求价值）。一个耗时0.1ms的缓存命中请求和一个耗时500ms的复杂计算请求，不能因为前者"等得久"就优先处理它。

OAM把这两维信息合成一个评分：路由相关性（QK^T，谁跟我有关）加上信号幅值（log‖V_j‖₂，有关的信息量有多大），用β=0.2的系数调节两者权重。β不能太高——超过0.3，幅值噪声会压过语义相关性，就像调度系统里给所有低优先级任务都加了权重，反而把关键路径上的资源挤占了。

加上OAM之后，消融实验从31.43%提升到31.64%。看起来提升不大，但这是在TPD已经拿走大头之后再加的——架构师知道，优先级排序解决80%的问题，价值评估解决剩下的20%，两者叠加才是完整的决策框架。

25%算力不是"压缩"，是"取舍"

把TPD和OAM放在一起看，Stem用25%的算力逼近稠密注意力精度（LongBench上仅差0.54%），这件事的本质不是压缩，是取舍——架构师最核心的能力。

压缩是"用更少资源做同样的事"，暗示有损但尽量无痛。取舍是"明确什么重要什么不重要，把资源押在重要的地方"。TPD回答了"谁重要"（初始token），OAM回答了"怎么判断重要"（相关性+信息量），两者加在一起，25%的预算就够了。

这跟缓存淘汰策略的思路完全一致：LRU/LFU不是在"压缩"缓存，是在判断哪些数据值得保留。判断对了，小缓存也能扛住大流量；判断错了，缓存再大也是浪费。

把加速比兑现成毫秒

"道"讲的是Stem怎么判断谁重要，"术"要讲的是——判断完了，怎么让GPU真的快起来。

前面说过，稀疏注意力跑不通的两堵墙是内存访问不规则和计算粒度不匹配。Stem的解法不是绕过这两堵墙，而是在算法设计阶段就把它们当约束条件，让算法本身天然适配GPU的执行模型。具体来说，是两个算子的事。

HPC-Stem：评估选块，快几十倍

稀疏注意力推理分两步：先评估哪些块值得算（选块），再只算选中的块（执行）。选块阶段往往是隐藏的性能杀手——你得先算一遍评分，才能决定跳过谁，这个"先算一遍"本身就有开销。

传统方法的评分依赖softmax归一化，这带来两个问题：一是额外的gather操作，二是FP8精度下softmax的数值误差。更致命的是，128K上下文下中间张量可达数十GB——光评估就先把显存吃满了。

OAM在这里露出了一手精心设计的工程优势。它的评分公式是纯加法结构（路由相关性+信号幅值），不依赖softmax和gather。这个看似微小的数学性质，在GPU上产生了一个巨大的工程效果：原本需要对采样后的Q/K做全量矩阵乘法产生的中间张量，可以等价转化为先预计算Q和K各自的紧凑block级别表示，再用一次标准GEMM直接得到全部block评分——计算量大幅降低，中间张量完全消除。

这不是事后优化——OAM的纯加法结构在算法设计阶段就决定了它能享受这个优化。换个说法：如果OAM的评分公式里多一个softmax，这些优化就都不成立。

TPD的选块逻辑也做了类似的算子融合：把预算生成和选块排序融合为一个算子，省掉中间的显存读写。这些优化叠加起来，HPC-Stem的评估速度比朴素实现快数十倍。

HPC-BSA：跳块零开销的秘密

选块搞定了，接下来是执行——只算选中的块。这是"纸面加速30%实测纹丝不动"的重灾区。

MIT的原版BSA算子是这样工作的：内层循环里，每遇到一个KV分块，先判断"这个块要不要跳过"，如果要跳过就跳过，如果不要就计算。听起来很合理，但问题是——逐次判断本身就是开销，而且跳过和不跳过的代码路径不同，GPU的流水线没法高效调度。

HPC-BSA换了一个思路：kernel处理每个Q分块时，先将block mask缓存到片上高速存储，在线构建一个"需要计算的KV分块列表"。内层循环只遍历这个列表里的有效分块，完全避免逐次判断。被跳过的分块在数学上等价于注意力分数全为负无穷，不影响softmax正确性。

关键来了：内层计算路径与稠密Attention几乎一致。GPU跑稀疏注意力和跑稠密注意力的代码路径是一样的，只是遍历的分块少了。没有分支，没有判断，没有不规则访问。稀疏带来的加速是"白赚"的——跳块额外开销控制在极低水平。

数据说话：HPC-BSA的延迟与计算密度呈近乎完美的线性关系——50%稀疏度下延迟约为稠密基线一半，80%稀疏度下仅约五分之一。相比MIT原版BSA算子，全稀疏度范围内稳定保持显著加速（FP8计算吞吐加Hopper架构优化叠加），从短序列到超长序列优势保持稳定。

还有一个容易被忽略但极其重要的工程决策：HPC-BSA原生支持vLLM的Paged KV Cache。Paged KV Cache是vLLM的显存管理机制，把KV Cache分成固定大小的页，按需加载——它本身就是非连续的、逐页访问的。块级稀疏的"不规则访问"在Paged KV Cache的语境下反而变成了"天然适配"——你本来就不是连续读的，跳过几页有什么关系？这不是巧合，是对部署环境的深度理解。

答案不在算法里

回到开头的问题：为什么大多数稀疏注意力算法停留在论文里，而Stem跑通了全栈？

一个直觉性的回答是"算法更好"。但仔细看消融实验的数据：TPD比均匀分配提升了2.02个百分点，OAM在TPD基础上又提升了0.22个百分点。算法层面的改进是实的，但幅度并不夸张。真正拉开差距的，不是TPD或OAM某个具体算法，而是一种贯穿始终的决策方式。

把它摊开来看：

OAM用纯加法结构做评分，不是因为它比softmax更优雅，而是因为纯加法结构天然适配GEMM优化——计算量大幅降低，中间张量完全消除。如果评分公式里多一个softmax，这些优化全都不成立。换句话说，OAM在设计之初就把"GPU能不能高效算出来"当成了约束条件。

TPD的线性衰减预算，不只是"初始token更重要"这个洞察的数学表达。线性衰减意味着预算生成逻辑极其简单——一个公式就能算出每个位置的预算，可以跟选块排序融合成一个算子。如果衰减函数是非线性的，这个融合就没那么干净。

HPC-BSA把稀疏判断前置到循环外，让内层计算路径与稠密Attention一致——跳块额外开销极低。这不是算子工程师的灵光一现，而是算法层面的块级稀疏设计（固定块大小B=128）给了算子工程师"前置判断"的空间。如果稀疏粒度是token级的，每个token单独判断，前置就无从谈起。

直接在W8A8-FP8量化模型上验证，而不是在BF16学术基线上报数字——这意味着算法设计时就把量化精度损失考虑进去了，而不是先在BF16上跑出漂亮结果，再发现量化后精度崩塌。

这些决策有一个共同模式：

从设计第一天就把"能落地"当约束条件，而不是先算出漂亮数字再去想怎么部署。

传统路径是算法团队先出方案，HPC团队再想办法实现。两个团队各自优化自己的目标函数——算法团队优化精度和稀疏比，HPC团队优化算子性能。但"理论加速比在GPU上蒸发"这个问题，恰好落在两个团队的交界处，谁都不负责。

Stem的路径是算法和算子协同设计。OAM的纯加法结构、TPD的线性衰减、块级稀疏的固定粒度——每一个算法决策背后都有工程约束的影子。不是"先设计再优化"，而是"设计本身就是为工程环境量身定制的"。

这不是Stem独有的发现，而是架构设计里一个古老的原则：约束不是敌人，是设计输入。把部署约束写进算法设计，就像把性能要求写进架构评审——不是限制创造力，是让创造力花在对的地方。

架构师的判断框架

所以，下次你看到一个AI推理优化方案宣称"N倍加速"，怎么判断它是真能落地还是只活在论文里？

三个问题就够了。

第一，它报的数据在什么条件下测的？

BF16学术基线和W8A8-FP8工业部署是两个世界。如果只报BF16数据，要么是还没做到量化适配，要么是量化后精度崩塌了不敢报。Stem直接在FP8量化模型上验证，LongBench精度损失仅0.54%——这个数字的可信度远高于BF16下的任何数字。

第二，理论加速比和实测加速比的差距有多大？

如果论文只报"跳过了70%的计算"，不报端到端延迟，大概率是因为端到端数字不好看。Stem的HPC-BSA跳块额外开销极低，延迟与计算密度呈近乎完美线性关系——50%稀疏度下延迟约为稠密基线一半，这才是"理论兑现成实测"的样子。

第三，算法设计和工程实现是两拨人各干各的，还是从一开始就协同？

各干各的，算法团队优化精度和稀疏比，HPC团队优化算子性能，中间的缝隙就是理论加速比蒸发的地方。协同设计，算法的每一个决策都带着工程约束的影子——OAM的纯加法、TPD的线性衰减、块级稀疏的固定粒度，都是这种协同的产物。

三个问题背后是一个更根本的判断原则：

看这个方案是先设计再想部署，还是从第一天就把部署约束写进设计。前者大概率停在论文里，后者才有机会跑在生产环境。

Stem不是唯一走通这条路的方向，但它提供了一个足够清晰的范例——用25%的算力、3.7倍的首字延迟降低、0.54%的精度损失，证明了稀疏注意力从论文到生产之间那道鸿沟，不是不可逾越的。而逾越它的关键，不在算法本身，在架构思维。

开源地址：github.com/Tencent/AngelSlim