SeerAttention: Self-distilled Attention Gating for Efficient Long-context Prefilling

We sincerely thank the reviewer for the constructive feedback and valuable suggestions. We address the main concerns below:

1. Lack of differences from other method

We respectfully disagree with the assessment that our method lacks novelty. SeerAttention introduces learnable sparse attention specifically for post-training, requiring only a small number of additional gate parameters that are distilled during this stage. In contrast, NSA and MoBA adopt a similar idea but focus on the pre-training phase. As shown in the comparison of trainable parameters and training tokens, NSA and MoBA require training all model weights with a large number of tokens and extensive GPU hours. Our approach, by contrast, is a lightweight and cost-effective self-distillation method.

We have discussed this point with another reviewer, as seen in this comment (https://openreview.net/forum?id=Nf8yfPDFTl&noteId=LV1xhwZYfg). In fact, both of the two works have recognized acknowledged the novelty of SeerAttention. (ACs could help verify this if needed.) Therefore, we believe our work is indeed novel and offers a new direction for developing sparse attention methods.

2. Fairness of comparisons (trainable parameters and self-distillation)

We would like to clarify that the “self‑distillation” in our method is different from conventional knowledge distillation as understood in the review. Our self‑distillation only trains the newly introduced gate parameters, while keeping the original model weights entirely frozen. The additional parameters are minimal—for example, in LLaMA‑3.1‑8B, the gate introduces approximately 101M parameters (about 101M/8B = 1.3% of the original model size); see also our response to Reviewer DFJm (https://openreview.net/forum?id=Nf8yfPDFTl&noteId=yPnYUXTMey).

As illustrated in Figure 1, this distillation process simply guides the gate to mimic the attention pattern of the frozen full‑attention model, here we call it "self‑distillation". Hence, no additional knowledge distillation is applied to the backbone model itself. For this reason, it is not meaningful to “train the baselines with similar distillation methods,” since our approach does not perform distillation on the original model’s weights.

3. Different settings of sparsity and block size.

We appreciate the suggestion to include experiments with different sparsity and block size settings.

For sparsity, SeerAttention allows runtime adjustment via the threshold T, providing flexibility to trade off efficiency and accuracy. The table below shows that as T increases, sparsity rises while accuracy decreases slightly, and SeerAttention consistently outperforms DuoAttention at comparable sparsity.

LongBench performance & efficiency results on Llama-3.1-8B-Instruct

* 50% streaming heads, the real sparsity < 50%

In RULER, we employ a threshold of 5e-4 under 4k-128k test benches.

RULER performance & efficiency results on Llama-3.1-8B-Instruct

For block size, we appreciate the reviewer’s interest in this factor, which was one of the key considerations in our system and algorithm design. Intuitively, using a larger block size reduces memory allocation overhead because the block‑wise attention map has shape $[bs, nh, seq/block, seq/block]$. However, at the algorithmic level, a larger block size leads to more coarse‑grained distillation. We have already conducted ablation studies in this direction. As shown below, larger block sizes generally improve perplexity at longer sequences, with block size = 64 providing the best overall trade‑off. Nevertheless, the optimal choice may depend on the downstream task, and we will explore additional experiments in the revision if this factor is deemed critical.

We hope our responses have addressed your concerns and clarified the contributions of our work. Please let us know if there are further questions or suggestions, as we are happy to provide additional information.

We sincerely thank the reviewer for your thoughtful and constructive feedback. We are glad that the reviewer recognized the effectiveness, practical speedup, and scalability of SeerAttention. The comment have helped us identify ways to improve the clarity and fairness of our presentation, and we have carefully addressed each point below.

1. Comparisons setting

Thank you for pointing out this issue. In our experiments, we directly used the original open-source implementations of MoA and DuoAttention. We agree that a more controlled comparison (using dense decoding for all methods) would be fairer. We are re-implement MoA and DuoAttention frameworks and doing the following the experiments settings: variable input length (prefilling) + variable output length (decoding) to fairly compare. Due to the rebuttal time limit, we could not re-implement and rerun all experiments with modified decoding settings.

2. SeerAttention Decoding

We appreciate the suggestion to discuss decode‑stage evaluation. In this work, we primarily focus on accelerating the prefill stage and leave decode‑stage optimization for future exploration. Prefill and decode are two distinct phases, and several PD‑disaggregation serving systems, such as DistServe [1] and Mooncake [2], have explored separating prefill and decode to improve TTFT/TPOT. Many prior works also focus on only one phase: for instance, MInference mainly targets prefill, while Quest [3] and H2O [4] focus solely on decode.

Our approach can naturally extend to decode sparsity with a minor modification of the AttnGate design. Specifically, we remove the Q‑compression branch (pooling + linear) and instead distill using a 1D max‑pooling ground truth of the attention map (along the sequence dimension) from the base model. In other words, during distillation for decode, we directly use the original Q without sequence‑level compression, enabling each token’s Q to select important KV blocks at inference.

We evaluated this design on the DeepSeek‑R1‑Distill‑Qwen‑14B reasoning model using two challenging math tasks and compared it against Quest. Below we provide a brief summary of the results (block size = 64). A more complete version can be included in the appendix upon acceptance. The results show that SeerAttention achieves higher accuracy than Quest, even with larger KV blocks and sparser attention.

Model: DeepSeek-R1-Distill-Qwen14B

3. Clarity of illustrations and explanations

3.1 Thank you for the suggestion. We will update the notation to use different superscripts to avoid confusion with the original model weights. The gate weights, $W_q^{\text{gate}}$ and $W_k^{\text{gate}}$, are newly introduced learnable parameters, and only the gate is trained during self-distillation. For LLaMA‑3.1‑8B, the gate contains 101M parameters (~1.3% of the total model parameters).

3.2 We appreciate the comment on the wording (dynamic vs. heterogeneous) and will revise it accordingly.

3.3 For mask generation, we illustrate the process using an example within a single head of one layer. Suppose the sequence length is $S$ and the block size is $B$. The output of AttnGate in this head is $o = [S/B, S/B]$.

• For the top‑$k$ method, we apply a torch.topk operation on each row, obtaining indices with shape $[S/B, k]$.

• For the threshold‑based method, we compute a binary mask as $mask = o > \text{threshold}$. Blocks marked as True in this mask are directly activated in our kernel.

4. Novelty relative to prior work

Block-sparse attention indeed dates back to much earlier works such as BigBird [5] and BlockSparse-BERT [6], and MInference is also a more recent approach that adopts block-sparse attention with pooling-based block estimation. However, our method is fundamentally different: pooling-based selection itself is not the core contribution here. The key novelty lies in introducing a learnable gating mechanism for sparse attention, similar to how a trainable router in MoE outperforms static pruning or heuristic compression methods. Our approach enables dynamic and data-dependent sparsity decisions guided by full-attention supervision, which sets it apart from heuristic or static approaches like MInference. Also, there are some following works like NSA [7] and MoBA [8] adopt this approach to pre-training stage.

5. Question 2: performance-efficiency trade-offs

Thank you for this helpful suggestions. We agree that our presentation of the performance–efficiency trade-off could be improved. While the paper already contains the necessary data (e.g., accuracy in Table 2 and speedup in Figure 6), we did not present them together in a unified figure or table, which may have reduced clarity.

To better illustrate this trade-off, we provide additional results on the 128k RULER benchmark below, showing how accuracy varies with sparsity under different thresholds (T = threshold):

For the kernel-level efficiency, we measured different sequence length latency under different sparsity setting:

For tradeoffs between speedup and accuracy, we employ a threshold of 5e-4 for all AttnGates under 4k-128k test benches and count the average sparsity and accuracy to present RULER performance & efficiency results on Llama-3.1-8B-Instruct.

These results demonstrate that SeerAttention achieves the best overall accuracy–efficiency trade‑off, offering higher sparsity and greater speedup without degrading accuracy compared to prior methods.

Reference:

[1] DistServe: Disaggregating Prefill and Decoding for Goodput-optimized Large Language Model Serving, https://www.usenix.org/system/files/osdi24-zhong-yinmin.pdf

[2] Mooncake: A KVCache-centric Disaggregated Architecture for LLM Serving, https://arxiv.org/abs/2407.00079

[3] Quest: Query-aware sparsity for efficient long-context llm inference, https://arxiv.org/abs/2406.10774

[4] H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models, https://arxiv.org/abs/2306.14048

[5] Big Bird: Transformers for Longer Sequences, https://arxiv.org/abs/2007.14062

[6] Blockwise Self-Attention for Long Document Understanding, https://arxiv.org/abs/1911.02972

[7] Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention, https://arxiv.org/abs/2502.11089

[8] MoBA: Mixture of Block Attention for Long-Context LLMs, https://arxiv.org/abs/2502.13189

Thank you very much for your helpful feedback and for raising your score. We appreciate your recognition of the improvements of SeerAttention and your suggestions regarding clarity on novelty, which have been valuable in refining the paper. Thanks again for your constructive comments!

We thank the reviewer for recognizing that our method is well-justified and thoughtfully designed, especially the motivation around dynamic sparsity in attention, the use of self-distillation with our efficient max-pooling kernel from both system and algorithm prospective. This aligns well with our goal to address the long-context prefilling challenge.

1. Training cost vs accuracy gain

We appreciate your suggestion to clarify the training cost and parameter overhead. In Llama-3-8B-Instruct model,

• SeerAttention: ~40 A100 hours in total (5 hours on 8×A100)

• DuoAttention: “several hours on 8×A100”

• MInference: ~2 hours on 1×A100 for sparsity profiling/calibration but is less accurate and slower in inference

SeerAttention requires a training budget comparable to or lower than DuoAttention. Both methods train only the gating module without updating the base model weights, so the cost mainly comes from forward passes. Compared to post-training approaches such as SFT or RL finetuning, this cost is negligible.

For accuracy, the gain should be considered together with sparsity. DuoAttention has a fixed sparsity ratio (~50%), whereas SeerAttention achieves similar accuracy at much higher sparsity (80–90% on 128k RULER). MInference is faster to profile since it only calibrates a sparsity head on a given dataset, but it is less accurate. Moreover, MInference often suffers from runtime overhead caused by head-by-head index search and execution flow, as reported in GitHub issues (e.g., https://github.com/microsoft/MInference/issues/43), where users observed that MInference can even run slower than the dense baseline.

As shown in Fig. 9 (Appendix A.1), we design customized training kernels to distill the gate, and the memory/latency overhead during distillation is < 5% compared to native FlashAttention. For LLaMA‑3.1‑8B, the gate adds only 101 M parameters (~1.3% of the model), making the method lightweight and scalable.

About training data:

Thank you for this question. Selecting training data is less critical in our case as the gate only learns the distribution instead of the real end-to-end knowledge from the data. Regarding training data, our gate learns to approximate distribution patterns rather than full knowledge, so dataset choice is less critical. We use RedPajama [1] for its popularity in long-context training. MInference and DuoAttention rely on downstream or BookSum [2] datasets, and none of these works provide ablation on dataset choice. This remains an open question but does not undermine our method’s core merit.

2. Use GT directly in inference

We appreciate the reviewer’s interest in this question, which was a key consideration in our system–algorithm co-design. Our experiments show that directly using GT to select sparse blocks results in virtually no accuracy loss, even under high sparsity. However, computing GT requires evaluating the full quadratic softmax term, $\text{softmax}(QK^\top)$. If this computation is already performed, there is little performance benefit from only applying sparsity to the $\text{AttentionScore} \cdot V$ part.

In summary, GT is accurate but expensive, which motivates the need for a fast, low-cost AttnGate.

Regarding the question “How much accuracy is lost compared to using GT?”, intuitively:

• Full attention provides the upper bound for GT-based sparsity.

• GT-based sparsity provides the upper bound for our AttnGate-based sparsity.

As noted above, GT computation is even more expensive than full attention because we must offload all attention scores to CPU to avoid out-of-memory issues before using them as GT. Therefore, we use full attention as the baseline when comparing accuracy across different sparsity settings (see Tables 1–3 in the paper and additional results in the rebuttal to another reviewer: https://openreview.net/forum?id=Nf8yfPDFTl&noteId=yPnYUXTMey).

3. Pooling selection

Our pooling experiments are consistent across all tested GQA-based models. The best pooling strategy generalizes well within this setting. However, pooling selection remains an open question, especially for other attention architectures like MQA/MLA. Our main contribution is introducing the gated attention concept, predating NSA [3] and MoBA [4] (another reviewer asked). Future work may explore alternative gate structures such as strided MLPs or convolutions, inspired by MoE router literature.

4. Tunable parameters for each layer

We would like to clarify to the reviewer that by “tunable parameters” we consider two categories: (1) trainable parameters and (2) hyperparameters. Below, we list and compare these parameters for SeerAttention and DuoAttention accordingly.

For trainable parameters

• DuoAttention: a scale vector multiplied on streaming head and dense head.
• SeerAttention: a linear layer from compressed QK to gate hidden vector.

If we compare the number of trainable parameters, DuoAttention indeed has fewer parameters. However, even for SeerAttention, the additional parameters are minimal—less than 1% of the total. For instance, in LLaMA-3.1-8B, the gating module introduces roughly 101M parameters, which is about 101M / 8B ≈ 1.3% of the model size. This overhead is very small and keeps the method efficient to train. As shown in Figure 9 of Appendix A.1, we design customized training kernels to distill the gate, and the memory and latency overhead during this distillation stage are minimal compared to native FlashAttention.

For hyperparameters

SeerAttention introduces an additional hyperparameter, block size, as well as the number of training steps, in addition to the usual hyperparameters required for model fine-tuning. DuoAttention, on the other hand, requires selecting the percentage of streaming heads and full-attention heads, along with the number of training steps and other standard fine-tuning hyperparameters.

5. Modify results organization

Thank you for the suggestion since we just present the overall sparsity for different sequence length inside the sentences. We will revise the presentation if accepted.

6. About decoding

We appreciate the suggestion to discuss decode‑stage evaluation. In this work, we primarily focus on accelerating the prefill stage and leave decode‑stage optimization for future exploration. Prefill and decode are two distinct phases, and several PD‑disaggregation serving systems, such as DistServe [5] and Mooncake [6], have explored separating prefill and decode to improve TTFT/TPOT. Many prior works also focus on only one phase: for instance, MInference mainly targets prefill, while Quest [7] and H2O [8] focus solely on decode.

Our approach can naturally extend to decode sparsity with a minor modification of the AttnGate design. Specifically, we remove the Q‑compression branch (pooling + linear) and instead distill using a 1D max‑pooling ground truth of the attention map (along the sequence dimension) from the base model. In other words, during distillation for decode, we directly use the original Q without sequence‑level compression, enabling each token’s Q to select important KV blocks at inference.

We evaluated this design on the DeepSeek‑R1‑Distill‑Qwen‑14B reasoning model and compared it against Quest. Below we provide a brief summary of the results (we choose block size = 64 for K block). A more complete version can be included in the appendix upon acceptance.

7.About batch inference

We appreciate the reviewer’s interest in this question about batching inference. In fact, unlike MoE, batch inference is less challenging for attention. During prefill, variable-length inputs can be packed along the sequence dimension, hiding batch size within sequence length. For decoding, modern infrastructures (like vllm or sglang) support paged attention, allowing our block-sparse method to selectively activate pages naturally.

Reference:

[1] RedPajama, https://github.com/togethercomputer/RedPajama-Data

[2] BookSum: A Collection of Datasets for Long-form Narrative Summarization, https://arxiv.org/abs/2105.08209

[3] Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention, https://arxiv.org/abs/2502.11089

[4] MoBA: Mixture of Block Attention for Long-Context LLMs, https://arxiv.org/abs/2502.13189

[5] DistServe: Disaggregating Prefill and Decoding for Goodput-optimized Large Language Model Serving, https://www.usenix.org/system/files/osdi24-zhong-yinmin.pdf

[6] Mooncake: A KVCache-centric Disaggregated Architecture for LLM Serving, https://arxiv.org/abs/2407.00079

[7] Quest: Query-aware sparsity for efficient long-context llm inference, https://arxiv.org/abs/2406.10774

[8] H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models, https://arxiv.org/abs/2306.14048

We thank the reviewer for their thoughtful feedback and for recognizing the novelty.

1. Weaknesses 1 & 2: Sparsity Ratio and Speedup

We evaluated speedup–accuracy trade‑offs under varying sparsity. The table below shows kernel‑level speedups for different sparsity and sequence lengths:

In the end-to-end, by adjusting threshold T, SeerAttention flexibly balances sparsity and accuracy. RULER‑128k accuracy results:

Tradeoffs between Speedup and Accuracy

To evaluate the tradeoffs, we count the average sparsity and accuracy to present.In LongBench, SeerAttention employs a threshold of 2e-3 for all AttnGates. With the same threshold, different attention gates can exhibit varying sparsity ratios, and longer context data tends to be sparser.

LongBench performance & efficiency results on Llama-3.1-8B-Instruct

* 50% streaming heads, the real sparsity < 50%

In RULER, we employ a threshold of 5e-4 under 4k-128k test benches.

RULER performance & efficiency results on Llama-3.1-8B-Instruct

These results demonstrate that SeerAttention achieves the best overall accuracy–efficiency trade‑off, offering higher sparsity and greater speedup without degrading accuracy compared to prior methods.

2. Weakness 3 & Question 3: Training Cost and Parameter Overhead

We appreciate your suggestion to clarify the training cost and parameter overhead. In Llama-3-8B-Instruct model,

• SeerAttention: ~40 A100 hours in total (5 hours on 8×A100)

• DuoAttention: “several hours on 8×A100”

• MInference: ~2 hours on 1×A100 for sparsity profiling/calibration but is less accurate and slower in inference

The training cost of SeerAttention is comparable to or lower than DuoAttention. Although MInference calibrates faster, it sacrifices inference accuracy and speed. SeerAttention introduces only a small number of trainable parameters. As shown in Fig. 9 (Appendix A.1), we design customized training kernels to distill the gate, and the memory/latency overhead during distillation is < 5% compared to native FlashAttention. For LLaMA‑3.1‑8B, the gate adds only 101 M parameters (~1.3% of the model), making the method lightweight and scalable.

3. Weakness 4: Decoding-Stage Inference:

We appreciate the suggestion to discuss decode‑stage evaluation. In this work, we primarily focus on accelerating the prefill stage and leave decode‑stage optimization for future exploration. Prefill and decode are two distinct phases, and several PD‑disaggregation serving systems, such as DistServe [1] and Mooncake [2], have explored separating prefill and decode to improve TTFT/TPOT. Many prior works also focus on only one phase: for instance, MInference mainly targets prefill, while Quest [3] and H2O [4] focus solely on decode.

Our approach can naturally extend to decode sparsity with a minor modification of the AttnGate design. Specifically, we remove the Q‑compression branch (pooling + linear) and instead distill using a 1D max‑pooling ground truth of the attention map (along the sequence dimension) from the base model. In other words, during distillation for decode, we directly use the original Q without sequence‑level compression, enabling each token’s Q to select important KV blocks at inference.

We evaluated this design on the DeepSeek‑R1‑Distill‑Qwen‑14B reasoning model using two challenging math tasks and compared it against Quest. Below we provide a brief summary of the results (block size = 64). A more complete version can be included in the appendix upon acceptance. The results show that SeerAttention achieves higher accuracy than Quest, even with larger KV blocks and sparser attention.

Model: DeepSeek-R1-Distill-Qwen14B

4. Question 2: Short Context:

We appreciate the suggestion to clarify why certain short‑context tasks (e.g., MMLU, ARC‑c) exhibit lower sparsity. First, as the input sequence length decreases, the attention map — Softmax(QKᵀ) — naturally becomes denser: shorter sequences lead to more concentrated information, and the softmax distribution assigns higher average attention scores to each token. This behavior is inherent to attention computation itself and not specific to our gate predictor design. In SeerAttention, we can freely control sparsity by selecting different thresholds or Top‑K ratios. For a fixed threshold, the resulting sparsity naturally decreases for short contexts and increases for long contexts because the gating mechanism’s softmax rows sum to 1, so tokens in shorter sequences inherently receive larger attention scores. Moreover, attention computation is rarely the bottleneck for short contexts, whereas sparse attention provides significant efficiency gains for long‑context inference. Attention only becomes a major bottleneck when the sequence is long, where sparsity leads to meaningful speedups. We thank the reviewer for pointing this out — we have reorganized Table 3 accordingly. As sequence length decreases, the observed sparsity drops (e.g., 3.4% and 26%), which directly reflects this inherent behavior.

5. Question 4: Rope in KV cache

The rope design in decoding version is similar to prefill case. In decoding, SeerAttention maintains a small compressed K cache (after RoPE) specifically for gating. This avoids loading the full K cache at every step and allows selective V-loading, maximizing decode-stage speedup. Without this, sparsity-based selection would be limited to a ~2× upper-bound speedup due to full K-cache loading.

Once again, we sincerely thank the reviewer for their positive assessment of SeerAttention’s novelty, technical soundness, and practical impact. We will incorporate your suggestions to further strengthen the final version of the paper.

[1] DistServe: Disaggregating Prefill and Decoding for Goodput-optimized Large Language Model Serving, https://www.usenix.org/system/files/osdi24-zhong-yinmin.pdf

[2] Mooncake: A KVCache-centric Disaggregated Architecture for LLM Serving, https://arxiv.org/abs/2407.00079

[3] Quest: Query-aware sparsity for efficient long-context llm inference, https://arxiv.org/abs/2406.10774

[4] H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models, https://arxiv.org/abs/2306.14048

Thank you very much for your continued support and for highlighting the novelty and practical relevance of SeerAttention! We’re glad to hear that you find the method sound and well-evaluated.

Regarding point (2) on the careful selection of thresholds—we’d like to clarify that SeerAttention also supports TopK-based masking, which enables precise control over the sparsity level. This is shown in the right plot of Figure 1. While we chose threshold-based evaluation in the main experiments due to its ability to dynamically adapt sparsity across different sequence lengths without notable performance loss, users can also specify a desired sparsity level directly through TopK during inference.

We sincerely appreciate your thoughtful feedback and your recommendation!