ParallelComp: Parallel Long-Context Compressor for Length Extrapolation

Thank for your suggestions!

Q1：Key differences

• We focus on the memory-bound encountered during the implementation of length extrapolation.

• ParallelComp introduces a two-stage compression strategy: chunk-level kv cache eviction followed by intra-chunk compression. Attention calibration is applied to reduce the performance loss caused by compression. This enables both context window extension and improved inference throughput.

• APE consists of a shared prefix to reduce distributional differences, a low-temperature mechanism to sharpen attention, and a scaling factor to adjust for temperature-induced changes. It aims to align the attention patterns of parallel and sequential encoding.

• StarAttention targets models with native long-context support. It recalculates attention for each chunk at every step, while our method computes attention only once at the prefill stage and reuses the compressed KV cache for generation—greatly improving efficiency.

• We conducted a comparative experiment, with the code implementation based on [1,2]. However, APE lacks prompt design, Star is missing the Longbench experiment, and hyperparameter configurations are absent in both. For APE, we set the temperature to 0.5 and the scaling factor to 0.8. For Star, the chunk size was set to 2K. In our approach, we reuse these 6K position encodings for length extrapolation. Base model is Llama-3.1-8B-Instruct.

Q2: Supplementary baseline

• The KV cache management mechanism of Gemma2 is significantly different from that of LLaMA and Qwen models, so we have not yet fully implemented it. If we have more time, we promise to submit the test results of EXAONE 3.5 and Gemma2. We promise to provide it in a future version of the paper.

• We conduct experiments on Llama3.1 and Qwen2.5 using Longbench and InfiniteBench under the same 24K KV cache budget. While other models use 24K position encodings directly, our approach leverages extrapolation techniques to reuse 6K position encoding. Others use the original length.In the experiment, it is found that whether the prompt word(such as “<|begin_of_text|>” for Llama3.1) is used has a great influence on the performance, and we make a comparison in the case of the prompt word. It can be seen that, under the same 24K KV cache size budget, our approach still outperforms other models that do not implement extrapolation, especially on the InfiniteBench dataset.

Q3: Necessity to sparsify attention mechanisms

Based on our observations, certain tokens at specific layers do indeed contribute to improving the model's performance:

• In early layers, attention biases are crucial for context understanding—removing them harms performance.
• For code tasks, biases in middle layers matter most—removal also leads to performance drops.
• The attention sink boosts retrieval in early layers but hurts performance in later layers.

However, in other layers, these attention biases significantly hinder the model's ability to process long contexts. Therefore, removing these anomalous tokens can improve the model's performance.

Q4: Others

• We will consider your suggestions and incorporate additional supplementary materials into future versions of the paper, especially the release of the open-source code.

Reference:

[1]https://github.com/Infini-AI-Lab/APE

[2]https://github.com/NVIDIA/Star-Attention

We would like to express our sincere gratitude for your fruitful suggestions, and for all your questions below!

Q1: The authors conduct a thorough comparison against existing length extrapolation methods for tasks such as single/multi-document QA. However, these tasks could potentially benefit from RAG without requiring long-context extrapolation in LLMs. Would it be possible to compare the proposed approach with leading RAG methods such as [1,2]?

• Since Llama3-ChatQA-2-8B produces "Not Available" results on InfiniteBench when using RAG (w/ RAG), we only compared the results of Llama3-ChatQA-2-8B in Table 2.
• The above shows the comparison results between our method and the representative RAG method, ChatQA 2 [2]. It is clear that our method performs well compared to the strong RAG baseline, especially on the Longbench dataset.
• The performance collapse of the RAG method on InfiniteBench is an unknown phenomenon, which also demonstrates the robustness of our method.

Q2: Does the paper provide a detailed description of the calibration data used? Have they analyzed stability across different context sources (e.g., code vs. narrative text)?

• We calibrated the attention scores based on Equation 7, using an online method to evict anomalous tokens, so there is no need to sample data in advance for calibration.
• We found that in code tasks, the attention distribution of the intermediate layers are more sensitive to calibration, while in text tasks, such as in-context learning, the attention distribution of the earlier layers are more sensitive to calibration.

Reference:
[1] RAPTOR: recursive abstractive processing for tree-organized retrieval.
[2] ChatQA 2: Bridging the Gap to Proprietary LLMs in Long Context and RAG Capabilities.

We would like to express our sincere gratitude for your fruitful suggestions, and for all your questions below!

Q1: Minimal Performance Difference Between "Ours" and Other Methods in Token Eviction Settings.

• "Ours" implements the eviction scheme for the entire chunk's KV cache.
• "Ours-compression" refers to further compressing the KV cache within the chunk.
• "Ours-calibration-compression" attempts to recover performance further through calibration.

"Ours" and "Ours-compression" refer to the methods of adding chunk eviction and parallel KV cache eviction, respectively. Generally, parallel KV cache eviction tends to degrade the performance of many models, so we hope to recover the performance lost during compression by using an attention distribution calibration mechanism, namely "Ours-calibration-compression."

Q2: Difference Between "Ours" and InfLLM: Global Attention with Full vs. Subset of Chunks.
We need to clarify your misunderstanding of our method.

• "Ours" only applies attention to chunks in the priority queue, unlike InfLLM, which retains all chunks in the CPU and causes high I/O overhead and slow prefill speeds by swapping the KV cache between the CPU and GPU. In fact, it can attend to all global chunks simultaneously.

• InfLLM addresses the memory-bound issue during inference by retaining only the chunk representations, while "Ours" uses a chunk eviction strategy and parallel compresses the KV cache to increase chunk process throughput. We mainly focus on the KV cache compression strategy in parallel chunk processing

Q3: Analyzing the Benefits of Calibration and Compression vs. "Ours".

We will start our analysis from a core motivation, which is the memory-bound issue in the parallel encoding and the attention bias issue aggravated by compression..

• “Ours“
  Due to GPU memory limits, it's not feasible to store KV caches for many chunks at once. Offloading chunks to the CPU during parallel processing also introduces high I/O latency, making it impractical. We propose a more general approach based on two principles: Simplicity: We use negative log-likelihood (NLL) to model the relevance between queries and chunks without designing extra representations. This reduces communication overhead—especially important in multi-GPU setups where I/O is the main bottleneck. Efficiency: Sorting in the priority queue only involves scalar comparisons with O(log n) time complexity, unlike heavier methods like InfLLM that use O(n²) retrieval operations.

• “Ours-compression“
  Parallel KV cache eviction helps address memory limitations in parallel computation, but our analysis shows it worsens attention bias. To mitigate this, we introduced a simple attention calibration mechanism “ours-compression-calibration” for near-lossless compression, making our design effective.

Q4: Typo of equation 6.

• In Eq. 6, [·] represents the eviction operation, which means removing the indices in [·] from K_x.

Q5: Clarifying the Link Between Figure 3 Patterns and the proposed method.

• Observation1: Based on our observations, due to using almost the entire window length for parallel encoding, the model always exhibits a strong recency bias in the query chunk region.

• Design1 : It indicates that the model relies on the tokens in the query chunk to make next token predictions, which inspired us to use the negative log-likelihood of the tokens in this region to measure the importance of the chunk.

• Observation2: Another phenomenon is that the KV cache compression operation exacerbates attention bias as shown in Figure 3, as shown in the first image of Figure 5.

• Design2 : Therefore, designing a calibration mechanism is necessary to mitigate the performance loss caused by exacerbated attention bias.

Q6: The performance discrepancy for Llama-3-8B-Instruct (8k) on LongBench may be due to different settings, while results for InfiniteBench align with InfLLM's Table 1.

• LongBench: For LongBench, we do not use the original hyperparameters reported in the InfLLM paper. Instead, we re-evaluated all baselines under a unified testing framework to ensure fair comparison. Specifically, we adopted the evaluation setup from the open-source repo[1], and ported that framework to the InfLLM codebase for consistency. This was done to eliminate inconsistencies in evaluation scripts, precision settings, and other hyperparameters. The inference details temperature=0, and num_beams=1, float16 precision.
• For InfLLM, we use the default hyperparameters as provided in their official code repository to evaluate on longbench.

Reference:
[1] https://github.com/Zefan-Cai/KVCache-Factory.

Thank for your suggestions!

Q1:For general inquiries such as typos and others.

• We promise to release our github and will correct typo error in future paper version.
• Regarding citation issues, the NTK-by-parts position interpolation method was not proposed in a formal paper. It was also cited from a GitHub repository in other papers, with reference [7] in the reference paper [1].

Q2:Fig 1, is that one selected attention head in one layer, or the average over attention heads and all layers, or what else is it exactly? And for what model? Eq 3, how is this self-information log prob defined? Is this the attention score A? Why do you write P? Sec 3.1 "concatenate them with the input query for parallel encoding" - what exactly does this mean? What is the input query for parallel encoding? Where do I see the input query in the formula?

• The attention head we show here is the distribution change of the 21st head of layer 1, and the model is meta-llama/Llama-2-7b-chat-hf.
• Self-information is not attention score; essentially, it is the negative log-likelihood of the question predicted based on the chunk input to the LLM. It reflects the model's confidence in predicting the query based on the context.
• We concatenate the input question(X^q in the formula) with each chunk and encode them in parallel. The query serves two purposes: first, to assess the importance of each chunk and decide if its KV should be stored in the GPU’s priority queue to manage memory; second, to evict tokens from the chunk’s KV cache, optimizing memory usage and throughput. In short, the query is primarily used for compression.

Q3:Sec 3.1 Chunk Eviction: What is this about? "retaining only the most relevant chunks" - what does this mean? how is this different to the KV Cache Eviction? What exactly is parallel about it? The previous chunk eviction is not parallel

• Chunk eviction means removing the entire KV cache of a chunk from a priority queue based on self-information scores, while KV Cache Eviction refers to removing tokens from a specific layer and head within a chunk.
• The token-level KV cache eviction process is parallel across chunks, as it depends on the cumulative attention distribution of the query across layers and heads without dependencies between chunks. In contrast, traditional KV cache eviction methods such as H2O [2] have no extrapolation capability and can only handle KV caches with a fixed context length.
• The eviction of chunks depends on whether the priority queue has reached its maximum size, usually determined by the GPU memory. When the size is reached, chunks completing self-information calculation and KV cache eviction compare their self-information size with those in the queue to decide retention or eviction. This comparison is typically done sequentially.

Q4:Why is Flash Attention relevant here? Flash attention is just a specific implementation.

• We want to emphasize that flash attention provides faster computation and memory utilization. FlashAttention does not return the complete attention distribution matrix for each layer of the model[3], thus making it unable to perform KV cache eviction. To address this issue, we calculate the inter-chunk attention matrix for each chunk and its corresponding query block at each layer separately, based on formula (5).

Q5:Since it is hard to obtain the complete attention distribution from Flash Attention" - what do you mean? Sec 3.1 Attention Calibration: "we evict tokens with excessively high attention scores": Evict means, you only use those values with high att scores, or you exclude those and only use the others? What exactly is "Ours" (without the calibration R_h and compression R_l)?

• According to [3], FlashAttention doesn't return the attention distribution, preventing us from evicting the KV cache based on attention scores. To solve this, we use the last 8 tokens from the query chunk and multiply it with the chunk's value matrix to quickly estimate a cumulative attention distribution for evicting tokens using Pytorch.
• Equation 7 represents that we evict tokens with abnormally high scores. At the same time, we also evict tokens with low attention scores according to Equation 6. In summary, the eviction strategy removes tokens with either abnormally high or the lowest attention scores.
• We calculate the full attention distribution for each head using flash-attention. To prevent abnormal patterns, we first approximate the attention distribution with Equation 5, allowing us to evict tokens that would gather excessive attention.
• 'Ours' represents the case without any KV cache compression or attention calibration, only with the chunk eviction mechanism.

Reference:
[1] YaRN: Efficient Context Window Extension of Large Language Models.
[2] H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models.
[3] https://github.com/Dao-AILab/flash-attention/issues/1357