腾讯混元 2026-06-05 19:23 北京

告别延迟

将一篇数万字的超长文档输入大语言模型并点击发送后，你是否经历过光标闪烁、模型迟迟不输出首个字符的漫长等待？这段“等待第一个字”的耗时阶段，正是业内常说的 预填充（Prefill）阶段。

其核心瓶颈源于 Transformer 自注意力机制的二次方计算复杂度：输入上下文越长，预填充速度越慢，且呈现平方级增长。

稀疏注意力（Sparse Attention） 当前被视为最主流的突破方向：策略性地只计算“关键”token，从而节省冗余算力。

然而，现有方案在算法与算子层面均存在明显短板。算法上：一是对所有位置“一刀切”地分配同等稀疏预算，忽视了因果架构中初始token的递归依赖特性；二是仅依赖注意力分数筛选token，忽略了Value向量本身承载的信息量。算子层还横亘着一道隐形障碍：无论稀疏模式设计得多巧妙，若底层算子无法高效跳过被稀疏丢弃的块，加速比终将大打折扣。

针对这些痛点，腾讯混元提出 Stem 稀疏注意力算法，该成果已被机器学习顶会 ICML-26 收录——它从“因果信息流”的视角重构块级稀疏，借助 Token 位置衰减（TPD） 和 输出感知度量（OAM） 两大创新，仅用 25% 算力即可逼近稠密注意力的精度。配套的 HPC 算子库则成功将理论加速比落地为端到端的实测性能。

• Stem论文链接：https://arxiv.org/abs/2603.06274Stem

• 开源地址：https://github.com/Tencent/AngelSlimHPC

• 算子开源地址：https://github.com/Tencent/hpc-ops

Stem 算法：从“信息流”维度重新定义稀疏注意力

1. 核心洞察：初始 token 是信息流的“树干”

Stem 的名称源于“树干”这一比喻——在因果注意力架构中，序列初始位置的 token 好比一棵树的树干，支撑着后续所有信息的传递与汇聚。

为何初始 token 如此关键？ 在因果自注意力机制中，每层计算严格遵守因果约束：第 $l$ 层、第 $n$ 个位置的输出，由该层前 $n$ 个位置的Value向量经加权聚合而来（$O_n^{(l)} = sum_{j=1}^{n} P_{n,j}^{(l)} V_j^{(l)}$），而这些Value向量又源于上一层输出经线性映射得到。这意味着：

·位置1的Value $V_1^{(l)}$：会出现在本层每一个位置的输出计算中（$O_1^{(l)}, O_2^{(l)}, dots, O_N^{(l)}$）

·位置N的Value $V_N^{(l)}$：仅参与本层最后一个位置 $O_N^{(l)}$ 的计算

这种不对称性在多层级Transformer网络中会被递归放大。初始 token 的信息经由 $V^{(l)} \to O^{(l)} \to V^{(l+1)}$ 的路径逐层传递，最终渗透至深层网络的每一个表示中。如下图所示，一旦在某一层稀疏化操作丢弃了初始 token $V_1$，误差将沿红色虚线路径扩散至下一层的所有 token；而若丢弃末尾 token $V_N$，其影响则仅局限于最后一个位置（红色实线）。

2. 实战验证：Stem 全栈加速效果

剖析完“初始 token 为何至关重要”后，你可能会追问：这个发现，在生产环境中究竟能跑出怎样的性能数据？

我们并未止步于学术 Benchmark，而是直接将 Stem 集成进腾讯混元 Hy3 preview（W8A8-FP8）的 vLLM 推理框架，并搭配 HPC 团队优化的 Stem 算子，端到端实测首字延迟（TTFT）与模型精度——这要求 Stem 不仅能在 BF16 学术基线上顺利运行，更能在量化后的工业级模型上保持稳定。

至于 Stem 与其他稀疏算法在开源模型上的速度/精度对比，论文中已给出完整结果，此处不再赘述。

下面重点展示 Stem 在Hy3 preview（W8A8-FP8）上、更贴近真实生产环境的落地数据。

2.1 首字加速比

2.2 模型精度

3. 揭秘：Stem 如何同时实现“快”与“准”？

性能数据已然摆出，关键问题随之而来——Stem 如何在削减 75% 计算量的同时，依然保持精度不降？

答案隐藏在两个看似简单、却被以往研究长期忽视的细节里：预算分配应偏向谁？挑选 token 时该关注什么？Stem 的两大核心创新——Token 位置衰减（TPD） 和 输出感知度量（OAM）——分别精准回应了这两个问题。

3.1 Token 位置衰减策略（TPD）：重新分配预算，而非增加预算

现行方法为所有位置施加统一的Top-$k$ 预算——这是一种“一刀切”策略，完全忽略了前述的位置敏感性差异。

Stem提出线性衰减的预算分配方案：

·初始位置：分配较大预算 $k_{\text{start}}$，充分保留递归依赖链上的关键token

·末尾位置：预算线性衰减至 $k_{\text{end}} = \mu \times k_{\text{start}}$（默认 $\mu=0.7$），对冗余信息实施激进剪枝

关键在于：TPD并未增加总计算量，而是在相同预算下重新调配资源，将计算集中投放到信息流的关键节点上。

Stem整体流程图。（a）为传统Uniform Top-$k$ 的均匀预算分配，（b）为TPD的衰减式分配，直观展示了预算如何向初始位置倾斜。

3.2 输出感知度量（OAM）：不只考察“路由分数”，更衡量“信息贡献”

预算分配确定后，下一个问题是：在预算范围内，到底选择哪些 token？

传统方法仅使用注意力分数（Query 与 Key 的点积）来筛选 token，但 Stem 指出一个常被忽视的真相：分数高不等同于实际贡献大——一个 token 可能与 Query 高度相关，但其 Value 向量幅值接近零，对输出几乎无贡献；相反，分数中等但 Value 幅值较大的 token，才是真正的“信息富矿”。

注意力的本质是加权求和，一个 token 的真实贡献实际上等于 路由概率 × 信号幅值。Stem 据此提出输出感知度量（OAM）：

通过对数变换将“乘法”转为“加法”，排序结果保持不变，同时能直接套用标准 Top-$k$ 算子，实现近乎零开销的融合。

至此，Stem 算法层面的全貌已清晰呈现：TPD 决定“给谁更多预算”，OAM 决定“在那部分预算里挑选谁”，两者互为补充，将稀疏注意力从“一刀切”升级为“依据信息流动态决定稀疏预算”。

然而，算法选得再精准，最终能否在 GPU 上真正跑出 3.6× 加速比，关键还在于底层算子能否充分配合。下一节，我们将深入介绍HPC 算子。

算子优化：HPC-Stem + HPC-BSA

现有算子实现的挑战

块级稀疏注意力的计算流程分为两个阶段：评估选块（评估每个block的重要性并记录结果）和稀疏执行（依据评估结果跳过不重要的block，仅对被选中的块执行注意力计算）。在现有生态中，这两个阶段均存在瓶颈。在评估选块阶段，当前主流方法依赖softmax归一化，这会引入额外的gather操作和FP8 GEMM精度误差，并需维护较大的中间张量（128K下可达16 GB）。在稀疏注意力计算过程中，已有的开源BSA算子需要动态判断哪些块应该跳过，由此带来较为显著的跳块开销。

HPC优化：我们采取了哪些措施

针对上述两大瓶颈，我们分别设计了两款核心算子：HPC-Stem将评估选块流程加速数十倍，HPC-BSA面向Hopper架构将稀疏算法的理论加速比真正落地。

1）Stem算子优化：

HPC-Stem的评估流程分为OAM评分和TPD选块两步。针对OAM评分，我们发现其纯加法的度量结构（不依赖softmax和gather）允许一个关键的数学简化：原始实现中对采样后的Q/K执行全量矩阵乘法所产生的中间张量，可以等价转化为先预计算Q和K各自的紧凑block级别表示，再通过一次标准GEMM直接获得全部block评分，计算量降低约64倍，中间张量被完全消除。针对TPD选块，我们将预算的生成与选块排序的逻辑融合为单一算子，显著提升了评估速度。

2）BSA算子优化：

HPC-BSA面向Hopper架构从零设计，采用数据搬运与计算并行的流水线架构，原生支持vLLM的Paged KV Cache和FP8量化（计算吞吐量翻倍）。

在此基础上，块级稀疏的支持被自然融入。如上图所示，经过分页的KV Cache本身就是非连续的、逐页加载的。针对稀疏场景，kernel在处理每个Q分块时，首先将对应的block mask缓存到片上高速存储，并在线构建出需要计算的KV分块列表。内层循环仅遍历这些有效分块，完全避免了逐次判断跳过带来的额外开销。被跳过的分块在数学上等价于注意力分数全为负无穷，不影响softmax的正确性，计算逻辑因而无需任何修改。

如上图所示，相比MIT原版算子中逐步查询索引、计算偏移的跳块流程，HPC-BSA将稀疏的判断和筛选完全前置到循环之外，内层计算路径与稠密Attention几乎一致，实现了近零开销的块级跳过。

Benchmark：算子级性能

我们对HPC-BSA算子进行了性能测试，以HPC-Dense (FP8)和FlashAttention V3 (FP8)作为稠密基线，以MIT-BSA (BF16)和FlashPrefill-BSA (BF16)两个主流开源实现作为稀疏对照。

测试结果揭示了三个关键发现。第一，HPC-BSA的延迟与计算密度呈现近乎完美的线性关系：50%稀疏度下延迟约为稠密基线的一半，80%稀疏度下仅为约五分之一，跳块机制的额外开销控制在2.5%以内，算法的理论加速比被几乎完整地转化为实测性能。第二，HPC-BSA相比MIT开源的BSA算子（MIT-BSA）在全稀疏度范围内稳定保持约3倍加速，该收益源于FP8计算吞吐优势与Hopper架构优化的叠加。第三，上述优势在8K到256K的全序列长度范围内保持稳定，展现出良好的长序列扩展性。

不同稀疏度 (block sparse ratio) 下，BSA算子延时的变化：

不同稀疏度下 HPC-BSA (FP8) 的延迟与稠密基线对比。标注百分比为 HPC-BSA 延迟占 HPC-Dense 的比例。淡红虚线为理论最快速度上限（HPC-Dense × density）。

64K输入下，不同稀疏度 (block sparse ratio) 下HPC-BSA算子与主流开源BSA算子加速比，HPC-BSA相比MIT-BSA(BF16)在全稀疏度范围内稳定保持约3.1倍加速，相比于FlashPrefill (BF16) 稳定~2.1倍加速

75% 稀疏度下，不同序列长度下 HPC-BSA相比主流开源BSA算子的加速比，相较于两个基线同样保持了较稳定的效果

总结与展望

本文系统介绍了 Stem算法 × HPC算子 的全栈加速方案：算法层面，Stem通过Token位置衰减（TPD）和输出感知度量（OAM）在25%预算下实现近无损精度；算子层面，HPC开源的Stem+BSA算子将稀疏收益转化为真实硬件加速，128K上下文下首字延迟降低3.7倍。

算法决定了“省去哪些计算”，算子决定了“省下的计算能快多少”——两者协同，构成了从理论推演到部署落地的完整闭环。

伴随大模型上下文窗口持续向百万级别扩展，高效的长文本推理将成为刚需。我们期待这一“算法+算子”的协同思路能为长文本高效推理提供新的范式参考。

来源：互联网