A Training-Free Sub-quadratic Cost Transformer Model Serving Framework with Hierarchically Pruned Attention

W3. The authors have already shown improvements for certain sequence lengths, while it is expected to thoroughly explore how HiP scales to extremely long sequences or very large models. The current configuration on the sequence length (128k) and the model size (8B) are relatively small for a paper studies on the efficient decoding. Alternatively speaking, the scalability of this method deserves more experiments to support.

This is a misunderstanding; we have already performed extremely long (512k) latency reports, as shown in Table 6. We already showed the long context extendability in Table 10. To sum up, with our method, every LLM can extend their context window without being bottlenecked by memory and latency.

Our method not only gains speedup in extremely long contexts, but also in shorter contexts such as 32k and 128k which are used in LongBench and other benchmarks. In Tables 1, 2, and 3, we already use appropriate context length to measure attention speedup (Longbench: T=32k, Passkey, RULER: T=128k). Therefore, we could show the speedup with the tested dataset context length. We think the 128k context window size of Llama3.1 is quite long enough to show the speedup of HiP.

Moreover, in the KV cache offloading experiment of Table 6, we try to extend the context length up to a 512k context window. Therefore, we believe we have already shown the effectiveness of the HiP framework in the long context in terms of latency and scalability.

Regarding performance scalability, we have already verified that the original LLM's long-context ability was recovered using RULER. Therefore, we can expect HiP to be scalable in an extremely long context model in the future, too (if you scale the cost budget appropriately). We hope this resolves the concerns about HiP's performance scalability.

To show the context extendability of HiP, we used Gemma 2 in Appendix Section C.4 by extending its pretrained context window 16 times longer using HiP with SelfExtend. This shows that we think our method is also scalable when we extend the context window extremely.

Table 1. Updated Latency Results.

Table 2. Passkey Result (Llama 3.1 8B, Pretrained on 128k)

Table 3. InfiniteBench Result (Llama 3.1 8B, Pretrained on 128k).

Table 4. InfiniteBench Result (Exaone 3 7.8B, Pretrained on 4k).

To further support context extendability, in our latest research after submission, we aim to enhance the ability of HiP to extend the context while maintaining the overall framework and latency, as shown in Table 1. We increased flexibility in our framework, so now the users freely control each iteration (how many chunks should be selected, how to break down each chunk, and how to select the representative tokens). We could increase long-context performance by searching appropriate iteration configurations. However, since the latest development is still under construction, we leave this for future work (of course, we will go to open source everything, so you can check our source code in real-time if you are wondering how we built it). We could achieve 99.0 percent passkey accuracy result in 512k context length with Llama3.1 8B (Table 2) and even outperform FA2 in InfiniteBench 384k (Table 3). As shown in the table (Table 4), we could increase performance improvement by increasing context length more than 32 times than the pre-trained.

Therefore, we think a hierarchically pruned attention framework is a promising research direction and has great value to share with the research community to stimulate the following creative studies by community members who are trying to build extremely long context models using various approaches.

Thank you once again for taking the time to review our work and provide insightful feedback.

We hope to address and resolve your concerns during the discussion period. If our rebuttal successfully clarifies your questions, we kindly ask you to consider updating your score accordingly.

Thank you for taking the time to review our work and for providing insightful and supportive feedback. We have addressed your concerns and questions below.

Weaknesses:

W1. The effectiveness of HiP is contingent upon the presence of "attention localities," which might vary across different LLM architectures or tasks. The robustness of this operation deserves more discussions.

We tested attention locality in various tasks and pretrained models. Due to space layout, we omitted the empirical results from various models, but we updated the revision to have them now, so please refer to Figure 11-16 in Appendix A. The updated figures in Appendix A now include empirical data showing that attention locality is prevalent across various LLM architectures. As suggested, we revised the section more, so please give us more comments if you have any concerns.

W2. While HiP reduces latency in the decoding phase, the iterative pruning process might introduce overhead in the initial stages. How does the proposed method balance these two parts.

This is a misunderstanding. Our method does not cause latency degradation in both the prefill (initial) stage and the decoding stage. We always report the prefill and decode attention latencies (including masking and sparse attention) in Tables 1, 2, and 3 in the bottom row. The sorting overhead is negligible compared to Flash Attention, as shown in Figure 8. As we reported in Figure 7, our method outperforms FlashAttention2 in terms of prefilling and decoding attention latency.

Please be aware our method provides latency speedup on every phase of the Transformer model in a plug-and-play manner without any architecture modification or further training.

Continue to the following comment.

Thanks for the response by the authors, which makes the performance reported by this paper more convincing. Given the original score I gave, I decide to keep my score.

Q3. How well does HiP integrate with other efficiency techniques, such as quantization or model pruning, to further improve serving efficiency?

Our method is completely orthogonal to many other efficiency techniques, as HiP is a simple drop-in replacement for the Flash Attention call in whatever framework the user is operating in, such as vLLM, SGlang, and Huggingface Transformers.

We have already provided vLLM and SGlang integration in the supplementary files that are built with interoperability with other efficiency techniques in mind, such as Paged Attention, model quantization (GPTQ, AWQ), speculative decoding, prompt caching, structured generation, and more. Furthermore, our method is compatible with KV cache compression or eviction methods as long as they are still calling the regular attention mechanism internally.

We are looking forward to contributing to the open source, especially vLLM, and SGlang, so our method can be easily integrated with future efficiency techniques that are compatible with vLLM and SGLang, which is the most widely used LLM inference framework.

Q4. Can HiP be effectively applied to other Transformer architectures beyond the specific LLM model used in the paper?

Yes, HiP can be effectively applied everywhere Flash Attention is used and attention locality exists. We already applied HiP to various Transformer models such as LLM (Llama3.1, Llama2, Gemma2) and LMM (Llava1.6, section C.1) in the manuscript and the appendices.

Thank you once again for taking the time to review our work and provide insightful feedback.

We hope to address and resolve your concerns during the discussion period. If our rebuttal successfully clarifies your questions, we kindly ask you to consider updating your score accordingly.

Thank you for taking the time to review our work and for providing insightful feedback. We have addressed your concerns and questions below.

Weaknesses:

W1. HiP's effectiveness relies on the assumption of attention locality. While this assumption generally holds, there might be cases where it's violated, potentially impacting performance.

Please see Appendix A for more detailed discussions about attention locality with more theoretical analysis and empirical results. We have already tested various tasks (please refer to Fig. 11) to show that attention locality exists in general for a wide range of tasks. Plus, we have added empirical results from additional models, namely Qwen2 7B and Exaone 3.0 7.8B, to the figures in Appendix A.3 in the revision. (Please check the updated Fig. 11-16).

W2. HiP enforces the same sparsity across all rows of the attention matrix. Exploring mechanisms for dynamic sparsity, where sparsity varies based on the input, could further enhance efficiency and performance.

In this study, we are focused on proposing and validating the HiP (Hierarchically Pruned attention) framework. However, as we discussed in Appendix E.9, we agree that dynamic sparsity is a promising approach for future research.

In fact, our implementation of HiP already has support for tweaking the amount of sparsity on-the-fly, so future researchers can easily utilize our implementation for their own dynamic sparsity version of HiP. The block sparse attention kernel receives the non-zero entries' indices and the number of non-zero entries (commonly referred to as nnz in many sparse matrix data structure implementations) for each row. By controlling the nnz, the user can control the cost of the block sparse attention very easily. However, in this study, we opt for the simplicity of design since dynamically determining the sparsity based on the input is out of the scope of this study.

W3. The implementation and optimization of HiP, particularly the KV cache offloading, are tailored for specific hardware platforms (e.g., RTX 4090). Performance and optimal configurations might vary across different hardware accelerators.

This is a factual misunderstanding -- We have not tailored our method to specific hardware such as the 4090. In Table 6, we benchmark our KV Offloading on both 4090 and A100. Also, in Appendix E.10, we compare HiP's latency on H100 and analyze its hardware resource demand patterns. As discussed there, HiP utilizes the GPU's thread resources more heavily instead of floating point arithmetic than the matmul-heavy FlashAttention2. However, this does not mean that our kernel only works efficiently in a specific HW.

We understand that a HW difference may lead to non-optimal performance with the default hyperparameters. To alleviate this problem, in our implementation, we have added a built-in autotuner of kernel hyperparameters (e.g., block sizes (not $b_k$, $b_q$, but the memory tile block size), number of warps, number of registers, etc.). Since we built an autotuner using grid search, we will add a pre-built library of our HiP kernels in future code releases.

W4. The paper doesn't specifically address potential LLM alignment issues that might arise from applying HiP. Further investigation is needed to ensure HiP's safety and robustness in practical deployments.

We believe that this is out of the scope of our research. Although we agree that LLM safety is also a serious problem, we are proposing an unrelated orthogonal approach, an efficient attention mechanism, in this work. Since our extensive empirical analysis shows we can recover most LLM behavior, we expect the HiP applied model also to be aligned with guidelines well like the original LLM.

Questions:

Q1. How does the choice of chunk sizes in HiP's hierarchical pruning affect the trade-off between accuracy and efficiency?

We will assume that the term "chunk size" you mentioned refers to the block sizes $b_k$ and $b_q$. In Appendix E.4, we perform an ablation study on block sizes in HiP and discuss the results. We find that the perplexity gets better as $b_q$ increases, whereas it gets worse as $b_k$ increases. Please refer to Appendix E.4 for more details.

Q2. What strategies can be employed to further optimize KV cache offloading, such as using different memory tiers or compression techniques?

Please see Appendix E.7 and E.10 for a discussion of the possibility of additional memory tiers as a future direction. To summarize the discussion in that section, we believe we can add additional cache offloading tiers to the two-level hierarchy between GPU VRAM and CPU DRAM in order to further relax the memory pressure under extremely long sequences. We are looking forward to extending our offloading system to contain various levels, such as NVMe local storage and network storage, in future work.

Regarding compression techniques, we will answer this in more detail in Q3.

Continue on the following comment.

I thank the authors for the explanations. The explanations given here are much better and I suggest the authors to make revisions in the main manuscript accordingly. I have raised my score.

Thank you for taking the time to review our work and for providing insightful and supportive feedback. We have addressed your concerns and questions below.

Weaknesses:

W1. It would be great to have the long context benchmarks also for different models -- gemma and mistral are both open source and around similar sizes.

Please refer to section C.4. You can find additional Transformer models such as Llama-2 and Llava1.6 in our appendix sections E.4, E.4, and C.1. We already have results with Gemma 2 with a training-free context extrapolation method to test the 128k context window, which is 16 times larger than its pretrained context window. Here is a summarized result of Gemma 2 from Appendix Table 10.

Table 1. Gemma 2 Results from Appendix Table 10

As for our model selection, we choose to use Llama3.1 8B to test our method in various tasks and baseline methods. Since our benchmark suite is quite huge (Passkey, Perplexity, RULER, LongBench, LMMS-Eval, MMLU), we could not perform all of the experiments on as many models as we would have liked.

In our experience, our method generally shows similar performance with Llama 3.1 in other well-known LLMs because the attention locality is prevalent across various LLM architectures, such as Qwen2 and Exaone3. We have updated the figures in Appendix A.3 to include empirical results across various LLM models, which shows that attention locality exists in other various models as well. (Please check the updated Figures 11-16).

Continue on to the next comment.

Questions:

Q1. Gemma 2 uses this mix between sliding window and normal attention. it would be great to understand if there are any degradation on non-llama architectures.

We also think Gemma 2 is quite an interesting architecture to test, so we include a performance evaluation on Gemma 2 in Appendix C.4. In that section, we tested the effect of sliding window and HiP attention up to 128k context window using HiP and SelfExtend. Our method showed good performance recovery (Please check Table 2 below); sometimes, HiP even improves upon the original performance due to longer context being available on context-extended Gemma 2 (Please check Appendix Table 10).

Table 2. Gemma 2 Sliding Window Ablation.

Moreover, we observe noticeable performance improvements with sliding window attentions in longer contexts. We think the sliding window layer helps stabilize the HiP-replaced layer. HiP specializes in retrieving long-range information, and the sliding window specializes in retrieving local information. Therefore, we think that there is some synergy effect in combination of two types of attention mechanisms.

Table 3. Updated Latency Results.

Table 4. InifniteBench with the latest HiP

Table 5. Passkey with the latest HiP

In our latest HiP research, after submission, we could engineer HiP further to enhance the ability to extend the context window while maintaining overall framework and latency, as shown in Table 3. We increased flexibility in our framework, so now the users freely control each iteration (how many chunks should be selected, how to break down each chunk, and how to select the representative tokens). We could increase long-context performance by searching appropriate iteration configurations. However, since the latest development is still under construction, we leave this for future work (of course, we will go to open source everything, so you can check our source code in real-time if you are wondering how we built it). With the latest version, we could extend the Gemma 2's context window from 8k to 960k on passkey evaluation (above Table 4).

Therefore, we think a hierarchically pruned attention framework is a promising research direction and has great value to share with the research community to stimulate the following creative studies by community members who are trying to build extremely long context models using various approaches.

Q2. The method does show that MMLU is not degraded when using this method. It would be interesting to see this for a broader set of metrics if possible.

Please specify a set of metrics you would like us to evaluate our method in, if possible. However, please note that we already have various sets of metrics that examine context utilization (Passkey, RULER), long-context reasoning performance (LongBench), and multi-modality functions of LMM (LMMS-eval). We hope our extensive large-scale empirical evaluation helps verify the performance of our work.

Thank you once again for taking the time to review our work and provide insightful feedback.

We hope to address and resolve your concerns during the discussion period. If our rebuttal successfully clarifies your questions, we kindly ask you to consider updating your score accordingly.

Thank you for taking the time to review our work and for providing insightful and supportive feedback. We have addressed your concerns and questions below.

Questions:

Q1. It would be easier to understand to have a figure illustration showing the tree search (or improve figure 2, the figure 2 step 1 is a bit confusing) for section 3.1.

Thank you for your advice. In this discussion period, we have added a detailed figure on tree search in the appendix (please check Figure 19: Detailed Flowchart of Hierarchically Pruned Attention.), referenced from Figure 2 in the main paper due to space limitations in the main paper. Please review the new figures and let us know if you have any further comments.

Q2. The algorithm divides the sequence into k segments, and one token will be selected in each segment. Is it correct? If so, then the top k tokens must be distributed in the sequence, and cannot be concentrated on certain regions. Why do you make this design choice?

We would like to respectfully clarify a factual misunderstanding. We select top-k tokens from the whole sequence rather than select top-1 from each segment (chunk) individually. We believe this may be a similar misunderstanding to Q1. Our method performs each iteration as follows (We omit block approximation here for clarity.):

1. (Initialization step) Uniformly split the sequence of length $T$ into $k$ chunks.
2. Split each chunk in half, resulting in $2k$ chunks.
3. For each chunk, select a representative key token (the center of the chunk) and compute the dot product between the query and key token. This dot product score serves as the importance score for each chunk.
4. Select the top-$k$ most important chunks and return to step 2. If the chunk size is equal to or smaller than $b_k$, terminate the process.

Since we globally select the top-$k$ chunks from the $2k$ chunks in each iteration (step 4), our method can concentrate the non-zero entries in specific, high-importance sections of the sequence. Please refer to the Appendix Figure 27 for the attention mask visualization on each iteration step.

Q3. I am aware of some literatures that also use iterative refinement to dynamically calculate attention for efficiency. Have the authors tried to compare to these literatures?

Can you provide specific references to the related literature? We can try to compare them with our method.

Thank you once again for taking the time to review our work and provide insightful feedback.

We hope to address and resolve your concerns during the discussion period. If our rebuttal successfully clarifies your questions, we kindly ask you to consider updating your score accordingly.