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：

层段	alpha	策略
前 N 层（warmup）	1.0	保留更多 block，保证底层特征提取完整
后续层	0.7	更激进的稀疏化，加速高层推理

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

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，无需训练）

模型	方法	CC	FSL	MD1	MD2	SUM	SYN	AVG
Qwen3-8B	Dense	19.09	63.10	11.23	15.06	20.63	62.92	32.01
	MINF	18.91	61.87	10.78	14.46	20.26	55.32	30.27
	FLEX	19.75	55.31	10.76	13.49	21.18	50.83	28.55
	XATTN	21.03	62.30	11.42	14.64	20.47	52.92	30.46
	STEM	19.43	61.84	11.22	14.97	20.21	62.19	31.64
Llama-3.1-8B	Dense	36.08	64.29	27.54	31.19	25.10	67.92	42.02
	STEM	35.86	62.89	26.33	30.53	24.93	68.32	41.48

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 关键参数

参数	默认值	说明
`block_size`	128	attention block 大小
`stride`	8	scoring 下采样步长
`chunk_size`	2048	scoring 分块宽度
`initial_blocks`	4	始终保留的 sink token block
`window_size`	4	sliding window 保留的尾部 block
`stem_alpha`	1.0	每层衰减因子，可传 list 分层控制

6.3 支持矩阵

维度	支持情况
后端	PyTorch + Triton / HPC C++ 扩展
HPC 精度	bf16（dense prefill）、fp8（block-sparse prefill）
序列长度	无上限，建议 4K+ 以体现加速
模型	Qwen3、Llama-3.1、Hunyuan-Dense 等

---

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

方法	类型	训练需求	核心机制	局限
MINF	动态	无	基于注意力模式的启发式稀疏	精度损失较大
FlexPrefill	动态	无	基于 token 重要性的动态选择	开销较高
XAttention	动态	无	局部-全局混合 attention	实现复杂
STEM	动态	无	TPD + OAM，位置+贡献双感知	需要 block-sparse kernel 获得最大加速

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 #训练无关 #记忆 #小凯