Multipole Attention for Efficient Long Context Reasoning
Authors: Coleman Richard Charles Hooper
Sebastian Zhao
Luca Manolache
Sehoon Kim
Michael W. Mahoney
Sophia Shao
Kurt Keutzer
Amir Gholami
Venue: NeurIPS 2025 poster ()
Keywords: Long Context
Inference
LLM
Reasoning
Efficiency
ML Systems
Attention
Primary area: deep_learning
TL;DR: Accelerating attention for long-context reasoning by identifying and loading important tokens and by approximating attention to less important tokens

Abstract: Large Reasoning Models (LRMs) have shown promising accuracy improvements on complex problem-solving tasks. While these models have attained high accuracy by leveraging additional computation at test time, they need to generate long chain-of-thought reasoning in order to think before answering, which requires generating thousands of tokens.
While sparse attention methods can help reduce the KV cache pressure induced by this long autoregressive reasoning, these methods can introduce errors which disrupt the reasoning process.
Our work addresses these challenges by introducing Multipole Attention, which accelerates autoregressive reasoning by only computing exact attention for the most important tokens, while maintaining approximate representations for the remaining tokens. 
Our method first performs clustering to group together semantically similar key vectors, and then uses the cluster centroids both to identify important key vectors and to approximate the remaining key vectors in order to retain high accuracy.
Additionally, in order to accelerate long generation tasks, we design a fast cluster update process to quickly re-cluster the input and previously generated tokens, thereby allowing for  accelerating attention to the previous output tokens.
We evaluate our method using emerging LRMs such as Qwen-8B and Deepseek-R1-Distil-Qwen2.5-14B, demonstrating that our approach can maintain accuracy on complex reasoning tasks even with aggressive attention sparsity settings.
We also provide kernel implementations to demonstrate the practical efficiency gains from our method, achieving up to 4.5$\times$ speedup for attention in long-context reasoning applications.

Multipole Attention for Efficient Long Context Reasoning
Coleman Hooper
Sebastian Zhao*1
Luca Manolache'
Sehoon Kim Michael W. Mahoney 1.23 Yakun Sophia Shao
Kurt Keutzer' Amir Gholami
University of California, Berkeley 2ICSI LBBN
{chooper, sebbyzhao, luca. manollche,, sehoonkim, mahoneymı, yashao, keutzer, amirgh)@berkelley. .du
Abstract
Large Reasoning Models (LRMs) have shown promising accuracy improvements for complex problem-solving tasks. While these models have attained high ac- curacy by leveraging additional computation at test time, they need to generate long chain-of-thought reasoning in order think before answering, which requires generating thousands of tokens. While sparse attention methods can help reduce the KV cache pressure induced by this long autoregressive reasoning, these methods can introduce errors which disrupt the reasoning process. Additionally, prior meth- ods often pre- preeprocesed the input to make it easier to identify the important prompt tokens when computing attention during generation, and this pre-processt processing is challenging to perform online for newly generated reasoning tokens. Our work addresses these challenges by introducing MULTIPOLE ATTENTION, which ac- celerates autoregressive reasoning by only computing exact attention for the most important tokens, while maintaining approximate representations for the remain- ing tokens. Our method first performs clustering to group together semantically similar key vectors, and then uses the cluster centroids both identify important key vectors and to approximate the remaining key vectors in order to retain high accuracy. Additionally, we design fast cluster update process quickly re-cluster the input and previously generated tokens, thereby allowing for accelerating at- tention the previous output tokens. We evaluate our method using emerging LRMs such as Qwen- 8B and Deepscek-RI- Di-Disiil-Qwn2..55 11BB demonstrating that our approach can maintain accuracy on complex reasoning tasks even with aggressive attention sparsity settings. We also provide kernel implementations to demonstrate the practical efficiency gains from our method, achieving up to 4.5x speedup for attention in long-conted reasoning applications. Our code is available at http://hh....ooeeeeeeeeaaaaaaaaa
1 Introduction
Test-time compute has emerged as a new dimension for scaling up Large Language Model (LLM) performance [Snell al. 2024]. By scaling the amount of computation used at test time, we can allow LLMs to think longer on harder problems in order to attain higher accuracy. Large Reasoning Models (LRMs), which are LLMs post-trained to have strong reasoning capabilities, typically reason by generating long chain-of -hhiuhh,- where they produce a step-by-step chain to help guide themselves
to the correct solution (Wei et al. 2022
Guo et al. 2025
Owen 2025
This has been shown to
$ Equual Contribution
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

* Query
O Key
+ Centroid
Figure I: A visualization of the MULTIPOLE ATTENTION algorithm. For important keys (which close to the current query), we compute exact attention. For keys which are farther away from the current query, we instead approximate the attention these keys using the attention score with a representative cluster centroid. We use progressively coarser grained centroids as we get further away from the query. This allows us to maintain important contextual information from all tokens in the sequence, while only requiring loading a small number of keys for exact attention computation.
provide substantial accuracy gains in domains such math and coding which require problem solving capabilities [Guo et al.) 2025 Qwen
2025).
However, the long autoregressive generation required for reasoning also brings substantial inference efficiency challenges. Notably, the reasoning process leads to large Key/Value (KV) cache memory footprint [Hooper et al. 2024bb. One common method to reduce the KV cache footprint is through sparse attention methods that only load important KV cache tokens at each generation step [Tang et al. 2024 Hooper et al.) 2224a]. However, while these methods are promising for reducing the KV cache memory bottleneck, they lead to unacceptable accuracy loss. Prior work has also shown that problem solving tasks which require complex reasoning experience greater accuracy degradation when pruning the KV cache [Liu et al. 2025). Additionally, existing sparse attention methods often pre-process the KV cache from the prompt make it easier retrieve important prompt tokens during generation, and this pre- pre-processing is challenging to perform online during generation for the previously generated reasoning tokens (Hooper al. 2024a]
To address these challenges, we build on top of (Hooper et al., 2024a] and introduce MULTIPOLE ATTENTION, which leverages sparse attention in order to only load the important KV cache tokens, while still approximating attention to the rest of the context to maintain high accuracy. Figure [ provides intuitive visualization of our algorithm, which first clusters the keys and computes representative centroid for each cluster. When computing attention, our method compares the current query with the centroids determine the importance of key tokens. For the highest scoring keys, compute exact attention, For the remaining keys, we leverage the cluster centroid to approximate the attention score for all keys in that cluster, thereby retaining contextual information from the full KV cache. We also present a hierarchical generalization of our method which leverages progressively coarser grained centroids we get further away from the query. To accelerate attention to the previously generated reasoning tokens, we also design a fast cluster update strategy that can be applied during generation, thereby allowing our method accelerate attention to the previous output tokens.
Specifically, our work makes the following contributions:
Multipole Approximation: Our algorithm clusters keys based on semantic similarity and rep- resents all keys within a cluster with representative key centroid (as in [Hooper et al., (024aa) ). When computing attention, we then compare the current query with the key centroids to identify which keys are important for the current query; we only compute exact attention for these important keys, and we approximate the attention to the remaining keys using the attention scores with the cluster centroids. This allows our algorithm to maintain contextually relevant information from the full sequence, even with low KV cache budget. We also extend this to a hierarchical multipole approximation to improve the efficiency of the centroid comparison.
In order to accelerate attention to the KV cache entries that are appended during generation, we design a fast cluster update algorithm. Our algorithm initially assigns the new tokens to clusters, and then performs small number of refinement steps over the full keys to ensure high quality clustering. Additionally, to improve the scalability of our algorithm,
Fast Online Clustering:

---

leverage blockwise k-means clustering method (which clusters the input sequence in segments), and we design a shifting window method ensure that the final block always has sufficient tokens to perform clustering. These optimizations together facilitate efficient clustering for the generated output tokens.
System Implementation: We present a prototype system implementation with custom Triton kernels for performing multipole approximate attention. Our system implementation consists of a multi-stage kernel implementation that compares the incoming user query with the key centroids, then computes exact attention for the important keys and approximate attention for the less important keys. Our methodology achieves up 4.5x x attention speedup during decode for long context reasoning applications.
2 Related Work
2.1 Reasoning Models
Reasoning models have emerged as a new paradigm for solving complex problem-solving tasks [Guo et al. 2023 Qwen 2025|. These models are trained to generate a chain-of-thougle reasoning trajectory, which contains step-by-stey reasoning to help guide the model to the correct answer (Wei et al. 2022). The DeepSeek-Re model [Guo al.] [0255] pioneered a new reinforcement learning- based methodology for training reasoning models for problem solving tasks. Their model release also included smaller open-source models (distilled from their base model) based on the Llama3 [Grtttafirii et al. [024) and Qwen2.5 model series (Yang et al. [2024). There have also been several powerful reasoning models released with agentic tool-calling capabilites such as the QwQ-32B model [Qwen] 2025]. the Qwen3 model series Yang et al. 2024], and the Llama- nema--o model series [Bercovich et al. 2025
2.2 Long Context Inference
Long-context length LLMs. which can support context lengths greater than 100K tokens, have seen widespread use for applications such document question answering, summarization, literature review, and processing multi-turn conversations. There have been several closed-source models which support long context capabilities [Achiam al. 2023, Anthropic 2023), including Gemini (Google 2023] which supports greater than million context length. Newer open-source model families such as Llama3 [Meta] 2024) and Qwen2.5 (Yang al.) 2024) also support greater than 100K context length. Another line of work has focused on extending the context windows of pretrained models beyond the original supported context length through scaling or adjusting the positional embeddings [Chen et al.
2023
Peng et al. 2023
2.3
KV Cache Compression
KV cache compression has emerged as a key algorithmic tool for enabling long context length inference. For long context inference scenarios, the KV cache becomes the main memory bottleneck, making KV cache compression crucial for efficient inference (Hooper et al. 2024b]. Common approaches for compressing the KV cache include quantization [Hooper et al. 2024b, Liu et al. 2024b). which aims to compress the representation for each token, and sparsification [Tang et al.
2024 Hooper et al. 2024a Zhang al. 2024b), which aims to only load important KV cache tokens.
One previous approach for inducing sparsity to reduce KV cache size during inference is KV cache eviction, where less important tokens are evicted during the generation process to reduce memory consumption. Previous works have leveraged metrics such as attention score contribution (Zhang et al. 2024b) to select tokens to evict from the KV cache. One challenge with evicting tokens from the KV cache is that it irreversibly discards less important tokens if these tokens become important later during generation, they cannot be recovered. KV cache merging [Wang et al. 2024 Yuan et al. 2025) was proposed as solution to compensate for the error induced by KV cache pruning. KV cache merging involves merging evicted KV cache entries (based on key similarity (Wang et al. 2024) or historical attention scores (Yuan et al.) 2025) instead of discarding them completely. In contrast with these works, our algorithm retains the full original KV cache (thereby allowing the model refer back to any prior KV cache state exactly if they become important during generation), while approximating attention to less important clusters using the key centroids.
3

---

* Query
O Key + Centroid
Low Scoring
High Scoring
Exact
Approximate
a) Clustering using K-Means
b) Centroid Lookup
c) Multipole Attention
Figure 2: A diagram outlining how MULTIPOLE ATTENTION identifies and retains important keys, while approximating the attention to the remaining keys. a) We first construct Key index by performing k k-means clustering (thereby obtaining representative centroid for each cluster). b) When computing attention, we compare the current query with the centroids to estimate the importance of keys in each cluster ("Centroid Lookup"). c) For the most important ("High Scoring") key centroids, we then compute exact attention with the keys in the corresponding clusters. For the remaining keys, approximate the attention to each key cluster using the attention the representative centroid for that cluster. Overall, MULTIPOLE ATTENTION attains the efficiency of sparse attention while preserving important contextual information from the full sequence. Our method can also be generalized to a hierarchical approach, as outlined in Section 3.4
Another common approach for leveraging sparsity for the KV cache to retain the full KV cache and only load in the required KV cache entries for each decoding step. One prior work, QUEST [Tang et al. 2024). grouped consecutive KV cache entries and derived representative vectors for the keys within each group, and compared the query with these representative vectors to decide which KV cache entries to load. Whereas QUEST clustered keys based on positional proximity, Squeezed Attention [Hooper al. 2024a] instead clustered keys based semantic similarity, thereby allowing for precise identification of tokens which were likely to be high-scoring. Another related work, Tactic [Zhu al. 2025 performed clustering based semantic similarity and then used distribution fitting to select critical KV cache entries. Other previous works have framed the query-key comparison a retrieval problem in order accelerate attention using vector search methods [Zhang al.) 2024a Liu et al. 2024a| He et al. 2025). In contrast with existing methods, our approach aims to identify and load only small subset of KV cache entries for exact attention computation, while approximating the attention to the remaining KV cache entries. Our algorithm clusters keys based on semantic similarity and then performs fast centroid comparison to identify important keys, while also leveraging these representative centroids to approximate the attention to the keys which are not retrieved by the centroid lookup. A more detailed discussion of methods that sparsely load KV cache entries is provided in Appendix
3 Algorithm
3.1 Retrieving Important KV Cache Tokens
In order to be able to identify and retrieve important KV cache entries, we first cluster the keys based on semantic similarity using k-means clustering. We then derive representative centroids for each cluster by taking the mean of all vectors in each cluster, which can then be used to quickly estimate the importance of each key cluster (as in [Hooper et al.) [024a]). To identify whether the keys in cluster are important for given query token 9. compare the query with the key centroid for that cluster to estimate the attention score for the corresponding keys:
exp (9Ko))
5 N₁ exp (9K.7)
where Kej is the key centroid for cluster j and N; is the number of keys in cluster j. We can then use the centroid attention score to identify the clusters which are likely to be high scoring, and only load in the corresponding keys from these clusters for exact attention computation. To decide which

---

keys to retain, we sort the clusters based on the centroid scores, and then retain clusters until we hit a target token budget. Note that for models which support grouped query attention, we aggregate importance by averaging the estimated attention across query heads which share keys and values Aiissie et al. 2023]
One challenge when clustering the keys is the impact of the Rotary Positional Embeddings (RoPE), which rotates key vectors by different amounts depending on their position in the sequence. This means that semantically similar keys at different positions in the sequence may not be clustered together. To address this, we employ the Windowed RoPE strategy from (He al., 2025| when computing the centroids (which assumes a fixed relative positional difference between the query and the key vectors), thereby improving the clusterability of the key vectors. We then use a query vector rotated at a fixed position when performing the centroid lookup (as in [He et al., 2025[).
3.2 Importance-Aware Multipole Approximation
After using our centroid lookup to identify the important tokens, we then aim to approximate the attention to the less important tokens in order to retain necessary contextual information. We directly use the attention score to the cluster centroid from Equation [ as the estimated attention score for the cluster. We also cache representative value centroid Ve for each cluster (which is the average of the corresponding value vectors for a given key cluster). The attention contribution for cluster i (excluding the Softmax denominator) can be represented follows:
exp (gK.)) Vei
where Nj. Kej- and Ve; are the number of keys for cluster i, the key centroid for cluster i, and the value centroid for cluster i, respectively. For the less important key clusters, we compute the attention We then merge the attention output from exact attention (for the important tokens) and from the centroids (for the less important tokens) in order to obtain the final attention output. Our algorithm is visualized in Figure 2 which shows how we leverage the key centroids both to identify important keys for the current query, as well as to approximate the attention for the less important keys.
contribution of each cluster using Equation 2
3.3 Efficient Online Clustering
A key challenge when employing our approach for long generation applications is the need to approximate attention to the generated tokens. This requires updating the clusters and corresponding centroids as new tokens are appended to the KV cache; however, this is computationally expensive if we need to re-run clustering for the entire sequence. To quickly update clusters and corresponding centroids when appending newly generated tokens, we incorporate two approaches: blockwise clustering methodology (with a sliding window to manage the transition between blocks when appending new tokens), as well as fast initial cluster assignment for newly appended tokens followed by a small number of cluster refinement steps.
Figure 3 outlines our blockwise clustering methodology to ensure that when we append to the centroids, we do not need to recluster the full sequence. We perform clustering blocks of W tokens, which means that when we append new tokens, we only need to recluster the final block. Additionally, if the final block was limited W tokens, after exceeding this limit it would be challenging to cluster the subsequent tokens until there were sufficient number of new tokens appended. We therefore allow the final block to extend to + W tokens before removing the first W tokens from this block and clustering them separately. This ensures that the final block never has fewer than tokens, meaning that the final block always contains a sufficient number of keys for clustering. Note that the blockwise clustering approach also improves clustering runtimes during prefill, improving the scalability of our algorithm.
Additionally, to accelerate the update process, we sample additional initial centroids randomly from the appended tokens, and then leverage fast initial assignment of the appended tokens to the nearest centroids and corresponding update of the centroids that these tokens map to (this is analogous to a batched version of the sequential k-means algorithm [MaQuueen 1967). This process only needs to be run for the newly appended tokens, which makes its runtime overhead minimal. After running this single-shot update, to ensure high clustering accuracy we run a small number of iterations of full k-means over the final block to refine the centroids and centroid assignment.
5

---

Untouched when updating clusters
Refined when updating clusters
Keys
W Tokens
W Tokens
Final & to W x Tokens
Centroids
When final window hits W x tokens, split into W and < tokens
W+ x Tokens
W Tokens
cc Tokens
Figure Visualization our blockwise clustering method and sliding window approach for updating clusters during generation. We cluster the keys in blocks of W tokens, except for the final block, which contains the last & to W + a tokens. When we update the clusters, we only need to update the final window. Once the final window reaches W + a tokens, we split it into the first W tokens (which becomes its own block) and the final & tokens (which becomes the new final block that we append new keys to).
3.4 Hierarchical Multipole Attention
Although our method allows for pruning large portion of KV cache entries while maintaining accuracy, the granularity of clustering needs to be sufficient for accurate retrieval and approximation. However, this can lead to the centroid lookup becoming bottleneck if the number of centroids that we need to compare with grows with the context length. In particular, as we push to more aggressive sparsity regimes, the overhead of the centroid lookup becomes more prominent.
To alleviate the overhead of the centroid lookup, while still providing an accurate approximation for the full KV cache, we extend the hierarchical centroid lookup method from (Hooper al. 2024a] to a hierarchical multipole approximation for attention operation. We perform hierarchical K-means clustering in order to derive progressively coarser grained cluster centroids. This hierarchical clustering is compatible with the fast online clustering update described in Section 3.3] as the blockwise clustering method can be applied each level of the hierarchy. At each level of the centroid lookup, we identify a small number of promising key centroids (for which we need to compare with the next level of refined centroids), and then approximate the less important keys using the attention to the key centroids.
For the case where we have two levels of hierarchy, we first compare the query with the coarse-grained (first-level) centroids to identify which regions of the keys are potentially important. This initial lookup allows us to only perform fine-grained (second-level) centroid comparisons with centroids that are likely to be high scoring. Additionally, for the low-importance first-level centroids, we can approximate the attention all of the tokens in the corresponding clusters using the attention to the first-level centroids. We then perform the second level centroid lookup only for the promising fine-gre fine-grained centroids. Here, identify the most important tokens (which need to be loaded for exact attention computation). and we approximate the attention to the remaining tokens which were only important enough to reach the second-level centroid lookup. This progressive refinement process allows us to employ more accurate approximation for keys depending on their importance. Overall, our hierarchical multipole algorithm reduces the overhead of the centroid comparison, while still providing accurate approximation for the less important keys at each step in the lookup.
4 System Implementation
We implement custom Triton kernels for computing the centroid lookup, sparse Flash Decoding, and centroid replacement operations. These kernels follow the parallelization strategies from FlashAttention-2 [Dao) 2023), including splitting computation across attention heads and along

---

the query sequence length dimension, and FlashDecoding [Da et al. split across the KV dimension and the intermediate results are reduced. We also build upon existing work on sparse FltehAttention kernels (Pagliardini et al. 2023 Hooper et al. [024a]
2023 where computation is
The first stage of our kernel implementation (the centroid lookup) is designed compare the query with the key centroids in order to identify the important tokens (as in [Hooper et al. (2024a). D. After obtaining the attention scores the centroids, we gather the centroid scores per-token and apply lop-K to identify which top-scoring clusters of keys can be retained under a fixed token budget. After identifying the important and less important clusters, the next phase of our kernel implementation (sparse FlashDecoding) calculates the attention output from keys belonging to selected (important) clusters. This kernel accepts as input a tensor containing the indices of keys, and only loads the keys corresponding to these indices for exact attention computation.
Finally, in the third stage of our kernel implementation (centroid replacement), we approximate the attention to the less important keys using the attention to the corresponding centroids. Based upon Equation we need exp (gK;;] ]). i.e. the result of comparing the query with the key centroids. Given that these values are computed during centroid lookup, we modify the centroid lookup kernels to save these values and then we later load them in the replacement kernel (to avoid re-loading the key centroids during the replacement stage). We also return a boolean mask from the centroid lookup that identifies which centroids are less important in order to only perform replacement for these centroids. The centroid replacement kernel loads the query-key centroid dot product results from the centroid lookup and then computes attention with the value centroids. The attention output is then merged with the output from the sparse FlashDecoding kernel.
For our fast clustering update, we implement this using PyTorch-level primitives (wrapped in torch.compile to reduce kernel launch overheads). We use fixed buffer of L to 2L local tokens which are not clustered and are loaded exactly at each decoding step, and we append the oldest tokens from the buffer to the clusters every L decoding steps. After performing clustering, we cache the key and value centroids in order to avoid recomputing these values for each generation step.
5 Results
5.1 Experimental Setup
We evaluate MULTIPOLE ATTENTION on two long context reasoning datasets: LongBench V2 Bai et al. 2024) and GSM-Infinite (Zhou et al. 2025). LongBench contains complex real- long- context questions across tasks such as document QA, long in-context learning, dialogue understanding. and code repository understanding. GSM-Infinite contains synthetic mathematical reasoning tasks that require multiple arithmetic operations. We leverage two open-source reasoning models for our evaluation: Qwen3-8B (Yang et al. 2024] and DeepSeek-R1 ---isss-----sssle--eeeeen 14B. We use YaRN scaling across our experiments to enable up to 128K context length [Peng et al. 2023]. We use the recommended decoding settings for both models, which is to temperature of 0.6 and a top-p value of 0.95, and to set the top-k value to 20 for Qwen3-8B. We report average results across trials for all experiments.
We evaluate our approach using token budgets of 128 and 512. We fix the block size for clustering as W- = 8K tokens throughout our evaluation and set a W , and we set L = 128 (meaning that we keep the last 128 to 256 tokens unclustered and we update the clusters every 128 decoding steps). We set the number of k-means iterations as 10 for prefill and for refinement with our fast cluster update, and we use random initialization for the centroids. We also leave the first 10 "Attention Sink" tokens untouched due to their disproportionate importance (Xiao ct al. 2023
We compare our method with Squeezed Attention [Hooper et al. 2024a and QUEST (Tang et al. 2024] as representative baselines for sparse attention. Note that for Squeezed Attention, we report results that leverage our optimized online clustering implementation in order to identify important KV cache entries among the generated output tokens (as well as windowed RoPE for improved key clustering), but without our method for approximating attention to clusters of tokens using the attention to the centroids. The token budgets for the baselines are adjusted to ensure that they retain the same number of tokens when factoring in the local buffer used by MULTIPOLE ATTENTION. For MULTIPOLE ATTENTION, we use one centroid per 16 tokens, and for QUEST we use a page size of 16. For Squeezed Attention, we use one centroid per 8 tokens for fair comparison, since it only
7

---

42.0
H H
38.0
Accuracy
36.0
34.0
32.0
30.0
28.0
26.0
Overall
Easy
Hard
Short
Medium
Long
Baseline
Squeezed Attention
Multipole Attention
Figure 4: LongBench V2 evaluation for DeepSeek-R1 -ieeee-----issiii-e-------- 14B. We report accuracy on the full dataset, as well easy/hard difficulty splits, and for short/medium/long (<32K/32K- (A////2KK11KK/I11 128KK122KK splits. We report accuracy for a token budget of 512, where we observe that MULTIPOLE ATTENTION (MpAttn) can achieve higher accuracy than Squeezed Attention (SqAttn) for the same token budget, and we observe no accuracy drop on the overall benchmark score.
Table 1: Accuracy LongBench- V2 for Qwen3-8B 3-B with different token budgets. We report accuracy on the "Short" (<32K) split from the dataset. We include baseline comparisons with QUEST [Tagg al. 2024 as well as Squeezed Attention [Hooper et al. 2024a
Budget 128
Budget
Raseline
QUEST
13.7
22.2
Squeezed Attention
36.1
40.2
MULTIPOLE
40.4
has store key centroids (whereas method stores key and value centroids). For evaluation with the hierarchical variant of MULTIPOLE ATTENTION, we use a two-level lookup with one centroid per 64 tokens for the first level, and with centroid per 8 tokens for the second level (and when comparing with the hierarchical centroid lookup from Squeezed Attention, we run their method with twice many centroids each level for fair comparison with our approach).
For our kernel benchmarking experiments, we use split size 2048 in our FlashDecoding baseline and sparse FlashDeccding kernels, well as the centroid lookup and replacement kernels. Note that this suboptimal for a batch size of 1 for the lookup and replacement kernels, as the number of centroids is too small saturate the GPU; however, fixed this across all batch sizes for simplicity. We benchmark our kernel implementations using triton. testing. _o__ddo_bbeeeee bench with 500 warmup runs and 500 measurement runs.
5.2 Eyaluation
Figure presents evaluation on the LongBenchV2 dataset for the Reepseek-R1-Distil-oe model. These results show how MULTIPOLE ATTENTION is able to preserve the accuracy of the baseline, even with aggressive sparsity settings. Additionally, these results highlight how our method outperforms existing sparse attention methods, demonstrating how our multipole attention approximation is able to retain closer accuracy to the baseline by leveraging the scores to the centroids. Table [ also presents evaluation LongBenchV2 for the Qwen3-8B model. We include baseline comparisons with QUEST (Tang et al. 2024) and Squeezed Attention, demonstrating how MULTIPOLE ATTENTION outperforms existing sparse attention methods different token budgets. Note that we only report accuracy on the "Short" (<32K) split from the dataset, since the Qwen3- 8B model achieves lower than 25% accuracy on the medium long splits (which worse than random since the questions are multiple choice with 4 answers), primarily due to issues with endless repetitions. Appendix B provides additional analysis on Long [Bai et al., 2023) using non-reasoning long context model, demonstrating how our method is generalizable for long context tasks outside of reasoning. Appendix C also provides ablations for the block size and number of cluster centroids hyperparameters.
14B

---

Table 2: GSM-Infinite performance of Qwen3-8B 8K and 16K context lengths. We measure the accuracy for 1-operation and 2-operation splits with the symbolic task. We report results for MULTIPOLE ATTENTION ("MpAttn") as well as Squeezed Attention [Hooper et al,, [024a] (t"n'Atn").. We also include results using the hierarchical extension of our methodology, as well as the hierarchical lookup approach from Squeezed Attention (hierarchical configurations are marked with "-H"). We report the portion of memory operations required relative to the baseline (accounting for metadata overhead and assuming 8K/16K context lengths), MULTIPOLE ATTENTION achieves higher accuracy for the same memory overhead relative to Squeezed Attention across all token budgets and context lengths.
(a) Context length 8K
(b) Context length 16K
Method
Memory Operations
Method
Token Budget
Memory Operations
Acc.
Budget
op.)
op.)
op.)
(2 op.)
Baseline
Baseline
0.65
0.28
AqAttn
128
0.11
0.58
0.16
SqAttn
128
0.09
0.28
0.13
MpArt
0.11
0.72
0.19
MpArt
128
0.09
0,40
0.17
SqAltn-H
128
0.08
0.65
0.17
SqAmn-H
128
0.05
0.28
0.09
MpAtin
MpAitm
128
SqAttn
0.16
0.17
SqAttn MpArt
512
0.11
0.41
0.18
MpAtto
0.16
0.82
0.22
512
0.11
0.61
0.30
Batch size 1
Batch size 4
Batch size 16
1.0
Normalized Latency
0.8
2.4X
2.5X
3.3X
4.1X
3.7X
4.5X
0.6
0.4
0.2
0.0
Baseline
90% Prune
95% Prune
Baseline
90% Prune
95% Prune
Baseline
90% Prune
95% Prune
Cluster Update
Centroid Replacement
Centroid Lookup
Sparse FlashDecoding
Figure 5: Attention runtime evaluation an A6000 GPU for the Qwen3-8B model. We report results normalized to the FlashDecoding baseline. We show results for 90% and 95% sparsity (with one centroid per 16 tokens). We can attain up to 4.5x speedup relative to the baseline FlashDecoding kernels by leveraging MULTIPOLE ATTENTION.
Table 2 provides evaluation for the Qwen3-8B model GSM-Infinite. We include results for the single-leved and hierarchical variants of our method and the Squeezed Attention baseline (Hooper al. [0244]] We report results for and operation difficulty levels (referring the number of operations required to get to the final answer). Our method retains closer to baseline accuracy than sparse attention methods like Squeezed Attention. The advantages of MULTIPOLE ATTENTION particularly pronounced for aggressive sparsity regimes since there is a greater chance that sparse attention methods like Squeezed Attention will discard tokens which are important for retaining accuracy, whereas our approach always ensures that we retain at least an approximate representation for these tokens. Additionally, our hierarchical multipole algorithm attains higher accuracy for the same number of memory operations required relative to the one-level approach by reducing the overhead of the centroid lookup and centroid replacement operations.
5.3 System Benchmarking
We benchmarked our custom kernel implementations on both A6000 and A A100 GPU platforms. Figure 5 shows the observed attention speedups on an A6000 GPU with our method, and kernel benchmarking evaluation an A100 GPU is also provided in Appendix E We observe up to

---

2.5x, x,1x, and 4.5x speedups for batch sizes of 1, 4, and 16, respectively. We also show the breakdown of the portion of the runtime spent on our initial centroid comparison ("Centroid Lookup"), attention approximation using the centroids ("Centroid Replacement") sparse Flash Decoding for exact attention with the important keys ("Sparse FlashDDecoiing", as well as the median overhead from the fast cluster update ("Cluster Update"). For larger batch sizes, the centroid lookup and sparse FlashDecoding runtimes are substantial relative to the attention approximation kernel and cluster update, demonstrating the low overhead of our centroid replacement method relative to sparse attention, and how our fast cluster update enables quickly re-building an index over the generated tokens. Appendix D compares the attention speedups from MULTIPOLE ATTENTION with Squeezed Attention, demonstrating that we achieve comparable speedups for the same memory footprint. Appendix F provides additional amalysis of the clustering overheads during prefill as well as decode.
6
Conclusion
Reasoning has emerged as a key enabler of complex problem solving capabilities in LLMs. While this reasoning through long chain-of of-uhh decoding is critical for achieving high accuracy on complex tasks, it comes with substantial efficiency costs due to the need to generate thousands of tokens before producing an answer. While sparse attention methods can accelerate decoding by reducing the memory band width requirements for loading the KV cache, these approaches lead to substantial accuracy loss with aggressive sparsity regimes. Our work addresses these challenges by introducing MULTIPOLE ATTENTION, which reduces the KV cache memory requirements by only loading a small number of entries for exact attention computation, while maintaining approximate representations for the rest of the tokens. Our algorithm first clusters the keys based on semantic similarity, and then uses the corresponding cluster centroids both to select important key vectors and to approximate the attention to less important keys, thereby retaining important contextual information from the full sequence. We also design a fast cluster update strategy to facilitate quickly re-clustering the input to incorporate the newly generated tokens, which allows to accelerate attention to the previous output tokens. We evaluate our method on emerging reasoning models and hard long-context reasoning tasks, demonstrating that our approach can maintain accuracy on complex reasoning tasks with aggressive KV cache sparsity settings. We also present efficient kernel implementations, demonstrating how our algorithm can accelerate attention by up to 4.5x for long context reasoning applications.
7 Limitations
One limitation of our method is that it focuses on speeding up generation, and it does not accelerate attention during the prefill phase. One potential avenue for future work would be extending our method to support accelerating prefill. Another limitation is that additional implementation effort would be required to support our method in existing LLM serving frameworks. A third limitation is that our method requires additional memory to store the centroids (although this overhead is small relative to the size of the full KV cache).
Acknowledgements
We acknowledge gracious support from the FuriosaAI team including Jihoon Yoon, Suyeol Lee, and Hyung 1l Koo, as well as from Intel, Apple, NVIDIA, and Mozilla. We also appreciate the support from Microsoft through their Accelcrating Foundation Model Research, including great support from Sean Kuno. Furthermore, we appreciate support from Google Cloud, the Google TRC team, and specifically Jonathan Caton, and Prof. David Patterson. Prof. Keutzer's lab is sponsored by the Intel corporation, UC Berkeley APeAPI Center of Excellence, Intel VLAB team, as well as funding through BDD and BAIR. We appreciate great feedback and support from Ellick Chan, Saurabh Tangri, Andres Rodriguez, and Kittur Ganesh. Sehoon Kim would like to acknowledge the support from the Korea Foundation for Advanced Studies (KFAS). Michael W. Mahoney would also like to acknowledge a J. P. Morgan Chase Faculty Research Award as well as the DOE, NSF, and ONR. This work was supported by the Director, Office of Science, Office of Advanced Scientific Computing Research, of the U.S. Department of Energy under Contract No. DE-AC02-05CHI 1231. Our conclusions do not necessarily reflect the position or the policy of our sponsors, and no official endorsement should be inferred.
10

---

References
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint 2rXiv:2303.08774 2023.
Joshua Ainslie, James Lee-Thorp. Michiel De Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Training generalized multi-query transformer models from multi-head checkpoints. arXiv preprint ar arXiv:2305. 2023.
Sanghai. Gqa:
Anthropic. Claude 2:
http://c.ccooeeeeccclccaaia clandes 2023.
Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, et al. Longbench: A bilingual, multitask benchmark for long context understanding. arXiv preprint arXiv:2308. 14508, 2023.
Yushi Bai, Shangqing Tu, Jiajie Zhang, Hao Peng, Xiaozhi Wang. Xin Lv, Shulin Cao. Jiazheng Xu, Lei Hou, Yuxiao Dong, et al. Longbench v2: Towards deeper understanding and reasoning on realistic long- congecontext multitasks. arXiv preprint arXiv:2412.15200 2024.
Akhiad Bercovich, Itay Levy, Izik Golan, Mohammad Dabbah, Ran EI- El-Ya,iv, Omri Puny, Ido Galil, Zach Moshe, Tomer Ronen, Najeeb Nabwani, Ido Shahaf, Oren Tropp, Ehud Karpas, Ran Zilberstein, Jiaqi Zeng, Soumye Singhal, Alexander Bukharin, Yian Zhang, Tugrul Komuk, Gerald Shen, Ameya Sunil Mahabaleshwarkar, Bilal Kartal, Yoshi Suhara, Olivier Delalleau, Zijia Chen, Zhilin Wang. David Mosallanezhad, Adi Renduchintala, Haifeng Qian, Dima Rekesh, Fei Jia, Somshubra Majumdar, Vahid Noroozi, Wasi Uddin Ahmad, Sean Narenthiran, Aleksander Ficek, Mehrzad Samadi, Jocelyn Huang. Siddhartha Jain, Igor Gitman, Ivan Moshkov, Wei Du, Shubham Toshniwal, George Armstrong, Branislav Kisacanin, Matvei Novikov, Daria Gitman, Evelina Bakhturina, Jane Polak Scowcroft, John Kamalu, Dan Su, Kezhi Kong, Markus Kliegl, Rabeeh Karimi, Ying Lin, Sanjeev Satheesh, Jupinder Parmar, Pritam Gundecha, Brandon Norick, Joseph Jennings, Shrimai Prabhumoye, Syeda Nahida Akter, Mostofa Patwary, Abhinav Khattar, Deepak Narayanan, Roger Waleffe, Jimmy Zhang, Bor- BiinG Su, Guyue Huang, Terry Kong, Parth Chadha, Sahil Jain, Christine Harvey, Elad Segal, Jining Huang, Sergey Kashirsky, Robert McQueen, Izzy Putterman, George Lam, Arun Venkatesan, Sherry Wu, Vinh Nguyen, Manoj Kilaru, Andrew Wang, Anna Wamo,, Abhilash Somasamudramath, Sandip Bhaskar, Maka Dong, Nave Assaf, Shahar Mor, Omer Ullman Argov, Scot Junkin, Oleksandr Romanenko. Pedro Larroy, Monika Katariya, Marco Rovinelli, Viji Balas, Nicholas Edelman, Anahita Bhiwandiwalla. Muthu Subramaniam. Smita Ithape Karthik Ramamoorthy, Yuting Wu, Suguna Varshini Velury, Omri Almog, Joyjit Daw, Denys Fridman, Erick Galinkin, Michael Evans, Katherine Luna, Leon Derczynski, Nikki Pope, Fileen Long, Seth Schneider, Guillermo Siman, Tomasz Grzegorzek, Pablo Ribalta, Monika Katariya, Joey Conway, Trisha Saar, Ann Guan, Krzysztof Pawelec, Shyamala Prayaga, Oleksii Kuchaiev, Boris Ginsburg, Oluwatobi Olabiyi, Kari Briski, Jonathan Cohen, Bryan Catanzaro, Jonah Alben, Yonatan Geifman, Eric Chung, and Chris Alexiuk. Llama-nemotron: Efficient reasoning models, 2025. URL http::tts:///arat org/ 255g//55.2500..... 2505. 00949
Shouyuan Chen, Sherman Wong, Liangjian Chen, and Yuandong Tian. Extending context window of large language models via positional interpolation. arXiv preprint S.Ki::2306./5599 2023.
Ronald Coifman, Vladimir Rokhlin, and Stephen Wandzura. The fast multipole method for the wave equation: A pedestrian prescription. IEEE Antennas and Propagation magazine, 35(3):7-12, 1993.
Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. arXiv preprint arXiv:2307.0869) 2023.
Tri Dao, Daniel Haziza, Francisco Massa, and Grigory Sisov. Flash-decoding for long-context inference: /attss://rfft: http:e////ccc....ttaa...eeee//// wtt.a//cwc.....///////////aaaaaa
2023.
Google. Gemini 1.5 https://blog... htt.a///gggggge.eeeooo////////// generration-model--eeraaaa //tp://ooa///////////////////o//
-enrini--eett:
2023.
Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The llaaa herd of models. arXiv preprint arXiv:2407.21771, 2024.

---

Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song. Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, et Deepseek-rl: Incentivizing reasoning capability in llms via reinforcement learning. arXiv preprini arXiv:2501. 2at.:22001...44,,. 2025.
Junhui He, Junna Xing, Nan Wang, Rui Xu, Shangyu Wu, Peng Zhou, Qiang Liu, Chun Jason Xue, and Qingan Li. A ATS: Retrieval-based kv cache reduction via windowed rotary position embedding and query-aware vector quantization. arXiv preprint arXiv:2502.12660 2025.
Coleman Hooper, Sehoon Kim, Hiva Mohammadzadeh, Monishwaran Maheswaran, June Paik, Michael W Mahoney, Kurt Keutzer, and Amir Gholami. Squeezed attention: Accelerating long context length Ilm inference. arXiv preprint arXiv:2411.0900 1.09688, 2024a.
Coleman Hooper, Sehoon Kim, Hiva Mohammadzadeh, Michael W Mahoney, Yakun Sophia Shao, Kurt Keutzer, and Amir Gholami. Kvquant: Towards 10 million context length Ilm inference with kv cache quantization. Advances in Neural Information Processing Systems, 2024b.
Yanming Kang. Giang Tran, and Hans De Sterek. Fast multipole attention: A divide-and and-conquer attention mechanism for long sequences. arXiv preprint arXiv:2310.11960, 2023.
Di Liu, Meng Chen, Baotong Lu, Huiqiang Jiang, Zhenhua Han, Qianxi Zhang, Qi Chen, Chen- gruidong Zhang, Bailu Ding, Kai Zhang, et al. Retrievalattention: Accelerating long-context Ilm inference via vector retrieval. arXiv preprint arXiv:2409. 105 16, 2024a.
Xiang Liu, Zhenheng Tang, Hong Chen, Peijie Dong, Zeyu Li, Xiuze Zhou, Bo Li, Xuming Hu, and Xiaowen Chu. Can lims maintain fundamental abilities under kv cache compression? arXiv preprint arXiv:2502.011 01944,,, 2025.
Zirui Liu, Jiayi Yuan, Hongye Jin, Shaochen Zhong, Zhaozhuo Xu, Vladimir Braverman, Beidi Chen, and Xia Hu. Kivi: A tuning-free asymmetric 2bit quantization for kv cache. arXiv preprint 202.:22402.027500 2024b.
James MacQueen. Some methods for classification and analysis of multivariate observations. In Proceedings of the Fifth Berkeley Symposium on Mathematical Stutistics and Probubility Volume 1: Statistics, volume 5, pages 281-298. University of California press, 1967.
Meta. Llama 3.1: https://ai meta.
httts::/aa.metaa.oomm// comboooommooo//neeta.aa 1lama-3---- 3-mma-3-111 2024.
Matanel Oren, Michael Hassid, Nir Yarden, Yossi Adi, and Roy Schwartz. Transformers are multi- state mns. arXiv preprint arXiv:2401.061044 2024.
Matteo Pagliardini, Daniele Paliotta, Martin Jaggi, and François Fleuret. Fast attention over long sequences with dynamic sparse flash attention. Advances in Neural Information Processing Systems, 36:59808-59831 2023.
Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. Yarn: Efficient context window extension of large language models. arXiv preprint arXiv:2309.0000 2023.
Qwen. Qwq-32b: Embracing the power of reinforcement learning, March 2025. URL https //qwenlm.. github. 10/ ht/iio//lllle hhtp//llogg/g//220
Charlie Snell, Jaehoon Lee, Kelvin Xu, and Aviral Kumar. Scaling Ilm test-time compute optimally can more effective than scaling model parameters. arXiv preprint M.Xiv:2408.03314 2024.
Jiaming Tang, Yilong Zhao, Kan Zhu, Guangxuan Xiao, Baris Kasikci, and Song Han. Quest: Query-aware sparsity for efficient long-context I1m inference. arXiv preprint arXiv:2406.10774 2024.
Zheng Wang. Boxiao Jin, Zhongzhi Yu, and Minjia Zhang. Model tells you where merge: Adaptive kv cache merging for llms on long-context tasks. arXiv preprint AXi::2407.084555 2024.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoe V Le, Denny Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models, Advances in neural information processing systems, 25:24824-24837, 2022.
12

---

Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. arXiv preprint arXiv:2309. 17955333. 2023.
Guangxuan Xiao, Jiaming Tang, Jingwei Zuo, Junxian Guo, Shang Yang, Haotian Tang, Yao Fu, and Song Han. Duoattention: Efficient long-context IIm inference with retrieval and streaming heads. arXiv preprint arXiv:2410.1055 2024.
An Yang, Baosong Yang. Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chengyuan Li, Dayiheng Liu, Fei Huang, Haoran Wei, Huan Lin, Jian Yang, Jianhong Tu, Jianwei Zhang, Jianxin Yang. Jiaxi Yang. Jingren Zhou, Junyang Lin, Kai Dang, Keming Lu, Keqin Bao, Kexin Yang, Le Yu, Mei Li, Mingfeng Xue, Pei Zhang, Qin Zhu, Rui Men, Runji Lin, Tianhao Li, Tingyu Xia, Xingzhang Ren, Xuancheng Ren, Yang Fan, Yang Su, Yichang Zhang, Yu Wan, Yugiong Liu, Zeyu Cui, Zhenru Zhang, and Zihan Qiu. Qwen2.5 technical report. arXiv preprint 1411 2024.
Jian Yuan, Ziwei He, Haoli Bai, Jingwen Leng, and Bo Jiang. Weightedkv: Attention scores weighted key-value cache merging for large language models. ICASSP 2025-2025 IEEE International Conference Acoustics, Speech and Signal Processing (ICASSP), pages 1-5. IEEF, 2025.
Hailin Zhang, Xiaodong Ji, Yilin Chen, Fangcheng Fu, Xupeng Miao, Xiaonan Nie, Weipeng Chen, and Bin Cui. Pqcache: Product quantization based kvcache for long context Ilm inference. arXiv preprint 1rXiv2240..12820, 2024a.
Zhenyu Zhang. Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, et al. H2o: Heavy-hitter oracle for efficient generative inference of large language models. Advances in Neural Information Processing Systems, 36, 2024b.
Yang Zhou, Hongyi Liu, Zhuoming Chen, Yuandong Tian, and Beidi Chen. Gsm-infinite: How do your Ilms behave over infinitely increasing context length and reasoning complexity? arXiv preprint M.Xiv:2502.0525555 2025.
Kan Zhu, Tian Tang, Qinyu Xu, Yile Gu, Zhichen Zeng, Rohan Kadekodi, Liangyu Zhao, Ang Li, Arvind Krishnamurthy, and Baris Kasikci. Tactic: Adaptive sparse attention with clustering and distribution fitting for long-context Ilms. arXiv preprint arXiv:2502. 12216, 2025.
13

---

A Extended Related Work: Only Loading Important KV Cache Entries
One previous approach for sparse attention has aimed to retain the full KV cache, but only load in the required KV cache entries for each decoding step. A common method that these approaches have used to identify important KV cache entries is clustering the keys. QUEST [Tang et al.) 2024] grouped consecutive KV cache entries and derived representative vectors for the keys within each group. Their approach performs retrieval by comparing the current query with the representative vectors in order to identify whether the group is important and needs to be loaded for exact attention computation. While this method grouped consecutive keys, positional proximity in the sequence is not necessarily indicative of the similarity of the corresponding key vectors. Squeezed Attention Hooper et al. 2024a] instead clustered keys based on semantic similarity, thereby allowing for precise identification of tokens which were likely to be high-scoring. Tactic (Zhu et al. [2025] also clustered the keys based on semantic similarity and then performed distribution fitting to help predict important KV cache entries.
Multiple prior works have framed the query-key comparison as retrieval problem in order to leverage vector search methods to accelerate attention. One such method, PQCache (Zhang et al., 2024a]. performed product quantization in order to perform vector search to select important keys, and then only computed exact attention using the selected keys. A2ATS (He et al., 2025 performed query-aware vector quantization in order to allow for accurate retrieval of important KV cache entries. RetrievalAAtentiion [Liu al., [0244] constructs a vector index over the keys in order to facilitate retrieval of important keys, and offloads the full KV cache and the vector index to the CPU to perform
the lookup.
In contrast with prior work on accelerating decoding through sparsely loading the KV cache, our method aims to select and load small subset of important tokens for exact attention computation, while maintaining an approximate representation for the remaining KV cache entries. Our approach uses clustering group semantically similar keys, followed by a fast centroid lookup that identifies important keys. Notably, our key centroids also serve as representative vectors for the keys in each of the clusters, and can therefore be used approximate the attention the keys which are not retrieved by the centroid lookup. This allows our method to maintain high accuracy even with aggressive sparsity settings (and for tasks which are sensitive to KV cache pruning).
Another related prior work is (Kang et al., 2023]. which utilized a Fast Multipole Method (FMM)- inspired approach for accelerating the attention computation in Transformers. FMM [Coifman et al. 1993) is a related technique for accelerating simulations with N-Body problems which approximates groups of point masses a single point mass and performs coarser grained approximations for masses which are further away from the current point of interest. This approach is analogous to our algorithm, which approximates keys using progressively coarser-grained key centroids as get further from the current query token. However, unlike (Kang et al. [0033 which grouped KV cache entries based on positional proximity in the sequence, our method clusters keys based on semantic similarity, and therefore determines the granularity of approximation based on the importance of KV cache entries for the current query.
B Additional Long Context Experiments
To demonstrate the broader applicability of our methodology on long context tasks other than reason- ing, we evaluate our method on LongBenchVV [Bai al., 2023] using a non-reasoning long context model --ama-3.1-8B--1 Instrrcct [Meta] (024)). ) with results provided in Table 3 We also provide additional baseline comparisons with Squeezed Attention [Hooper et al. 2024a]. DuoAttention (Xiao et al. 2024) and TOVA [Oren et al. 2024]. We use the evaluation setup from the DuoAAtention open-source code, which simulates decoding for the last 50 prompt tokens, and we use a temperature of O. We configure MULTIPOLE ATTENTION to use 95% sparsity and one centroid per 16 tokens, and then adjust the other methods to have equivalent KV cache budgets (88.75% sparsity when accounting for metadata). Our results demonstrate how our method substantially outperforms Squeezed Attention, AttAntention, and TOVA on range of long context evaluation tasks, and also how our method generalizes to long context tasks outside of reasoning.
14

---

Table 3: MULTIPOLE ATTENTION ("MpAttn") evaluation non-synthetic tasks from LongBench
Bai et al. 2023 for the Llama-3. 1-ama--1--B-Insttvee Instruuucctt model. We include baseline comparisons with Squeezed Attention ("SqAAtt") (Hooper et al. 2024a well as DuoAttention ("DuoAttt")) [Xiao et al. 2024) and TOVA [Oren et al. 2024] (configured to have the same KV cache budget as our method). We also provide estimates for the number of memory operations required for KV cache loading, including metadata (normalized to the number of KV cache memory operations required for the baseline). Note that the TOVA algorithm allows each query head to select different KV entries and would therefore have higher KV footprint for GQA models even for the same sparsity settings. MULTIPOLE ATTENTION provides noticeable accuracy improvements for the same KV cache budget relative to existing methods.
Single-Doc
Multi-Doc.
Summarization
Few-shot Learning
Code
Config
Mem. Ops.
NQA
Qspr
MFQA
HIQQA
2Wikd
MSQ
GRep
QMSM
MNews
TREC
TQ4
58mm
LCC
Avg.
RB
Baseline
1.00
30.1
45.3
55.5
$6.0
44.8
30.5
35.2
25.3
27.3
72.5
91.2
43.2
49.0
52.5
47.0
DuaAut
0.1125
13.2
16.8
21.3
29.1
17.4
8.2
20.5
17.3
20.1
32,0
52.4
36.8
47.0
49.0
27.2
TOVA
0.1125
24.1
14.0
23.2
41.5
21.1
15.1
25.5
19.9
20.9
35.0
86.
39.9
49.1
48.6
33.1
SqArtn
0.1125
29.3
24.2
38.6
50.1
31.5
24.5
55.0
41.5
49.9
47.4
39.6
MpAtta
0.1125
31.6
32,9
42.0
54.3
31.3
30.3
32.8
24.7
22.9
67.5
86.7
42.1
46.8
48.7
42.5
C Hyperparameter Ablations
Tables 4 5 provide ablations for the number of centroids as well as block size, respectively. We report LongBenchV2 accuracy for the Qwen3-8B model on the "short" (<32K context length) split with token budget of 128. We find that if the number of KV tokens per centroid is increased or if the block size decreased, the accuracy degraded. However, increasing these parameters leads to additional inference costs. We therefore used centroid per 16 tokens and a block size of 8K to retain model accuracy without introducing substantial memory and latency overheads.
Table 4: LongBench accuracy on the "Short" (<32K) split versus ratio of centroids KV tokens for the Qwen3-8B model. We report results for a token budget of 128. The configuration used throughout the paper is bolded.
Baseline
Ratio=1/128
Ratio=1/64
Ratio=1/32
Ratio=1/16
Ratio=1/8
41.5
29.3
35.4
38.5
38.1
37.0
Table 5: LongBench V2 accuracy the "Short" (<32K) split versus clustering block size W for the Qwen3-8B model. We report results for a token budget 128. The configuration used throughout the paper is bolded.
Baseline
2K
4K
8K
16K
41.5
35.0
38.0
38.1
36.9
D Runtime Comparison with Squeezed Attention
In Table 6 we compare the latency for MULTIPOLE ATTENTION versus Squeezed Attention Hooper al.) [024aa] for the same memory footprint (using one centroid per 16 tokens for MULTIPOLE ATTENTION and one centroid per 8 tokens for Squeezed Attention, as outlined in Section 5.I.. Note that the Squeezed Attention baseline we compare with in our experiments also contains the clustering improvements from MULTIPOLE ATTENTION, which are required to allow it to be applied for tasks where the prefill only known runtime and during the generation process to accelerate attention to newly generated tokens (as the baseline Squeezed Attention method cannot be run online). We report the speedups for Squeezed Attention and MULTIPOLE ATTENTION on an A6000 GPU with batch sizes of 1/4/16 (assuming 90% sparsity). These results outline how MULTIPOLE ATTENTION provides similar latency for the same sparsity level as Squeezed Attention, while providing substantial accuracy improvements (as highlighted in Section 5.2
15

---

Table 6: Attention speedup for MULTIPOLE ATTENTION and Squeezed Attention Hooper al. 2024a for the Qwen3-88 8B model on an A6000 GPU, relative to the full attention baseline.
Method
Batch Size
Batch Size
Batch Size 16
Squeezed Attention
2.8x
3.3x
366 3.3x
MULTIPOLE ATTENTION
24x
3.7x
Batch size
Batch size 4
Batch size 16
1.0
1.2X
1.3X
1.9X
2.1X
2.4X
2.8X
Normalized Latency
0.8
0.6
0.4
0.2
0.0
Baseline
90% Prune
95% Prune
Baseline
90% Prune
95% Prune
Baseline
90% Prune
95% Prune
Cluster Update
Centroid Replacement
Centroid Lookup
Sparse FlashDecoding
Figure 6: Attention runtime evaluation A100 GPU for the Qwen3-8B model. We report results normalized to the FlashDecoding baseline. We show results for 90% and 95% pruning, and for batch sizes of 1, and 16 (and assuming one centroid per 16 tokens). We attain up to 2.8 speedup relative to the baseline FlashDecoding kernels by leveraging MULTIPOLE ATTENTION.
E A100 Kernel Benchmarking
Figure shows the observed attention speedups A100 GPU with our method. We observe up 1.3x, 2.1 x, and 2.8x speedups for batch sizes of 1, 4, and 16, respectively. Note that there are reduced speedups for batch size of due to there being insufficient work to saturate the GPU.
F Clustering Runtime
Tables and g provide measured runtime for the fast cluster update during generation on A6000 and A100 GPUs, respectively. We find that the overhead of the fast cluster update is typically 3-4% of the runtime of the FlashDecoding baseline on an A6000 GPU, and 5% of the runtime of the FlashDecoding baseline which is relatively low (between 13-15% on A6000/A 100 GPUs).
A100 GPU. Table 9 also reports the clustering overhead during prefill,
Table 7: Runtime (in milliseconds) for running our fast clustering update for the Qwen3-8B model on A6000 GPU. The clustering runtimes for 8K and 12K tokens serve the median and maximum clustering runtimes, respectively, since we W 8K and W throughout Since we update the clusters once every 128 decoding steps (and need to run FlashDecoding at each step), clustering adds median of 3-4% overhead and maximum of 6% overhead during the decoding process.
evaluation.
Batch Size
FlashDecodiing Runtime (128K context length)
Clustering Runtime (8K)
Clustering Runtime (12K)
2.9
12.4
23.1
16

---

Table 8: Runtime (in milliseconds) for running our fast clustering update for the Qwen3-8B model on an A100 GPU. The clustering runtimes for 8K and 12K tokens serve the median and maximum clustering runtimes, respectively, since we set W = 8K and = 1 W throughout Since we update the clusters once every 128 decoding steps (and need to run FlashDecoding at each step), clustering adds median of 5% overhead and maximum of 8-9% overhead during the decoding process.
evaluation.
Barch Size
FlashDsooding Runtime (128K context length)
Clustoring Runtime (8K)
Clustering Runtime (12K)
0.4
4.2
1.3
8.4
14.8
Table Runtime overhead (in milliseconds) for running our blockwise clustering algorithm during prefill for the Qwen3-8B model. We report overheads on both A100 and A6000 GPUs for batch size 1, assuming 128K context length prefill (divided into 8K blocks). Clustering adds a relatively low overhead of between 3315%% when processing the input prompt.
GPU
FlashAttention Runtime (128K context longth)
Clustering Runtime blocks of size 8K)
Clustering Overhead (%).
A6000
1506
192
A100
787
118
17

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
The abstract and introduction clearly state the contributions of the work
(namely, optimizing long-context reasoning through a novel algorithm and system imple- mentation).
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We included a limitations section in the appendix.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-speciiiatiinn asymptotic approximations only holding locally). ). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
18

---

Answer:
[NA]
Justification:
We did not include any theoretical results requiring proofs.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification: We include detailed description of the method as well hyperparameter settings and system configuration.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, descríbing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e-g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
5. Open access to data and code
19

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes)
Justification: Yes, we will submit our code as part of the supplemental material
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips... cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
/ "cippscc/puulli//guu /utee//ggudddeeeedSuuuubmiibii ssionPolicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: All hyperparameter settings are specified.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification:
Our experimental results are for individual LLM evaluation runs.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
20

---

The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We specify the hardware platforms used for benchmarking results.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental rums as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics htttp:://eeurip:: http://wwwwp..ccoooa/ll/cccccaaa
Answer: [Yes]
Justification: Our work does not violate the code of ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our work is focused on efficiency for long context reasoning. There are hence no direct links applications with positive or negative societal impacts.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
21

---

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
We do not release models or datasets in our work.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We cite the creators of existing models and datasets that we use, as well as previous code implementations that build upon.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the Datasets.papersstt paperswi thcode. cww/ datasets package should be provided. For popular datasets,
has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
22

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: We provide documentation and instructions with our code to understand and run it.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Question:
Answer: [NA]
Justification:
We did not perform research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation,
or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification:
We did not perform research with human subjects or crowdsoureing.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such the institution conducting the review.
23

---

16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
Justification: We did not use LLMs in non-standard manner when producing this work.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to L.LM policy http:p/s/eaaaeesssss////oo/eeeea for what should or should not be described.
24


Reviews:
[{"review_id": "iClXtZicdD", "reviewer": "Reviewer_zAWi", "summary": "The paper introduces Multipole Attention, a sparse attention mechanism which selectively computes exact attention only for the most relevant tokens, while using approximate attention based on semantically clustered key-value centroids for the rest of the sequence. To achieve this, the method first performs semantic k-means clustering over key vectors and derives representative centroids. At inference time, the model compares the query with these centroids to determine which tokens to retrieve for exact attention and which to approximate. To ensure scalability during generation, the authors propose a fast online clustering algorithm with blockwise updates and a sliding window to re-cluster newly generated tokens efficiently. Additionally, the paper extends this to a hierarchical multipole formulation, enabling progressively coarser approximations for less important tokens. The authors implement this method using custom Triton kernels and demonstrate that it achieves up to 4.5× decoding speedup on reasoning benchmarks like LongBenchV2 and GSM-Infinite, while maintaining accuracy significantly better than prior sparse attention baselines such as QUEST and Squeezed Attention.", "questions": "Can multipole attention be applied to tasks that are not reasoning related?", "limitations": "Yes", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 4, "originality": 4, "strengths_and_weaknesses": "Strengths:\n- introduces a novel sparse attention mechanism which can be used to substitute attention during inference\n- experimental results on long reasoning benchmarks demonstrates that proposed method can retain strong reasoning performance, even with aggressive KV cache sparsity.\n- proposed technique is designed to be efficient and scalable (authors propose using fast online clustering and have a custom triton kernel implementation)\n- writing is clear\n\nWeaknesses:\n- generalization outside of reasoning- while the paper is focused on the reasoning task, I don't think the idea of being able to perform well on long contexts is restricted to reasoning.  It would be interesting to see how multipole attention performs in broader use cases such as summarizing/qa for long documents\n- missing ablations for hyperparameters such as number of centroids and block size", "was_revised": true, "final_justification": "Overall, I thought this was an interesting work that provides rigorous experimentation to demonstrate that the proposed multipole attention helps improve long context reasoning performance.  I appreciate that the authors performed additional experiments on other long context tasks outside of reasoning and additional hyperparameter ablations.  Because of this I recommend accept.", "rebuttal": "> R1-1: Generalization outside of reasoning- while the paper is focused on the reasoning task, I don't think the idea of being able to perform well on long contexts is restricted to reasoning. It would be interesting to see how multipole attention performs in broader use cases such as summarizing/qa for long documents\n\nWe completely agree and appreciate the reviewer’s feedback. We focused our evaluation on reasoning models since newer reasoning models typically have much longer output generation lengths, which makes accelerating the generation phase particularly critical, and since reasoning tasks require broad understanding of the input context and experience greater accuracy loss with existing sparse attention methods [1]. However, note that we evaluated these reasoning models on a range of tasks requiring complex reasoning capabilities; LongBenchV2 includes questions from a range of downstream tasks, including single and multi-document QA, in-context learning, and code repository understanding, and GSM-Infinite contains arithmetic reasoning questions. \n\nAdditionally, we would like to highlight that we have evaluated Multipole Attention on non-reasoning models (see response R3-1) on various long context length tasks, including the LongBench dataset (which contains single and multi-document QA, summarization, in-context learning, and code repository understanding) and passkey retrieval. These results demonstrate the generalizability of our method.\n\n[1] Can LLMs maintain fundamental abilities under KV cache compression?\n\n\n> R1-2: Missing ablations for hyperparameters such as number of centroids and block size\n\nWe have included ablations here for the number of centroids as well as block size.  We report LongBenchV2 accuracy for the Qwen3-8B model on the “short” (<32K context length) split with a token budget of 128. For the number of centroids ablation, these results demonstrate that if we use coarser grained clustering than one centroid per 16 tokens (the main configuration used in the paper), the accuracy is noticeably degraded - we therefore used 1 centroid per 16 tokens as the main configuration throughout our work. For the block size ablation, we find that if we use less than 8K block size, the accuracy is degraded. If we use greater than 8K block size,  the efficiency overhead of clustering will be noticeably increased; we therefore selected 8K as the setting throughout our experiments.\n\n| Centroid / Token Ratio | Baseline | 1/128 | 1/64 | 1/32 | 1/16 | 1/8  |\n| ---------------------- | -------- | ----- | ---- | ---- | ---- | ---- |\n| LongBenchV2 Accuracy   | 44.4     | 31.1  | 35   | 36.1 | **40.6** | 38.3 |\n\n*LongBenchV2 accuracy versus ratio of centroids to KV tokens for Multipole Attention with the Qwen3-8B model (bold - the config used in the paper).*\n\n| Block Size           | Baseline | 2K   | 4K   | 8K   | 16K  |\n| -------------------- | ------- | ---- | ---- | ---- | ---- |\n| LongBenchV2 Accuracy | 44.4     | 35.0 | 38.3 | **40.6** | 36.7 |\n\n*LongBenchV2 accuracy versus clustering block size for Multipole Attention with the Qwen3-8B model (bold - the config used in the paper).*"}, {"review_id": "syikL6V67a", "reviewer": "Reviewer_zdAu", "summary": "This paper proposes a method to reduce KV cache requirements during long-context inference by performing exact attention only for tokens identified as important, while using approximate cluster representations for the remaining tokens. Attention for these other tokens is computed with respect to the cluster centroids. This approach reduces KV cache demands while maintaining a reasonably good representation over the sequence and preserving final accuracy fairly well.", "questions": "How does Multipole Attention compare to Squeezed Attention when both accuracy and runtime are jointly taken into account? In particular, is it possible to report results in a way that clarifies the trade-off between these two dimensions?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "## Strengths\n\nThe paper introduces a potentially useful method to speed up long-context inference while preserving accuracy to a reasonable degree. The algorithmic design and its description appear sound. The system-level implementation make likely to be practical in real-world scenarios.\n\n## Weaknesses\n\nThe contribution is conceptually an incremental improvement over Squeezed Attention. In practice, while it appears to preserve accuracy slightly better than squeezed attention, it also seems more computationally expensive, since it requires computing approximate attention over cluster centroids.\n\nIn general, the comparison with Squeezed Attention could be performed along two dimensions simultaneously: accuracy preservation and runtime. This would allow us to see whether the proposed approach is significantly better for a given runtime budget, or significantly faster for a given target accuracy. \n\nRegarding Figure 4: “We report accuracy for a token budget of 512, where we observe that MULTIPOLE ATTENTION (MpAttn) can achieve higher accuracy than Squeezed Attention (SqAttn) for the same token budget.” This choice of token budget appears somewhat cherry-picked.\n\nOverall, it would be more convincing if the experimental evaluation was more thorough and systematic.\n\nFinally, there is no measure of uncertainty reported in the tables and figures (e.g., standard deviation or confidence intervals).\n\nMinor: The citations are not correctly formatted (\\citet, \\citep, etc.)", "was_revised": true, "final_justification": "Clarifications from the authors about the evaluation setup convinced me to update the overall assessment.", "rebuttal": "> R2-1: The contribution is conceptually an incremental improvement over Squeezed Attention. In practice, while it appears to preserve accuracy slightly better than squeezed attention, it also seems more computationally expensive, since it requires computing approximate attention over cluster centroids. In general, the comparison with Squeezed Attention could be performed along two dimensions simultaneously: accuracy preservation and runtime. This would allow us to see whether the proposed approach is significantly better for a given runtime budget, or significantly faster for a given target accuracy. \n> How does Multipole Attention compare to Squeezed Attention when both accuracy and runtime are jointly taken into account? In particular, is it possible to report results in a way that clarifies the trade-off between these two dimensions?\n\nWe appreciate the feedback. **As highlighted below, we can achieve 43-54% less accuracy degradation from the baseline than SqueezeAttention when comparing configurations with the same runtime.** To attain these benefits, we made multiple improvements relative to Squeezed Attention:\n\n- We incorporate algorithmic improvements by approximating attention to less important KV cache tokens, and we develop a systems implementation to compute this approximate attention to less important regions with minimal overhead \n- We design an efficient clustering strategy such that we can apply clustering online during inference (whereas Squeezed Attention could only be applied for tasks where the prefill was available offline ahead of time), and we also design a fast cluster update procedure such that we can also apply clustering to the newly generated tokens during long generation\n\nIn the paper, we provided comparisons which used memory footprint as a proxy for runtime, and demonstrated how Multipole Attention provides improved accuracy relative to Squeezed Attention for the same memory footprint. **For example, Table 2 highlights how the best Multipole Attention configuration with token budget 128 achieves 56-96% less accuracy degradation from the baseline than the best Squeezed Attention configurations with similar memory footprint (across both context lengths and 1 and 2-operation tasks).** Here, we also provide additional analysis to clarify how the reduced memory footprint correlates with a reduction in runtime, which we will incorporate into the final version of our paper. \n\nWe compare the accuracy with Multipole Attention and Squeezed Attention for the Qwen3-8B model on GSM-Infinite (16K context length) and LongBenchV2 with a 128-token budget and using one centroid per 16 tokens (and using 2x more centroids for Squeezed Attention to ensure the same memory footprint, since Multipole Attention requires computing centroids for the keys and values). Note that the Squeezed Attention baseline we compared with in our experiments also contains the clustering improvements from Multipole Attention, which are required to allow it to be applied for tasks where the prefill is only known at runtime and during the generation process to accelerate attention to newly generated tokens (as the clustering from the original Squeezed Attention method cannot be run online). We also report the speedups for Squeezed Attention and Multipole Attention on an A6000 GPU with batch sizes of 1/4/16 (assuming 90% sparsity). These results outline how Multipole Attention provides substantial accuracy improvements for a similar runtime speedup relative to Squeezed Attention (with 43-54% less accuracy degradation from the baseline across GSM-Infinite and LongBenchV2 datasets).\n\n|                     | GSM-Inf - 1-Op Accuracy | GSM-Inf - 2-Op Accuracy | LongBenchV2 Accuracy | Speedup (Batch size 1) | Speedup (Batch size 4) | Speedup (Batch size 16) |\n| ------------------- | ----------------------- | ----------------------- | -------------------- | ---------------------- | ---------------------- | ----------------------- |\n| Baseline            | 0.62                    | 0.28                    | 44.4                 | 1                      | 1                      | 1                       |\n| Squeezed Attention  | 0.27                    | 0.11                    | 36.1                 | 2.8                    | 3.3                    | 3.6                     |\n| **Multipole Attention** | **0.42**                    | **0.19**                    | **40.6**                 | **2.4**                    | **3.3**                    | **3.7**                     |\n\n*GSM-Infinite 1-op/2-op accuracy and LongBenchV2 accuracy versus attention speedup for Squeezed Attention and Multipole Attention with the Qwen3-8B model.*\n\n\n\n> R2-2: Regarding Figure 4: “We report accuracy for a token budget of 512, where we observe that MULTIPOLE ATTENTION (MpAttn) can achieve higher accuracy than Squeezed Attention (SqAttn) for the same token budget.” This choice of token budget appears somewhat cherry-picked.\n\nWe would like to kindly note that we reported accuracy results for token budgets of 128 and 512 across both datasets (LongBenchV2 as well as GSM-Infinite) and for both models, and that Figure 4 contains results for a token budget of 128 and 512. Please see Table 1 and Figure 4 for LongBenchV2 results, and Table 2 for GSM-Infinite results.\n\n\n> R2-3: Finally, there is no measure of uncertainty reported in the tables and figures (e.g., standard deviation or confidence intervals).\n\nWe appreciate the reviewer’s feedback. In the final version, we will incorporate experimental results with multiple samples with error bars.\n\n\n> R2-4: Minor: The citations are not correctly formatted (\\citet, \\citep, etc.)\n\nThank you for the formatting correction - we will fix this in the final version of our paper."}, {"review_id": "byUzGfXpB8", "reviewer": "Reviewer_vJGV", "summary": "This paper introduces a method that uses hierarchical key clustering and centroid-based attention to reduce the KV cache while preserving contextual information. To accelerate attention computation, the authors propose a Fast Online Clustering algorithm and implement system-level optimizations using a Triton kernel.", "questions": "How does the method perform across different context lengths? An ablation study or detailed analysis would be helpful.", "limitations": "See Weakness", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. The idea of using clustering to compress the KV cache is well-motivated, and the hierarchical clustering design aligns well with modern memory access patterns;\n2. The proposed Efficient Online Clustering algorithm makes KV cache compression practically feasible;\n3. The system implementation based on Triton kernels demonstrates the efficiency of Efficient Online Clustering convincingly.\n\nWeaknesses:\n1. The experiments are limited to only two benchmarks: LongBenchV2 and GSM-Infinite. The evaluation lacks diversity in tasks and datasets.\n2. The comparison is also restricted, focusing solely on Squeezed Attention and QUEST.\n3. It would strengthen the paper to include additional benchmarks, such as Needle-in-a-Haystack [1], and comparisons with more recent methods like DuoAttention [2] and TOVA [3].\n\nReferences:\n\n[1] Needle-in-a-Haystack: https://github.com/gkamradt/LLMTest_NeedleInAHaystack\n\n[2] DuoAttention: Efficient Long-Context LLM Inference with Retrieval and Streaming Heads\n\n[3] Transformers are Multi-State RNNs", "was_revised": true, "final_justification": "This paper is well-motivated and provides solid experiments. All my concerns have been addressed; therefore, I will maintain my positive score.", "rebuttal": "> R3-1: The experiments are limited to only two benchmarks: LongBenchV2 and GSM-Infinite. The evaluation lacks diversity in tasks and datasets. The comparison is also restricted, focusing solely on Squeezed Attention and QUEST. It would strengthen the paper to include additional benchmarks, such as Needle-in-a-Haystack [1], and comparisons with more recent methods like DuoAttention [2] and TOVA [3].\n\nWe appreciate the reviewer’s feedback. In our submission, we evaluated on a range of tasks requiring complex reasoning capabilities; LongBenchV2 includes questions from a range of downstream tasks, including single and multi-document QA, in-context learning, and code repository understanding, and GSM-Infinite contains arithmetic reasoning questions. \n\nTo bolster the evaluation from our paper, we have conducted additional evaluation on the LongBenchV1 benchmark suite and passkey retrieval task to demonstrate the generalizability of our method. We have included additional comparisons with DuoAttention and TOVA below, which we will include in the final version of our paper. We compare accuracy on LongBenchV1 for the Llama-3.1-8B-Instruct model using the evaluation setup from the open-source DuoAttention code which simulates decoding for the final 50 prompt tokens. We configure Multipole Attention to use 95% sparsity and 1 centroid per 16 tokens, and then adjust the other methods to have equivalent KV cache budgets (around 11.25% of the full KV cache). Our results demonstrate how our method substantially outperforms DuoAttention and TOVA, and also how our method generalizes to non-reasoning models and other long context tasks. Note that the TOVA algorithm allows each query head to select different keys and values and therefore has a larger KV footprint for GQA models even for the same sparsity settings.\n\n| Method                     | Baseline Accuracy | DuoAttention | TOVA | Squeezed Attention | **Multipole Attention** |\n| -------------------------- | ----------------- | ------------ | ---- | ------------------ | ------------------- |\n| Average LongBench Accuracy | 47.8              | 26.2         | 34.3 | 41.0               | **43.7**                |\n\n\n*Average LongBenchV1 accuracy (Llama-3.1-8B-Instruct, evaluated on English tasks) for DuoAttention, TOVA, Squeezed Attention, and Multipole Attention.*\n\nAdditionally, we have provided needle-in-a-haystack evaluation for our method using the evaluation code from DuoAttention (sweeping across different context lengths and depths). We ran evaluation for the baseline as well as Squeezed Attention and Multipole Attention with the Llama-3.1-8B-Instruct model using 128 token budget. We found that the accuracy for Multipole Attention / Squeezed Attention were equal / nearly equal to the baseline accuracy, even with aggressive KV cache compression. Note that while Multipole Attention retains the accuracy of the baseline model, the benefits relative to Squeezed Attention (and other prior sparse attention methods) are more pronounced for tasks which require more broad contextual information from across the input context and previously generated tokens (e.g. complex reasoning tasks). We will include the full visualizations across different context lengths and depths in the final version of our paper.\n\n| Method              | Overall Accuracy |\n| ------------------- | ---------------- |\n| Baseline            | 0.912            |\n| Squeezed Attention  | 0.911            |\n| **Multipole Attention** | **0.912**            |\n\n*Passkey retrieval accuracy for Llama-3.1-8B-Instruct model (token budget 128).*\n\n\n> R3-3: How does the method perform across different context lengths? An ablation study or detailed analysis would be helpful.\n\nWe show results in Figure 4 which outlines how Multipole Attention performs across different context length splits for document processing tasks in LongBenchV2 (short/medium/long, i.e. <32K/32K-128K/128K+ context lengths)). Looking at the configuration with token budget of 512, we observe comparable performance to the baseline across all 3 context length splits. Additionally, Table 2 highlights the performance of Multipole Attention for both 8K and 16K context length inputs with arithmetic computation tasks (GSM-Infinite). We find that Multipole Attention provides consistent accuracy advantages relative to Squeezed Attention across different context lengths across both document processing tasks (LongBenchV2) as well as arithmetic reasoning tasks (GSM-Infinite), and retains close to the baseline accuracy even with aggressive KV cache compression budget."}, {"review_id": "OGWUuYtaNI", "reviewer": "Reviewer_TWnq", "summary": "This paper proposes Multipole Attention, a method to speed up long-chain reasoning in Large Reasoning Models (LRMs) by computing exact attention only for key tokens and approximating the rest. It clusters semantically similar key vectors and uses their centroids to guide attention, balancing efficiency and accuracy. A fast cluster update mechanism enables efficient handling of long sequences during generation. Evaluations on models like Qwen-8B and Deepseek-R1 show up to 4.5× speedup with minimal accuracy loss on complex reasoning tasks.", "questions": "N/A", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. The motivation of this paper is apparent and addresses a practical need, especially given the growing use of LRMs in scenarios where they need to generate very long outputs. The proposed method has high application value in such contexts.\n2. This paper introduces the idea of applying clustering techniques to the attention mechanism, which is a relatively novel and innovative approach.\n3. The experiments in this paper are comprehensive and effectively support the paper’s main claims.\n\nWeaknesses:\n1.  The paper is not well-written; it contains a large amount of textual description but lacks formal mathematical formulations, which makes it difficult to follow and understand.\n2. In line 179 of the \"Efficient Online Clustering\" section, the term \"BS Tokens\" is used without any prior explanation or background, which makes the paper less reader-friendly and harder to follow.\n3. As shown in Figure 4, the proposed method shows the largest performance gap compared to the baseline on medium-length sequences. The authors attribute this to the issue of endless repetition. However, they do not provide further analysis or validate this explanation on datasets that do not suffer from repetition issues.", "was_revised": true, "final_justification": "Recommended score: 4\n\n Justification: Considering the novelty of the paper and the sufficiency of the experiments, I believe the minor writing issues do not undermine the strength of the work. Therefore, I recommend a score of 4.", "rebuttal": "> R4-1: The paper is not well-written; it contains a large amount of textual description but lacks formal mathematical formulations, which makes it difficult to follow and understand.\n\nWe appreciate the feedback from the reviewer. For the final version of the paper, we will ensure that we include relevant mathematical formulations in addition to textual descriptions in order to ensure the methodology is described clearly.\n\n\n> R4-2: In line 179 of the \"Efficient Online Clustering\" section, the term \"BS Tokens\" is used without any prior explanation or background, which makes the paper less reader-friendly and harder to follow.\n\nWe appreciate the reviewer’s feedback with respect to the writing. Line 179 (`We perform clustering in blocks of BS tokens`) was intended to define BS as a variable corresponding to the number of tokens in each block. We will rework this sentence to enhance clarity.\n\n\n> R4-3: As shown in Figure 4, the proposed method shows the largest performance gap compared to the baseline on medium-length sequences. The authors attribute this to the issue of endless repetition. However, they do not provide further analysis or validate this explanation on datasets that do not suffer from repetition issues.\n\nWe would like to kindly note that we did not attribute the performance gap to endless repetitions. The repetition issue was observed with the Qwen3-8B model when running on medium / long context length inputs for the LongBenchV2 dataset. This is a known issue with the model when not using recommended decoding settings (as mentioned in the model page [1]); however, for long context inputs, we observed this even when using recommended decoding settings. As such, we did not include Qwen3-8B evaluation for medium/long context length splits in the paper.\n\n[1] https://huggingface.co/Qwen/Qwen3-8B"}]

Author rebuttal:

Meta-review: This paper proposes a method to reduce KV cache requirements during long-context inference by performing exact attention only for tokens identified as important, while using approximate cluster representations for the remaining tokens. 

Strengths:

1. The paper introduces a novel sparse attention mechanism, which has high application value in LRMs for long-contexts.
2. The algorithmic design and its description appear sound, being efficient and scalable (i.e., fast online clustering and a custom triton kernel implementation).
3. The experiments in this paper are comprehensive and effectively support the paper’s main claims.

Weakness:

1. generalization outside of reasoning tasks and some hyperparameters analysis missing (` zAWi `)
2. lack of detailed comparison with baseline Squeezed Attention (`zdAu`)
3. more benchmarks and baselines (`vJGV`)

All weakness have been addressed during rebuttal, and all reviewers consider this is a good paper. Please do revise it accordingly.
Decision: Accept (poster)