Mixture of Attention Spans: Optimizing LLM Inference Efficiency with Heterogeneous Sliding-Window Lengths

W: Terminology

I strongly object to the terms that the authors use to describe their work -- in particular the words "sparse" and "compression", which are right in the title. I have seen this misuse a couple of times just recently in other papers as well, and it needs to stop.

We thank the reviewer for pointing out the misleading use of sparse and compression. We will fully adopt your recommendation and revise the manuscript as follows:

• New technical wording.

  • We will describe our method simply as heterogeneous sliding-window attention and as an efficiency optimization technique.
  • We will reserve sparse attention exclusively for truly irregular, fine-grained sparsity. For block-pattern methods such as Sparse Transformers [1], we will explicitly write block-sparse attention.

• New title.

  • Mixture of Attention Spans: Optimizing LLM Inference Efficiency with Heterogeneous Sliding-Window Lengths

• Global edits.

  • All occurrences of sparse and compression in the abstract, main text, and figures will be replaced with the above terminology.

We appreciate your careful feedback. Our original terminology was never intended as false advertising; we viewed sliding-window attention as a structured, block-sparse variant of dense attention. Nevertheless, we recognise that its contiguous, coarse-grained pattern differs from the fine-grained sparsity typically implied by sparse attention, and we have adopted the revised phrasing to avoid any ambiguity. We welcome any further suggestions. The modified abstract is attached below.

[1] Child et al., “Generating Long Sequences with Sparse Transformers,” NeurIPS 2019.

The modified abstract

Sliding-window attention is a hardware-friendly way to mitigate the significant memory and throughput demands of Large Language Models (LLMs) in long contexts. Existing methods typically employ a single window length across different attention heads and input lengths. However, this uniform approach fails to capture the diverse attention patterns inherent in LLMs, ignoring their distinct accuracy-latency trade-offs. To address this challenge, we propose the Mixture of Attention Spans (MoA), which automatically tailors distinct sliding-window length configurations to different heads and layers. MoA constructs and navigates a search space of various window lengths and their scaling rules relative to input sequence lengths. It profiles the model, evaluates potential configurations, and pinpoints the optimal length for each head. MoA adapts to varying input sizes, revealing that some attention heads expand their focus to accommodate longer sequences, while other heads consistently concentrate on fixed-length local contexts. Experiments show that MoA increases the effective context length by 3.9x with the same average sliding-window length, boosting retrieval accuracy by 1.5-7.1x over the uniform window baseline across Vicuna-{7B,13B}, and Llama3-{8B,70B} models. Moreover, MoA narrows the capability gaps to full attention, reducing the maximum relative performance drop from 9%-36% to within 5% across three long-context understanding benchmarks. MoA achieves a 1.2-1.4x GPU memory reduction, boosting decode throughput by 6.6-8.2x and 1.7-1.9x over FlashAttention2 and vLLM, with minimal performance impact.

We hope these revisions address your concern and are happy to refine further if needed.

W1: SnapKV & PyramidKV comparison

While the paper compares MoA against StreamingLLM and H2O, it overlooks critical baselines that define state-of-the-art sparse attention (e.g., [1] SnapKV, [2] PyramidKV). These methods address similar efficiency-accuracy trade-offs but employ distinct strategies (e.g., token prioritization, hybrid sparse-dense blocks).

We provide a direct comparison with the official SnapKV and PyramidKV implementations in Fig. 4. The table reports retrieval accuracy on LongEval using Vicuna-7B with an 8k-token input on a single A100-80GB GPU, with all methods using 25% average attention density:

MoA achieves both higher accuracy and significantly higher throughput, due to its effective offline optimization and the static, lightweight sliding-window during inference.

W2: Exposition of Sections 2-3

Sections 2–3 suffer from disjointed exposition:

• Section 2 introduces heterogeneous patterns but lacks a visual roadmap (e.g., a flowchart) of how elastic rules are derived from attention matrices.
• Section 3.1’s attention influence derivation (Eq. 2) is abrupt; a step-by-step intuition (e.g., "masking token j redistributes attention mass to k ≠ j") would aid readability.
• The interplay between profiling (Section 3.1) and optimization (Section 3.2) is under-explained. A unified algorithmic pseudocode (combining both stages) would clarify the pipeline.

We thank the reviewer for the valuable suggestions. We plan to streamline as follows:

1. Visual roadmap: Revise Figs. 2 and 3(a) to link the attention matrix, desired attention mask, and elastic rule. Update Fig. 3(b) to a clearer flowchart illustrating the derivation of elastic rules.

2. Intuition before formalism: Add a brief intuitive explanation before Eq. 2 as suggested, explaining that masking token $j$ redistributes its attention mass to tokens k$\neq$j.

3. Pseudocode: Introduce a pseudocode block and cross-reference it from both Sections 3.1 and 3.2, clearly connecting both stages, as well as Figure 3.

We believe these changes will create a more coherent visual and algorithmic pipeline. We are happy to further refine the writing if there are additional suggestions.

W3.1: Barriers of custom CUDA kernel

The reliance on custom CUDA kernels introduces practical barriers:

Compatibility issues with existing frameworks (e.g., HuggingFace, vLLM) requiring significant engineering effort.

Our CUDA kernel is exposed through a lightweight Python wrapper that mirrors torch.nn.functional.scaled_dot_product_attention, adding only a single tensor to specify each head’s attention span. Integrating MoA into HuggingFace Transformers requires only a few modifications to the AttentionModule, after which it can be enabled with just two lines of user-level code (see Appendix D.4). Since MoA employs fixed sliding-window masks at inference, integrating the kernel into frameworks like vLLM should involve only moderate engineering effort.

W3.2: Energy efficiency and FLOPs

No discussion of energy efficiency—a critical metric for real-world deployment. For instance, does MoA’s sparsity reduce FLOPs sufficiently to offset kernel overhead?

We thank the reviewer for the suggestion. We have added energy efficiency results under the same setup as Table 5, using Vicuna-7B on a single A100-80GB GPU:

Energy: MoA achieves a 8.7–10x reduction in energy per output token, driven by slightly lower GPU power and, more importantly, much higher throughput. We will add these results to the paper.

FLOPs: During decoding, each head accesses only its (shortened) KV cache, with an O(#heads) offset computation that is negligible compared to sequence length. Since attention FLOPs scale linearly with KV length, MoA reduces attention FLOPs per decode step by approximately 50%, much more than offsetting the minimal addressing overhead.

W1: Limited evaluation to retrieval tasks

Limited evaluation: the paper doesn't provide evaluation other than retrieval tasks such as coding, math Q&A and etc.

We evaluate MoA on a broad range of tasks beyond retrieval, including code completion, single- and multi-document QA, summarization, and few-shot learning, using 11 and 13 sub-datasets from LV-Eval and LongBench (L626). Averaged results are reported in Table 4, with detailed per-task results provided in Appendix B.1.3 (Table 8 and Figure 10). MoA consistently matches the original dense model across tasks, whereas StreamingLLM and InfLLM show greater variance: sometimes exceeding the original model, but with significant drops in other tasks. We will further emphasize the comprehensiveness of our evaluation in the revised manuscript.

W2: Dynamic vs. static attention mask

Lack of dynamic sparse attention mechanisms: it doesn’t incorporate dynamic sparse attention mechanisms that adapt to varying input patterns in real-time, not context-dependent

Dynamic sparsity can, in theory, remove more FLOPs, yet on commodity GPUs its fine-grained control flow often reduces those gains. MoA chooses a different point on this trade-off curve:

• Design Choice. We adopt static, per-head heterogeneous windows because they execute in a single branch-free kernel on standard GPUs.

• Effectiveness. Although static at run time, the window lengths are optimized from gradients and differ across heads, layers, and input lengths. At 50 % density MoA is 6–8x faster than FlashAttention-2 (Tab. 5) while showing smaller accuracy drops than dynamic-sparsity baselines (Tab. 4).

• Complementarity. MoA can simply replace any homogeneous local window inside a dynamic scheme. As a demonstration, we swap the 0.15$\times$N uniform local window in H2O [1] to MoA heterogeneous spans with the same average density, keeping the 0.15$\times$N dynamic heavy-hitter unchanged. At the same 30% density, our version better restores dense-level accuracy:

• Future work. MoA pipeline can also be adept to optimize hyper-parameters inside dynamic attentions (e.g., the H2 rate of H2O). Exploring these methods, especially on custom accelerator hardware, is a promising future work.

W3: Novelty over MInference and SpargeAttn

Limited novelty: similar work like MInference, already addresses the prefill inefficiency in long-context LLM inference by using a more practical and implementation-ready approach. It dynamically constructs sparse indices and leverages optimized GPU kernels for real-time acceleration. The paper doesn't show its advantages on the existing methods like Minference and SpargeAttn, which solve the similar problem.

MInference and SpargeAttn reduce prefill with dynamic sparsity, but during decoding they still access the full KV Cache. MoA tackles the complementary bottleneck of KV-cache reduction and decode acceleration.

MoA is also practical and implementation-ready, providing a two-line Hugging Face interface (Appendix D.4). Empirically, it advances the accuracy-throughput frontier over six decoding baselines (Fig. 4).

We will cite and discuss MInference and SpargeAttn in the final paper.

Q1: Clarifying attention span and elastic rule

It is unclear how the attention span is explored by the pipeline. It will be better to provide more details on the relationship between attention span function and elastic rules. The current explanation is limited to fully understand what are the elastic rules and how they are explored.

The attention span $S$ is the sliding-window size plus the first 64 always-visible tokens (L129). The elastic rule $(\alpha,\beta)$ specifies how $S$ scales with input length $N$.

Logic

1. Eq. (1): elastic rule $(\alpha,\beta)$ $\overset{\text{length }N}{\longrightarrow}$ attention span $S$
2. Sec. 2.2:  $S \longrightarrow$ attention mask $M$
3. Eq. (2–3): $M$ $\overset{\text{calibration data}}{\longrightarrow}$ prediction loss $L$

Optimization Pipeline

Eq. 4 then minimizes $L$ across lengths by searching over elastic rules of each head with multi-objective mixed-integer programming (MIP).

We will add this step-by-step explanation and refine the terminology in the final version.

Q2: Reading Figure 3(a)

Span to mask

It will be better to give more details on Figure 3(a) and tie back to your explanation in the paragraph. How does the attention span related to the masks, by colors or by length?

Figure 3(a) shows examples of attention masks from the search space. Each $N\times N$ grid is a candidate attention mask for a single head.

Two vertically placed masks: same length, two elastic rules $\rightarrow$ two masks.

Two horizontally placed masks: same elastic rule, two lengths $\rightarrow$ two masks.

How masks are applied

Are masks all overlapped together on the same input sequence or sequentially applied?

All heads apply their masks in parallel; they are neither overlapped nor sequential. Each mask gates both the attention computation and the KV Cache entries for its head.

Elastic rule to mask

Does elastic rules decide the range of masks?

Yes. An elastic rule $(\alpha_h,\beta_h)$, with input length $N$, determines attention span $S_h$ (Eq. 1); the resulting window width $S_h-64$ fixes the mask.

Selecting span/window

How are the sizes of attention span and attention windows selected separately/related.

MoA searches elastic rules ($\alpha_h$, $\beta_h$) for each head offline. At inference time, $N$ determines each head’s span and window via Eq. 1.

We will incorporate these explanations in the camera-ready version.

Q3: Perplexity length range

For the experiments, could you please explain the reason that perplexity is only reported with length 8k - 12k.

Apologies for the confusion. Perplexity is calculated over all response tokens for every test sample. The label “8k–12k” is not a cut-off we imposed—it simply marks the prompt + response length of data items. Detailed perplexity setups are discussed in Appendix A.1. We will clarify this more explicitly in the final version.

Q4: Data efficiency & task agnosticism

What is the data efficiency of finding the Pareto point, is this task agnostic?

• Data Efficient. The entire optimization only uses 150 calibration sequences (2k-8k tokens each) for profiling, and 250 sequences (12k tokens each) for validation (Appendix. A.1).

• Task-agnostic rule. The optimized elastic rules are used across all downstream benchmarks—code, QA, retrieval, etc.—without any per-task re-profiling or tuning.

• Generalization empirical evidence. As discussed in Section 5.3, MoA offline optimized elastic rules generalize well to different tasks and unseen lengths.

Q5: Shared patterns across heads

Do you observe any shared pattern between heads?

Yes, layer position strongly influences spans:

• Early layers: heads keep a broad, almost uniform view of the sequence, matching the wide-span example in Fig. 2 (right).
• Last layers: most heads focus on the first few tokens and diagonal, as in Fig. 2 (left); only a few still use wide spans.

MoA naturally recovers this trend. As detailed in Appendix C.2 and Figure 11,

• Early layer heads are assigned larger spans,
• Last layer heads mostly receive shorter spans.

This aligns with known attention behaviors. We will add examples in the Appendix to fully visualize how MoA’s learned masks correspond to existing attention patterns at different layers.

Q6: LongBench accuracy averaged across tasks?

Is the accuracy for LongBench averaged across tasks?

Yes. The reported accuracy is averaged over all data items across LongBench. Per-task scores are provided in Figure 10 (Appendix B.1.3) for full transparency.

I appreciate authors' answer to my questions and further clarifications. I will raise my score by 1