HShare: Fast LLM Decoding by Hierarchical Key-Value Sharing

Thank you for acknowledging our work. Your rating has been a great source of encouragement for us. Below is our detailed response.

W1: The paper should have some figures that discuss the tradeoff between accuracy and efficiency.

Thank you for the valuable suggestion. Here we provide the attention latency and GSM8K accuracy of our proposed HShare under different sharing ratios in Table 1 to reveal the tradeoff between accuracy and efficiency within our method HShare. We have also included a figure to present the average score on Longbench and end-to-end throughput of various methods, highlighting the trade-off between accuracy and efficiency across all methods. Please refer to Figure 7 in the revised paper.

In addition, we provide system efficiency comparisons with other baselines, including two KV cache eviction algorithms (StreamingLLM and H2O) and two query-aware token sparsity algorithms (Quest and DS). From the results, it can be observed that KV cache eviction algorithms have a clear advantage in system efficiency compared to query-aware token sparsity algorithms. Specifically, StreamingLLM achieves the fastest speed as it applies a fixed sparsity pattern. However, both StreamingLLM and H2O fall short in terms of accuracy. In contrast, HShare not only preserves accuracy but also achieves significant speedup compared to other query-aware dynamic token sparsity methods. For example, in terms of attention latency when batch size is 16 and sequence length is 2k, our method delivers up to a 2.80x speedup over Quest and a 2.21x speedup over DS. We have also added these tables into our revised manuscript (Appendix).

Tabel 1 Attention latency ($\downarrow$) and GSM8K accuracy of HShare under different sharing ratios.

Table 2 Attention latency ($\downarrow$) of different methods across various batch sizes and sequence lengths.

Table 3 End-to-end throughput ($\uparrow$) of different methods across various batch sizes and sequence lengths.

Q1: What would be your suspected reason for the performance gain by enabling KV sharing in Figure 6?

Thank you for your question. We have also observed that a certain degree of sharing may improve performance on some datasets, especially when we only perform head-level sharing. We suspect that this may be due to the model architecture and the sharing of critical tokens plays a role in regularization. Specifically, LLaMA2-7b-chat uses multi-head attention with 32 heads, some of these heads show high similarity in certain layers. When the sharing rate is relatively small, the regularization would improve performance to a certain extent.

Thank you for acknowledging our work. Your rating has been a great source of encouragement for us. We truly appreciate the time and effort you've dedicated to reviewing our work, as well as your thoughtful and incredibly detailed feedback.

W1: The language is a big issue. I had to read several sentences multiple times to make sense of them. So I suggest to improve the writing, so that language reads well.

Based on your valuable comments, we have revised the language to improve clarity and readability, with the hope that it will now be easier to follow. Your suggestions are very helpful in refining the manuscript, and we are grateful for your contribution to improving its quality.

W2: The cache hierarchy is not discussed, i.e. the order in which the three dimensions are applied, which data they might share and so on. The individual levels are discussed in how they are constructed and their details, but not how they are applied together apart from that they are independent from each other (in their memory benefits), which contradicts the initial hierarchy claim.

Thank you for your suggestions. Here, we provide an illustration of the hierarchical sharing architecture (https://postimg.cc/w1v0qpnm). In details, HShare first evaluates whether the current query should share critical token indices with those in the previous query. If this condition is met, there’s no need to proceed with layer-level and head-level sharing. If not, HShare then checks whether the current layer should share critical token indices with those in its previous layer. If layer-level sharing is applied, head-level sharing will not be performed. Only when query sharing and layer sharing are not performed, the head-level sharing will be considered according to the calculated sharing config.

We have reported the results of the three levels working together as well as their individual effects in our paper. However, they can also be combined in any combination. Here, we report the results of pairwise combination sharing in Table 1, with the sharing ratio for each level set to 1/2.

Table 1 The results of pairwise combination sharing.

W3: some of the equations do not add much (eq. 4 and 5)

Thanks, we have revised the corresponding part in the new manuscript.

W4: The performance of the original models without modification is kind of missing as a baseline for the evaluation.

Thank you for your valuable suggestion. We appreciate your insight and follow your recommendation, we have already added the performance of the unmodified models as a baseline for the evaluation in the revised manuscript. We hope this could help provide a clearer context for the results.

Specific Issues:

line 64: "Although these two methods only load critical tokens, they retain all the KV cache" That seems like a contradiction.

These two methods retain all the KV cache, but they only select critical tokens to participate in the attention output calculation. Since these methods use query-aware token sparsity, meaning that the critical tokens selected change each time, they need to retain all the KV cache instead of performing KV cache eviction.

Table 3: I question the validity of computing an average from different types of scores (similarity, F1 and rouge).

Averaging the scores of different metrics doesn't have intrinsic meaning, but we do it because we believe it can somewhat reflect the overall performance of the algorithm across different datasets. And here we humbly point out that some other works[1,2] also report the average score on Longbench.

[1] InfLLM: Training-Free Long-Context Extrapolation for LLMs with an Efficient Context Memory. NeurIPS, 2024.

[2] Palu: Compressing KV-Cache with Low-Rank Projection. Arxiv, 2024.

Thank you for acknowledging our work and your suggestions are very helpful. Below is our detailed response.

W1: While HShare achieves substantial speedup, the layer- and head-level sharing configurations are tested only on a few model architectures, leaving broader applicability unexplored. Additionally, its performance in tasks that demand highly dynamic attention, such as document summarization with varying context lengths, could be tested to reveal any potential trade-offs. Minor computational overhead remains from online calculation in the prefill phase, though this appears manageable.

Thank you for the valuable suggestions. We have discussed potential adaptations of our method to other model architectures. For more details, please refer to the answer to Q1.

To test our method in document summarization with varying context lengths, we conduct experiments on the MultiNews dataset, which belongs to the document summarization category. We evaluate the proposed HShare across different context lengths and sharing ratios. The results are presented in Table 1. It should be noted that regardless of the length of the text, we consistently selected 256 critical tokens. The results show that the score decreases as the context length increases, and similarly, higher degrees of sharing also lead to a decline in performance. This suggests that when dealing with long context lengths, a more moderate sharing strategy is needed to maintain accuracy, whereas for shorter texts, a higher degree of sharing can be applied. We have added the above description and experiment into the revised manuscript (Appendix).

Table 1. Performance of HShare on MultiNews.

Q1: Can HShare effectively support transformer variants such as convolutional or graph neural networks? A brief discussion on potential adaptations would help generalize its usability.

Yes, we believe HShare can effectively support transformer variants. HShare aims to optimize long sequence attention operations by sharing critical key and value indices between layers, heads, and queries. As such, convolutional or graph neural networks with attention can also benefit from the proposed HShare. Below are two examples of potential applications:

CNN with attention: [1] proposes to adopt CNNs to extract image features and use spatial-temporal attention to identify crucial frames. To apply HShare in this network, we can predict the indices of critical features (similar to critical tokens), which can then be shared across heads and layers. The potential adaptation involves using only the critical features for attention computation. Since video input exhibits temporal redundancy across nearby frames, critical feature indices can also be shared across frames. HShare can further be applied to other CNN-related works, such as VQA [2].

Graph neural network with attention: [3] adopts multi-head attention to extract features from a set of input nodes. HShare can be seamlessly applied to identify critical nodes and share their indices across heads and layers (if multiple layers are used). Depending on the problem, if the graph nodes are close and similar, the key indices of these nodes can also be shared across nodes, similar to query sharing in HShare. Similar works [4, 5] can also benefit from applying HShare to reduce computational load.

The discussion will be added to the revised manuscript.

References

[1] CSTA: CNN-based Spatiotemporal Attention for Video Summarization. CVPR, 2024.

[2] Bottom-up and top-down attention for image captioning and visual question answering. CVPR, 2018.

[3] Graph attention networks, ICLR, 2018.

[4] Attention-based graph neural network for semi-supervised learning. arXiv,2018.

[5] Representing long-range context for graph neural networks with global attention. NeurIPS, 2021.

We hope this response could answer your questions and address your concerns, looking forward to receiving your further feedback soon.

Thank you for acknowledging our work. Your rating has provided us with great encouragement, and your detailed suggestions are really helpful. Below is our response.

W1. I think the author should report the corresponding system efficiency with the same set of baselines in Section 5.1 to fully reveal the trade-offs introduced in Table 1.

Table 1 Attention latency ($\downarrow$) of different methods across various batch sizes and sequence lengths.

Table 2 End-to-end throughput ($\uparrow$) of different methods across various batch sizes and sequence lengths.

Thank you for the constructive advice. Here, we provide system efficiency comparisons as the same settings in Section 5.1 (including two KV cache eviction algorithms StreamingLLM and H2O, two query-aware token sparsity algorithms Quest and DS). The results are shown in Table 1 (attention latency) and Table 2 (end-to-end throughput).

From the results, it can be observed that KV cache eviction algorithms have a clear advantage in system efficiency compared to query-aware token sparsity algorithms. Specifically, StreamingLLM achieves the fastest speed as it applies a fixed sparsity pattern. However, both StreamingLLM and H2O fall short in terms of accuracy. In contrast, HShare not only preserves accuracy but also achieves significant speedup compared to other query-aware dynamic token sparsity methods. For example, in terms of attention latency when batch size is 16 and sequence length is 2k, it delivers up to a 2.80x speedup over Quest and a 2.21x speedup over DS.

We have added the above description and experiment to the revised manuscript (Appendix).

W2. There are some trivial presentation issues. For example, on page 6, line 270, there is no space between "across" and "layers".

Thanks. We have taken note of the issues you mentioned and have made the necessary revisions to the wording in the article. Specifically, the spacing issue between "across" and "layers" on line 270 of page 6 has been corrected. We really appreciate your careful review.

We hope this response could answer your questions and address your concerns, looking forward to receiving your further feedback soon.

Thanks for your review and feedback, below is our detailed response.

W1: The experiments only selected 6 out of the 16 commonly used English datasets from Longbench for evaluation, which may not be sufficient to draw solid conclusions. For example, the performance of H2O and Quest on these 6 datasets in Table 3 appears similar, whereas in the Quest paper, H2O showed significantly lower accuracy. This discrepancy may suggest that specific methods perform better on these six datasets. To clarify, the authors are encouraged to provide results across all datasets.

Table 1 Evaluation of different methods on sixteen datasets in Longbench.

Thanks for the suggestions. The results across the 16 commonly used English datasets from Longbench are shown in Table 1. It can be seen that in the 7/16 dataset, HShare achieves the highest score. On average, HShare outperforms StreamingLLM, H2O, and Quest. Although HShare performs slightly worse than DS, this is a negligible performance decline considering the speed advantage of HShare over DS.

Here we humbly point out that the original H2O paper only applies token sparse attention for the decoding stage, with dense attention used in the prefill stage. In our previous manuscript, we did not take this into account, which may have led to some unfair comparisons, here we apologize for any potential misunderstandings. In the revised manuscript, for all methods, we apply token-sparse attention to all layers across both the decoding and prefill stages to ensure a fair comparison.

W2: In terms of speedup, the authors analyzed the theoretical complexity of KV cache selection in Table 2. However, this analysis alone is insufficient to demonstrate that the proposed method significantly reduces computational overhead, as the selection cost may only account for a small portion of the entire decoding process. To strengthen this claim, the authors could provide direct evidence of decoding speed in experiments. Additionally, an experimental study on the trade-off between accuracy and speed would be valuable.

Tabel 2 Attention latency ($\downarrow$) and GSM8K accuracy of HShare under different sharing ratios.

Thank you for the valuable suggestions. Here we provide the attention latency and GSM8K accuracy of our proposed HShare under different sharing ratios in Table 2 to reveal the tradeoff between accuracy and efficiency in our method.

As you mentioned, the theoretical complexity of KV cache selection is insufficient, and a comparison of decoding speed is needed. We have provided system efficiency comparisons with other methods, please refer to the answer to W3 for further details.

Dear Reviewer ZRXS,

Thank you for actively participating in the discussions during the rebuttal period and for your constructive feedback. We have made revisions to our manuscript based on your suggestions.

As the rebuttal phase is nearing its end, we look forward to hearing from you, should our modifications and clarifications in the revised manuscript meet your expectations.

We greatly appreciate your time and effort.

Best regards, Authors