大语言模型在长上下文推理时，KV 缓存的注意力成本成为瓶颈。稀疏注意力通过只访问缓存的一小部分来降低成本，但现有方法在确定性检索（捕获高分 token）和采样（提升覆盖率但方差大）之间存在权衡。混合方案将 top‑k 检索与均匀采样结合，却在高稀疏度下仍与密集注意力有显著差距。
我们提出 Stratified Sparse Attention (SSA)，在统一框架下融合确定性检索和基于分数的采样。SSA 首先确定性地关注得分最高的 token，然后对剩余部分使用代理注意力分数进行采样，并根据包含概率进行校正，从而恢复检索集之外的注意力质量。SSA 包含 top‑k、top‑p 以及 top‑k+采样等特例。
在 RULER、LongBench、LOFT、AIME 四套基准以及 3B‑27B 参数的五个模型上，SSA 在相同稀疏度下始终优于基线。在 RULER‑HARD 上使用 Ministral‑3‑8B，SSA 在 32K 长度下将准确率从 81.8% 提升至 84.2%，在 128K 长度下从 63.8% 提升至 69.3%，对应稀疏度 50× 时分别恢复了密集注意力的 94% 与 85%。配合优化的 CUDA kernel，SSA 在 1M 上下文下相较 FlashInfer 实现最高 2.96× 的解码加速，证明精度提升可转化为实际推理加速。代码已在 https://github.com/skylight-org/SSA-Paper-Code 开源。
点评：SSA 通过分层采样与概率校正，在极端稀疏条件下显著缩小了与密集注意力的性能差距，并兼具实现简洁和加速潜力，值得在长上下文 LLM 推理中推广。