More Value per Key: Asymmetric Sparse Attention for Faster LLM Decoding
_{Official PyTorch implementation of SAGA (Sparse Asymmetric Group-Query Attention)}

Noam Elata*, Itay Lamprecht*, Mikey Shechter*, Daniel Ohayon, Itay Hubara, Daniel Soudry
NeurIPS 2026

* Equal contribution

Autoregressive generation in Large Language Models (LLMs) is constrained by the memory and computational demands of attention mechanisms. Sparse attention methods mitigate this cost by selecting only high-probability entries of the attention matrix. We observe that in many such methods, this renders the probability-value multiplication negligible, shifting the bottleneck to the query-key step. Key heads can therefore be reduced to accelerate inference, while retaining more value heads preserves capacity with limited additional decoding cost. We introduce Sparse Asymmetric Group-Query Attention (SAGA), which decouples key and value head counts to exploit this principle, and pair it with approximate top-N (Atop-N) attention, a simple sparse attention method designed to study the interaction between sparsity and head-count asymmetry. We formalize the benefits of this asymmetry theoretically and validate them empirically through latency measurements and quality evaluations on models up to 1.5B parameters. Together, SAGA and Atop-N achieve end-to-end decoding speedups exceeding 2× over our full-attention GQA baseline at long contexts. Models trained from scratch with SAGA nearly match the quality of comparable GQA variants on the evaluated benchmarks. To facilitate adoption, we introduce an efficient fine-tuning method that converts pretrained models to the SAGA architecture, enabling practitioners to benefit from our approach without costly retraining.

Throughout the code and CLI flags, Atop-N is written atopn.

Install the dependencies:

pip install -r requirements.txt

The repository supports three workflows. Each has a one-line entry point below and a detailed guide.

Trains SmolLM2-360M with fewer key heads than value heads:

bash train_scratch.sh

Builds the data caches, initializes the student by merging key heads, then distills it from the original model in three stages of increasing context length:

bash run_prepare_data.sh   # build data caches (once per model)
bash finetune.sh           # student init + stages 1-3

Runs accuracy (HellaSwag, WikiText, RULER) and decoding-latency benchmarks:

EVAL_MODEL=./experiments/qwen_merged/final_model bash run_eval.sh

@inproceedings{elata2026saga,
  title     = {More Value per Key: Asymmetric Sparse Attention for Faster {LLM} Decoding},
  author    = {Elata, Noam and Lamprecht, Itay and Shechter, Mikey and Ohayon, Daniel and Hubara, Itay and Soudry, Daniel},
  booktitle = {Advances in Neural Information Processing Systems},
  year      = {2026}
}