引言

想象一下这个场景：你把一篇数万字的文档丢给大语言模型，点击发送，然后……光标开始闪烁，模型迟迟不吐出第一个字。这段 "等待第一个字" 的过程，就是预填充（Prefill）阶段。它背后的瓶颈，其实来自 Transformer 自注意力机制的二次方计算复杂度——输入越长，预填充越慢，而且是平方级增长。

为了解决这个问题，稀疏注意力（Sparse Attention）成为了当下最主流的破局方向：只计算那些"重要"的 token，把不必要的算力省下来。不过，从算法到算子，现有的方案都存在着明显短板。

算法层面主要有两个问题：一是对所有位置"一刀切"地分配相同的稀疏预算，忽视了因果架构中初始 token 的递归依赖特性；二是只看注意力分数来挑选 token，忽略了 Value 向量本身携带的信息量。算子层面还有一道隐形门槛——再聪明的稀疏模式，如果底层算子无法真正高效地跳过被丢弃的块，加速比也会大打折扣。

今天要聊的，是腾讯混元 AI Infra 团队提出的 Stem，它已被机器学习顶会 ICML-26 收录。Stem 从"因果信息流"的角度重新审视块级稀疏，用两大创新——Token 位置衰减（TPD）和输出感知度量（OAM）——仅用 25% 的算力就逼近了稠密注意力的精度。同时，配套的 HPC 算子库将这份理论加速真正转化为了端到端的实测性能。

一、Stem 算法：从"信息流"重新理解稀疏注意力

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

Stem 的名字来源于"树干"的隐喻——在因果注意力架构中，初始位置的 token 如同一棵树的主干，支撑着所有后续信息的传递。

为什么初始 token 如此重要？在因果自注意力中，每一层的计算遵循严格的因果约束：第 N 个位置的输出，由该层前 N 个 token 共同决定，而这些 Value 向量又源自上一层输出的映射。这意味着：

• 位置 1 的 Value 参与了后续所有位置的注意力计算
• 位置 N 的 Value 只参与了位置 N 的注意力计算

这种不对称性在多层 Transformer 中被递归放大。初始 token 的信息经由层层传递，影响着整个网络的输出。如果稀疏掉初始 token，误差会沿因果路径扩散至下一层的所有 token；而稀疏掉末尾的 token，影响则相对有限。这个看似简单的观察，恰恰是 Stem 的核心洞察所在。

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

讲完了"为什么初始 token 重要"这个理论问题，你可能会问：这个发现，最终在生产环境里到底能跑出什么样的数字？

Stem 团队没有停留在学术 benchmark 上，而是直接将其集成进腾讯混元 Hy3 preview（W8A8-FP8）的 vLLM 推理框架，搭配 HPC 团队优化的 Stem 算子，端到端测量了首字延迟（TTFT）与模型精度。这意味着 Stem 不仅要在 BF16 学术基线上跑得通，还要在量化后的工业级模型上稳得住。

至于 Stem 与其他稀疏算法在开源模型上的速度/精度对比，论文中已给出完整结果，下面呈现 Stem 在 Hy3 preview（W8A8-FP8）上更贴近生产环境的真实落地数据。

2.1 首字加速比

……（相关数据与图示）

2.2 模型精度

……（相关数据与图示）

二、揭秘：Stem 凭什么又快又准？

数字摆在这里，问题就来了——Stem 是怎么在砍掉 75% 计算的同时，还把精度保住的？答案藏在两个看似简单、却被以往工作长期忽视的细节里：给谁多一点预算？挑 token 时该看什么？

Stem 的两大核心创新——Token 位置衰减（TPD）和输出感知度量（OAM）——分别回答了这两个问题。

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

现有方法对所有位置施加统一的 Top- b 稀疏预算，默认每个位置的注意力有相同的信息需求。但 Stem 从因果信息流的角度指出：初始 token 的稀疏预算应该更宽裕，因为它们在后续计算中被依赖的次数更多。

基于此，TPD 策略设计了非均匀的预算分配：

• 初始位置：分配更多预算，充分保留递归依赖链上的关键 token
• 末尾位置：预算线性衰减，因为它们在计算中参与度较低

这意味着在总预算不变的情况下，TPD 将计算资源倾斜到了真正重要的位置上。Stem 整体流程图（见下图）清晰展示了这一过程。

2.2 输出感知度量（OAM）：挑 token 时，该看什么？

传统稀疏方法只通过注意力分数（Query 和 Key 的点积）来判定 token 重要性，但忽略了 Value 向量的信息量。OAM 将视角扩展到了"输出"维度，同时考虑注意力分数和 Value 向量的重要性。

在实际实现中，OAM 通过对数变换把"乘法"转成"加法"，排序结果不变，又能直接套用标准 Top- b 算法，大幅降低了计算复杂度。

三、HPC-BSA 算子：把稀疏理论翻译成底层 "快" 代码

算法层面算出了"该省哪些计算"，但落地时还需要一套高效的算子把省下的计算真正变快。HPC 团队开源的块级稀疏注意力算子正是为此设计的。

3.1 核心设计理念：把智能判断前置到循环之外

块级稀疏的核心思路是：将 KV Cache 划分为固定大小的块（比如 64×128 或 128×128），效果好的块整个参与计算，无关的块整块跳过。然而，在底层 GPU 实现中，每次判断是否跳过都需要访问全局内存，这个开销在极端稀疏场景下可能吃掉收益。

HPC-BSA 的关键创新在于：将稀疏的判断和筛选完全前置到循环之外。内层循环只遍历预先筛选出的有效分块，完全避免了逐次判断跳过带来的额外开销。被跳过的分块在数学上等价于注意力分数全为负无穷，不影响 softmax 的正确性，计算逻辑无需任何修改。

3.2 Benchmark：算子级性能

性能测试以 HPC-Dense (FP8) 和 FlashAttention V3 (FP8) 作为稠密基线，以 MIT-BSA (BF16) 和 FlashPrefill-BSA (BF16) 作为稀疏对照。

结果揭示了三个关键发现：

• HPC-BSA 的延迟与计算密度近乎完美线性关系——50% 稀疏度下延迟约为稠密基线的一半，80% 稀疏度下仅为约五分之一，跳块机制的额外开销控制在 2.5% 以内。
• 相比 MIT 开源的 BSA 算子（MIT-BSA），HPC-BSA 在全稀疏度范围内稳定保持约 3 倍加速。
• 上述优势在 8K 到 256K 的全序列长度范围内保持稳定，展现出良好的长序列扩展性。

不同稀疏度下，BSA 算子延时的变化示意图：

四、总结与展望

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

算法决定"省哪些计算"，算子决定"省下的计算能快多少"——两者协同，构成了从理论到部署的完整闭环。随着大模型上下文窗口持续扩展至百万级别，高效的长文本推理将成为刚需。这一"算法 + 算子"的思路，有望为长文本高效推理提供新的范式参考。