HShare: Fast LLM Decoding by Hierarchical Key-Value Sharing

Authors: Huaijin Wu, Lianqiang Li, Hantao Huang, Tu Yi, Jihang Zhang, Minghui Yu, Junchi Yan
Primary area: foundation or frontier models, including LLMs
Keywords: Large Language Model, Decode, Key-Value Sharing, Critical Token, Hierarchical

Abstract:
The frequent retrieval of Key-Value (KV) cache data has emerged as a significant factor contributing to the inefficiency of the inference process in large language models. Previous research has demonstrated that a small subset of critical KV cache tokens largely influences attention outcomes, leading to methods that either employ fixed sparsity patterns or dynamically select critical tokens based on the query. While dynamic sparse patterns have proven to be more effective, they introduce significant computational overhead, as critical tokens must be reselected for each self-attention computation. In this paper, we reveal substantial similarities in KV cache token criticality across neighboring queries, layers, and heads. Motivated by this insight, we propose HShare, a hierarchical KV sharing framework. HShare facilitates the sharing of critical KV cache token indices across layers, heads, and queries, which significantly reduces the computational overhead associated with query-aware dynamic token sparsity. In addition, we introduce a greedy algorithm that dynamically determines the optimal layer-level and head-level sharing configuration for the decoding phase. We evaluate the effectiveness and efficiency of HShare across various tasks using three models: LLaMA2-7b, LLaMA3-70b, and Mistral-7b. Experimental results demonstrate that HShare achieves competitive accuracy with different sharing ratios, while delivering up to an $8.6\times$ speedup in self-attention operations and a $2.7\times$ improvement in end-to-end throughput compared with FlashAttention2 and GPT-fast respectively. The source code is publicly available at ~\url{https://github.com/wuhuaijin/HShare}.

Decision: accept (poster)
Reviewer scores: 6, 6, 6, 8, 8

Reviews and discussion:
<thread>
creation date:
1729129223330

thread_type:
official_review

<thread_object>
creator: reviewer_bqih

<content>
summary:
This paper discusses HShare, a hierarchical framework that facilitates the sharing of critical KV cache token indices across layers, heads, and queries. From observed similarities in KV token criticality across different levels, HShare enables a more efficient sharing strategy. The paper also introduces a greedy algorithm to dynamically determine the optimal sharing configuration for the decoding phase. The experimental results demonstrate that HShare can maintain similar accuracy and improve the end-to-end throughput.

strengths:
1. The problem is very practical and of high interest to the community, especially with limited computation resources
2. A very comprehensive set of experiments to showcase the performance (latency and accuracy) on different datasets.
3. Well-organized paper with adequate discussions on other existing methods.

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

questions:
1. What would be your suspected reason for the performance gain by enabling KV sharing in Figure 6?
2. The sharing ratio can be more finely tuned. Do you think the sharing ratios would be dependent on the dataset and the task?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730389270922

thread_type:
official_review

<thread_object>
creator: reviewer_rqy3

<content>
summary:
HShare aims to reduce computation complexity of key value caches of transformer-based LLMs by sharing theses caches between attention heads within a layer, the layers themselves and queries within a batch. They identify critical tokens from three areas: initial/sink tokens, most recent window and significant tokens in the middle. For each of three dimensions (head, layer, query) they identify elements in these dimensions that are similar (based on a self-defined metric: overlap of cached key-value pairs) and only compute the respective key-value cache once, which is then shared between the respective elements. The resulting reduction in computation improves the latency and throughput for most configurations while maintaining the accuracy.

strengths:
* A big strength is the evaluation. There is a wide variety of datasets used with a reasonable number of models, which show some variance in model size and architecture. Four additional framework are chosen to serve as baseline for comparison. The evaluation not only focuses on accuracy, but on compute performance (latency, throughput) as well. The provided ablation study is also fair.
* The main idea is presented clearly for the most part and is reasonable motivated with examples/microbenchmarks.
* The reduction in computation complexity with a subsequent performance improvement seems like a relevant contribution.

weaknesses:
* 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.
* 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.
* some of the equations do not add much (eq. 4 and 5)
* The performance of the original models without modification is kind of missing as a baseline for the evaluation.

more specific issues:
* line 64: "Although these two methods only load critical tokens, they retain all the KV cache"
  * That seems like a contradiction.
* 2.1: recent efforts are from 2023 and 2024, but their downsides are discussed in a paper from 2022 - seems counterinituitive
* line 248: "Fig. 1(c) and Fig. 1(d)" - I think these are not the correct subfigures: it should be 1e and 1f. I think "Fig. 1(e) and Fig. 1(f)" in line 250 is similarly wrong.
* line 404: first time that the full name "DoubleSparse (DS)" was introduced; all previous mentions just refer to DS
* table 3: I question the validity of computing an average from different types of scores (similarity, F1 and rouge).
* I think the table and figure placement for the evaluation can be improved, i.e. moved closer to where they are discussed.

detailed/minor issues:
* language:
  * line 40/41: "since the context length expands" - probably better "since as the context length expands"
  * line 43: "to addressing this issue" - I think it should either be "to address this issue" or "for addressing this issue".
  * line 45: "many works manage to only load these critical KV cache tokens" - I think "manage" is the wrong word here, maybe better: "many works choose to only load these critical tokens into the KV cache"
  * line 135: "make efforts" - should probably be past tense: "made efforts"
  * line 136: "Jiang et al. (2023b) trains" - train (multiple authors as subject), or even better "trained"
  * line 186: "also the same operation for value matrix and form a new value" - missing verb and general improvement: "also apply the same operation to the value matrix to form a new value matrix ..."
  * line 265: "layers of attention blocks" - I think it should be "layers of the attention blocks"
* typos:
  * line 46: "tokens(also" - missing whitespace
  * line 85: "GSM8K Cobbe et al. (2021)" - missing brackets around the reference, probably the wrong cite command is used
  * line 120: "scenarios.Decoder-only" - missing whitespace
  * line 143: "MInference(Jiang et al., 2024)" - missing whitespace
  * line 182: "i-th head(layer) and the j-th head(layer)" - missing whitespace between "head(layer)", two times; alternatively you could also use "head/layer" or "head|layer" (also applies to other occurences of the same phrase)
  * line 186: "S}corresponding" - missing whitespace
  * line 210: "HSahre" -> "HShare"
  * line 363: "Sec.4.1" - missing whitespace
  * line 408: "We use LM-Eval framework" - missing "the" before "LM-Eval"
  * Figure 6: x-axis label for TrivialQA is "Sparsity Level", when "Sharing ratio" is used for all other subfigures
* references:
  * "Anze Xie Ying Sheng Lianmin Zheng Joseph E. Gonzalez Ion Stoica Xuezhe Ma Dacheng Li*, Rulin Shao* and Hao Zhang" - author list is broken: wrong order, missing commas between the authors and please also remove the asterisks
  * consider using title case for the titles

questions:
* line 49: What are "heavy tokens"?
* line 137: "However, such a compression strategy ... and also increases the overhead during inference." - In what regard? Wasn't the 
 compression of prompts applied to reduce certain types of overheads, for example memory complexity?
* Figure 2: On what query sequence is the figure based?
* 5.2.2 Where does the query data come from?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730443576814

thread_type:
official_review

<thread_object>
creator: reviewer_ddxb

<content>
summary:
The paper introduces HShare, a hierarchical key-value (KV) sharing framework to accelerate inference in large language models (LLMs). HShare exploits observed similarities in critical KV cache token usage across layers, heads, and adjacent queries to reduce the computational overhead of selecting query-specific tokens repeatedly. It introduces a method to share critical KV cache token indices within and across layers, heads, and queries and proposes a greedy algorithm to optimize layer- and head-level sharing configurations. Experiments demonstrate that HShare maintains accuracy while providing up to 8.6× speedup in self-attention and a 2.7× improvement in end-to-end throughput across various LLM architectures, including LLaMA and Mistral models.

strengths:
The hierarchical sharing of KV cache tokens across heads, layers, and queries is an innovative approach to reducing latency in self-attention. HShare's observation of token reuse across different levels of attention blocks offers a new angle to address computational efficiency in LLM decoding. HShare's ability to maintain high accuracy while delivering substantial improvements in self-attention latency and throughput makes it significant for practical deployment scenarios. The method's compatibility with several model types suggests its value for LLM inference optimization.

weaknesses:
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.

questions:
Can HShare effectively support transformer variants such as convolutional or graph neural networks? A brief discussion on potential adaptations would help generalize its usability.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730626424266

thread_type:
official_review

<thread_object>
creator: reviewer_zrxs

<content>
summary:
The authors conducted systematic experiments to analyze the similarity of KV cache selection among queries, heads, and layers. Capitalizing on these findings, they proposed a method that shares KV selection across the dimensions of query, head, and layer, thereby reducing the overhead of KV selection.

strengths:
The sharing-based method proposed in this paper seems novel.

weaknesses:
1.	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.

2.	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 entired 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.

3.	The authors provided only a speed comparison with FlashAttention. To better demonstrate the effectiveness of their method, comparisons with other query-aware dynamic token sparsity methods are necessary. 

4.	The authors mention using an approximate computation approach from Double Sparse ('Following the approach in Yang et al. (2024)') to calculate approximate attention weights in their implementation. As a result, it is unclear how much of the observed speed gain is attributable to HShare specifically versus the query-aware dynamic token sparsity method itself.

questions:
Please refer to weakness points 1-4
</content>
</thread_object>

</thread>
<thread>
creation date:
1730634619631

thread_type:
official_review

<thread_object>
creator: reviewer_mexh

<content>
summary:
This paper explores the similarities in the KV cache tokens across neighboring queries, layers, and heads. Building on this finding, the paper introduces an advanced solution named HShare, which leverage a hierarchical framework for KV cache sharing. HShare enables the sharing of critical KV cache token indices across layers, heads, and queries, substantially reducing the computational load associated with query-aware dynamic token sparsity. Additionally, we propose a greedy algorithm that dynamically selects the optimal layer-level and head-level sharing configurations during the decoding phase. Emprical study was conducted to assess HShare's effectiveness and efficiency on various tasks using three models—LLaMA2-7b, LLaMA3-70b, and Mistral-7b. Experimental results reveal that HShare preserves accuracy while enhancing system efficiency.

strengths:
- S1. This paper addresses a crucial and cutting-edge problem in the field, tackling the computational challenges associated with managing KV cache in long-context LLMs. 

- S2.  The techique solution is well-desgined. By optimizing the reuse and sharing of critical KV cache token indices across queries, layers, and heads, the proposed solution significantly reduces computational overhead without compromising accuracy.

weaknesses:
- W1. The main concern I have is about the system-efficiency evaluation in Section 5.2. Concretely, I think the baseline is a little weak; we know FlashAttention could be the state-of-the-art implementation, but it is not equipped with advanced algorithm design for selective KV-cache usage. As motivated in the introduction, other methods also attempt to reduce the computation budget based on sparsity; 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. 

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

questions:
Please address the corresponding problem listed in the Weaknesses Section.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732114522827

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response by Authors

comment:
We express our gratitude to all the reviewers for dedicating their time and providing valuable comments. They acknowledged that our work is meaningful (mExh,Bqih), well-designed (mExh,RqY3,Bqih), contributive (mExh,DdXB,Bqih), effective (RqY3,Bqih), and presents a novel approach (ZRXS, DdXB). While the overall feedback from the reviewers is positive, reviewer ZRXS has some concerns regarding the evaluation of LongBench and the efficiency comparison with other baselines. In the following response, we provide detailed answers to all the questions and comments point-by-point. In particular, we have included results from 16 commonly used English datasets from LongBench, as well as system efficiency comparisons with other methods, in the rebuttal. Additionally, we have revised the manuscript based on the reviewers' suggestions, with all revisions and additions clearly marked in blue.

We deeply appreciate the suggestions to improve this paper. If you have any further questions, please let us know so that we can provide a timely follow-up response.
</content>
</thread_object>

</thread>


Metareview:
metareview:
Summary:
This paper introduces HShare, a hierarchical framework for sharing KV cache token indices across layers, heads, and queries in LLM inference. The authors exploit observed similarities in critical token usage patterns to reduce computational overhead while maintaining accuracy, proposing a greedy algorithm for determining optimal sharing configurations that demonstrates significant speedups across multiple models and tasks.

Main Strengths:
The technical merit of this work is evident in its hierarchical sharing approach, which is built on theoretical foundations and employs an effective greedy configuration algorithm. The empirical validation is sufficient, featuring evaluation across multiple models and tasks, demonstrating good accuracy maintenance alongside significant speedup, and including ablation studies and comparisons. The practical impact of this work is significant as it addresses an important efficiency challenge in LLM inference, demonstrates compatibility with multiple architectures, and shows clear potential for real-world deployment.

Main Weaknesses:
The initial evaluation had several limitations, including restricted dataset coverage in LongBench evaluation, missing baselines in efficiency comparisons, and inconsistency in the prefill stage sparsity approach. The presentation suffered from clarity issues, with some unclear technical descriptions, language problems in several sections, and missing details on hierarchical integration. There were also practical concerns regarding limited discussion of sharing ratio selection, inadequately explored trade-offs, and questions about broader applicability.