SAS：端到端优化上下文排名的稀疏注意力机制 · 干货攻略

• 链接: https://x.com/_akhaliq/status/2099511334987522298 | https://arxiv.org/abs/2609.13141
• 分类: x-tips
• 来源: X @_akhaliq
• 作者: Jay
• 更新: 2026-09-28
• 仓库: Tencent-Hunyuan/Simple-Attention-Sparsification

这是什么

SAS（Simple Attention Sparsification）是腾讯混元团队（Tencent Hunyuan）提出的一种门控稀疏注意力机制，2026 年 9 月 11 日发布于 arXiv（arXiv:2609.13141），代码仓库位于 Tencent-Hunyuan/Simple-Attention-Sparsification，同步开源于 HuggingFace tencent/Simple-Attention-Sparsification。

核心解决的问题：后训练（post-training）稀疏化注意力时，传统的 Hard Top-K 选取会阻断语言模型损失向 selector 的梯度传播，导致 selector 只能通过蒸馏原始密集注意力分布来间接训练，存在排名不对齐的问题。SAS 用连续门控（continuous gating）替代离散 Top-K 选取，将 selector 的连续得分注入注意力 logits，使语言建模损失能直接通过标准反向传播优化上下文选择器。

作者来自腾讯混元 + 香港科技大学（广州），接受采访的模型规模为 Qwen3-4B / 8B / 14B，训练数据集为 OpenR1-Math-220k，资源需求为 8 张 NVIDIA H20 GPU。

为什么值得关注

长上下文推理是 LLM 部署的核心效率瓶颈。自回归生成每个新 token 时都需要对全部历史上下文 token 做密集注意力，累计注意力成本随上下文长度二次增长。

业界已有多种稀疏注意力方案：

• 训练无关方法（如 H2O、StreamingLLM）：用启发式规则选 token，够快但无法适应模型自身预测行为
• 可训练的稀疏方法（如 SeerAttention-R）：引入 learnable selector 对上下文块打分，再经 Hard Top-K 选取性能更高——但 Hard Top-K 不可导，selector 只能靠蒸馏密集注意力分布来训练，这带来了排名不对齐：selector 学到的是"原始模型在哪里注意"，而非"在有限注意力预算下哪些上下文对最终预测最重要"

SAS 的核心贡献是把 selector 的梯度通路打通，让语言建模损失直接决定"哪些上下文块值得被选中"，而不是被蒸馏信号绕道。论文标题里的 "End-to-End Optimization of Context Ranking" 就是这个意思。

核验过程

本攻略基于以下来源编写，关键说法均经交叉验证：

关键说法核验结论：

• ✅ "连续门控替代 Hard Top-K selection"：论文原文 + GitHub README 一致确认
• ✅ "Triton kernel 集成 FlashAttention"：原文 + README 一致确认，原文描述为 "memory-efficient Triton kernel that integrates our method into FlashAttention-style computation"
• ✅ "解决注意力稀疏化梯度阻断问题"：论文 Introduction Figure 1 明确描述这一 motivation
• ✅ Qwen3-4B 在 2048 token 预算下 MATH500=93.47%（SeerAttention-R=91.85%）：AlphaXiv 来源的 benchmark 数字，与 arXiv HTML 页面数据一致
• ✅ 资源需求 8×H20 GPU：来自 GitHub README，原文标注
• ⚠️ "特别适用于 tight budget"：原文主张：SAS 在小预算下增益最大（如 1024 token 时比 SeerAttention-R 高 6.0-7.7 分 on MATH500），但这一相对优势是否在所有任务类型普遍成立，需视具体场景

上手步骤

环境安装

需要先安装 pixi（Rust 版包管理器）：

curl -fsSL https://pixi.sh/install.sh | bash  # 安装 pixi（如已安装跳过）
pixi install && git submodule update --init --recursive

下载模型和数据集

SAS 当前开源的训练配方支持 Qwen3-4B、Qwen3-8B、Qwen3-14B。下载基座模型和训练数据：

export MODEL_PATH=/path/to/Qwen3-4B      # 或 8B、14B
export DATA_PATH=/path/to/OpenR1-Math-220k/data

数据集托管在 HuggingFace：open-r1/OpenR1-Math-220k

启动训练

bash scripts/train/simple_sparse_attention_Qwen3-4B.sh
# 对应 8B：scripts/train/simple_sparse_attention_Qwen3-8B.sh
# 对应 14B：scripts/train/simple_sparse_attention_Qwen3-14B.sh

推理评估（需要 SGLang 后端）

先构建推理依赖：

cd third_party/sglang-blocksparse && pixi install && cd -

推理任务 benchmark（数学 + 长上下文）

export GATES=/path/to/AttnGates        # 或用 dense 基线
export MODE=dense_4b BASE_MODEL=/path/to/Qwen3-4B
export BUDGET=2048                     # 注意力 token 预算
export TP=1 DP=8                       # TP=1，DP=8 得 8 个完整副本
export TASK=math,gpqa,aime24,aime25

bash scripts/eval/run_reasoning.sh     # MATH500 / GPQA / AIME
bash scripts/eval/run_longbench.sh      # LongBench 长上下文理解

函数调用 benchmark（BFCL）

git submodule update --init third_party/gorilla
python -m venv third_party/gorilla/.venv
third_party/gorilla/.venv/bin/pip install -e third_party/gorilla soundfile
git -C third_party/gorilla apply "scripts/eval/patches/bfcl.patch"

export GATES=/path/to/AttnGates
export BUDGET=2048 TP=1 DP=8
export CATEGORIES=multi_turn
bash scripts/eval/run_bfcl.sh

多域工具调用 benchmark（VitaBench）

git submodule update --init third_party/vitabench
python -m venv third_party/vitabench/.venv
third_party/vitabench/.venv/bin/pip install -e third_party/vitabench
git -C third_party/vitabench apply "scripts/eval/patches/vitabench.patch"

export GATES=/path/to/AttnGates
export BUDGET=2048 TP=1 DP=8
export DEEPSEEK_API_KEY=sk-...         # 需要 DeepSeek-V4-Pro 做用户模拟和评估
export DOMAIN=delivery,instore,ota
bash scripts/eval/run_vitabench.sh

推理时如何生效

• Prefill 阶段：使用密集注意力（标准 FlashAttention）
• Decode 阶段：启用块选择（block selection），结合 SGLang + Paged KV Cache + FlashInfer 稀疏解码
• 核心数据结构：AttnGates 目录保存训练产出的门控权重，评估时通过 GATES 环境变量指定路径

坑与适用边界

适用场景

• 推理效率优化：长序列场景（>8K tokens），预算紧张（tight budget）时收益最大
• Agentic 任务：函数调用、多步工具调用（BFCL、VitaBench 等 benchmark 验证）
• 长上下文理解：LongBench 验证有效，适合 RAG + 长文档分析管线
• 快速训练适配：基于 Qwen3 系列后训练，无需改动模型结构

已知限制

• 基座模型锁定：当前开源配方仅支持 Qwen3-4B/8B/14B，不包含其他模型架构
• 硬件要求：需 8×H20 GPU，单卡或小显存设备无法本地复现训练
• 评估依赖 SGLang fork：推理后端基于 rayleizhu/sglang 分支（非官方 SGLang），部署有额外维护成本
• Tight budget 以外优势递减：原文数据显示，budget 越大 SAS 相对基线优势越小——极端宽松预算下可考虑直接用密集注意力
• 训练稳定性：门控 log-space 注入需要正确的归一化处理，README 强调了几个关键设计选择，细节不到位可能训练不稳定

踩坑提示

• 评估脚本会自动启动 sglang server，运行完后会清理；如果 port 冲突，检查是否有残留进程
• VitaBench 需要真实 API key 且调用量大，生产使用前先确认 DEEPSEEK API 配额
• submodule patch 文件依赖仓库根目录执行，切换目录后 patch 路径会失效

一句话结论

SAS 通过在注意力 softmax 内部注入连续门控（而非靠蒸馏间接训练），打通了语言建模损失到 selector 的梯度通路，在推理效率与精度之间实现了更优的上下文排名，尤其适合 tight budget 下的长上下文与 Agent 场景——是目前稀疏注意力领域「端到端可训练」设计的最新 State of the Art。