tencent/Simple-Attention-Sparsification

Simple Attention Sparsification (SAS)

SAS learns to rank and select KV blocks for each query. Unlike methods that train a sparse-attention selector by distilling dense attention scores, SAS adds continuous gates to the selected blocks so that the language-modeling loss can optimize context ranking end to end.

Overview of SAS. Continuous gates provide a differentiable path through discrete Top-K block selection, enabling end-to-end context-ranking optimization with the language modeling loss.

Released checkpoints

Directory	Base model	Gate parameters	Gate checkpoint size
Qwen3-4B-AttnGates	Qwen3-4B	33.0M	64 MiB
Qwen3-8B-AttnGates	Qwen3-8B	33.0M	64 MiB
Qwen3-14B-AttnGates	Qwen3-14B	42.0M	81 MiB

Each directory is an SGLang-compatible AttnGates package containing:

• attn_gate_weights.pth: the learned gate-router weights;
• config.json: the base-model reference and seerattn_* configuration;
• tokenizer and chat-template files copied from the corresponding Qwen3 model.

These are router-only checkpoints, not standalone language models. The Qwen3 backbone was frozen during training and is not included. Standard transformers.AutoModelForCausalLM.from_pretrained() cannot directly run these directories. Inference requires the corresponding Qwen3 base model and the seer_attn backend in our sglang-blocksparse fork.

Checkpoint configuration

All three checkpoints use the same sparse-attention setup:

Setting	Value
KV block size	64 tokens
Training Top-K	31 historical blocks
Gate hidden size	128
Query projection	Qproj
Key block pooling	max + min + average
Q/K normalization	enabled
Gate RoPE	enabled
Training sequence length	32,768 tokens
Training data	OpenR1-Math-220k

Only the gate routers were optimized; all base-model parameters remained frozen. The released setup uses a 2,048-token sparse decode budget by default. The same checkpoints can be evaluated with 1,024-, 2,048-, or 4,096-token budgets without retraining.

Inference

1. Set up SAS and the sparse SGLang backend

git clone https://github.com/Tencent-Hunyuan/Simple-Attention-Sparsification.git
cd Simple-Attention-Sparsification

git submodule update --init --recursive

# Install the SAS environment.
curl -fsSL https://pixi.sh/install.sh | bash  # skip if pixi is already installed
pixi install

# Build the environment used by the sparse inference server.
cd third_party/sglang-blocksparse
pixi install
cd ../..

Download this Hugging Face repository and point GATES to the directory that matches the desired base-model size.

2. Launch a server directly

The following example serves the Qwen3-4B gate checkpoint:

export GATES=/path/to/downloaded-repo/Qwen3-4B-AttnGates
export SGLANG_SEER_TOKEN_BUDGET=2048
export CUDA_VISIBLE_DEVICES=0

cd third_party/sglang-blocksparse
pixi run python -m sglang.launch_server \
  --model-path "$GATES" \
  --served-model-name qwen3-sas-4b \
  --attention-backend seer_attn \
  --tool-call-parser qwen \
  --trust-remote-code \
  --tp-size 1 \
  --port 30000

config.json points to the corresponding public Qwen3 repository, so the base model is downloaded separately when needed. For offline use, download the base model in advance and replace the base_model value in the gate checkpoint's config.json with its local path.

To use another checkpoint, change GATES and the served model name:

# Qwen3-8B
export GATES=/path/to/downloaded-repo/Qwen3-8B-AttnGates

# Qwen3-14B
export GATES=/path/to/downloaded-repo/Qwen3-14B-AttnGates

The server exposes an OpenAI-compatible API at http://127.0.0.1:30000/v1.

3. Run the official evaluation wrappers

The code repository also provides wrappers that start the server, run a benchmark, and stop the server automatically:

cd /path/to/Simple-Attention-Sparsification

export GATES=/path/to/downloaded-repo/Qwen3-4B-AttnGates
export MODE=seer_4b
export BUDGET=2048
export TP=1
export DP=8

# MATH, GPQA-Diamond, and AIME 2024/2025
export TASK=math,gpqa,aime24,aime25
bash scripts/eval/run_reasoning.sh

# LongBench-E
bash scripts/eval/run_longbench.sh

Use MODE=seer_8b or MODE=seer_14b with the matching gate directory for the other model sizes.

Evaluating on BFCL and VitaBench requires additional benchmark-specific environments and dependencies. Please follow the setup instructions in the code repository before running scripts/eval/run_bfcl.sh or scripts/eval/run_vitabench.sh.

Citation

If you find SAS or these checkpoints useful, please cite:

@misc{li2026sassimpleattentionsparsification,
  title         = {SAS: Simple Attention Sparsification via End-to-End Optimization of Context Ranking},
  author        = {Zhiwei Li and Lei Zhu and Hao Gu and Xiang Hu and Yan Wang and Haitao Mi and Sirui Han and Leo Liang and Zhijiang Guo},
  year          = {2026},
  eprint        = {2609.13141},
  archivePrefix = {arXiv},
  primaryClass  = {cs.CL},
  url           = {https://arxiv.org/abs/2609.13141}
}

Acknowledgements

This release builds on Qwen3, SeerAttention-R, VeOmni, and sglang-blocksparse, our fork of SGLang providing the seer_attn inference backend.