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

## Abstract

Abstract Large language models (LLMs) have driven significant advancements across diverse NLP tasks, with long-context models gaining prominence for handling extended inputs. However, the expanding key-value (KV) cache size required by Transformer architectures intensifies the memory constraints, particularly during the decoding phase, creating a significant bottleneck.
Existing sparse attention mechanisms designed to address this bottleneck have two limitations: (1) they often fail to reliably identify the most relevant tokens for attention, and (2) they overlook the spatial coherence of token selection across consecutive Transformer layers, which can lead to performance degradation and substantial overhead in token selection.
This paper introduces TidalDecode, a simple yet effective algorithm and system for fast and accurate LLM decoding through position persistent sparse attention. TidalDecode leverages the spatial coherence of tokens selected by existing sparse attention methods and introduces a few token selection layers that perform full attention to identify the tokens with the highest attention scores, while all other layers perform sparse attention with the pre-selected tokens.
This design enables TidalDecode to substantially reduce the overhead of token selection for sparse
attention without sacrificing the quality of the generated results.
Evaluation on a diverse set of LLMs and tasks shows that TidalDecode closely matches the generative performance of full attention methods while reducing the LLM decoding latency by up to 2.1 × 2.1\times 2.1 × 1 1 1 The codebase to reproduce performance and efficiency results for TidalDecode included in this paper can be found at https://github.com/DerrickYLJ/TidalDecode .

## 1 Introduction

Large language models (LLMs) have revolutionized natural language processing (NLP) by achieving state-of-the-art performance on various applications.
As LLMs evolve, they are increasingly being adapted to manage tasks with long contexts, such as Chain-of-Thought reasoning (Wei et al., ), document summarization (Huang et al., ), and retrieval-augmented generation (Ram et al., ; Zhang et al., ).
However, quickly and efficiently serving long-context LLMs is challenging due to the inherent memory and compute bottlenecks in the Transformer architectures (Vaswani et al., ).

LLM inference involves two separate stages: prefilling and decoding.
The prefilling stage computes the activations for all input tokens and stores the keys and values for all tokens in the key-value (KV) cache, allowing the LLM to reuse these keys and values to compute attention for future tokens.
In each decoding stage, the LLM decodes one new token using all input tokens and previously generated tokens.
The KV cache size grows linearly in the sequence length (Kwon et al., ).
For instance, with a context length of 128K tokens, the KV cache of LLama2-7B with half-precision can easily reach 64 GB(^2^22The KV cache size is computed as: $\text{Layers}\times\text{KV Heads}\times\text{Head Dim}\times\text{Seq Len} \times\text{FP16 Size}\times 2\ (\text{for K+V})=32\times 32\times 128\times 1 28K\times 2\ \text{bytes}\times 2=64\ \text{GB}$.), creating substantial memory pressure for LLM serving.
In addition, the LLM decoding stage is memory-bounded since decoding one new token requires accessing all previous tokens in the KV cache, making KV cache access the primary bottleneck for long-context LLM decoding.
This memory-bound nature severely limits the scalability and efficiency of LLM serving.

To address this problem, recent work has introduced sparse attention, which approximates full attention using a small portion of tokens with the highest attention scores. Compared to full attention, sparse attention reduces computation cost and memory access while preserving the LLM’s generative performance (Ge et al., ; Zhang et al., ).
Existing sparse attention techniques can be classified into two categories: eviction- and selection-based methods.

First, eviction-based sparse attention reduces memory usage for the KV cache by selectively discarding less relevant tokens from the KV cache, therefore reducing the number of tokens computed in attention mechanisms (Xiao et al., ; Zhang et al., ).
While these methods decrease the size of the KV cache, they can be inadequate for tasks where critical information is carried by tokens that are prematurely evicted, such as the needle-in-the-haystack tasks (Peng et al., ).
On the other hand, selection-based sparse attention maintains all tokens in the KV cache, estimates their attention scores, and selects a small subset of tokens to participate in each LLM decoding step.
This approach is prone to issues related to distribution shifts caused by appending sparsely attended, biased KV representations back into the cache.

Figure: Figure 1: The heatmap for one decoding step of Llama3-8B-Instruct (AI, ), where columns and rows indicate different Transformer layers and tokens in the KV cache, respectively. For each layer, the 5 tokens (10% sparsity) with the highest attention scores of the first attention head are highlighted in yellow, which are the tokens used for sparse attention. We feed an input prompt “Use only the provided search results to write a high-quality, concise answer to the question.\n<—begin_of_text—>\n The magic number is: 15213. \n\n\n Question: What is the magic number? Keep the response short and direct. Answer: ”, and the LLM outputs “15213”. The results show strong spatial coherence of tokens chosen for sparse attention in the decoding step.
Refer to caption: x1.png

This paper presents TidalDecode, an algorithm and system for fast and precise LLM decoding, utilizing position persistent sparse attention (PPSA).
A key insight behind TidalDecode is the observation that tokens chosen for sparse attention — based on their highest attention scores — exhibit significant overlap across consecutive Transformer layers within each decoding phase.
Figure 1 illustrates this overlap in a single decoding step of LLaMA-3-8B instruct AI () with an input of 51 tokens.
Each column in the figure corresponds to a Transformer layer, and each row indicates one token in the KV cache.
Selection-based sparse attention methods select the 5 tokens with the highest attention scores (highlighted in yellow) for attention computation in each head.
As the figure depicts, there is a recurring pattern where consecutive layers consistently focus on the same set of tokens, indicating a spatial coherence in the selection of tokens for sparse attention.

Instead of independently selecting tokens for sparse attention at each layer, TidalDecode introduces a few token selection layers, which perform full attention to identify the tokens with the highest attention scores.
All remaining layers implement position persistent sparse attention, where only the tokens selected by the token selection layers are retrieved from the KV cache for attention.
Consequently, all other layers between two token selection layers operate on the same set of tokens, reducing the overhead of token selection.
Experiments across a diverse set of LLMs and datasets demonstrate that using just two token selection layers — one at the beginning and one in the middle — is sufficient to achieve high generative performance while minimizing computation and memory overheads.

This design enables TidalDecode to substantially reduce the overhead of token selection for sparse attention without sacrificing the quality of the generated results. Additionally, to address the KV cache distribution shift, TidalDecode introduces a cache-correction mechanism that periodically refills the KV cache using full attention for all sparsely decoded tokens to mitigate bias in the KV representations.

Comprehensive evaluation with the LongChat-7b-v1.5-32k, Llama-3-8B, Llama-3-70B, and Llama-3.1-8B models on the Needle-in-the-Haystack, PG-19, and LongBench tasks demonstrates that TidalDecode can consistently achieve the best performance efficiency trade-off compared with the best existing sparse attention methods.
We have implemented custom GPU kernels for PPSA and an end-to-end system for TidalDecode. Compared with existing full and sparse attention implementations, our system reduced the end-to-end inference latency by up to $2.1\times$ and $1.2\times$, respectively.
In conclusion, our contributions are:

- •
We propose TidalDecode, a streamlined and efficient algorithm and system for fast and high-quality LLM decoding, utilizing position persistent sparse attention.
- •
To address KV cache distribution shifts, we introduce a cache-correction mechanism that periodically refills the KV cache with using full attention for sparsely decoded tokens.
- •
Empirically, we demonstrate the effectiveness and efficiency of TidalDecode through comprehensive evaluation, showing that TidalDecode significantly outperforms existing sparse attention methods.

## 2 Related Work

##### Long-context model.

Efficiently handling long-context inputs is essential for various LLM tasks in real-world applications such as document summarization, question answering, and dialogue systems (Wang et al., ).
Recent advancements, including rotary positional encoding (RoPE) (Su et al., ), have enabled models to manage extended context lengths effectively.
The LLaMA-3 model series supports up to 8K tokens, with enhanced versions such as Gradient-AI-Llama3 (AI, ) and LLaMA 3.1 (AI, ) extending this limit to 128K tokens.
Additionally, proprietary LLMs such as GPT-4 Turbo and GPT-4o (OpenAI, ) support up to 128K tokens, and Claude 3.5 Sonnet allows up to 200K tokens (Anthropic, ).
While recent work has introduced efficient attention kernel implementation (Dao et al., ; Dao, ), processing long-context inputs continues to be constrained by significant memory usage and computational costs from the extended KV cache.
TidalDecode is designed to mitigate these challenges by reducing latency and memory overhead through an efficient strategy for selecting tokens with the highest attention scores and one-time intermediate re-calibration, ensuring both efficiency and high-quality output.

To alleviate the intrinsic computational and memory bottleneck in long-context LLM inference, recent works on sparse attention have approached this problem from two main perspectives: eviction- and selection-based methods.

##### Eviction-based sparse attention.

Xiao et al. (); Zhang et al. () propose to reduce KV cache memory usage by evicting tokens that are considered less relevant during inference.
These suffer from potential performance degradation, especially in tasks where every token may carry crucial information (e.g., needle-in-the-haystack tasks), since tokens with high importance for a future decoding step can be mistakenly evicted as the generation proceeds, which makes selection-based methods more popular choices in latest sparse attention works.

##### Selection-based sparse attention.

Instead of evicting past tokens in the KV cache, Child et al. (); Kitaev et al. (); Choromanski et al. (); Tang et al. (); Ribar et al. () preserve the full KV cache and only select important tokens to attend with the attention module on the fly. More specifically, Child et al. () leverages a fixed attention mask to select tokens while Tang et al. (); Ribar et al. (); Choromanski et al. (); Kitaev et al. () aim to identify and retain the most relevant tokens at each layer by approximating attention scores.
Although these methods are more selective, they operate independently at each layer and are not guaranteed to obtain the ground-truth tokens with the highest attention scores, failing to capture token relevance patterns that persist across layers.
Moreover, attention score estimation algorithms sometimes introduce unnecessary complexity, diminishing the practical efficiency gains they are designed to achieve.
Improving upon prior works, TidalDecode leverages a shared pattern of most important tokens across consecutive layers to further reduce the computational overhead and memory access required for token selection.

## 3 Methodology

Figure: Figure 2: An overview of the decoding step in TidalDecode, which performs full attention for the first two layers, full attention with token selection for the third layer and a middle layer, and position persistent sparse attention for all other layers.
Refer to caption: x2.png

This section introduces TidalDecode, an efficient algorithm and system for fast LLM decoding using position persistent sparse attention and KV cache correction. Figure 2 shows an overview of TidalDecode.
TidalDecode uses the same prefilling mechanism as existing systems and performs full attention to compute the key-value (KV) cache for all prompt tokens.
In each decoding step, TidalDecode uses three types of attention layers: full attention, full attention with token selection, and position persistent sparse attention.
First, TidalDecode performs full attention for the initial Transformer layers to avoid early performance degradation as identified by prior work (Tang et al., ).
Second, the layer immediately after full attention and a single middle layer (e.g., layer 2 and 13 in Figure 2) perform full attention with token selection, where TidalDecode stores the inner product(^3^33We don’t store attention score as state-of-the-art attention kernels don’t materialize the attention score. Since the softmax operation is ordering invariant, we store the inner product value instead.) between the current query and key vectors of all tokens in KV cache during full attention and then selects $k$ tokens contributing to the highest attention scores.
Third, all other layers perform position persistent sparse attention, where only tokens selected from the previous token selection layer are loaded from the KV cache to perform attention computation.

### 3.1 Position Persistent Sparse Attention (PPSA)

Attention mechanisms have been widely used in today’s LLMs. For each attention head, the output is computed via scaled multiplicative formulation as follows.

$$ $A_{i}=Q_{i}K_{i}/\sqrt{d},\quad H_{i}={\rm softmax}(A_{i})V_{i}$ (1) $$

where $Q_{i}$, $K_{i}$, and $V_{i}$ are the query, key, and value tensors for the $i$-th attention head. $A_{i}$ is a matrix representing the attention scores between tokens, and $H_{i}$ is the output of the $i$-th attention head.
Instead of attending to all input tokens, existing sparse attention methods approximate attention computation by attending the query $Q_{i}$ to a subset of previous tokens the highest attention scores.
Prior work generally performs token selection for individual attention heads and Transformer layers, introducing high runtime overhead.
For example, selecting the tokens with highest attention scores using top-k can take longer than computing full attention (see Figure 7), thus diminishing the benefits of performing sparse attention.

The key insight behind TidalDecode’s position persistent sparse attention is an observation that tokens with highest attention scores for consecutive Transformer layers highly overlap.
We use the LLaMA-3-8B model and the needle-in-the-haystack test on PG-19-mini dataset with a context length of 100K tokens to quantify this observation.
We randomly select 100 requests from the dataset, compute full attention, and analyze the top 256 tokens with the highest attention scores for each Transformer layer.
Figure 3(a) shows the overlap ratios for all pairs of transformer layers, where an overlap ratio of 1 indicates that the tokens with highest attention scores are always identical in these layers, and an overlap ratio of 0 means the top tokens do not overlap in the two layers. Note that we select top 256 tokens from 100K tokens in the KV cache, so randomly selected tokens hardly overlap.
In Figure 3(b), we compute average recall rates of selected tokens by choosing different re-selection layers. We observe that without re-selection layers, where all layers possess a low overlap ratio with Layer 3 shown in the purple box in Figure 3(a), the average recall rates are less than 20%. When we choose Layer 13 to perform re-selection, the average recall rates boost to almost 40% due to higher overlap ratios between Layer 13 and its subsequent layers, shown by red boxes in Figure 3(a).

Figure: (a) Overlap of Tokens with the Highest Attention Scores between Layers
Refer to caption: x3.png

Figure: Algorithm 1 TidalDecode

Based on this observation, we design position persistent sparse attention to maximally leverage the token overlaps between consecutive Transformer layers to reduce the computation cost for token selection while achieving high predictive performance. Algorithm 1 shows the TidalDecode algorithm for interleaving full attention and PPSA layers.
After the initial full attention layers, TidalDecode uses a token selection layer that computes full attention and selects tokens with the highest attention scores.
To select tokens, TidalDecode stores the inner product $\langle Q,K\rangle$ on the fly together with full attention calculation. TidalDecode then selects the top $k$ tokens with the highest inner product values to form a token set $\mathcal{T}$. Note that using the inner product to select top-$k$ is equivalent to the post-softmax attention score as the softmax operator is ordering invariant.
All PPSA layers after a token selection layer computes sparse attention by only loading the keys and values for tokens in $\mathcal{T}$, thus limiting the number of tokens participating in attention computations and reducing memory access.

A straightforward approach to designing TidalDecode is to select the tokens $\mathcal{T}$ once after full attention and perform PPSA using the same set of tokens for all subsequent layers.
However, our preliminary experimentation shows that using a single token set for all Transformer layers reduces the LLM’s predictive performance by a large margin since distant Transformer layers are less correlated compared to consecutive layers, as shown in Figure 3(a).
To address this issue, TidalDecode performs token re-selection in a middle layer, where TidalDecode recalibrates the selected tokens with the highest attention scores by applying full attention and re-selecting top-$k$ token to update $\mathcal{T}$, ensuring that token selection remains optimal for the remaining layers.
This re-selection mechanism significantly boosts the model performance and promotes accurate and efficient PPSA throughout the model.

Extensive evaluation on both small and large models on a wide range of datasets shows that using a single middle layer for token re-selection is sufficient to preserve the LLM’s generative performance, while introducing small runtime overhead.
However, deciding which layer to perform token re-selection is critical to model performance.
As shown in Figure 3, choosing different layers for token re-selection results in different recall rates, where layer 11 and 13 achieve optimal performance.
Introducing a one-time token re-selection at an optimal layer ensures the selected tokens are re-calibrated, effectively mitigating the drift in token importance and elevates accuracy from 15% (without re-selection) to almost 40%.

### 3.2 KV Cache Correction

For tokens decoded by sparse attention methods, their key/value representations can deviate from the original representation of full attention decoded ones, which we refer to as polluted tokens. The problem can be further exacerbated as their KV pairs are added to the KV cache, resulting in the error accumulation or distribution shift of the KV cache. This can lead to model performance drop in scenarios where the generation length is fairly long. To this end, TidalDecode uses a cache-correction mechanism as shown in Figure 4 to periodically correct the polluted tokens in the KV cache. For every $T$ decoding step performed by TidalDecode, there will be a cache correction step through a prefill over all polluted tokens to update their KV representations in the cache. The choice of $T$ can be at the level of thousands of decoding steps but also depend on different models and tasks. Notice that the cache correction step can be performed concurrently with the sparse decoding step. Nevertheless, we haven’t used cache correction in our evaluations to make it a fair comparison against existing methods.

Figure: Figure 4: Cache Correction
Refer to caption: x5.png

## 4 Experiments

### 4.1 Experiment Setting

In this section, we conduct extensive experiments to assess both the performance and efficiency of TidalDecode.
Our evaluations are performed on widely used open-source models, including Llama-2-7B Touvron et al. () and Llama-3-8/70B.
Both models are pretrained decoder-only transformers, exhibiting similar yet distinct architectural features.
For instance, Llama 3-8B incorporates group query attention (GQA), a feature not present in Llama 2-7B. In Section 4.2, we evaluate TidalDecode’s performance on various tasks, including needle-in-the-haystack, language modeling on PG-19, and LongBench. In Section 4.3, we write customized attention kernels and compare TidalDecode’s kernel efficiency against existing state-of-the-art sparse attention methods. Finally, in Section 4.4, we conclude our evaluations with a detailed sensitivity analysis on the choice of different token selection layers. We use TD+LX to denote TidalDecode with layer X selected as the token re-selection layer throughout this section.

### 4.2 Performance Evaluation

To evaluate the effectiveness of TidalDecode, we conduct two key downstream NLP experiments: the needle-in-the-haystack test and perplexity evaluation on the PG-19 dataset (Rae et al., ). These tasks provide robust benchmarks for measuring both sparse attention models’ ability to retrieve critical information in challenging scenarios and their performance on long-context language modeling tasks.

#### 4.2.1 Needle-in-the-Haystack

**Table 1: Results of 10k-context-length Needle-in-the-Haystack test on LongChat-7b-v1.5-32k. TidalDecode achieves the same or better results than Quest and significantly better results than cache eviction algorithms such as H2O, TOVA, and StreamingLLM. TidalDecode achieves full accuracy with only a 512 token budget.**
| Method / Budget | K=32 | K=64 | K=128 | K=256 | K=512 |
| --- | --- | --- | --- | --- | --- |
| H2O | 0% | 1% | 1% | 1% | 3% |
| TOVA | 0% | 1% | 1% | 3% | 8% |
| StreamingLLM | 1% | 1% | 1% | 3% | 5% |
| Quest | 65% | 99% | 99% | 99% | 100% |
| TD+L7(Ours) | 73% | 92% | 98% | 99% | 100% |

**Table 2: Comprehensive results of 10K-, 32K-, and 100K-context-length Needle-in-the-Haystack test on Llama-3-8B-Instruct-Gradient-1048k, Llama-3.1-8B-Instruct, and Llama-3-70B-Instruct-Gradient-1048k with PG-19-mini dataset. Across all models, TidalDecode consistently outperforms Quest, showing that TidalDecode with only two token selection layers can effectively retain critical information. TidalDecode achieves full accuracy with 64, 64, and 128 tokens in 10K-, 32K-, and 100K-context-length tests, which is only 0.6%, 0.2%, and 0.1% of total input lengths, respectively.**
| Model (context length) | Method / Budget | K=32 | K=64 | K=128 | K=256 | K=512 |
| --- | --- | --- | --- | --- | --- | --- |
| LLaMA-3-8B<br>(10K) | Quest | 74% | 84% | 99% | 98% | 100% |
| TD+L13(Ours) | 88% | 98% | 100% | 100% | 100% |  |
| TD+L15(Ours) | 92% | 88% | 94% | 94% | 100% |  |
| LLaMA-3-8B<br>(100K) | Quest | 38% | 50% | 65% | 87% | 98% |
| TD+L13(Ours) | 86% | 92% | 100% | 100% | 100% |  |
| TD+L15(Ours) | 84% | 90% | 92% | 98% | 100% |  |
| LLaMA-3.1-8B<br>(10K) | Quest | 74% | 86% | 94% | 100% | 98% |
| TD+L13(Ours) | 100% | 100% | 100% | 100% | 100% |  |
| TD+L14(Ours) | 98% | 100% | 100% | 100% | 100% |  |
| LLaMA-3.1-8B<br>(32K) | Quest | 78% | 88% | 92% | 100% | 100% |
| TD+L13(Ours) | 98% | 100% | 100% | 100% | 100% |  |
| TD+L14(Ours) | 80% | 98% | 100% | 100% | 100% |  |
| LLaMA-3-70B<br>(10K) | Quest | 68% | 72% | 90% | 98% | 100% |
| TD+L14(Ours) | 87% | 93% | 100% | 100% | 100% |  |
| TD+L31(Ours) | 90% | 97% | 100% | 100% | 100% |  |
| LLaMA-3-70B<br>(32K) | Quest | 50% | 80% | 88% | 92% | 78% |
| TD+L14(Ours) | 82% | 98% | 98% | 100% | 100% |  |
| TD+L31(Ours) | 80% | 82% | 92% | 98% | 100% |  |

The Needle-in-the-Haystack test assesses LLMs’ ability to handle long-dependency tasks, which is particularly critical for sparse attention algorithms.
Eviction-based methods Xiao et al. (); Zhang et al. () may discard essential tokens, while selection-based approaches often fail to consistently identify the ground-truth tokens with the highest attention scores in long contexts.
Since Quest is the current state-of-the-art approach on this task, we first run TidalDecode on the same test as Quest on the LongChat-7b-v1.5-32k model and obtained Table 1 with competitive performance.
To demonstrate the effectiveness of TidalDecode on long-dependency tasks, we further evaluate TidalDecode on tasks with 10K-, 32K-, and 100K-context-window lengths with the LLaMA-3-70B, LLaMA-3-8B, LLaMA-3.1-8B model using the PG-19-mini dataset, shown in Table 2.
To ensure fairness, both TidalDecode and Quest use dense attention in the first two layers. In each test, we inserted a random password within the text and tested whether the specific method could retrieve the password correctly.

From Table 2, TidalDecode consistently outperforms Quest and achieves full accuracy with an extremely low sparsity (about 0.5% across all context lengths and models). These results demonstrate TidalDecode can achieve state-of-the-art performance with only two token selection layers.
While Quest relies on page-level importance estimation for token selection, TidalDecode’s exact selection with token reuse approach proves more effective for this task. Also, note that TidalDecode can reduce the token budget by up to $8\times$ when achieving a 100% accuracy compared with Quest. This further demonstrates that TidalDecode’s exact token selection layer can obtain more relevant tokens than Quest.

#### 4.2.2 Language Modeling

Figure: (a) Token Budget = 2048
Refer to caption: x6.png

Perplexity measures the negative likelihood of how well a model predicts the next word in a sequence, with lower values indicating better performance. We evaluate TidalDecode on Llama-3-8B-Instruct-Gradient-1048k with the PG-19 dataset, which includes up to 100 books, providing a comprehensive long-context benchmark.

As shown in Figure 5, TidalDecode+L9/13/15 consistently achieves lower perplexity than Quest across all token budget options (2048, 4096). This indicates that TidalDecode’s position persistent sparse attention mechanism can effectively retain critical information without significantly sacrificing model accuracy, even as the sequence length grows, demonstrating its robustness for long-context inputs.

#### 4.2.3 LongBench Experiment

**Table 3: Performance comparison on eight LongBench datasets evaluating single/multi-document QA, summarization, and retrieval tasks using Llama-3-8B-Instruct-Gradient-1048k. TidalDecode outperforms Quest at a 4096 token budget and achieves an average score higher than full-weight attention. The maximum F1-score for each task is in bold.**
| Method (K)/Task | MFQA | NrtQA | Qasp | 2Wiki | HotQA | QMSm | TrQA | PRe | Avg |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Full | 30.76 | 5.52 | 14.56 | 13.32 | 11.50 | 19.43 | 86.56 | 77.00 | 32.33 |
| Quest (1024) | 26.21 | 4.08 | 12.19 | 12.61 | 10.75 | 19.56 | 83.47 | 63.84 | 29.09 |
| TD+L13 (1024) | 28.57 | 7.63 | 11.11 | 13.56 | 9.82 | 20.37 | 79.78 | 75.17 | 30.75 |
| Quest (4096) | 28.92 | 3.74 | 13.63 | 12.83 | 12.15 | 19.36 | 85.91 | 72.50 | 31.13 |
| TD+L13 (4096) | 30.94 | 6.19 | 13.85 | 14.40 | 13.71 | 19.48 | 86.30 | 78.00 | 32.86 |

We also evaluate TidalDecode on LongBench, a benchmark designed to test LLMs on long-context tasks across diverse NLP domains (Bai et al., ). We focus on eight tasks: MultiFieldQA (MFQA), NarrativeQA (NrtQA), Qasper (Qasp), 2WikiMQA (2Wiki), HotpotQA (HotQA), QMSum (QMSm), TriviaQA (TrQA), and Passage Retrieval (PRe), which collectively composite a comprehensive evaluation benchmark in single/multi-document QA, summarization, and retrieval.

We evaluate all methods with LLaMA-3-8B-Instruct-Gradient-1048k. TidalDecode is compared against full-weight attention and Quest at token budgets of 1024 and 4096. As shown in LABEL:tab:longbench, TidalDecode consistently outperforms Quest on all tasks at $K=4096$ and on five tasks at $K=1024$. Surprisingly, TidalDecode, in most cases, matches or exceeds full attention baseline with notable sparsity: 14% on NrtQA, 50% on MFQA, 80% on Qasp, 50% on 2WikiMQA, 32% on HotQA, 29% on QMSm, 35% on TrQA, and 33% on PRe. We hypothesize this is because TidalDecode’s token selection process can filter out irrelevant information, thus leading to higher performance.

These results demonstrate TidalDecode’s generic ability to select tokens with the highest attention scores, achieving competitive or superior performance while significantly reducing token usage, making it ideal for long-context scenarios.

### 4.3 Efficiency Evaluation

Figure: Figure 6: End-to-end latency results on LLaMA-2-7B model for Full attention baseline(Full), Quest, and TidalDecode(TD) when context length is 10K, 32K, and 100K, respectively.
Refer to caption: x8.png

To show the efficiency of TidalDecode, we write customized kernels for our approach and measure the end-to-end decoding latency. We conduct evaluation under the configuration of Llama-2-7B on one Nvidia A100 (80 GB HBM, SXM4) with CUDA 12.2. We compare TidalDecode with state-of-the-art full attention serving library FlashInfer (Ye et al., ) and also the Quest implementation. As shown in Figure 6, we can observe that TidalDecode can consistently outperform full attention baseline and Quest by a large margin under all token budgets and context lengths. TidalDecode achieves this through token pattern reuse to minimize the token selection overhead. Notice that the latest LLaMA-3 model shares the same architecture as LLaMA-2, except it uses Group-Query-Attention instead of Multi-Head-Attention. However, this does not affect the relative efficiency comparison against Quest and full attention.

Refer to caption: x9.png

In Figure 7, we compare the overall attention latency between different methods on the LLaMA model with 32/64 layers. For the 32-layer LLaMA model, we have 2 full attention layers + 2 token selection layers + 28 sparse attention layers, while Quest has 2 full attention layers + 30 Quest attention layers. For the 64-layer LLaMA model, we have 2 full attention layers + 2 token selection layers + 60 sparse attention layers, while Quest has 2 full attention layers + 62 Quest attention layers. Thus, by completely removing the token estimation overhead in the sparse attention layers, for the 32-layer and 64-layer LLaMA model under all context lengths, TidalDecode can consistently achieve the lowest serving latency while bringing up to $5.56\times$ speed-up ratio against the full attention baseline and $2.17\times$ speed-up ratio against Quest. When the context length is 10K, Quest has a higher latency due to the token selection overhead, which aligns with the end-to-end results in Figure 6. In contrast, TidalDecode still achieves significant speed-up by utilizing the position persistent sparse attention mechanism.

Figure: Figure 8: The breakdown latency results for the full attention, token selection attention, sparse attention, and Quest attention kernels over 10K, 32K, and 100K context length. We use full attention latency as a reference and report other kernels’ relative latency ratio. We use a token budget of K=512 for TidalDecode and Quest across all evaluations.
Refer to caption: x12.png

In Figure 8, we further break down the latency comparison for different attention modules to show why TidalDecode can bring significant speed-up consistently. We compare different attention modules, namely, the full attention layer, the token selection layer, TidalDecode’s sparse attention layer, and the Quest attention layer over the 10K, 32K, and 100K context length. We can observe that, as TidalDecode’s sparse attention kernel can directly reuse previous token patterns, it completely removes the important token estimation overhead in the Quest attention kernel, resulting in up to $3.36\times$ speed-up compared with the Quest implementation. On the other hand, even though TidalDecode’s token selection layer has a slightly higher latency, we only have two token selection layers even in the 70B LLaMA model that has 64 layers in total.

### 4.4 Sensitivity analysis on token re-selection layer

Figure: (a) LLaMA-2 Model
Refer to caption: x13.png

In this section, we conduct sensitivity studies for different choices of the token re-selection layer. As TidalDecode only has one token re-selection layer in the middle, it is critical to choose the best-performed one. As shown in Figure 9, we have two interesting findings: (1). Different choices of token re-selection layers can significantly affect the accuracy of the results (2). For models within the same model family, the optimal token re-selection layer is consistent over different tasks. In our setup, the optimal token re-selection layer for the LLaMA-2-7B model is layer 7, while for the LLaMA-3-8B/LLaMA-3.1-8B model is layer 13. A concurrent KV cache compression work also identifies that layer 13 is surprisingly important for their approach as well (Shi et al., ). For a more detailed sensitivity results on the choice of different token re-selection layers, please refer to the appendix for more results.

## 5 Conclusion

To conclude, we introduce TidalDecode, an efficient LLM decoding framework with sparse attention. On observing the correlation of the pattern of tokens with the highest attention scores across different consecutive layers, TidalDecode proposes only to select tokens twice: once at the beginning layers and once in the middle layer to serve as a token re-selection layer. We find that using two token selection layers is necessary and sufficient to achieve high-generation quality. Additionally, by reusing the token patterns throughout the sparse attention layer, TidalDecode greatly reduces the token selection overhead, resulting in a significant end-to-end speed-up ratio against existing methods. More interestingly, the optimal choice of the token re-selection layer is consistent across different tasks if the model is in the same model family.

## Acknowledgment

This research is supported by NSF awards CNS-2147909, CNS-2211882, CNS-2239351, and research awards from Amazon, Cisco, Google, Meta, Oracle, Qualcomm, and Samsung.

## Appendix A Appendix

### A.1 LongBench for LLaMA-3.1-8B-Instruct

**Table 4: Performance comparison on eight LongBench datasets evaluating single/multi-document QA, summarization, and retrieval tasks using Llama-3.1-8B-Instruct. The maximum F1-score for each task is in bold.**
| Method | MFQA | NrtQA | Qasp | 2Wiki | HotQA | QSm | TrQA | Pre | Avg |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Full | 27.02 | 25.59 | 13.05 | 16.64 | 16.86 | 23.88 | 91.48 | 97.67 | 39.02 |
| Quest (1024) | 22.35 | 14.89 | 12.44 | 14.24 | 14.12 | 23.86 | 81.71 | 95.73 | 34.92 |
| TD+13 (1024) | 23.70 | 23.25 | 11.14 | 13.53 | 13.72 | 22.69 | 92.35 | 92.15 | 36.57 |
| Quest (4096) | 26.34 | 21.17 | 11.99 | 15.61 | 16.26 | 23.61 | 90.73 | 96.35 | 37.76 |
| TD+13 (4096) | 25.89 | 26.29 | 12.65 | 16.86 | 15.94 | 23.27 | 90.22 | 95.47 | 38.32 |

TidalDecode and full-weight attention share the maximum F1 scores for all tasks, achieving the best scores in three tasks (NrtQA, 2Wiki, and TrQA). TidalDecode significantly outperforms Quest in 4/8 tasks (NrtQA, Qasp, 2Wiki, and TrQA) and full-attention in 3/8 tasks (NrtQA, 2Wiki, and TrQA). For other tasks, we stay close to the full attention and also obtains a higher average score than Quest.

### A.2 End-to-end efficiency evaluation results

**Table 5: TidalDecode end-to-end efficiency results on LLaMA-2-7B**
| Context Length | Full Attention<br>(ms) | TidalDecode(ms) |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| K=32 | K=64 | K=128 | K=256 | K=512 | K=1024 | K=2048 | K=4096 |  |  |
| 10K | 19.22 | 16.94 | 17.18 | 17.15 | 16.98 | 16.96 | 17.32 | 17.19 | 17.63 |
| 32K | 25.71 | 17.89 | 17.92 | 17.64 | 17.70 | 17.91 | 17.97 | 18.48 | 18.98 |
| 100K | 45.70 | 21.26 | 21.09 | 21.38 | 21.19 | 21.13 | 21.38 | 21.65 | 22.34 |

**Table 6: Quest end-to-end efficiency results on LLaMA-2-7B**
| Context Length | Full Attention<br>(ms) | TidalDecode(ms) |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| K=32 | K=64 | K=128 | K=256 | K=512 | K=1024 | K=2048 | K=4096 |  |  |
| 10K | 19.22 | 20.39 | 19.86 | 19.44 | 19.35 | 20.18 | 19.91 | 20.23 | 21.09 |
| 32K | 25.71 | 20.47 | 20.85 | 20.73 | 21.06 | 20.62 | 20.94 | 21.35 | 22.11 |
| 100K | 45.70 | 24.93 | 25.18 | 24.77 | 24.90 | 24.84 | 25.10 | 25.77 | 26.17 |

### A.3 Full sensitivity studies on different token re-selection layer

**Table 7: Sensitivity study of re-selection layer (RL) on 10k-context-length Needle-in-the-Haystack test for LLaMA-3.2-3B-Instruct with TidalDecode. The best accuracy for each token budget (K) is in bold.**
| RL/K | 32 | 64 | 128 | 256 | 512 |
| --- | --- | --- | --- | --- | --- |
| TidalDecode+L2 | 2% | 12% | 14% | 16% | 30% |
| TidalDecode+L3 | 2% | 6% | 8% | 10% | 24% |
| TidalDecode+L4 | 6% | 14% | 16% | 20% | 28% |
| TidalDecode+L5 | 2% | 10% | 22% | 26% | 36% |
| TidalDecode+L6 | 18% | 26% | 26% | 32% | 46% |
| TidalDecode+L7 | 6% | 10% | 16% | 18% | 28% |
| TidalDecode+L8 | 20% | 20% | 46% | 60% | 84% |
| TidalDecode+L9 | 6% | 12% | 32% | 58% | 66% |
| TidalDecode+L10 | 44% | 58% | 50% | 60% | 64% |
| TidalDecode+L11 | 4% | 12% | 16% | 22% | 28% |
| TidalDecode+L12 | 50% | 84% | 96% | 98% | 98% |
| TidalDecode+L13 | 42% | 80% | 94% | 98% | 100% |
| TidalDecode+L14 | 28% | 44% | 54% | 60% | 72% |
| TidalDecode+L15 | 2% | 8% | 16% | 22% | 36% |
| TidalDecode+L16 | 4% | 16% | 12% | 22% | 34% |
| TidalDecode+L17 | 2% | 6% | 16% | 20% | 32% |
| TidalDecode+L18 | 2% | 10% | 12% | 18% | 28% |
| TidalDecode+L19 | 2% | 6% | 10% | 18% | 32% |
| TidalDecode+L20 | 6% | 10% | 12% | 18% | 24% |
| TidalDecode+L21 | 6% | 8% | 10% | 16% | 26% |
| TidalDecode+L22 | 6% | 0% | 12% | 12% | 26% |
| TidalDecode+L23 | 4% | 14% | 14% | 18% | 26% |
| TidalDecode+L24 | 2% | 10% | 16% | 20% | 28% |
| TidalDecode+L25 | 4% | 8% | 14% | 16% | 22% |
| TidalDecode+L26 | 2% | 10% | 10% | 22% | 26% |
| TidalDecode+L27 | 0% | 10% | 12% | 22% | 26% |
| Quest | 46% | 56% | 72% | 88% | 96% |

**Table 8: Sensitivity study of re-selection layer (RL) on 10k-context-length Needle-in-the-Haystack test for LLaMA-3.1-8B-Instruct with TidalDecode. The best accuracy for each token budget (K) is in bold. Layer 13 and Layer 14 are the best two re-selection layers for accuracy.**
| RL/K | 32 | 64 | 128 | 256 |
| --- | --- | --- | --- | --- |
| TidalDecode+L2 | 36% | 38% | 46% | 58% |
| TidalDecode+L3 | 8% | 10% | 14% | 34% |
| TidalDecode+L4 | 0% | 10% | 16% | 34% |
| TidalDecode+L5 | 14% | 30% | 52% | 52% |
| TidalDecode+L6 | 8% | 12% | 28% | 40% |
| TidalDecode+L7 | 6% | 10% | 10% | 18% |
| TidalDecode+L8 | 34% | 44% | 50% | 66% |
| TidalDecode+L9 | 64% | 78% | 82% | 90% |
| TidalDecode+L10 | 56% | 74% | 84% | 94% |
| TidalDecode+L11 | 52% | 76% | 82% | 86% |
| TidalDecode+L12 | 8% | 10% | 28% | 40% |
| TidalDecode+L13 | 100% | 100% | 100% | 100% |
| TidalDecode+L14 | 98% | 100% | 100% | 100% |
| TidalDecode+L15 | 56% | 78% | 88% | 96% |
| TidalDecode+L16 | 18% | 46% | 54% | 72% |
| TidalDecode+L17 | 64% | 74% | 86% | 98% |
| TidalDecode+L18 | 64% | 70% | 74% | 84% |
| TidalDecode+L19 | 58% | 50% | 60% | 68% |
| TidalDecode+L20 | 68% | 60% | 62% | 76% |
| TidalDecode+L21 | 40% | 48% | 48% | 62% |
| TidalDecode+L22 | 28% | 38% | 46% | 56% |
| TidalDecode+L23 | 40% | 46% | 52% | 64% |
| TidalDecode+L24 | 30% | 46% | 54% | 66% |
| TidalDecode+L25 | 40% | 54% | 50% | 66% |
| TidalDecode+L26 | 34% | 48% | 62% | 64% |
| TidalDecode+L27 | 38% | 50% | 54% | 70% |
| TidalDecode+L28 | 30% | 40% | 56% | 58% |
| TidalDecode+L29 | 32% | 48% | 56% | 68% |
| TidalDecode+L30 | 36% | 48% | 52% | 70% |
| TidalDecode+L31 | 30% | 36% | 42% | 56% |

**Table 9: Sensitivity study of re-selection layer (RL) on 10k-context-length Needle-in-the-Haystack test for LLaMA-3-70B-Instruct-Gradient-1048k; we first run top_k = 512 and filter out those layers that do not achieve full accuracy with TidalDecode. The best accuracy for each token budget (K) is in bold. Layer 14 and Layer 31 are the best two Re-selection layers for accuracy.**
| RL/K | 32 | 64 | 128 | 256 | 512 |
| --- | --- | --- | --- | --- | --- |
| L2 | - | - | - | - | 6% |
| L3 | - | - | - | - | 37% |
| L4 | - | - | - | - | 23% |
| L5 | - | - | - | - | 63% |
| L6 | - | - | - | - | 70% |
| L7 | - | - | - | - | 90% |
| L8 | - | - | - | - | 70% |
| L9 | - | - | - | - | 30% |
| L10 | - | - | - | - | 83% |
| L11 | - | - | - | - | 70% |
| L12 | - | - | - | - | 63% |
| L13 | - | - | - | - | 50% |
| L14 | 87% | 93% | 100% | 100% | 100% |
| L15 | - | - | - | - | 97% |
| L16 | - | - | - | - | 63% |
| L17 | - | - | - | - | 87% |
| L18 | 50% | 70% | 83% | 97% | 100% |
| L19 | - | - | - | - | 93% |
| L20 | - | - | - | - | 87% |
| L21 | 53% | 80% | 93% | 97% | 100% |
| L22 | - | - | - | - | 97% |
| L23 | 53% | 93% | 97% | 100% | 100% |
| L24 | 33% | 60% | 77% | 93% | 100% |
| L25 | - | - | - | - | 80% |
| L26 | - | - | - | - | 87% |
| L27 | 50% | 87% | 93% | 93% | 100% |
| L28 | 80% | 83% | 93% | 87% | 100% |
| L29 | - | - | - | - | 97% |
| L30 | 33% | 67% | 80% | 90% | 100% |
| L31 | 90% | 97% | 100% | 100% | 100% |
| L32 | 27% | 73% | 80% | 97% | 100% |

**Table 10: Sensitivity study of re-selection layer (RL) on 10k-context-length Needle-in-the-Haystack test for Llama-3-8B-Instruct-Gradient-1048k with TidalDecode. The best accuracy for each token budget (K) is in bold. Layer 9, Layer 13, and Layer 14 are the best three re-selection layers for accuracy.**
| RL/K | 16 | 32 | 64 | 128 | 256 | 512 |
| --- | --- | --- | --- | --- | --- | --- |
| TidalDecode+L2 | 78% | 84% | 76% | 94% | 88% | 98% |
| TidalDecode+L3 | 0% | 6% | 10% | 16% | 28% | 64% |
| TidalDecode+L4 | 2% | 10% | 16% | 28% | 68% | 84% |
| TidalDecode+L5 | 10% | 12% | 32% | 52% | 72% | 80% |
| TidalDecode+L6 | 4% | 6% | 10% | 14% | 16% | 24% |
| TidalDecode+L7 | 2% | 10% | 10% | 10% | 14% | 28% |
| TidalDecode+L8 | 26% | 64% | 80% | 90% | 92% | 96% |
| TidalDecode+L9 | 52% | 90% | 96% | 100% | 98% | 100% |
| TidalDecode+L10 | 72% | 76% | 86% | 94% | 96% | 100% |
| TidalDecode+L11 | 56% | 74% | 94% | 100% | 98% | 98% |
| TidalDecode+L12 | 8% | 14% | 22% | 44% | 66% | 94% |
| TidalDecode+L13 | 92% | 92% | 96% | 100% | 100% | 100% |
| TidalDecode+L14 | 74% | 68% | 88% | 98% | 100% | 100% |
| TidalDecode+L15 | 74% | 94% | 92% | 88% | 100% | 100% |
| TidalDecode+L16 | 44% | 50% | 72% | 66% | 82% | 94% |
| TidalDecode+L17 | 42% | 60% | 74% | 82% | 96% | 96% |
| TidalDecode+L18 | 60% | 72% | 74% | 74% | 88% | 98% |
| TidalDecode+L19 | 58% | 74% | 82% | 84% | 98% | 96% |
| TidalDecode+L20 | 64% | 74% | 96% | 78% | 90% | 98% |
| TidalDecode+L21 | 10% | 38% | 60% | 66% | 90% | 94% |
| TidalDecode+L22 | 60% | 70% | 68% | 72% | 82% | 98% |
| TidalDecode+L23 | 58% | 78% | 70% | 86% | 88% | 98% |
| TidalDecode+L24 | 62% | 58% | 76% | 70% | 78% | 92% |
| TidalDecode+L25 | 66% | 86% | 84% | 82% | 92% | 100% |
| TidalDecode+L26 | 54% | 64% | 66% | 80% | 90% | 94% |
| TidalDecode+L27 | 84% | 80% | 94% | 96% | 88% | 100% |
| TidalDecode+L28 | 66% | 66% | 76% | 84% | 94% | 94% |
| TidalDecode+L29 | 72% | 80% | 88% | 80% | 90% | 96% |
| TidalDecode+L30 | 80% | 90% | 86% | 88% | 96% | 100% |
| TidalDecode+L31 | 74% | 90% | 88% | 84% | 90% | 96% |

**Table 11: Sensitivity study of re-selection layer (RL) on 3k-context-length Needle-in-the-Haystack test for LongChat-7b-v1.5-32k with TidalDecode. The best accuracy for each token budget (K) is in bold. Layer 7 serves the best re-selection layer for accuracy.**
| RL/K | 32 | 64 | 128 | 256 |
| --- | --- | --- | --- | --- |
| TidalDecode+L2 | 2% | 2% | 6% | 54% |
| TidalDecode+L3 | 10% | 52% | 67% | 78% |
| TidalDecode+L4 | 4% | 36% | 65% | 76% |
| TidalDecode+L5 | 17% | 87% | 94% | 99% |
| TidalDecode+L6 | 70% | 96% | 99% | 99% |
| TidalDecode+L7 | 80% | 98% | 100% | 100% |
| TidalDecode+L8 | 58% | 82% | 96% | 96% |
| TidalDecode+L9 | 7% | 31% | 59% | 71% |
| TidalDecode+L10 | 16% | 59% | 71% | 78% |
| TidalDecode+L11 | 34% | 61% | 68% | 77% |
| TidalDecode+L12 | 17% | 32% | 53% | 77% |
| TidalDecode+L13 | 5% | 10% | 28% | 48% |
| TidalDecode+L14 | 24% | 41% | 57% | 64% |
| TidalDecode+L15 | 37% | 47% | 62% | 69% |
| TidalDecode+L16 | 16% | 24% | 28% | 46% |
| TidalDecode+L17 | 4% | 4% | 10% | 34% |
| TidalDecode+L18 | 2% | 3% | 8% | 15% |
| TidalDecode+L19 | 0% | 1% | 7% | 19% |
| TidalDecode+L20 | 0% | 3% | 6% | 20% |
| TidalDecode+L21 | 0% | 2% | 10% | 19% |
| TidalDecode+L22 | 0% | 4% | 4% | 18% |
| TidalDecode+L23 | 0% | 2% | 5% | 13% |
| TidalDecode+L24 | 0% | 2% | 6% | 21% |
| TidalDecode+L25 | 0% | 2% | 7% | 16% |
| TidalDecode+L26 | 0% | 1% | 7% | 19% |
| TidalDecode+L27 | 0% | 2% | 4% | 21% |
| TidalDecode+L28 | 1% | 2% | 10% | 17% |
| TidalDecode+L29 | 0% | 3% | 7% | 16% |
| TidalDecode+L30 | 1% | 1% | 9% | 15% |
| TidalDecode+L31 | 1% | 2% | 5% | 16% |

**Table 12: Sensitivity study of re-selection layer (RL) on 3k-context-length Needle-in-the-Haystack test for Yarn-Llama-2-7b-128k with TidalDecode. The best accuracy for each token budget (K) is in bold. Layer 7 serves the best re-selection layer for accuracy.**
| RL/K | 32 | 64 | 128 | 256 | 512 |
| --- | --- | --- | --- | --- | --- |
| TidalDecode+L2 | 0% | 0% | 0% | 3% | 25% |
| TidalDecode+L3 | 11% | 26% | 39% | 65% | 85% |
| TidalDecode+L4 | 5% | 17% | 34% | 60% | 92% |
| TidalDecode+L5 | 16% | 42% | 65% | 87% | 96% |
| TidalDecode+L6 | 73% | 83% | 89% | 95% | 100% |
| TidalDecode+L7 | 73% | 95% | 98% | 98% | 100% |
| TidalDecode+L8 | 87% | 92% | 97% | 94% | 99% |
| TidalDecode+L9 | 7% | 21% | 43% | 60% | 95% |
| TidalDecode+L10 | 12% | 31% | 58% | 69% | 93% |
| TidalDecode+L11 | 20% | 21% | 46% | 68% | 97% |
| TidalDecode+L12 | 2% | 15% | 28% | 51% | 92% |
| TidalDecode+L13 | 4% | 5% | 20% | 34% | 88% |
| TidalDecode+L14 | 16% | 20% | 49% | 53% | 91% |
| TidalDecode+L15 | 2% | 25% | 44% | 56% | 90% |
| TidalDecode+L16 | 10% | 13% | 21% | 43% | 86% |
| TidalDecode+L17 | 3% | 4% | 9% | 16% | 85% |
| TidalDecode+L18 | 0% | 1% | 2% | 15% | 84% |
| TidalDecode+L19 | 0% | 2% | 3% | 7% | 80% |
| TidalDecode+L20 | 0% | 2% | 0% | 7% | 79% |
| TidalDecode+L21 | 0% | 0% | 1% | 5% | 77% |
| TidalDecode+L22 | 0% | 0% | 0% | 8% | 76% |
| TidalDecode+L23 | 0% | 0% | 1% | 7% | 74% |
| TidalDecode+L24 | 0% | 2% | 0% | 9% | 73% |
| TidalDecode+L25 | 0% | 0% | 2% | 10% | 71% |
| TidalDecode+L26 | 0% | 1% | 1% | 5% | 70% |
| TidalDecode+L27 | 0% | 1% | 3% | 8% | 68% |
| TidalDecode+L28 | 0% | 1% | 0% | 9% | 67% |
| TidalDecode+L29 | 0% | 0% | 2% | 8% | 65% |
| TidalDecode+L30 | 0% | 1% | 1% | 8% | 64% |
| TidalDecode+L31 | 0% | 0% | 2% | 10% | 62% |

## Peer Reviews

confidence: 5
contribution: 3
initial_rating: 6
presentation: 4
questions: see cons.
soundness: 3
strengths: The advantages of this paper are as follows:

1) The logic is very clear, and the writing is well-structured. The figures and captions are precise, making it easy for readers to quickly understand.


2) The focus of the paper is excellent, and I completely agree with it. Particularly with group query attention, KV cache storage is no longer a major issue; instead, the I/O for the KV cache in attention deserves attention. And I really like the mechanism cache correction.


3) Most of the experiments are solid, although there is room for improvement. And the cuda kernel is good.
summary: This paper introduces TidalAttention, which uses a layer-wise selection approach to choose the top-k keys and values needed in top-k attention, thereby reducing the amount of KV cache requiring I/O.
weaknesses: I think the paper has the following limitations or raises some questions for me:

1) I believe there may be an issue with the dataset selection in the experiments. Since the authors’ method retains the full KV cache as opposed to KV cache eviction, it would be best to evaluate on datasets with longer outputs. If the dataset only outputs one token, as in NIAH, there should be no difference between eviction-based and selection-based sparse methods, which doesn't highlight the importance of maintaining the full KV cache. Additionally, some QA datasets in LongBench also have this issue. I recommend that the authors supplement more summarization datasets from LongBench (such as gov_report) and code completion datasets (like LCC and repo-bench).


2) The authors claim that their kernel stores dot product instead of attn score, so I just wondering how this method works for group query attention. How can the authors decrease the IO of gqa if the queries in the same group with different top-k?


3) The authors’ description of Quest is not very clear, and it would be beneficial to add citations for previous baselines in the tables.

---

confidence: 5
contribution: 2
initial_rating: 5
presentation: 2
questions: NA
soundness: 2
strengths: - The paper presents an optimization strategy to improve speed and latency by recognizing that the selected tokens in sparse attention remain consistent across layers, so that the selection process may only need to happen in lower layers.
- The paper presents evaluation to show the method is competitive to some baselines on latency evaluation
summary: The paper introduce a decoding method for sparse attention with pre-selected tokens. The method essentially performs full attention at lower layers and sparse selective attention at the upper layers to obtain the selected tokens, and then full-attention on reduced sets of tokens in the highest layers. The paper presents experiments to show that TidalDecode is competitive to full attention decoding while reducing latency.
weaknesses: - The writing and presentation of the paper significantly need improvement. There are too much unnecessary underlined texts.
- The motivation and introduction are written unclearly, with too much irrelevant introductional texts.
- The preliminary and methodology are written in confusing ways. There is a need to put clearer definition and better writing. 

Generally, I would lean towards acceptance if the paper is revised thoroughly.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: 1.Building on Weakness 1, an adaptive layer selection strategy seems necessary. Did the authors explore any strategies in this regard?
2.I am quite curious about the actual application effects of KV Cache Correction. Were corresponding experimental validations conducted?
3.In Table 3, it is interesting to note that on certain datasets like HotQA, the sparse attention mechanism sometimes outperforms the full attention mechanism. Is there a relevant explanation for this, or an empirical rationale behind it?
soundness: 3
strengths: 1.The article is written in a clear and comprehensible way, with strong consistency across empirical experiments, motivation, methodology, and the effectiveness of the proposed method.
2.The method is simple yet effective, supported by extensive experimental and theoretical validations. A wide range of diverse experiments, including multiple LLMs and tasks, demonstrates the excellent performance of the proposed method.
summary: This article introduces a new sparse attention mechanism, TidalDecode, which is an efficient selection-based attention computing mechanism for the second stage of LLM inference, specifically the decoding stage. The authors discovered through empirical experiments that there is a significant overlap among the Top-K tokens in the attention scores between consecutive Transformer layers. Thus, they propose performing full attention computation only in the initial and intermediate layers, while the other layers select the most relevant KV-cache to compute attention, achieving an efficient sparse attention mechanism. Experimental results across diverse LLMs and tasks demonstrate significant improvements in both decoding latency and the quality of generated results. The optimal choice of the token re-selection layer is consistent across different tasks if the model belongs to the same model family.
weaknesses: 1.The chosen token across different tasks for the same model family is almost consistent; however, the dramatic variation in performance between L12 and L13 in Table 7 presents a significant challenge for the method's generalization across different models and datasets. This variability could lead to catastrophic results when generalizing to other datasets.
2.Regarding Figure 2a, the overlap results do not seem significant based solely on the visual representation. Please provide a detailed explanation of the information that can be interpreted from the figure.

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 3
questions: Please see weaknesses
soundness: 3
strengths: The proposed method is effective, maintaining the advantages of both eviction and selection-based sparse attention methods according to experiments.
The paper presents robust evaluations across various LLMs and tasks, including LongBench and Needle-in-the-Haystack tests. Results indicate that TidalDecode consistently performs well compared to state-of-the-art methods.
summary: The paper presents TidalDecode to enhance the efficiency and accuracy of large language models (LLMs) through position persistent sparse attention. TidalDecode introduces token selection layers to identify and retain the most relevant tokens, with other layers using sparse attention to manage memory and reduce computational demands. Empirical evaluations indicate that TidalDecode can achieve up to a 2.1× reduction in latency while maintaining comparable generative performance to full attention mechanisms.
weaknesses: 1. Definition of `spatial coherence`. The spatial coherence seems just come from nowhere. I can mostly understand it refers to the consecutive layers have significant attention overlaps. But there is not a `scientific definition` in this paper.
2. Complexity of Implementation: TidalDecode's reliance on custom GPU kernels and specific layer configurations could limit its accessibility and adoption. Please provide pseudo code or GPU consumption data, it could make the results more clear.
3. The proposed method can be considered as a kind of hierarchical attention, more related baselines can be included. such as [1,2]
4. KV-cache correction seems to be crucial, no ablation study (or I just missed). Please provide the results without KV-cache correction.
5. training details, e.g, learning rate, training tokens, GPU hours, etc. It would help the reproduction.
6. Does `L` matter in different models? How to determine the optimal hyper-parameters. It would be appreciated if automatic prediction of `L` or a guideline for `L` selection across different model size, model architecture, etc.
7. There might be some conflicts in the packages used in latex.

[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.

Meta-review: The paper proposes a novel sparse attention mechanism, TidalDecode, to enhance the efficiency and accuracy of large language models (LLMs) during the decoding stage of inference. This mechanism leverages a layer-wise selection approach to optimize attention computation. By reducing the amount of KV cache requiring input/output, TidalDecode significantly decreases latency. Empirical evaluations demonstrate that TidalDecode achieves up to a 2.1× reduction in latency while maintaining comparable generative performance to full attention mechanisms. 

The paper is clear and easy to follow. Sufficient experiments and analyses make the results very convincing. In the rebuttal stage, the authors actively interact with reviewers to answer their questions. I am happy to recommend acceptance of the paper. However, please also consider the negative comments when revising the paper.

Decision: Accept (Poster)