TidalDecode: Fast and Accurate LLM Decoding with Position Persistent Sparse Attention

We thank the reviewer for their review and address their main concerns below.

W1. Spatial Coherence

We thank the reviewer for pointing out our use of spatial coherence and the need for a definition. We have modified accordingly in the fourth paragraph of the introduction section with a definition of spatial coherence that is highlighted in the underlines.

W2. Complexity of Implementation

For the GPU kernel implementation, we are building on top of the state-of-the-art attention kernel implementation library FlashInfer, which has been used in vLLM, SGLang, MLC-LLM, etc. Thus, we believe that our implementation can be easily integrated into popular LLM serving systems. In addition, our kernel implementation code has been anonymized and included in the supplementary materials for your reference.

W3. Hierarchical Attention Networks

We thank the reviewer for suggesting potential related works. We want to clarify that TidalDecode is a training-free method that can be directly applied to existing pretrained models like LLaMA-2 and LLaMA-3 during the inference phase, which is different from the hierarchical attention methods [1, 2] that requires pre-training.

W4. KV-cache correction

Even though we have introduced the mechanism of KV cache correction, we don't use the feature for all our current experiments to make it a fair comparison against other baseline models that haven’t applied cache correction. Additionally, for most of the tasks in our evaluations, the generation length is insufficient to see a degradation in model quality, so the model performance is still maintained even without KV cache correction. However, to show the effects of KV cache correction, we provide additional experimental results for the LongBench evaluation with KV cache correction on the original eight tasks and include results in Figure 9 in the revised paper. Albeit the improvement is not significant due to a short generation length, cache correction with a stride of 4 improves performance in five out of eight tasks, and a stride of 8 yields improvements in six out of eight tasks. These results highlight the efficacy of cache correction in mitigating the adverse effects of polluted KV representations, which stem from sparse attention mechanisms and can accumulate errors, ultimately degrading model performance.

W5. Training details

Please refer to the response to W3, as our method is a training-free method.

W6. How to choose the optimal token re-selection layer (L)

Our observations show that the optimal token re-selection layer choice is consistent across different task and decoding steps for models within the same model family. We hypothesize that this is an intrinsic property resulting from specific models' pretraining process. This phenomenon is an interesting observation but lacks theoretical analysis, which we leave as a future work. To sum up, we don’t have an automatic layer selection strategy at the current stage since the optimal choice is consistent within the same model family. We can thus directly identify the optimal choice from sample data only once and use it as a guideline for all later deployments. If the optimal layer choice within a model family does vary for different downstream tasks or even different decoding steps, there is certainly a need to design a flexible, adaptive layer selection strategy.

W7. Conflicts in the packages used in latex.

We have tried our best to resolve all the warnings in the compilation for the revised paper. Please let us know if the package conflicts still persist.

References:

[1] Chalkidis, Ilias, et al. "An exploration of hierarchical attention transformers for efficient long document classification." arXiv preprint arXiv:2210.05529 (2022).

[2] Yang, Zichao, et al. "Hierarchical attention networks for document classification." Proceedings of the 2016 conference of the North American chapter of the association for computational linguistics: human language technologies. 2016.

Thanks for the responses.

1. I'm confused by the response to W4. It doesn't make any sense to me that the authors introduced some methods, but didn't use them in the experiments. TBH the augments cannot convince me. It's just a useless method that would make the paper more "scientific" to me if the authors express such arguments: Even though we have introduced the mechanism of KV cache correction, we don't use the feature for all our current experiments to make it a fair comparison against other baseline models that haven’t applied cache correction.
2. Response to W6. As a summarization, the authors cannot provide any selection strategy, this will greatly reduce the real world application of the proposed method, as we can only try every L for optimal performance. Even on a small dev test, it would be very unreliable and time-consuming.

As a result, the author's responses did not address my concerns, I would like to keep the score unchanged.

We thank the reviewer for their review and address their main concerns below.

Q1/W1. Adaptive layer selection strategy

We agree with the reviewer that the model's performance is highly sensitive to the choice of the token re-selection layer. However, as you have suggested, our observations show that the optimal token re-selection layer choice is consistent across different task and decoding steps for models within the same model family. We hypothesize that this is an intrinsic property resulting from specific models' pretraining process. This phenomenon is an interesting observation but lacks theoretical analysis, which we leave as a future work. To sum up, we don’t have an adaptive layer selection strategy at the current stage since the optimal choice is consistent within the same model family. We can thus directly identify the optimal choice from sample data only once and use it for all later deployments. If the optimal layer choice within a model family does vary for different downstream tasks or even different decoding steps, there is certainly a need to design a flexible, adaptive layer selection strategy.

W2. Figure 2

We thank the reviewer for pointing out the confusion in our presentation for Figure 2. We have updated our descriptions and definitions for terms (e.g., recall rate) in the text and the caption, which have been highlighted in blue. More specifically, Figure 2a shows the overlap ratio of the top-k tokens selected between different layers, which is used to calculate the recall rate in Figure 2b. The recall rate is defined as the averaged overlap ratio in Figure 2a over the sparse attention layers between the third layer and the corresponding token re-selection layer, and over all the sparse layers after the chosen re-selection layer. For instance, the recall rate for choosing layer 13 as the token re-selection layer is the average of the overlap ratio over the two red bars in Figure 2a. Then, we can clearly see that adding a token re-selection layer around layer 13 can significantly boost the recall rate due to the improvement in the overlap ratio from the purple bar under layer 3 to the red bar under layer 13 in Figure 2a. Figure 2b quantitatively captures the improvement. We can observe that selecting layer 13 as the token re-selection layer is optimal if we only allow one additional token re-selection layer.

The reason why the overlap is not that significant is two-folded: 1. As the context length that we used is 100K and the token budget is 256, an overlap ratio of 40% is relatively significant compared to the sparsity ratio of 1/400. However, the heatmap we used for Figure 2a allows a maximum value of 1, making 0.4 less obvious. 2. Besides the token selection overlap ratio, we have added additional evidence that supports our claim in Figure 8 in the appendix. The figure captures the cosine similarity between adjacent layers’ attention scores, and we have provided our justifications along with the figure. The reason why we use cosine similarity is that the token overlap ratio only focuses on the number of shared tokens rather than the contribution of each token in terms of the attention score. For instance, two layers can have only a few tokens that are shared out of the 256 tokens we chose, but these several tokens can contribute a huge portion of the attention scores, which can be captured by the cosine similarity metrics we used in Figure 8. The resulting attention score similarity (correlation) for adjacent layers in Figure 8 is much higher than the token overlap ratio in Figure 2a, which provides further evidence for our design.

Q2. KV Cache Correction

For most of the tasks in our evaluations, the generation length is insufficient to see a degradation in model quality. On the other hand, to make it a fair comparison against other baseline models that haven’t applied cache correction, we don’t use the KV cache correction feature for all our current experiments. However, to provide some results on KV cache correction, we provide additional experimental results for the LongBench test with KV cache correction in Figure 9 in the appendix. Albeit the improvement is not significant due to a short generation length, the consistent improvement brought by cache correction indicates that cache correction effectively mitigates the adverse effects of polluted KV representations, which stem from sparse attention mechanisms and can accumulate errors, potentially reducing model performance.

Q3. Better performance than full attention baseline

Besides potential variances, we hypothesize that the rationale behind this phenomenon is that the sparse attention mechanism can focus more on the actual important tokens while removing noisy ones that are irrelevant to the question. We also identify it as a pretty interesting observation but lacks comprehensive analysis, which we leave as a future work.

We thank the reviewer for their review and address their main concerns below.

Q1. Writing - underlines

We thank the reviewer for pointing out the issues in our presentation. In our revised paper, we removed most of the phrases emphasized with underlines and only kept essential ones when they first appeared.

Q2. Writing - motivation and introduction

We have removed some unnecessary descriptions in the introduction section and rewritten some paragraphs (in blue) to make it more concise while preserving essential information. The introduction section now follows an order of 1. LLM and long context generation background 2. The bottleneck of long context LLM decoding lies in KV cache memory access (motivation) 3. A brief overview of existing methods 4. A brief introduction of our method (TidalDecode) 5. Evaluation results and summary of contributions.

Writing - preliminary and methodology

We have rewritten some descriptions in the methodology section that might cause confusion. Additionally, we have added the necessary explanations and definitions for certain figures and terms, like recall rate. All the modifications are highlighted in blue.

We hope we have addressed most of the writing concerns and misunderstandings raised by the reviewer. Please let us know if there is anything else we can further clarify.

We thank the reviewer for their review and address their main concerns below.

Q1. dataset selection in the experiments.

We agree with the reviewer that selecting a dataset with a longer generation length is important to distinguish between the eviction-based and selection-based methods. To this end, we further provide experimental results on the summarization task GovReport and code completion tasks RepoBench (RB) and LCC in Table 5 below and the revised paper. TidalDecode achieves the highest score on the added summarization task GovReport, where the model generates 512 tokens for each question. For code completion tasks, when selecting top-4096 tokens, TidalDecode outperforms full attention and Quest on RepoBench and stays close to full attention on LCC. We also include Table 4 below and in the revised paper to show the generation length for each task in LongBench.

Table 4:

Table 5:

Additionally, some other evidence that shows the improvement brought by TidalDecode: 1. In the original and extended LongBench results (Table 3 and Table 5), TidalDecode can both achieve a better score on average compared to the full-attention version across all eight or eleven tasks, where the average generation length is 145 tokens for a test. 2. In our sensitivity analysis for the choice of the token reselection layer (Table 9-14), choosing different token reselection layers can introduce a huge model performance degradation (e.g., from 100% to 0%), which indicates that the generation quality in the needle-in-the-haystack tasks also highly depends on the algorithm we used in the decoding stage.

Q2. I/O for group query attention.

As state-of-the-art kernel implementation for the Group Query Attention (GQA) only loads the keys and values from KVCache once within a group for all the query heads in the same group, if we choose to use a token budget of $k$, then the IO reduction ratio for KV loading would be: $\frac{\min{(\text{seqlen}, \bigcup_{i \in \text{q-head-group}}B_i)}}{\text{seqlen}}$ where $B_i$ is the top-k token indices selected by query head $i$ within the group. So, the reduction rate range spans from $\frac{k}{\text{seqlen}}$ to $\min{(1, \frac{\text{GroupSize}*k}{\text{seqlen}})}$. In practice, using a $k$ satisfying $\text{GroupSize}*k \ll \text{seqlen}$ can already achieve a pretty good performance, and the query head will usually share many top-k indices to reduce the I/O further.

For the dot product results that we saved, the I/O is negligible compared to the I/O for the KVs as the dot product is of shape (batch_size, seq_len, q_head_num) while the KVs are 2*(batch_size, seq_len, head_dim*kv_head_num). In the future, we plan to remove the materialization of the dot product by fusing the top-k operation with the attention kernel.

Q3. Presentation and references.

We thank the reviewer for pointing out the issues in our presentation. We have added all the citations for the baselines in the caption of Table 1 and a necessary description of Quest in section 4.2.1. All the modifications are highlighted in blue.