Paper deep dive
More Value per Key: Asymmetric Sparse Attention for Faster LLM Decoding
Noam Elata, Itay Lamprecht, Mikey Shechter, Daniel Ohayon, Itay Hubara, Daniel Soudry
Intelligence
Status: not_run | Model: - | Prompt: - | Confidence: 0%
Entities (0)
Relation Signals (0)
No relation signals yet.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:utoregressive 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\times$ 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.
Tags
Links
- Source: https://arxiv.org/abs/2610.04753v1
- Canonical: https://arxiv.org/abs/2610.04753v1
Trouble viewing inline? Open PDF directly →
Full Text
70,195 characters extracted from source content.
Expand or collapse full text
More Value per Key: Asymmetric Sparse Attention for Faster LLM Decoding Noam Elata †thanks: Equal contribution. Affiliation: Technion – Haifa, Israel Email: noamelata@technion.ac.il Itay Lamprecht11footnotemark: 1 †thanks: Work was partially done while the author was at Intel†thanks: Israel. Affiliation: Technion – Haifa, Israel Affiliation: Crusoe AI Email: itayl@campus.technion.ac.il Mikey Shechter11footnotemark: 1 Affiliation: Technion – Haifa, Israel Email: m.sachter@technion.ac.il Daniel Ohayon Affiliation: Corma Affiliation: Technion – Haifa, Israel Email: ohayon.daniel4@gmail.com Itay Hubara Affiliation: Stealth Startup Email: itayhubara@gmail.com Daniel Soudry Affiliation: Technion – Haifa, Israel Email: daniel.soudry@gmail.com Abstract 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×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.11 1 Code is available at https://github.com/noamelata/SAGA. 1 Introduction Attention-based architectures have become dominant across AI applications [1], with large language models such as GPT [2, 3], Llama [4], and Qwen [5, 6] demonstrating strong capabilities in natural language understanding [7], generation [8], and reasoning [9]. However, the quadratic scaling of attention with sequence length [10, 11] remains a significant bottleneck, constraining context window sizes and limiting the ability of models to reason over lengthy documents. For autoregressive generation, KV caching [12] provides substantial speedups by reusing key and value computations across generation steps. This shifts the bottleneck from compute capacity to memory capacity and bandwidth [13]. Group-Query Attention (GQA) [14] further exploits this setting by reducing the number of key-value heads, accelerating inference for long sequences. Yet even with these optimizations, bounded context windows remain a practical necessity. Sparse attention methods [15, 16, 17, 18, 19] offer a complementary direction, accelerating attention by computing only a subset of query-key interactions. Such methods span both the prefill and decoding phases, and range from zero-shot approximations to approaches integrated into pre-training. In this work, we focus on the decoding phase and use a generic sparsity method, as primary examples of the effects of sparsity on the computational profile of generation. To this end, we adopt top-N attention as a representative sparsity method. Building on previous observations [20, 21] that models assign significant attention weight to only a small subset of keys, top-N attention has the advantage of simplicity: it requires no training to apply and admits a portable implementation, a property uncommon among sparse attention methods. Our efficient approximate top-N attention (Atop-N) implementation gathers only the most relevant attention entries, reducing the probability-value multiplication to a fraction of its original cost. Analyzing the resulting computational profile reveals a fundamental asymmetry: under sparse attention, the bottleneck shifts from the probability-value multiplication to the query-key multiplication. This insight motivates Sparse Asymmetric Group-Query Attention (SAGA), an architectural modification that decouples the number of key heads hkh_k from value heads hvh_v. While GQA ties these together, SAGA reduces hkh_k independently, directly accelerating the query-key multiplication. The number of value heads remains large to preserve model capacity, incurring no additional cost since only the top-N values are accessed. We formalize this asymmetry theoretically and validate it empirically, demonstrating end-to-end speedups exceeding 2×2× in long contexts while maintaining the majority of baseline quality. To facilitate practical adoption, we introduce an efficient fine-tuning method for converting existing GQA models to the SAGA architecture. We merge teacher key heads using a closed-form initialization derived from minimum-weight perfect matching on attention logit reconstruction error (Sec. 4), followed by a short distillation phase. This enables the benefits of SAGA without costly retraining. To summarize, our contributions are: • We present a novel analysis of LLM inference trade-offs with accelerated sparse attention. • We introduce Atop-N, a simple and portable sparse attention algorithm for decoding. • Motivated by our finding, we propose SAGA, an extension of the GQA architecture tailored for sparse attention inference. We show that SAGA is competitive both when trained from scratch and when finetuned from existing LLM transformers, while demonstrating more than 2×2× runtime speedups end-to-end on long sequence generation. 2 Background and related work 2.1 Multi-head self-attention The Transformer architecture, introduced by Vaswani et al. [1], revolutionized natural language processing. Its core mechanism, multi-head self-attention, enables parallel processing of sequences and captures long-range dependencies by allowing the model to attend to all positions simultaneously. Formally, given a hidden state ∈ℝlq×DX ^l_q× D, we apply three projections Q,K,V∈ℝD×(h⋅d)W_Q,W_K,W_V ^D×(h· d) to obtain queries ∈ℝh×lq×dQ ^h× l_q× d, keys ∈ℝh×lkv×dK ^h× l_kv× d, and values ∈ℝh×lkv×dV ^h× l_kv× d: =Q,=K,=V,Q=XW_Q, K=XW_K, V=XW_V, (1) where h is the number of heads, D and d are the hidden and head dimension sizes respectively, and lql_q and lkvl_kv are the query and key-value sequence lengths. The attention output is computed as: =⊤dA= QK d (2) =softmax()P=softmax(A) (3) =,O=PV, (4) where ,∈ℝh×lq×lkvA,P ^h× l_q× l_kv denote the attention and probability matrix respectively. The output O is then multiplied by O∈ℝ(h⋅d)×DW_O ^(h· d)× D. Modern LLMs stack many such layers in a decoder-only design, modeling complex relationships between all sequence elements regardless of distance. 2.2 KV cache Key-Value (KV) caching [12] is an optimization technique that significantly improves the inference efficiency of autoregressive Transformer models. During the generation of each new token, Transformer models would typically recompute the key and value projections for all previous tokens in the sequence. The KV cache stores projections from previous steps, eliminating redundant computations. For an autoregressive model generating the lkvl_kv-th token, the attention computation with KV caching can be expressed as: lkv=softmax(lq[1:(lkv−1),lkv]⊤d)[1:(lkv−1),lkv],O_l_kv=softmax ( Q_l_q[K_1:(l_kv-1),K_l_kv] d )[V_1:(l_kv-1),V_l_kv], (5) where 1:(lkv−1)K_1:(l_kv-1) and 1:(lkv−1)V_1:(l_kv-1) are retrieved from the cache, and only lqQ_l_q, lkvK_l_kv, and lkvV_l_kv need to be computed for the new token. By caching key and value vectors from previous time steps, the model only needs to compute new projections for the most recently generated token, leading to substantial speedups during inference. This optimization is particularly valuable for long-form text generation, where the model must generate many tokens sequentially. The trade-off is increased memory usage proportional to sequence length, as the cache must store key-value pairs for all tokens in the context window. 2.3 Group-query attention and head reduction Group-Query Attention (GQA) [14] addresses the memory overhead of KV caching by recognizing an asymmetry in the attention computation during decoding: there are significantly more cached keys and values than queries (lkv≫lql_kv l_q). In this context, it is sensible to allocate more of the model’s capacity to the queries, which are computed fresh at each decoding step, rather than to the keys and values, which are retrieved from the cache. GQA achieves this by decoupling the number of key-value heads hkvh_kv from the number of query heads hqh_q, such that hq>hkvh_q>h_kv. Multiple query heads share a single key-value head, reducing the size of the KV cache while preserving the expressiveness of the query projections. This approach yields a significant improvement in inference efficiency with negligible degradation in model quality, making it a widely adopted technique in modern LLM architectures. Beyond reducing the total number of KV heads, recent work has explored making key and value head counts asymmetric. DHA [22] fuses similar attention heads in MHA architecture into decoupled per-layer key and value configurations to reduce memory footprint. Sigma [23] performs differential rescaling of query, key, and value dimensions during pretraining and finds that reducing key capacity is preferable to reducing value capacity—an observation consistent with our analysis. While conceptually related, SAGA differs in both motivation and applicability, exploiting the shifted computational bottleneck under sparse attention to accelerate decoding. Multi-head latent attention with DeepSeek Sparse Attention [24, 25] is related in spirit but must be trained from scratch with a shared latent KV and a learned indexer; it reduces the number of key positions scored per query, whereas SAGA reduces the number of key heads, so the two are orthogonal rather than competing. Figure 1: Diagram of SAGA. There are two significant decoding bottlenecks in attention mechanisms: key-query multiplications and probability-value multiplications (top). SAGA addresses the query-key bottleneck by reducing the number of key heads (left). The probability-value bottleneck is mitigated by utilizing Atop-N to select important values (right). 2.4 Attention sparsity and KV cache acceleration Recent efforts to mitigate the quadratic complexity of the attention mechanism have largely focused on attention sparsity and KV cache compression. Sparse attention methods, such as BigBird [17], Longformer [16], and Sparse Transformer [15], utilize predefined patterns like sliding windows or global tokens to restrict the number of query-key pairs computed. While effective for the prefill phase, these static patterns often struggle to capture the dynamic, data-dependent importance of tokens during autoregressive decoding. NSA [19], SpargeAttention [26], BLASST [27], and BSFA [28] also aim to bridge the gap between algorithmic sparsity and hardware efficiency, mostly focusing on the more compute-intensive pre-fill phase. Our work, SAGA, focuses instead on the decode phase, and offers a new framework by combining a simple attention sparsity method with a dedicated neural architecture. Other methods accelerate decoding by evicting low-importance tokens from the KV cache [29, 30] or selectively loading cache pages using query-aware metadata [31, 32]. Complementary techniques reduce cache memory via quantization [33, 34]. SAGA is orthogonal to these approaches; it optimizes the head architecture itself and can be readily combined with cache compression methods. 2.5 Top-N attention Top-N attention [20, 35] (also known as top-K attention or kkNN attention) is based on the observation that only a small subset of attention keys meaningfully influence the attention output. This sparsity arises from the softmax operation, which effectively learns to identify the best query-key matches while suppressing mismatched pairs. Formally, top-N attention identifies the indices of the N largest values in each row of A: j,k,1:N=TopN(j,k,1:lkv),i_j,k,1:N=TopN(A_j,k,1:l_kv), (6) and masks all other entries of A not in i before applying the softmax. This formulation has been shown to enable memory-efficient implementations [20, 35], though it does not yield computational speedups since the full attention matrix must still be computed. Several recent works have extended top-N attention in various directions [36, 37] or strengthened its theoretical foundations [38]. 3 Method We begin by introducing the attention sparsity method underlying SAGA. For clarity, we present an architecture-agnostic formulation of top-N attention to illustrate our approach and the trade-offs many online sparse attention method induce. In the following analysis, we show that under sparse attention, the computational bottleneck shifts from the probability-value multiplication to the query-key multiplication—the property that SAGA is designed to exploit. 3.1 Efficient approximate top-N attention for LLM decoding We propose an efficient framework for top-N attention designed for LLM decoding. Using the indices ∈ℤh×lq×Ni ^h× l_q× N from Eq. 6, we define the Top-N Attention Matrix []∈ℝh×lq×NA[i] ^h× l_q× N and corresponding Top-N Values []∈ℝh×lq×N×dV[i] ^h× l_q× N× d, and compute attention over only the gathered entries: =softmax([])[].O=softmax(A[i])V[i]. (7) In autoregressive generation (lq=1l_q=1), this reduces the probability-value cost from O(hlkvd)O(h\,l_kv\,d) to O(hNd)O(h\,N\,d), potentially saving nearly 50% of computation and bandwidth when N≪lkvN l_kv. However, naïvely finding the true top-N entries incurs lengthy hardware runtimes, negating the benefit. We address this with an approximate threshold approach: given the N-th largest attention value as a threshold, top-N entries can be identified by a simple per-element comparison — a parallel, portable operation much faster than sorting. Since attention distributions shift only marginally between consecutive decoding steps, we compute exact thresholds once every n tokens and reuse them, amortizing the cost over the full text generation. We term this method Atop-N attention; full implementation details are in App. C. For long sequences, Atop-N yields an overall speedup of nearly 2×2× for the full attention operation. Section 5 shows that for practical values of N, Atop-N retains 83–97% of baseline RULER accuracy. 3.2 Sparse Asymmetric Group-Query Attention The analysis in the previous section reveals that under Atop-N attention, the probability-value multiplication accesses only hq×lq×Nh_q× l_q× N value vectors, N per query head, regardless of the number of key-value heads hkvh_kv. Consequently, hkvh_kv affects computational complexity only through the query-key multiplication in Eq. 2. This asymmetry suggests that the number of key heads and value heads should be disentangled. We propose Sparse Asymmetric Group-Query Attention (SAGA), which decouples the key and value head counts into separate hyperparameters hkh_k and hvh_v. Reducing hkh_k directly reduces the memory bandwidth required to load keys during the query-key multiplication step (Eq. 2), the dominant latency bottleneck in LLM decoding under Atop-N attention. The number of value heads hvh_v can remain large to preserve model capacity, since only the top-N values are accessed regardless of hvh_v. A diagram of SAGA is shown in Fig. 1. Beyond the bandwidth reduction, fewer key heads yield a smaller key projection KW_K (reducing per-token compute) and a smaller KV cache footprint, enabling longer context lengths or larger batch sizes within the same memory budget. SAGA can be applied directly when training from scratch, as detailed in Sec. 5.2. To leverage existing pretrained models, we describe a procedure for converting GQA models to SAGA in Sec. 4 below. 4 Converting GQA models to SAGA We describe a practical procedure for converting a pretrained GQA model into a SAGA model. Our goal is to transform the attention module of an existing GQA model T, with equal numbers of key and value heads hkT=hvT,h_k^T=h_v^T, (8) into an attention module of a SAGA model S with fewer key heads than value heads, hkS<hvS,h_k^S<h_v^S, (9) while preserving the attention behavior as closely as possible. The conversion consists of two stages: 1. Choosing an initialization for the SAGA key projection matrices KSW_K^S that minimizes deviation from the original GQA attention, as described below. 2. Distillation-based fine-tuning, using the original GQA model T as a teacher and the SAGA model S as a student, as implemented in Sec. 5.3. Our GQA to SAGA conversion recipe naturally extends to multi-head attention (MHA) models. 4.1 Initializing SAGA from GQA We consider a GQA model with hkTh_k^T key heads, hvT=hkTh_v^T=h_k^T value heads, and hqh_q query heads. We convert it to a SAGA model with fewer key heads hkSh_k^S while keeping the same number of value heads hvS=hvTh_v^S=h_v^T. We assume the reduction factor is a positive integer m≥2m≥ 2 hkT=m⋅hkS(m∈ℕ),h_k^T=m· h_k^S (m ), (10) so that each SAGA key head is shared by exactly m value heads. For simplicity we describe the case m=2m=2, i.e. hkT=2hkSh_k^T=2h_k^S, and note that the construction extends directly to any m∈ℕm by replacing the minimum-weight perfect matching in Theorem 4.2 with a minimum-weight perfect m-partitioning of G. Let K1,…,K2hkSW_K^1,…,W_K^2h_k^S denote the key projection matrices, and Q1,…,QhqW_Q^1,…,W_Q^h_q the query projection matrices of an attention module of the teacher model T. Let ∈ℝL×DX ^L× D denote a matrix of hidden states sampled from a representative dataset. Our objective is to construct new key projections ^K1,…,^KhkS W_K^1,…, W_K^h_k^S for the SAGA model such that the resulting attention logits (Qj)(^Ki)⊤d (XW_Q^j)(X W_K^i) d (11) closely approximate the corresponding logits produced by the original GQA model for every query head j that originally shared one of the two merged teacher key heads. We choose to optimize only the merged key projections, and to match the attention logits while ignoring rotary embeddings because this is the broadest objective that enables us to derive a closed-form solution. Matching attention probabilities, or jointly optimizing the query and key projections, leads to a nonlinear problem with no closed-form minimizer. Reducing the number of key heads requires deciding which original key heads should be merged together. We formalize this choice using a permutation σ=(σ1,…,σ2hkS)σ=( _1,…, _2h_k^S) of 1,…,2hkS\1,…,2h_k^S\. The permutation specifies a partition of the original key heads into hkSh_k^S disjoint pairs: (σ2i−1,σ2i),i=1,…,hkS.( _2i-1, _2i), i=1,…,h_k^S. (12) Each pair is merged into a single SAGA key head. Once σ is determined, we apply the induced permutations to the query, value, and output projection matrices so that they remain aligned with the reordered key heads and each merged key is shared by its intended value-head pair. The induced permutations definitions are given in App. B. Now we explain how we choose the permutation and merge the chosen key heads. For each teacher key head r∈1,…,2hkSr∈\1,…,2h_k^S\, let rQ_r denote the query heads that share KrW_K^r, and let j:=QjQ^j:=XW_Q^j denote the realized query matrices. For any pair of teacher key heads a,b∈1,…,2hkSa,b∈\1,…,2h_k^S\ and candidate merged projection M, define the pairwise logit-matching loss ℒa,b():=∑j∈a _a,b(M):= _j _a ‖j()⊤−j(Ka)⊤‖F2 Q^j(XM) -Q^j(XW_K^a) _F^2 (13) +∑j∈b + _j _b ‖j()⊤−j(Kb)⊤‖F2 Q^j(XM) -Q^j(XW_K^b) _F^2 where ∥A∥F=∑i,jAij2 A _F= _i,jA_ij^2 denotes the Frobenius norm of a matrix A. We omit the attention scaling factor 1/d1/ d, since it multiplies the entire objective by the constant 1/d1/d and therefore does not affect the minimizer. Our objective is to jointly choose the pairing σ and the merged key projections [1:hkS]M^[1:h_k^S] so as to minimize the total squared logit reconstruction error over the dataset: minσmin[1:hkS]∑i=1hkSℒσ2i−1,σ2i(i). _σ _M^[1:h_k^S] _i=1^h_k^SL_ _2i-1, _2i(M^i). (14) Lemma 4.1. For any pair of teacher key heads a,b∈1,…,2hkSa,b∈\1,…,2h_k^S\, an optimal merged key projection minimizing ℒa,b()L_a,b(M) is ^K(a,b)=(Kaa+Kbb)(a+b)†. W_K(a,b)= (W_K^aS_a+W_K^bS_b ) (S_a+S_b ) . (15) where r:=∑j∈r(j)⊤(j)S_r:= _j _r(Q^j) (Q^j), and † denotes the Moore–Penrose pseudoinverse. The proof of Lemma 4.1 is given in App. A. Let G=(V,E,w)G=(V,E,w) be the complete weighted graph with vertex set V=1,…,2hkSV=\1,…,2h_k^S\, where for each pair of teacher key heads a,b\a,b\, E contains an edge with weight w(a,b):=ℒa,b(^K(a,b))=minℒa,b().w(a,b):=L_a,b( W_K(a,b))= _ML_a,b(M). (16) where the second equality follows from Lemma 4.1. Theorem 4.2. A minimum-weight perfect matching of G yields an optimal solution to Equation 14, with σ∗σ^* being a permutation whose consecutive pairs are the matched edges, and for each matched pair (σ2i−1∗,σ2i∗)( _2i-1^*, _2i^*) the corresponding optimal merged key projection is ^Ki=^K(σ2i−1∗,σ2i∗). W_K^i= W_K( _2i-1^*, _2i^*). (17) Proof. For a fixed permutation σ, we look at the inner objective in Eq. 14. This objective is a sum of hkSh_k^S independent pairwise terms, one for each pair (σ2i−1,σ2i)( _2i-1, _2i). Therefore min[1:hkS]∑i=1hkSℒσ2i−1,σ2i(i)=∑i=1hkSminiℒσ2i−1,σ2i(i)=∑i=1hkSw(σ2i−1,σ2i), _M^[1:h_k^S] _i=1^h_k^SL_ _2i-1, _2i(M^i)= _i=1^h_k^S _M^iL_ _2i-1, _2i(M^i)= _i=1^h_k^Sw( _2i-1, _2i), (18) Moreover, for each i, the minimizing choice is ^Ki=^K(σ2i−1,σ2i) W_K^i= W_K( _2i-1, _2i). Thus the value of Eq. 14 for a given permutation σ is exactly the total weight of the perfect matching encoded by its consecutive pairs. Any perfect matching of G can be written as such a permutation, and conversely any such permutation defines a perfect matching. Minimizing Eq. 14 over σ is therefore exactly the minimum-weight perfect matching problem on G. ∎ Figure 2: End-to-End speedup over baseline full attention as a function of sequence length, for Llama 3.2-1B (hk:8→4h_k:8→ 4, H200) and Qwen2.5-1.5B (hk:2→1h_k:2→ 1, B200) at various batch sizes. Merging alone (SAGA full) gives a consistent speedup across all configurations. Combining merging with Atop-N attention adds overhead at short sequences but provides substantial gains at long contexts. 5 Experiments 5.1 Inference decoding speedups We measure end-to-end per-token decode latency on Llama 3.2-1B [39] (hk:8→4h_k:8→ 4) and Qwen2.5-1.5B [6] (hk:2→1h_k:2→ 1) using compiled CUDA graphs. We compare the baseline model (original GQA head configuration) against the SAGA model (merged key heads), each evaluated with full attention and with Atop-N attention at N∈512,1024,2048N∈\512,1024,2048\. We sweep sequence lengths from 8K to 128K at different batch sizes, reporting mean per-token latency over 10 runs of 32 generated tokens each. Fig. 2 presents end-to-end speedup over baseline full attention. To showcase different GPU architectures, Llama is evaluated on H200 while Qwen is evaluated on the newer B200 (complementary results in App. D). Merging key heads alone (SAGA full) provides a modest but consistent speedup across all sequence lengths and batch sizes, reflecting the reduced query-key cost from fewer key heads (see App. D for a kernel-level profile, larger models and additional ablations). This benefit grows with batch size as queries compete for memory bandwidth: on Llama at batch size 16, SAGA full attention is up to 1.47×1.47× faster than baseline full at 128K context. Combining merged heads with Atop-N attention introduces overhead at short sequences due to the threshold calibration and sparse gather, but yields substantial speedups at long contexts. At 128K with batch size 1, SAGA Atop-N achieves up to a 2.2×2.2× end-to-end speedup on Qwen and 2.1×2.1× on Llama over baseline full attention. Even at batch size 16, where the Atop-N overhead is larger, SAGA Atop-N reaches 2.2×2.2× and 2.6×2.6× speedups at 128K on Qwen and Llama respectively. We note that flash attention [40] provides little benefit during autoregressive decode, where the attention computation is a matrix-vector product with no quadratic attention matrix to tile. Our full-attention baseline uses faster standard compiled attention, enabling portable support for the asymmetric GQA and ensuring a fair comparison when applying Atop-N attention. To assess the quality impact of the Atop-N approximation, we evaluate all configurations on the full RULER benchmark [41] at context lengths of 8K, 16K, and 32K in Sec. 5.3 below. Fig. 4 reports per-task accuracy. On Qwen2.5-1.5B, Atop-N retains 83–97% of full attention quality across all context lengths, with most tasks virtually unaffected. The degradation is concentrated in multi-key retrieval tasks, which require the model to simultaneously attend to many dispersed keys – a setting where aggressive sparsity discards relevant entries. 5.2 Training from scratch Figure 3: Training from scratch on SmolLM2-360M with hq=16h_q=16 query heads. Left: validation perplexity (↓ ). Right: validation accuracy (↑ ). SAGA (hk,hv)=(1,8)(h_k,h_v)=(1,8) closely tracks the high-capacity baseline (8,8)(8,8) while outperforming both the matched-key-budget (1,1)(1,1) baseline, and the parameter-matched baseline (4,4)(4,4). Solid curves are means of 3 seeds; dotted lines indicate ±1± 1 std. We train SAGA from scratch on the SmolLM2-360M [42] backbone, leaving the model unchanged except for the head counts. The number of query heads is fixed to hq=16h_q=16, while hkh_k and hvh_v are varied. SAGA with (hk,hv)=(1,8)(h_k,h_v)=(1,8) is compared against three GQA baselines: high-capacity with (hk,hv)=(8,8)(h_k,h_v)=(8,8), low-capacity with (hk,hv)=(1,1)(h_k,h_v)=(1,1) matching the key budget of SAGA, and parameter-matched with (hk,hv)=(4,4)(h_k,h_v)=(4,4) roughly matching the parameter count of SAGA. This isolates the contribution of decoupling hvh_v from hkh_k. Each configuration is run with three seeds. Sec. D.2 shows that a larger-scale Qwen3-0.6B experiment trained on 50B tokens exhibits the same trend. Fig. 3 shows that SAGA (1,8)(1,8) closely matches the quality of the (8,8)(8,8) baseline despite using a little over half the KV cache of (8,8)(8,8). It also clearly outperforms the (1,1)(1,1) baseline as well as the (4,4)(4,4) baseline despite having a similar number of parameters. This result suggests that the improvement comes from allocating capacity asymmetrically. The same holds for how the key budget itself is spent: a Sigma-style variant [23] that reduces the key-head dimension instead of reducing the number of heads, (hk,dk)=(2,32)(h_k,d_k)=(2,32), matches SAGA’s (1,64)(1,64) in KV-cache size and nearly in parameter count (both with hv=8h_v=8, dv=64d_v=64), yet reaches 41.7% final validation accuracy versus 42.7% for SAGA. This is consistent with the results of the SIGMA paper, who favored reducing the number of key heads while leaving the key-head dimension unchanged. Combined with Atop-N attention at inference, this translates to the substantial decoding speedups reported in Sec. 5.1 at a small cost in quality. Our from-scratch experiments are necessarily limited in scale, and we hope future work will validate SAGA’s benefits in a large-scale pretraining setting. 5.3 Application to pretrained models To evaluate SAGA on models with realistic capabilities, we convert pretrained models to the SAGA architecture via knowledge distillation, avoiding the compute and data requirements of training a competitive LLM from scratch. We start from Qwen2.5-1.5B [6], a highly capable model with only hk=2h_k=2 key heads, where the matching from Sec. 4 is trivial (a single pair) and reduces to averaging (hk:2→1h_k:2→ 1). We then follow up with Llama 3.2-1B [39], using the full matching-based initialization from Sec. 4 to merge 8 key heads into 4 (see Sec. B.1 for a comparison of merging strategies). Ideally, a fair comparison would increase the student’s capacity to compensate for the reduced key heads, for example by increasing the head dimension or adding value heads. However, as teacher and student representations diverge, a short distillation cannot effectively leverage the teacher’s signal. For simplicity we keep the student architecture identical to the teacher except for the reduced hkh_k, and accept the slight quality degradation from this conservative approach. Figure 4: RULER accuracy for full attention and Atop-N attention at N∈512,1024,2048N∈\512,1024,2048\, across context lengths 8K, 16K, and 32K on Qwen2.5-1.5B, for different task categories. SAGA with Atop-N remains near baseline quality, while alternative eviction-based methods degrade substantially at longer contexts. The student is initialized by merging the teacher’s key projection heads and trained with a combination of KL divergence from the teacher and hard-label cross-entropy. We train in several stages to expose the student to increasing context lengths while keeping total training short at only around 23K iterations (architecture dependent). We use a mix of SlimPajama [43], QA [44, 45] and long-context data [46, 47], as well as synthetic needle-in-a-haystack examples. We note that the student is heavily weighted to match the teacher’s output probabilities on these sequences, so its retained performance reflects architectural capacity rather than task-specific adaptation. Full hyperparameters and data mix details are provided in Sec. C.3. The distilled student retains 96.3% of the teacher’s aggregate RULER score across 8K, 16K, and 32K context lengths, with per-task results shown in Fig. 4. A complementary Llama 3.2-1B comparison is shown in App. D. As shown in Tab. 1, the SAGA student remains close to the teacher on general-capability benchmarks despite the reduced key head capacity. For comparison, we evaluate alternative decode-acceleration methods: sliding window with attention sinks (SW+S) [30], H2O [29], and KIVI [48]. These methods fall below 50% aggregate RULER accuracy at 8K context and degrade rapidly at longer contexts, despite the 4096-token budget for sliding window and H2O. We note that SAGA could readily combine with KV cache quantization and eviction methods, such as those used for comparison above. Furthermore, to experiment at a larger scale, we evaluate Llama3.1-70B-Instruct [39] with AtopN only. Since AtopN requires no retraining, these experiments remain feasible with limited computational resources. To highlight that the AtopN approximation holds over long generations, we evaluate on tasks that elicit extended outputs before a final answer: GPQA Main, limited to 8K output tokens, and GPQA Diamond, limited to 16K output tokens [49]. Both consist of graduate-level questions spanning multiple domains, however Diamond is a harder subset. In practice, the model produces long generations on these tasks, averaging over 1.2K output tokens on GPQA Main and 2.1K output tokens on GPQA Diamond. We experiment with N∈2048,1024,512N∈\2048,1024,512\ and find that GPQA Main accuracy is unharmed even at N=512N=512, whereas GPQA Diamond incurs some degradation. Results are available in Table 2 Table 1: General-capability benchmarks for the teacher and SAGA student models. Qwen2.5-1.5B (hk:2→1h_k:2→ 1) Llama 3.2-1B (hk:8→4h_k:8→ 4) Baseline SAGA Baseline SAGA HellaSwag (acc_norm ↑ ) 0.678 0.644 0.643 0.583 WikiText (word ppl ↓ ) 12.10 12.43 11.98 13.09 Table 2: Llama3.1-70B-Instruct accuracy on GPQA with AtopN for varying N. GPQA Main is limited to 8K output tokens and GPQA Diamond to 16K. N GPQA Diamond @16K (↑ ) GPQA Main @8K (↑ ) Baseline 0.4343 0.4241 2048 0.3939 0.4129 1024 0.3838 0.4330 512 0.3586 0.4241 6 Limitations, discussion and conclusion We introduced SAGA, an attention architecture designed around the computational trade-offs of sparse decoding. By decoupling key and value head counts, SAGA directly addresses the query-key bottleneck that dominates under sparse attention, while preserving model capacity through additional value heads. Our experiments show that SAGA nearly matches equivalent GQA models when trained from scratch and retains 96.3% of the teacher’s aggregate RULER score when distilled from pretrained checkpoints, while delivering substantial inference end-to-end speedups of more than 2×2× at long contexts. Several limitations point to promising future directions. First, our current implementation relies on the simple Atop-N method; integrating hardware-aware sparse attention kernels could yield further speedup gains. Second, our distilled SAGA models maintain the same architecture as the teacher aside from the reduced key projections — future work could compensate for this slight capacity reduction by reallocating parameters, for example by increasing the value head count or dimension. Finally, our quality experiments are limited to models of up to 1.5B parameters; validating SAGA’s benefits in a full-scale pretraining setting remains an important open question. More broadly, we believe that co-designing attention architectures with their inference-time sparsity patterns is a fertile direction, and hope that SAGA serves as a step toward making this a standard consideration in LLM design. Acknowledgments and Disclosure of Funding The research of DS was funded by the European Union (ERC, A-B-C-Deep, 101039436). Views and opinions expressed are however those of the author only and do not necessarily reflect those of the European Union or the European Research Council Executive Agency (ERCEA). Neither the European Union nor the granting authority can be held responsible for them. The research of DS was also supported by the Council for Higher Education in Israel under the Moonshot Project. References [1] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin (2017) Attention is all you need. Advances in neural information processing systems 30. Cited by: §1, §2.1. [2] A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, I. Sutskever, et al. (2019) Language models are unsupervised multitask learners. OpenAI blog 1 (8), p. 9. Cited by: §1. [3] T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, et al. (2020) Language models are few-shot learners. Advances in neural information processing systems 33, p. 1877–1901. Cited by: §1. [4] H. Touvron, T. Lavril, G. Izacard, X. Martinet, M. Lachaux, T. Lacroix, B. Rozière, N. Goyal, E. Hambro, F. Azhar, et al. (2023) Llama: open and efficient foundation language models. arXiv preprint arXiv:2302.13971. Cited by: §1. [5] A. Yang, B. Yang, B. Hui, B. Zheng, B. Yu, C. Zhou, C. Li, C. Li, D. Liu, F. Huang, G. Dong, H. Wei, H. Lin, J. Tang, J. Wang, J. Yang, J. Tu, J. Zhang, J. Ma, J. Xu, J. Zhou, J. Bai, J. He, J. Lin, K. Dang, K. Lu, K. Chen, K. Yang, M. Li, M. Xue, N. Ni, P. Zhang, P. Wang, R. Peng, R. Men, R. Gao, R. Lin, S. Wang, S. Bai, S. Tan, T. Zhu, T. Li, T. Liu, W. Ge, X. Deng, X. Zhou, X. Ren, X. Zhang, X. Wei, X. Ren, Y. Fan, Y. Yao, Y. Zhang, Y. Wan, Y. Chu, Y. Liu, Z. Cui, Z. Zhang, and Z. Fan (2024) Qwen2 technical report. arXiv preprint arXiv:2407.10671. Cited by: §1. [6] Q. Team (2024) Qwen2.5: a party of foundation models. External Links: Link Cited by: §1, §5.1, §5.3. [7] A. Wang, Y. Pruksachatkun, N. Nangia, A. Singh, J. Michael, F. Hill, O. Levy, and S. Bowman (2019) Superglue: a stickier benchmark for general-purpose language understanding systems. Advances in neural information processing systems 32. Cited by: §1. [8] K. Papineni, S. Roukos, T. Ward, and W. Zhu (2002) Bleu: a method for automatic evaluation of machine translation. In Proceedings of the 40th annual meeting of the Association for Computational Linguistics, p. 311–318. Cited by: §1. [9] J. Wei, X. Wang, D. Schuurmans, M. Bosma, F. Xia, E. Chi, Q. V. Le, D. Zhou, et al. (2022) Chain-of-thought prompting elicits reasoning in large language models. Advances in neural information processing systems 35, p. 24824–24837. Cited by: §1. [10] Z. Shen, M. Zhang, H. Zhao, S. Yi, and H. Li (2021) Efficient attention: attention with linear complexities. In Proceedings of the IEEE/CVF winter conference on applications of computer vision, p. 3531–3539. Cited by: §1. [11] F. D. Keles, P. M. Wijewardena, and C. Hegde (2023) On the computational complexity of self-attention. In International conference on algorithmic learning theory, p. 597–619. Cited by: §1. [12] R. Pope, S. Douglas, A. Chowdhery, J. Devlin, J. Bradbury, J. Heek, K. Xiao, S. Agrawal, and J. Dean (2023) Efficiently scaling transformer inference. Proceedings of Machine Learning and Systems 5, p. 606–624. Cited by: §1, §2.2. [13] M. Davies, N. Crago, K. Sankaralingam, and C. Kozyrakis (2025) LIMINAL: exploring the frontiers of llm decode performance. External Links: 2507.14397, Link Cited by: §1. [14] J. Ainslie, J. Lee-Thorp, M. De Jong, Y. Zemlyanskiy, F. Lebrón, and S. Sanghai (2023) Gqa: training generalized multi-query transformer models from multi-head checkpoints. arXiv preprint arXiv:2305.13245. Cited by: §1, §2.3. [15] R. Child (2019) Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509. Cited by: §1, §2.4. [16] I. Beltagy, M. E. Peters, and A. Cohan (2020) Longformer: the long-document transformer. arXiv preprint arXiv:2004.05150. Cited by: §1, §2.4. [17] M. Zaheer, G. Guruganesh, K. A. Dubey, J. Ainslie, C. Alberti, S. Ontanon, P. Pham, A. Ravula, Q. Wang, L. Yang, et al. (2020) Big bird: transformers for longer sequences. Advances in neural information processing systems 33, p. 17283–17297. Cited by: §1, §2.4. [18] N. Kitaev, Ł. Kaiser, and A. Levskaya (2020) Reformer: the efficient transformer. External Links: 2001.04451, Link Cited by: §1. [19] J. Yuan, H. Gao, D. Dai, J. Luo, L. Zhao, Z. Zhang, Z. Xie, Y. X. Wei, L. Wang, Z. Xiao, Y. Wang, C. Ruan, M. Zhang, W. Liang, and W. Zeng (2025) Native sparse attention: hardware-aligned and natively trainable sparse attention. External Links: 2502.11089, Link Cited by: §1, §2.4. [20] A. Gupta, G. Dar, S. Goodman, D. Ciprut, and J. Berant (2021) Memory-efficient transformers via top-k attention. arXiv preprint arXiv:2106.06899. Cited by: §1, §2.5, §2.5. [21] P. Wang, X. Wang, F. Wang, M. Lin, S. Chang, H. Li, and R. Jin (2022) Kvt: k-n attention for boosting vision transformers. In European conference on computer vision, p. 285–302. Cited by: §1. [22] Y. Chen, L. Zhang, J. Shang, Z. Zhang, T. Liu, S. Wang, and Y. Sun (2024) DHA: learning decoupled-head attention from transformer checkpoints via adaptive heads fusion. In Advances in Neural Information Processing Systems, Vol. 37. Cited by: §2.3. [23] Z. Lin, Z. Tang, X. Liu, Y. Gong, Y. Chang, J. Jia, J. Li, Y. Shen, C. Lin, N. Duan, and W. Chen (2025) Sigma: differential rescaling of query, key and value for efficient language models. External Links: 2501.13629, Link Cited by: §2.3, §5.2. [24] A. Liu, B. Feng, B. Wang, B. Wang, B. Liu, C. Zhao, C. Dengr, C. Ruan, D. Dai, D. Guo, et al. (2024) Deepseek-v2: a strong, economical, and efficient mixture-of-experts language model. arXiv preprint arXiv:2405.04434. Cited by: §2.3. [25] DeepSeek-AI (2025) DeepSeek-V3.2-Exp: boosting long-context efficiency with DeepSeek sparse attention. Note: https://github.com/deepseek-ai/DeepSeek-V3.2-Exp Cited by: §2.3. [26] J. Zhang, C. Xiang, H. Huang, H. Xi, J. Zhu, J. Chen, et al. (2025) SpargeAttention: accurate and training-free sparse attention accelerating any model inference. In Forty-second International Conference on Machine Learning, Cited by: §2.4. [27] J. Yuan, C. Shinn, K. Xu, J. Cui, G. Klimiashvili, G. Xiao, P. Zheng, B. Li, Y. Zhou, Z. Ye, et al. (2025) BLASST: dynamic blocked attention sparsity via softmax thresholding. arXiv preprint arXiv:2512.12087. Cited by: §2.4. [28] D. Ohayon, I. Lamprecht, I. Hubara, I. Cohen, D. Soudry, and N. Elata (2025) Block sparse flash attention. arXiv preprint arXiv:2512.07011. Cited by: §2.4. [29] Z. Zhang, Y. Sheng, T. Zhou, T. Chen, L. Zheng, R. Cai, Z. Song, Y. Tian, C. Ré, C. Barrett, et al. (2023) H2o: heavy-hitter oracle for efficient generative inference of large language models. Advances in Neural Information Processing Systems 36, p. 34661–34710. Cited by: §2.4, §5.3. [30] G. Xiao, Y. Tian, B. Chen, S. Han, and M. Lewis (2024) Efficient streaming language models with attention sinks. External Links: 2309.17453, Link Cited by: §2.4, §5.3. [31] Y. Li, Y. Huang, B. Yang, B. Venkitesh, A. Locatelli, H. Ye, T. Cai, P. Lewis, and D. Chen (2024) Snapkv: llm knows what you are looking for before generation. Advances in Neural Information Processing Systems 37, p. 22947–22970. Cited by: §2.4. [32] J. Tang, Y. Zhao, K. Zhu, G. Xiao, B. Kasikci, and S. Han (2024) Quest: query-aware sparsity for efficient long-context llm inference. arXiv preprint arXiv:2406.10774. Cited by: §2.4. [33] A. Liu, J. Liu, Z. Pan, Y. He, G. Haffari, and B. Zhuang (2024) Minicache: kv cache compression in depth dimension for large language models. Advances in Neural Information Processing Systems 37, p. 139997–140031. Cited by: §2.4. [34] Y. Sheng, L. Zheng, B. Yuan, Z. Li, M. Ryabinin, B. Chen, P. Liang, C. Ré, I. Stoica, and C. Zhang (2023) Flexgen: high-throughput generative inference of large language models with a single gpu. In International Conference on Machine Learning, p. 31094–31116. Cited by: §2.4. [35] Y. Wu, M. N. Rabe, D. Hutchins, and C. Szegedy (2022) Memorizing transformers. arXiv preprint arXiv:2203.08913. Cited by: §2.5, §2.5. [36] A. Bertsch, U. Alon, G. Neubig, and M. Gormley (2023) Unlimiformer: long-range transformers with unlimited length input. Advances in Neural Information Processing Systems 36, p. 35522–35543. Cited by: §2.5. [37] C. Lou, Z. Jia, Z. Zheng, and K. Tu (2024) Sparser is faster and less is more: efficient sparse attention for long-range transformers. arXiv preprint arXiv:2406.16747. Cited by: §2.5. [38] T. Haris (2024) k N attention demystified: a theoretical exploration for scalable transformers. arXiv preprint arXiv:2411.04013. Cited by: §2.5. [39] A. Grattafiori A. Dubey et al. (2024) The llama 3 herd of models. External Links: 2407.21783, Link Cited by: §5.1, §5.3, §5.3. [40] T. Dao, D. Y. Fu, S. Ermon, A. Rudra, and C. Ré (2022) FlashAttention: fast and memory-efficient exact attention with io-awareness. External Links: 2205.14135, Link Cited by: §5.1. [41] C. Hsieh, S. Sun, S. Kriman, S. Acharya, D. Rekesh, F. Jia, Y. Zhang, and B. Ginsburg (2024) RULER: what’s the real context size of your long-context language models?. arXiv preprint arXiv:2404.06654. Cited by: §5.1. [42] L. B. Allal, A. Lozhkov, E. Bakouch, G. Penedo, N. Tazi, L. Tunstall, et al. (2025) SmolLM2: when smol goes big – data-centric training of a small language model. arXiv preprint arXiv:2502.02737. Cited by: §5.2. [43] D. Soboleva, F. Al-Khateeb, R. Myers, J. R. Steeves, J. Hestness, and N. Dey (2023) SlimPajama: a 627b token cleaned and deduplicated version of redpajama. Note: https://w.cerebras.ai/blog/slimpajama-a-627b-token-cleaned-and-deduplicated-version-of-redpajama External Links: Link Cited by: §5.3. [44] Z. Yang, P. Qi, S. Zhang, Y. Bengio, W. W. Cohen, R. Salakhutdinov, and C. D. Manning (2018) HotpotQA: a dataset for diverse, explainable multi-hop question answering. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, p. 2369–2380. External Links: Link Cited by: §5.3. [45] P. Rajpurkar, R. Jia, and P. Liang (2018) Know what you don’t know: unanswerable questions for SQuAD. In Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics, p. 784–789. External Links: Link Cited by: §5.3. [46] Y. Bai, X. Lv, J. Zhang, Y. He, J. Qi, L. Hou, J. Tang, Y. Dong, and J. Li (2024) LongAlign: a recipe for long context alignment of large language models. External Links: 2401.18058, Link Cited by: §5.3. [47] Y. Chen, S. Qian, H. Tang, X. Lai, Z. Liu, S. Han, and J. Jia (2024) LongLoRA: efficient fine-tuning of long-context large language models. External Links: 2309.12307, Link Cited by: §5.3. [48] Z. Liu, J. Yuan, H. Jin, S. Zhong, Z. Xu, V. Braverman, B. Chen, and X. Hu (2024) KIVI: a tuning-free asymmetric 2bit quantization for KV cache. In Proceedings of the 41st International Conference on Machine Learning, External Links: 2402.02750 Cited by: §5.3. [49] D. Rein, B. L. Hou, A. C. Stickland, J. Petty, R. Y. Pang, J. Dirani, J. Michael, and S. R. Bowman (2023) GPQA: a graduate-level google-proof q&a benchmark. External Links: 2311.12022, Link Cited by: §5.3. [50] I. Loshchilov and F. Hutter (2019) Decoupled weight decay regularization. In International Conference on Learning Representations, Cited by: §C.2. Appendix A Proof of Lemma 4.1 Recall that rQ_r denotes the query heads that share KrW_K^r. Define the realized query matrices j:=QjQ^j:=XW_Q^j, and fix a pair of teacher key heads a,b∈1,…,2hkSa,b∈\1,…,2h_k^S\. Recall the pairwise loss ℒa,b():=∑j∈a _a,b(M):= _j _a ‖j()⊤−j(Ka)⊤‖F2 Q^j(XM) -Q^j(XW_K^a) _F^2 (19) +∑j∈b + _j _b ‖j()⊤−j(Kb)⊤‖F2. Q^j(XM) -Q^j(XW_K^b) _F^2. Proof of Lemma 4.1. For any fixed j∈aj _a, we have ‖j()⊤−j(Ka)⊤‖F2=‖j((−Ka))⊤‖F2 Q^j(XM) -Q^j(XW_K^a) _F^2= Q^j (X(M-W_K^a) ) _F^2 (20) =tr(j((−Ka))⊤(−Ka)(j)⊤) =tr\! (Q^j (X(M-W_K^a) ) X(M-W_K^a)(Q^j) ) =tr((−Ka)⊤⊤(−Ka)(j)⊤(j)), =tr\! ((M-W_K^a) X X(M-W_K^a)(Q^j) (Q^j) ), where we used the trace form of the Frobenius norm, ∥A∥F2=tr(AA⊤) A _F^2=tr(A ), and the cyclic property of the trace. Denoting a:=∑j∈a(j)⊤(j),b:=∑j∈b(j)⊤(j)S_a:= _j _a(Q^j) (Q^j), S_b:= _j _b(Q^j) (Q^j) (21) and summing over j∈aj _a gives ∑j∈a‖j()⊤−j(Ka)⊤‖F2=tr((−Ka)⊤⊤(−Ka)a). _j _a Q^j(XM) -Q^j(XW_K^a) _F^2=tr\! ((M-W_K^a) X X(M-W_K^a)S_a ). (22) Applying the same argument to the second summation term, ℒa,b()L_a,b(M) is equivalent to minimizing tr((−Ka)⊤⊤(−Ka)a)+tr((−Kb)⊤⊤(−Kb)b).tr\! ((M-W_K^a) X X(M-W_K^a)S_a )+tr\! ((M-W_K^b) X X(M-W_K^b)S_b ). (23) This is a convex quadratic function of M. We find the minimizer by taking the derivative with respect to M. Using the identity ∂tr(⊤)=2 ∂Mtr\! (M BMC )=2BMC for symmetric positive semidefinite B and C, and applying it with ←−KaM←M-W_K^a, ←⊤B←X X, and ←aC←S_a (and analogously for the second term), the gradient of the objective with respect to M is 2⊤(−Ka)a+2⊤(−Kb)b.2X X(M-W_K^a)S_a+2X X(M-W_K^b)S_b. (24) Setting this to zero gives ⊤((a+b)−Kaa−Kbb)=0.X X (M(S_a+S_b)-W_K^aS_a-W_K^bS_b )=0. (25) Hence any M satisfying (a+b)=Kaa+KbbM(S_a+S_b)=W_K^aS_a+W_K^bS_b (26) is optimal. One such solution is =(Kaa+Kbb)(a+b)†.M= (W_K^aS_a+W_K^bS_b ) (S_a+S_b ) . (27) Since aS_a and bS_b are positive semidefinite, the rows of Kaa+KbbW_K^aS_a+W_K^bS_b lie in the row space of a+bS_a+S_b, so this choice indeed satisfies the linear system above. ∎ Appendix B Induced permutations for query, value, and output projections Once a matching has been chosen, let σ=(σ1,…,σ2hkS)σ=( _1,…, _2h_k^S) be any permutation of 1,…,2hkS\1,…,2h_k^S\ whose consecutive pairs are exactly the matched pairs, i.e. (σ2i−1,σ2i)( _2i-1, _2i) for i=1,…,hkSi=1,…,h_k^S. The purpose of the induced permutations is to re-index the teacher heads before we merge its two key heads so that each matched pair becomes consecutive. We first permute the key and value heads by σ: ~K:=[Kσ1⋯Kσ2hkS],~V:=[Vσ1⋯Vσ2hkS]. W_K:= [W_K _1\;\;·s\;\;W_K _2h_k^S ], W_V:= [W_V _1\;\;·s\;\;W_V _2h_k^S ]. (28) For the query heads, let g:=hqhkT=hq2hkS,Br:=((r−1)g+1,…,rg),r=1,…,2hkS,g:= h_qh_k^T= h_q2h_k^S, B_r:= ((r-1)g+1,…,rg ), r=1,…,2h_k^S, (29) so BrB_r is the block of query heads that share the teacher key head r. The associated query-head permutation πQ _Q is the concatenation of these blocks into a single sequence of length hqh_q: πQ:=(Bσ1∥Bσ2∥⋯∥Bσ2hkS)∈1,…,hqhq, _Q:= (B_ _1\|B_ _2\|·s\|B_ _2h_k^S )∈\1,…,h_q\^h_q, (30) where ∥\| denotes concatenation, so that πQ(k) _Q(k) for k=1,…,hqk=1,…,h_q refers to the k-th entry of this flattened sequence. Writing the query projection in head blocks, Q=[Q1⋯Qhq],W_Q= [W_Q^1\;\;·s\;\;W_Q^h_q ], (31) the permuted query projection is ~Q:=[QπQ(1)⋯QπQ(hq)]. W_Q:= [W_Q _Q(1)\;\;·s\;\;W_Q _Q(h_q) ]. (32) For the output projection, write OW_O as a vertical stack of query-head input blocks, O=[O,1O,hq],O,r∈ℝd×D.W_O= bmatrixW_O,1\\ \\ W_O,h_q bmatrix, W_O,r ^d× D. (33) We permute these input blocks by the same query-head permutation: ~O:=[O,πQ(1)O,πQ(hq)]. W_O:= bmatrixW_O, _Q(1)\\ \\ W_O, _Q(h_q) bmatrix. (34) Before merging, this re-indexing leaves the attention module unchanged. After merging, the two value heads at positions 2i−12i-1 and 2i2i and their associated 2g2g query heads are exactly the heads that should share the merged SAGA key head i. B.1 Merging strategies ablation We compare three strategies for initializing the merged SAGA key projections on Llama 3.2-1B, measuring cross-entropy loss and per-token KL divergence from the teacher immediately after initialization (before distillation): • Sequential pairing: merges consecutive head pairs in their original tensor order (effectively arbitrary pairing) by averaging their key projection matrices. • Key matching: minimum-weight perfect matching with edge weight w(a,b):=∥Ka−Kb∥F2,w(a,b):= XW_K^a-XW_K^b _F^2, (35) i.e. the same matching procedure as ours, but the graph weights are based only on the realized key matrices instead of attention logits. Here, the minimizer is the average 12(Ka+Kb) 12(W_K^a+W_K^b). • Logit matching (ours): the minimum-weight perfect matching on attention logit reconstruction error described in Sec. 4.1. Table 3: Initialization quality before distillation (Llama 3.2-1B). Lower is better for both metrics. Reduction Method CE KL 8→48→ 4 Sequential pairing 7.231 4.754 8→48→ 4 Key matching 7.727 5.254 8→48→ 4 Logit matching (ours) 6.748 4.271 8→28→ 2 Sequential pairing 10.306 7.819 8→28→ 2 Key matching 9.077 6.596 8→28→ 2 Logit matching (ours) 7.597 5.125 Logit-matching initialization consistently achieves the lowest reconstruction error across both reduction ratios, with the advantage growing as the reduction becomes more aggressive (8→28→ 2). The two other merging strategies do not have a clear superior method. Sequential pairing is better at →48\!→\!4, and key matching at →28\!→\!2. Appendix C Implementation details Code implementation of our experiments is included in https://github.com/noamelata/SAGA. C.1 Atop-N attention Our Atop-N attention is an efficient implementation of the algorithm described in Sec. 3.1. Specifically, it uses a pre-calibrated per-head threshold to collect all entries that exceed it, greatly improving runtimes. This comparison operation uses a custom portable Triton22 2 https://github.com/triton-lang/triton, MIT license. kernel, replacing pytorch’s CUDA-optimized torch.topk used for true top-N values. During the first decode step after prefill, we compute exact attention probabilities via softmax and extract the N-th largest value per head using torch.topk. This value serves as the threshold for the subsequent n decode steps. Because the attention distribution changes slowly across consecutive decode steps, we use n=32n=32, thus amortizing the computational cost of torch.topk across 32 decode steps. Alongside with the reduced probability-value multiplication, computing softmax only on the top-N values is a significant source of speedup. Yet, we find that thresholds in logit space tend to perform poorly. Instead, we make use of a probability-space threshold τ transformed into the logit space, a much cheaper but mathematically equivalent operation. Given raw attention logits z, we approximate the max logit m m from a small tail of positions and compute the log-sum-exp denominator s^=∑iexp(zi−m^) s= _i (z_i- m). The logit-space threshold is then m^+log(τ⋅s^) m+ (τ· s), which requires only a single reduction pass over the full sequence rather than a full softmax. The selected indices are used to gather the corresponding value vectors, after which the logit-space weights are normalized via softmax (the −∞-∞ padding receives zero probability) and multiplied with the gathered values via regular batched matrix multiplication (for batch < 16) or sparse batched matrix multiplication (otherwise) to produce the attention output. Certain layers can be designated as exempt from Top-N selection and always use full attention. In our experiments, we exempt the first and 14th attention layers for Qwen2.5-1.5B33 3 https://huggingface.co/Qwen/Qwen2.5-1.5B, Apache 2.0 license. and the first layer for Llama 3.2-1B44 4 https://huggingface.co/meta-llama/Llama-3.2-1B, Llama 3.2 Community License., as these layers tend to have more diffuse attention patterns. The threshold comparison and index collection are implemented as a single Triton kernel. Each row of the attention weight matrix is partitioned into blocks (autotuned over 512, 1024, and 2048 elements). Within each block, elements exceeding the threshold are identified via a comparison mask, their positions computed via an intra-block cumulative sum, and results written to an output buffer using an atomic counter to coordinate across blocks. The kernel outputs both the selected values and their indices into a fixed-size buffer of capacity 2N2N, with unused slots initialized to −∞-∞. If the number of entries exceeding the threshold exceeds 2N2N, the atomic counter saturates and excess entries are silently dropped, preserving memory safety. The buffer is allocated once and reused across steps, ensuring compatibility with CUDA graph capture. C.2 Training SAGA from scratch We use the SmolLM2-360M55 5 https://huggingface.co/HuggingFaceTB/SmolLM2-360M, Apache 2.0 license. backbone. All models are trained on SlimPajama66 6 https://huggingface.co/datasets/cerebras/SlimPajama-627B, Apache 2.0 license. for 20K steps with batch size 256 and sequence length 1024, totaling ∼ 5.24B tokens. We use the AdamW optimizer [50] with (β1,β2)=(0.9,0.999)( _1, _2)=(0.9,0.999) and weight decay 0.1 and a cosine learning rate schedule with 500 warmup steps, peak learning rate 3×10−43× 10^-4, and minimum learning rate 7.5×10−57.5× 10^-5. Training is performed in bfloat16, and each configuration is run with 3 seeds. C.3 Distillation to SAGA The student model is initialized from the teacher by merging the key projection heads according to the procedure described in Section 4.1. We then reorder the teacher’s query, value, and output projection heads using the induced permutations described in Appendix B so that they remain aligned with the merged key heads, and copy all parameters unchanged from the teacher. The student is trained with a weighted combination of KL divergence from the teacher’s output distribution (weight α=0.7α=0.7, temperature T=3T=3) and hard-label cross-entropy (weight 1−α=0.31-α=0.3), using AdamW with fused implementation. Training proceeds in three stages with increasing context length. In Stage 1, the student trains for 15K steps at context length 4096 with a learning rate of 2×10−42× 10^-4 and global batch size 64. The learning rate schedule follows a warmup-stable-decay (WSD) pattern: 200 warmup steps, 9800 stable steps, and 5000 decay steps. The data mix is 75% SlimPajama (packed), 15% QA data (HotpotQA77 7 https://huggingface.co/datasets/hotpotqa/hotpot_qa, C BY-SA 4.0 license. and SQuAD88 8 https://huggingface.co/datasets/rajpurkar/squad_v2, C BY-SA 4.0 license., packed), and 10% synthetic needle-in-a-haystack examples. Stage 2 starts from the Stage 1 checkpoint with a fresh optimizer and trains for 5K steps at context length 20480 with a learning rate of 1×10−41× 10^-4 and global batch size 32. The WSD schedule uses 200 warmup, 3300 stable, and 1500 decay steps. The data mix shifts to 75% SlimPajama, 5% QA, 10% long-context documents (LongAlign-10k99 9 https://huggingface.co/datasets/THUDM/LongAlign-10k, LongAlpaca-12k1010 10 https://huggingface.co/datasets/Yukang/LongAlpaca-12k, individually padded), and 10% synthetic NIAH. Stage 3 continues from the Stage 2 checkpoint and trains at context length 36864 with a learning rate of 5×10−55× 10^-5 and global batch size 32. For Qwen2.5-1.5B: 2K steps, α=0.7α=0.7, T=3T=3, WSD schedule with 100 warmup and 500 decay steps. For Llama 3.2-1B: 4K steps, α=0.85α=0.85, T=4T=4, WSD schedule with 200 warmup and 800 decay steps. Stage 3 interleaves 20K-length microbatches (2 per gradient step) with the 36K-length data to prevent forgetting of shorter-range attention patterns. Training at 36864 tokens provides headroom over the 32K evaluation target. Total training time is approximately 60 hours on 8×H200 GPUs. C.4 Evaluation Quality is evaluated using the RULER benchmark1111 11 https://huggingface.co/datasets/THUDM/RULER, Apache 2.0 license. for long-context tasks, and HellaSwag1212 12 https://huggingface.co/datasets/Rowan/hellaswag, MIT license. and WikiText1313 13 https://huggingface.co/datasets/Salesforce/wikitext, C BY-SA 3.0 license. for general capabilities. All evaluations use the lm-eval-harness framework1414 14 https://github.com/EleutherAI/lm-evaluation-harness, MIT license.. Appendix D Additional results D.1 Qwen2.5-1.5B speedups on H200 Figure 2 reports Qwen2.5-1.5B latency measurements on B200. Fig. 5 shows the corresponding results on H200, confirming that SAGA’s speedup gains are consistent across GPU generations. Figure 5: Speedup over baseline full attention on Qwen2.5-1.5B (hk:2→1h_k:2→ 1) measured on H200 at various batch sizes and sequence lengths. D.2 Qwen3-0.6B trained from scratch on 50B tokens We train a Qwen3-0.6B-based architecture from scratch on 50B tokens. This model is approximately 1.7×1.7× larger than the SmolLM2-360M model evaluated in Sec. 5.2, and its training budget is 9.54×9.54× larger. We run a single seed for each configuration due to computational limitations. Tab. 4 shows that SAGA with (hk,hv)=(1,8)(h_k,h_v)=(1,8) is closer to the high-capacity (8,8)(8,8) baseline than to the matched-key-budget (1,1)(1,1) baseline on HellaSwag accuracy, WikiText perplexity, and validation accuracy. Thus, the same trend holds at this larger model and training scale. Table 4: From-scratch results for a Qwen3-0.6B-based architecture trained on 50B tokens. Configuration HellaSwag (acc_norm ↑ ) WikiText (word ppl ↓ ) Validation accuracy (↑ ) (hk,hv)=(1,1)(h_k,h_v)=(1,1) 0.3310±0.00470.3310± 0.0047 34.9129 0.440 (hk,hv)=(1,8)(h_k,h_v)=(1,8) 0.3343±0.00470.3343± 0.0047 33.8604 0.443 (hk,hv)=(8,8)(h_k,h_v)=(8,8) 0.3368±0.00470.3368± 0.0047 33.3382 0.444 D.3 Llama 3.2-1B RULER accuracy Figure 6 presents the per-task RULER accuracy breakdown for Llama 3.2-1B (hk:8→4h_k:8→ 4), complementing the Qwen2.5-1.5B results in the main text. The SAGA student retains the majority of baseline quality, with degradation again concentrated in multi-key retrieval tasks. Figure 6: Per-task RULER accuracy for Llama 3.2-1B (hk:8→4h_k:8→ 4) with full attention and Atop-N attention at N∈512,1024,2048N∈\512,1024,2048\, across context lengths 8K, 16K, and 32K. D.4 Kernel-level decode profile Tab. 5 profiles one decode step of Llama 3.2-1B at 64K context and batch size 16 on an H100. Under full attention the query-key and probability-value multiplications cost a comparable amount (870 vs. 1091 μ ). Merging key heads 8→48→ 4 halves the query-key term while leaving the probability-value term unchanged, and Atop-N reduces the probability-value term to 88 μ , leaving the query-key multiplication as the dominant remaining cost. The two components thus address the two bottlenecks of Fig. 1. Table 5: Decode-step profile of Llama 3.2-1B (hk:8→4h_k:8→ 4) at 64K context, batch size 16, on H100. Kernel times are μ per non-exempt layer per decode step; the last rows give end-to-end per-token latency. Kernel GQA (full) SAGA (full) SAGA + Atop-N 1024 Query-key multiplication 870.1 466.0 446.0 Probability-value multiplication 1090.9 1091.1 87.5 Softmax / threshold 249.8 257.3 107.7 Gather 1.2 1.2 88.8 Q/K/V/O projections 22.8 23.3 22.8 MLP 42.7 43.0 43.3 Layer total 2277.6 1881.9 796.0 End-to-end (ms / token) 40.57 31.08 18.99 Speedup 1.00×1.00× 1.31×1.31× 2.14×2.14× D.5 Decode latency at larger model scale Decode latency depends on tensor shapes and kernels rather than on weight values, so the efficiency of SAGA can be measured on larger architectures without training them. Tab. 6 reports per-token decode latency of the Qwen2.5-32B architecture (hk:8→4h_k:8→ 4) with untrained weights at 128K context on H100. The speedup of SAGA with Atop-N persists at this scale and grows with batch size, as in the smaller models, though at a slower rate. Quality at this scale is not evaluated. Table 6: Per-token decode latency (ms) of the Qwen2.5-32B architecture (hk:8→4h_k:8→ 4) at 128K context on H100. Batch 1 Batch 4 Batch 16 GQA + full attention 46.90 84.99 263.06 SAGA + Atop-N 1024 39.43 68.93 175.06 Speedup 1.19×1.19× 1.23×1.23× 1.50×1.50× D.6 Separating the head configuration from the selector Tab. 7 reports Qwen2.5-1.5B decode speedups for GQA and SAGA under full attention, Atop-N, and exact top-N, which recomputes the true top-N indices with torch.topk at every step. Comparing GQA with and without Atop-N isolates the selector; comparing SAGA against GQA under the same attention method isolates the head configuration. SAGA is faster under all three, e.g. 1.23×1.23× over GQA with Atop-N and 1.15×1.15× with exact top-N at batch size 4 and 128K, so the architectural gain does not depend on the threshold approximation of Atop-N. Exact top-N is slower than dense attention except at the longest contexts, which is what motivates Atop-N. Table 7: Decode speedup of Qwen2.5-1.5B relative to GQA with full attention at the same batch size (H100), N=1024N=1024. Batch Heads Attention 8K 16K 32K 64K 128K 1 GQA full 1.00 1.00 1.00 1.00 1.00 1 GQA Atop-N 0.93 1.00 1.17 1.43 1.88 1 GQA exact top-N 0.68 0.63 0.89 1.08 1.39 1 SAGA full 1.00 1.02 1.05 1.05 1.04 1 SAGA Atop-N 0.95 1.02 1.23 1.53 2.06 1 SAGA exact top-N 0.68 0.64 0.92 1.14 1.48 4 GQA full 1.00 1.00 1.00 1.00 1.00 4 GQA Atop-N 0.85 0.93 1.06 1.24 1.51 4 GQA exact top-N 0.67 0.73 0.82 0.96 1.10 4 SAGA full 1.05 1.05 1.08 1.11 1.14 4 SAGA Atop-N 0.89 0.98 1.16 1.41 1.86 4 SAGA exact top-N 0.69 0.75 0.87 1.05 1.27