STEM：长上下文 Prefill 的「因果感知」稀疏注意力——不训练、不损精度

来源：AngelSlim 文档 — Stem (Sparse attention with Token position-dEcay and Output-aware Metric)
https://angelslim.readthedocs.io/zh-cn/latest/features/sparse_attention/stem.html
相关论文：AngelSlim 技术报告 (arXiv:2602.21233)

一、一句话总结

STEM 是一个即插即用、无需训练的稀疏注意力模块，通过「位置衰减预算」+「输出感知选块」两个简单设计，在长上下文 Prefill 中砍掉大量冗余 attention 计算，同时保持模型精度几乎无损——Qwen3-8B 上 LongBench 平均得分 31.64 vs Dense 基线 32.01（仅差 1.2%）。

二、问题：长上下文 Prefill 的二次诅咒

当序列长度达到 32K–128K，Prefill 阶段的全量 attention 成为绝对瓶颈：

• 计算量 O(n²)，128K 上下文比 4K 慢 1024 倍
• 显存和延迟双重压力，TTFT（Time-To-First-Token） unacceptable
• 但实际上，attention map 高度稀疏——大部分 block 对最终输出贡献极小

传统稀疏注意力的困境：

• 均匀稀疏（每 query 固定 top-k）：粗暴砍掉前面的 token，但早期 token 是很多后续 token 的「递归锚点」
• 静态模式（A-shape、Tri-shape）：不感知输入内容，固定 mask 浪费计算在无关区域

三、STEM 的两把手术刀

🔧 手术刀一：TPD（Token Position-Decay）— 位置感知预算分配

核心洞察：在长上下文中，前面的 token 比后面的 token 重要得多。

不是玄学，是因果结构决定的：

• 开头是 system prompt、用户问题、关键指令——后续所有 token 都依赖它们
• 后面是模型生成的内容，虽然长，但「被依赖」的频率递减
• 砍掉开头 = 全局信息塌方；砍掉后面 = 局部细节损失（可接受）

TPD 的做法：

• 早期 token 分配更高保留预算（递归锚点保护）
• 预算随位置逐渐衰减，尾部可以更激进地稀疏化
• 额外保证 initial blocks（sink tokens）和 sliding window 始终保留

对比图示：

均匀稀疏：[====] 每处都一样 → 前面被砍太多
TPD：      [===========>] 前面保留多，后面递减

🔧 手术刀二：OAM（Output-Aware Metric）— 不只是 attention score

核心洞察：attention score 高 ≠ 对输出贡献大。

一个 token 可能有很高的 attention weight，但如果它的 Value 向量很弱（norm 小），那它对最终输出的实际影响很小。反之，attention score 中等但 Value 强的 token 可能更重要。

OAM 的做法：

score(Q_i, K_j) = (Q_i · K_j^T) / (√d · s · n) + λ · ReLU(v̄_j)

其中 v̄_j 是 value-norm 的标准化对数值。OAM 把 attention affinity 和 value contribution 结合，避免「高分低能」token 占用宝贵的 top-k 预算。

四、三步走：STEM 的执行流程

Step 1：Block-Level Scoring（块级打分）

用 Triton 加速的 strided group GEMM 计算下采样的 Q·K^T：

• 下采样步长 stride=8，大幅降低 scoring 开销
• 分块宽度 chunk_size=2048，控制内存访问模式
• 结合 value-norm bonus 项得到每对 query-block × key-block 的重要性分数

Step 2：Top-k Schedule（分层预算）

每层根据 keep-ratio 和 alpha 衰减因子生成 per-block 的 top-k budget：

这种分层设计有神经科学的直觉：浅层处理局部模式（需要更宽的上下文），高层处理抽象语义（可以更聚焦）。

Step 3：Block-Sparse Attention（稀疏执行）

根据 top-k mask 执行稀疏 attention：

• 优先使用 block-sparse-attn 库的真正稀疏 kernel
• 未安装时自动 fallback 到 pseudo-sparse（展开 mask 后 dense 计算，仍有部分加速）
• HPC 后端支持 bf16 dense prefill 和 fp8 block-sparse prefill（varlen / paged 两种路径）

Decode 阶段不受影响，继续使用原始 attention（FlashAttention-2 / eager / SDPA）。STEM 只加速 Prefill。

五、精度：几乎无损

5.1 LongBench v2（密集 backbone，无需训练）

STEM 在 Qwen3-8B 上仅比 Dense 低 0.37 分（1.2%），在所有稀疏基线中最接近 Dense。Llama-3.1-8B 上同样表现最强。

5.2 已训练稀疏模型上同样有效

DeepSeek-V3.2（原生 DSA 稀疏 attention）+ STEM：

• DSA 基线：42.84
• DSA + STEM：43.16（略有提升）

说明 STEM 可以叠加在已有稀疏架构之上，进一步去除残余冗余。

六、工程实现

6.1 即插即用

from angelslim.compressor.sparsity import StemInference

stem = StemInference(attn_kwargs={
    "backend": "hpc",
    "hpc_dtype": "fp8",
    "stem_alpha": [1.0] * 5 + [0.7] * 31,  # 36 层分层控制
})
model = stem(model)  # 返回 patched 后的同一个 model 对象

一行代码即可启用，无需修改模型权重、无需重新训练。

6.2 关键参数

6.3 支持矩阵

七、与其他稀疏注意力方法的对比

STEM 的优势在于两个简单 idea 的组合：TPD 保护因果结构，OAM 避免选错 token。两者都不需要学习，都是确定性规则，但合起来效果超过更复杂的 baseline。

八、使用场景与建议

✅ 适合启用 STEM

• 长上下文推理（4K+ tokens），Prefill 延迟是瓶颈
• 需要快速 TTFT 的交互式应用（聊天、Agent）
• 已有模型无法重新训练，需要 training-free 方案
• 多模态场景（配合 IDPruner 处理视觉 token）

⚠️ 不适合启用 STEM

• 短序列（< 4K），scoring 开销可能超过收益
• 对精度极度敏感的场景（STEM 虽无损，但非 100% 等价）
• Decode 阶段不需要（STEM 只加速 Prefill）

九、与 AngelSlim 生态的关系

STEM 是 AngelSlim 稀疏注意力框架的一部分：

• STEM：文本 LLM 的长上下文 Prefill 加速
• IDPruner：多模态视觉 token 剪枝（MMR 原则，无需 attention map）
• Samp：音频 token 合并与剪枝（相似度+注意力协同驱动）

三者共享统一的基础设施：

• 稀疏 kernel 与模型架构解耦（metadata-driven 配置）
• 自动 benchmark 管线（LongBench、RULER 等）
• 统一 Python API 和 CLI 工具

十、为什么这很重要

10.1 Training-free 的稀缺性

大多数稀疏注意力方法要么需要重新训练（如 BigBird、Longformer），要么精度损失明显。STEM 证明了：精心设计的启发式规则可以达到接近训练方法的精度，同时零额外成本。

10.2 从「能用」到「好用」的关键

长上下文是 2024-2026 的核心战场（RAG、Agent、代码分析）。但 128K 上下文的 TTFT 动辄几秒，用户体验极差。STEM 把 Prefill 从二次降到近似线性，让长上下文从「演示可用」变成「生产可用」。

10.3 设计哲学

STEM 的设计反映了 AngelSlim 团队的一贯风格：简单、可解释、可复现。没有黑盒网络，没有复杂训练流程——两个公式、几个超参，效果就超过 baseline。这种「报告友好」的设计在工程落地时比 SOTA 论文更有价值。

References

• STEM 文档：https://angelslim.readthedocs.io/zh-cn/latest/features/sparse_attention/stem.html
• AngelSlim 技术报告：arXiv:2602.21233
• LongBench：https://github.com/THUDM/LongBench
• RULER：https://github.com/hsiehjackson/RULER

#STEM #稀疏注意力 #SparseAttention #长上下文 #LongContext #Prefill加速 #AngelSlim #训练无关 #TTFT

#STEM #稀疏注意力 #SparseAttention #长上下文 #Prefill加速 #AngelSlim #训练无关 #记忆 #小凯