A Little Goes a Long Way: Efficient Long Context Training and Inference with Partial Contexts

## Abstract

Abstract Training and serving long-context large language models (LLMs) incurs substantial overhead.
To address this, two critical steps are often required: a pretrained LLM typically undergoes a separate stage for context length extension by training on long-context data, followed by architectural modifications to reduce the overhead of KV cache during serving.
This paper argues that integrating length extension with a GPU-friendly KV cache reduction architecture not only reduces training overhead during length extension,
but also achieves better long-context performance.
This leads to our proposed LongGen , which finetunes a pretrained LLM into an efficient architecture during length extension. LongGen builds on three key insights:
(1) Sparse attention patterns, such as window attention (attending to recent tokens), attention sink (initial ones), and blockwise sparse attention (strided token blocks) are well-suited for building efficient long-context models, primarily due to their GPU-friendly memory access patterns, enabling efficiency gains not just theoretically but in practice as well.
(2) It is essential for the model to have direct access to all tokens.
A hybrid architecture with 1/3 full attention layers and 2/3 efficient ones achieves a balanced trade-off between efficiency and long-context performance.
(3) Lightweight training on 5B long-context data is sufficient to extend the hybrid model’s context length from 4K to 128K. We evaluate LongGen on both Llama-2 7B and Llama-2 70B, demonstrating its effectiveness across different scales.
During training with 128K-long contexts, LongGen achieves 1.55x training speedup and reduces wall-clock time by 36%, compared to a full-attention baseline.
During inference, LongGen reduces KV cache memory by 62%, achieving 1.67x prefilling speedup and 1.41x decoding speedup.
Compared to baselines that apply KV-cache reduction techniques to full-attention long-context LLMs, LongGen achieves substantially stronger performance not only on the Needle-in-a-Haystack retrieval task, but also on more challenging long-context reasoning tasks, including BABILong and RULER.

## 1 Introduction

Transformer-based large language models (LLMs) capable of processing long contexts have unlocked new opportunities across a wide range of applications, such as processing long multi-turn dialogue, understanding code repositories, and answering complex queries that require synthesizing information from multiple documents.
However, their quadratic complexity incurs substantial overhead for both training and inference with long contexts (Jiang et al., ; Sun et al., ).
To address this, typical approaches involve a two-stage framework (Xiong et al., ; Dubey et al., ).

- •
To enhance the model long context capabilities without the substantial overhead of training with long contexts, an LLM is typically pretrained on large short-context datasets (e.g., 2TB of 4K tokens) and undergoes a separate context length extension phase, usually through continual pretraining on a smaller amount of long-context data (e.g., 5B of 128K tokens).
- •
Serving a long-context LLM is also challenging due to high memory requirements for caching key and value vectors (KV cache). For example, a 7B Llama-2 model in BF16 precision uses 14 GB for model weights, while the KV cache for a single 128K sequence adds 69 GB - exceeding an H100 GPU’s capacity. Recent works have sought to reduce KV cache memory overhead post-hoc by analyzing LLM attention patterns and leveraging these insights to create sparse attention mechanisms, usually without additional fine-tuning.

Despite their promising performance in language modeling perplexity, these methods underperform on long-context retrieval and complex reasoning tasks (§2).
This paper argues that integrating the two stages not only builds an efficient long-context LLM with improved long-context performance, but also reduces the cost of length extension training.

Figure: Figure 1: Overview of LongGen. Left: It uses a hybrid architecture, and applies KV-reduced attention in 2/3 layers at the top and bottom, while keeping the middle 1/3 layers full attention. Right: Two KV-reduced attention variants are explored.
Refer to caption: extracted/5895696/figures/model_architecture.png

Figure: Figure 2: Training and inference efficiency under different sparsity levels. Left: Training wall-clock speedup. Mid: KV memory reduction. Right: Inference speedup. We compare training wall-clock time with FlashAttention and benchmark inference on vLLM. All results are measured on Llama2-7B, which consists of 32 layers in total. “1/7”, “1/5”, “1/3 Full” and “All Full” indicate using 5, 7, 12, and 32 full layers, respectively.
Refer to caption: x1.png

Our proposed approach LongGen, as illustrated in Figure 1, is a simple and effective hybrid transformer architecture that could be built upon any pretrained transformer to extend its context length.
LongGen incorporates two key designs:
(1) It conducts context extension with various GPU-friendly KV cache-saving designs, including but not limited to window attention (Jiang et al., ), attention sink (Xiao et al., ), and blockwise strided attention (Zaheer et al., ).
Compared with delicately curated KV cache reduction methods, LongGen employs a simpler and more efficient KV design, which maintains uniform memory access patterns for attention heads and achieves load balance among token blocks.
Practically, our customized triton training kernel inherits from FlashAttention-2 (Dao, ), but achieves faster speed in sparse settings, and the inference kernel is well-suited to vLLM (Kwon et al., ) for high-throughput serving.
(2) LongGen uses this sparse attention in a hybrid architecture, where 2/3 of the attention layers use sparse attention while the remaining 1/3 retain full attention, which we find is essential for handling complex tasks requiring direct access to long-context information(§4.4).

Experiments are conducted on a Llama-2-7B base model (Touvron et al., ) and its 70B counterpart with group query attention.
We extend their context length from 4K to 128K and perform evaluations on multiple long-context benchmarks, including the needle-in-a-haystack (NIAH) retrieval, and two more distinguishing benchmarks, BABILong (Kuratov et al., ) and RULER (Hsieh et al., ), where long-context reasoning is also needed.
Through ablations on the position and number of sparse attention layers, we study how sparsity level affects long-context capability.
Results show that with full layers in the middle, LongGen could achieve comparable performances with standard transformers even with 2/3 of its remaining layers being with sparse attention.
Meanwhile, LongGen only requires lightweight training on 5B Slimpajama (Shen et al., ) tokens (less than 0.1% GPU hours of pretraining) and maintains its short context capability, as indicated by the MMLU (Hendrycks et al., ) score.
With our customized kernel, it achieves a 1.55x wall-clock training speedup, as shown in Figure 2.
During inference, it reduces 62% KV cache memory, bringing 1.67x prefilling speedup and 1.41x decoding speedup.

## 2 Background

#### LLM context length extension.

Due to the overhead of training transformers on long sequences, context extension is usually separated from standard pretraining as a dedicated post-training stage.
State-of-the-art models, such as Gemini (Team et al., ), Llama series (Touvron et al., ; Dubey et al., ) and Qwen (Bai et al., ) are typically pretrained on large short-sequence corpora, then undergo length extension on relatively smaller amounts of longer sequences.
For example, Llama-3 is pre-trained on 15T tokens within an 8K context length, then post-trained on an additional 800B tokens to extend to 128K length (Dubey et al., ; Xiong et al., ).
A related data engineering work (Fu et al., ) significantly reduce the tokens needed for context extension to only 5B by carefully balancing the data source to be similar with that of the pretraining corpus.
Typically, there is no architectural modifications during length extension.
As a result, these long-context models are challenging to serve due to the memory overhead incurred by KV cache (Jiang et al., ).
Therefore, techniques are required to reduce the inference-time memory overhead of these models.

#### Inference-time KV reduction fails to generalize to long contexts.

Figure: Figure 3: Inference time KV cache reduction methods fail on long context.
Refer to caption: extracted/5895696/figures/inference_fail.png

To reduce the inference cost of large transformer models, one straightforward idea is to apply inference time KV reduction methods.
By only saving the mostly attended KV Cache, they reduce inference memory consumption and forward computing FLOPs.
For efficient decoding, the key and value vectors over the context are kept in the GPU memory, usually called KV cache.
Its memory overhead has become the primary bottleneck of serving LLMs.
Various techniques have been developed to reduce KV cache overhead by storing only a subset of the KV and evicting the rest (Xiao et al., ; Liu et al., ; Ge et al., ).
However, as we show here, they often underperform in long-context scenarios, where KV cache reduction is needed the most.

We evaluate two established KV cache reduction methods, Attention Sink (Xiao et al., ), and PyramidKV (Zhang et al., ) on a context-extended Llama2-7B-128K model.(^1^11Details about this Llama2-7B-128K model can be found in §4.1.)
Attention Sink (AttnSink) only retains initial and local tokens, while PyramidKV identifies the unique attention distribution at each layer and allocates a dynamic KV budget.
Similarly to Heavy-Hitter (H2O) (Zhang et al., ), PyramidKV preserves tokens that have the largest cumulative attention scores.
Additionally, we include the original Llama2-7B-128K baseline and our LongGen, and report their performance in Figure 3.
For fair comparison, we keep the same KV cache budget for Attn Sink, PyramidKV, and LongGen, i.e., 60% of the original full transformer.
Figure 3 shows the results on the needle-in-a-haystack (NIAH) retrieval task.
Both the original Llama2-7B-128K base model and LongGen can pass the needle retrieval task with 100% accuracy, which we will analyze later in the experiment section §4.2.
However, both KV eviction methods underperform on NIAH, especially with $>32$K contexts.
Intuitively, their underperformance can be attributed to two factors:

- •
Lack of full context access. Due to KV reduction at all layers, the model loses *direct* access to certain positions, making it more challenging for the model to perform accurate retrieval.
- •
Lack of sparse context adaption.
The model parameters, especially position embeddings, remain fixed during cache reduction.
However, it has been observed that the widely used RoPE-style position embedding (Su et al., ) struggles to generalize to unseen position ranges (Wang et al., ).
After KV reduction, the position range becomes non-contiguous,
creating a mismatch with the training setup.
This can potentially hurt the model’s ability to adapt to longer contexts.

The lessons from the above experiment motivates LongGen’s key design choices:
(1) retaining full context at some layers, and (2) training the model to adapt to sparse attention at the others.

## 3 LongGen

LongGen improves long-context training and inference efficiency.
It finetunes a pretrained LLM into an efficient architecture during length extension.
LongGen uses a hybrid architecture,
with KV-reduced attention patterns in bottom and top layers, while retaining full contexts in the middle layers (Figure 1).
We will explore various strategies to initialize the sparse attention,
followed by a comprehensive analysis of FLOPs and memory savings to demonstrate the theoretical efficiency benefits of LongGen.

### 3.1 Context Extension with Efficient Architectures

#### A hybrid architecture

Previous work has found attention is relatively specialized in the middle layers.
Some middle-layer attention heads are shown to be strongly related to retrieval and reasoning, indicating refined information aggregation in those layers Wu et al. (); Fu (); Zhang et al. ().
Inspired by their findings, we design LongGen to be in an hourglass shape—sparse attention in both ends to efficiently reduce computing cost, and full attention in the middle to maximize long context performance.
We will show the necessity of preserving full contexts for middle layers in §4.4.
We post-train this hybrid architecture on on longer sequences with the same loss function as in pretraining.

#### Different KV-Reduced Attention Strategies

Efficient attention layers in LongGen can be initialized using various KV cache reduction strategies.
Instead of choosing attention patterns that adaptively allocate memory for each head, LongGen prefers uniformly structured attention patterns.
Specifically, we introduce two criteria for sparse attention patterns:
(1) Static access of position.
For each head, the position to be retained should be consistent and agnostic of input tokens.
This enables a uniform GPU memory access pattern in training, which allows for optimized sparse training kernel and memory IO.
During inference, it is also feasible to enforce the a fixed cache management strategy, thus avoiding extra computation and tensor replacement.
From both sides, static attention patterns lead to better efficiency.
(2) Block-wise context handling.
In GPU, threads are assigned to streaming multiprocessors in block granularity, where each block handles consequent memory space, which is usually significantly larger than the memory occupied by each token.
Consequently, token-wise KV design will create idling threads and lower the utilization rate of SRAM.
Therefore, we load context by blocks and set the block size to be similar with each thread block size.
Guided by the two principles, we equip LongGen with two existing KV cache reduction strategies and illustrate them in Figure 1.
Our first variant is Attention Sink, which combines initial tokens with local blocks of context window (Xiao et al., ).
This technique builds upon the proven effectiveness of sliding window attention in long context training, as demonstrated by Mistral (Jiang et al., ).
An alternative initialization strategy is block sparse attention, which divides the context into multiple blocks and selectively attends to specific block indices based on a predetermined stride (Zaheer et al., ; Qiu et al., ).

### 3.2 Kernel-Level Optimization to Improve Efficiency

Since our implementation differs from traditional transformers in sparse attention computation only, LongGen can benefit from the tools developed to improve transformer efficiency, such as flash attention.
In fact, we build our customized kernel upon flash attention.
Similarly to FlashAttention (Dao et al., ), LongGen takes an attention mask matrix as input.
The attention mask is represented in compressed sparse row format, which consumes $O(N)$ memory instead of the original $O(N^{2})$.
Note that the sparse mask is static for all heads, so there is no overhead in building the dynamic masks.
During forward and backward steps, we skip computing a position by avoiding loading it into HBM if its mask value is zero.
This implementation enables FLOPs saving, significantly accelerating training and inference.

During inference, the bottleneck is to load KV cache as fast as possible.
Instead of loading the entire head dimension, we split the loading of the head dimension across different thread blocks.
We empirically found this to be more efficient than the original FlashAttention, probably because of reduced SRAM usage per loading cycle.
For the attention sink initialization, we follow Mistral to use the rolling buffer cache for local sliding window tokens (Jiang et al., ).

### 3.3 Theoretical Training and Inference Advantage

Table: Table 1: Efficiency comparison between conventional full attention and the sparse attention in LongGen. Training FLOPs count both attention forward and backward operations. $L_{\text{full}}$ and $L_{\text{partial}}$ are the number of full context and partial context layers. $N$ and $D$ represent sequence length and head dimension, respectively. In LongGen’s attention sink variant, $S$ represents a fixed window block size. In the block sparse variant, $S$ represents a fixed total block size. In both variants, $S$ does not grow with context length, and $S\ll N$.

#### Training

Since each token needs to be stored to calculate the loss function, LongGen does not provide any memory savings during the training process.
The major efficiency gain comes from a substantial reduction in training FLOPs.
Table 1 demonstrates that when the context length $N$ significantly exceeds the block size $S$, the FLOPs required for sparse attention layers become negligible ($S\ll N$).
Consequently, LongGen effectively reduces the total FLOPs to a fraction of $L_{\text{full}}/(L_{\text{full}}+L_{\text{sparse}})$ compared to the original transformer model.
Empirical studies indicate that a ratio of $L_{\text{full}}:L_{\text{sparse}}=1:2$ provides an optimal balance, achieving significant efficiency gains without compromising the model’s long-context performance.

#### Inference

During the inference phase, LongGen demonstrates notable improvements in both prefilling time and KV cache memory usage, as shown in the last two columns of Table 1.
In models with full layers, KV cache typically grows linearly with sequence length, while prefilling time exhibits quadratic growth.
However, LongGen effectively mitigates these factors. Analogous to the training phase analysis, when $S\ll N$, the cost induced by sparse attention layers becomes negligible, reducing both metrics to $L_{\text{full}}/(L_{\text{full}}+L_{\text{sparse}})$ of the original values.

With our customized triton kernels, LongGen can translate the theoretical efficiency gain to wall-clock speedup and GPU memory saving, which we will further illustrate in § 4.3.

## 4 Experiments

### 4.1 Experimental Settings

In order to evaluate LongGen across different model scales and architecture, we experiment with Llama-2-7B-base (regular multihead attention) and Llama-2-70B-base (grouped-query attention).
Following Fu et al. (), we continue pretraining them on the same SlimPajama data blend, an open-source reproduction of Llama-2’s pretraining corpus (Shen et al., ).
We curate long sequence training data by concatenating short texts
in Slimpajama to 128K length and marking the document boundary with <bos> token.
We train models on 5B tokens, which only account for 0.2% of its 2.4T pretraining corpus.
With 32 Nvidia A100-80G GPUs, the post-training of 7B and 70B models with full attention consumes 74 and 387 GPU hours, respectively.
The learning rate is 2e-5, the global batch size is 32, and we modify the RoPE base to 5M.
Each kernel block handles 64 tokens.
For LongGen with attention sink, we retain the first block for sink tokens and the most recent 32 blocks for local context.
Similarly, for LongGen with block sparse, we set the stride length as 64 blocks, meaning we only retain one context block for every 64 blocks.
This ensures that when having the same numbers of sparse layers, LongGen w/ AttnSink and LongGen w/ BlockSparse share similar KV cache budgets, which are estimated as 2K retained tokens for a 128K sequence length.
We run inference on the vLLM benchmark (^2^22vllm/benchmark/benchmark_latebcg.py) and compare efficiency accordingly.

Table: Table 2: Averaged pass rate of different methods on the needle-in-a-haystack retrieval task. All methods are based on a Llama2-7B model. Evaluation context length ranges from 512 to 128K.

Table: Table 3: Evaluation results on the first five tasks of BABILong. Columns correspond to sequence lengths, rows denote models. Each number indicates the average accuracy of the model over 5 tasks at a given sequence length, calculated over 1000 samples.

We evaluate on three long-context benchmarks—Needle-in-a-Haystack (NIAH) retrieval, BABILong (Kuratov et al., ), and RULER (Hsieh et al., ).
The latter two require models to not only retireve multiple pieces of evidence from the long context, but also reason over them; each evidence sentence is randomly placed in a long sequence.
Following BABILong paper, we used 3 different random seeds to generate 2250 test samples for each task, and average model performance on the first five tasks.
For RULER, we only evaluate models on its two QA subtasks for two reasons:
(1) We observe huge performance variance on multi-needle retrieval and other synthetic tasks. When varying random seeds, the performance can fluctuate up to 40%.
(2) Compared with other synthetic counting tasks, the two multi-hop QA tasks are more practical and close to real-world application.
For short-context evaluation, we followed Llama2 to test on MMLU (Hendrycks et al., ), Math (Hendrycks et al., ), and BigBenchHard(BBH) (Suzgun et al., ).

### 4.2 Overall Performance

#### Baselines

We evaluate LongGen against various baselines, reporting performance on NIAH and BABILong datasets in Tables 2 and 3, respectively. The comparative methods span three categories:
(1) Full Attention: Serving as an upper-bound reference, this approach lacks efficiency optimizations and is computationally expensive for long contexts.
(2) KV Cache Eviction Methods: We incorporate two widely adopted inference-time KV reduction techniques: Attention Sink (Xiao et al., ) and Heavy Hitter (H2O) (Zhang et al., ).
Additionally, we include recently proposed methods such as RazorAttention (Tang et al., ) and PyramidKV (Zhang et al., ), which capture fine-grained attention distribution patterns and adaptively allocate KV budgets.
(3) Long Context Training Methods: This category includes sliding window attention, YOCO (Sun et al., ), and our proposed approach with both patterns.
Sliding window attention at all layers is employed by Mistral for long context training, while YOCO utilizes a cross-decoder design to reuse the KV cache from earlier layers.
To ensure a fair comparison, we maintain consistent KV cache budgets across all methods by adjusting the corresponding hyperparameters.
Due to variations in implementation, KV budgets may fluctuate by approximately 5% between different methods.

#### Performance Evaluation on Long-Context

Analysis of the NIAH results in Table 2 yields two key findings:
(1) Long context training methods consistently outperform inference-time KV eviction techniques.
As discussed in §2, ensuring full access to the context window and training models to adapt to longer context lengths is crucial.
(2) Among the long context training methods, LongGen uniquely achieves perfect needle retrieval results.
In contrast to window attention, LongGen retains the entire sequence.
Compared to YOCO, LongGen more closely resembles the original transformer architecture in pretraining.
Furthermore, our method takes advantage of the information aggregation capabilities in middle layers.

To assess advanced retrieval and reasoning capabilities, we present results on BABILong in Table 3, detailing accuracy at each sequence length.
While exhibiting similar trends to NIAH, BABILong provides better differentiation between methods.
Among efficient models, LongGen demonstrates superior performance, achieving an average accuracy of 0.27.
Notably, despite achieving 100% accuracy on NIAH, LongGen still slightly underperforms full attention on more challenging tasks.
These findings underscore the importance of evaluating models on challenging reasoning tasks.

Table: Table 4: Overall long context performance of LongGen compared with full attention baseline. Evaluation context length ranges from 512 to 128K. NIAH measures the averaged correct answer rate of the models on the needle-in-a-haystack retrieval task. BABILong refers to the averaged accuracy score on the first five tasks on the BABILong benchmark. RULER reports the averaged accuracy on two multi-hop question answering subtasks from RULER.

#### Comprehensive Comparison with Full Attention

We compare LongGen-7B and LongGen-70B with their full attention counterparts, which serves as the upper bound for our hybrid sparse architecture.
Our experiments revealed no significant performance disparities between LongGen w/ AttnSink and LongGen w/ BlockSparse since their accuracies exhibited minor fluctuations across different training iterations.
For simplicity, we present the evaluation results of LongGen w/ AttnSink on both short-context and long-context benchmarks in Table 4. It is worth noting that we do not express a preference for either strategy in terms of performance.

Although our proposed models demonstrate marginally lower performance compared to the full attention models, they achieve comparable results across various context scales while offering substantial efficiency improvements. Our analysis yielded two key findings:
(1) Their performance gap is more significant in long-context scenarios compared to short-context ones.
This can be attributed to the fact that most short context is fully captured by the local blocks in LongGen, whereas information from longer distances is only sparsely represented.
(2) The performance gap between ours and the upper bound narrows as the model scale increases.
This trend may be attributed to the higher degree of information sparsity exhibited by larger models, thus highlighting the potential for applying our approach to models with even larger parameter counts.
These findings underscore LongGen’s potential for scaling to more expansive parameter spaces.

### 4.3 Efficiency Improvement

#### Training

Similar to FlashAttention (Dao et al., ), we use wall-clock training speedup to evaluate training efficiency.
We measured the latency of training a Llama2-7b model on 128K sequence length with 256 A100-80G GPUS.
We set tensor parallel size as 8 and use distributed optimizer and checkpoint activation.
Results are presented in the leftmost part of Figure 2.
By adding sparse layers, LongGen significantly improves training throughput.
Our chosen “1/3 Full” setting could achieve 1.55 times speedup in total wall-clock training time.
While increasing the sparsity level will bring us more efficiency gain, we avoid doing this in pursuit of optimal long-context performance.
We will explain more details in §4.4.

#### Inference

To verify that LongGen will bring practical system-level improvement during serving, we pair our inference kernel to be compatible with vLLM.
Using vLLM’s official benchmarking tools, we report KV cache memory saving and inference latency reduction in Figure 2.
Both memory and latency are measured on one single 128K sequence with a tensor parallel size of 4.
From the figure, our chosen “1/3 Full (12 Full layers)” setting reduces memory consumption from 69.2 GB to 26.5 GB (a 62% reduction).
Meanwhile, it saves 40% profiling time and 29 % decoding time.
The results demonstrates LongGen’s potential for efficient long-context serving.
Similarly, we do not opt for sparser architecture due to performance consideration.

Table: Table 5: Ablation on the placement of the full attention layers. We train the base Llama2-7B model with 32 layers in total on 128K context length, and report results on BABILong. We fix the full attention layer budget as one-third of all layers (12 sparse layers in total) and test varied locations of full layers. Top, middle, and bottom settings stack all full layers together while interleaving setting inserts every full layer between two sparse ones.

Table: Table 6: Ablation on the number of full layers in LongGen w/ AttnSink. The base model is Llama2-7B-128K with 32 layers in total. We locate all full attention layers in the middle of the model for the best performance. All, 1/3, 1/5, and 1/16 Full indicate using 32, 12, 6, and 2 full layers, respectively.

### 4.4 Identifying Essential Factors for Long-Context

In this section, we investigate the impact of varying the number and location of full attention layers on long-context performance.

#### Position of Full Layers

We first examine the effect of full layer positioning by maintaining a constant number of 12 full layers and altering their placement.
We explore four configurations: (1) stacking in the top layers, (2) stacking in the middle layers, (3) stacking in the bottom layers, and (4) interleaving in the middle layers.
Table 5 presents the results on the BABILong dataset, demonstrating that the position of full layers significantly influences long-context performance.
Notably, positioning full layers in the middle of the network yields optimal results.
These findings align with previous research indicating that attention heads in the middle layers play a crucial role in information aggregation.

#### Number of Full Layers

Subsequently, we investigate the relationship between the number of full layers and model performance.
While increasing the number of full layers is expected to approach the performance upper bound, we aim to identify an optimal balance between computational efficiency and accuracy.
Table 6 illustrates the performance for various numbers of full layers.
Our analysis reveals that maintaining full layers between 1/5 (6 layers) and 1/3 (12 layers) of the total network depth adequately preserves satisfactory long-context accuracy.

These empirical observations inspire the final architectural design of LongGen, which incorporates 1/3 full layers positioned in the middle of the network.

## 5 Related Work

### 5.1 Efficient Long Context Training Architecture

To overcome the quadratic complexity of transformers, several model architecture alternatives have been introduced.
With recent key modifications (Gu & Dao, ) to enhance training efficiency, state-space models emerged as they only require linear computation complexity.
Gu et al. (; ) first mitigate the training scalability issue of sequential models with hippo matrix and parallelizable operands.
Mamba (Gu & Dao, ; Dao & Gu, ) introduces input-dependent state parametrization to SSMs and optimized hardware-aware implementations.
(Park et al., ; Zuo et al., ; Ma et al., ; Ren et al., ; Glorioso et al., ; Ren et al., ) propose hybrid architectures by integrating different variations of transformer layers into SSMs.
However, as recent studies suggests, SSMs are less competitive in long context capabilities compared to transformers (Wen et al., ; Arora et al., ).
Earlier, several works attempts to inject sparsity into transformers to improve efficiency in long context settings (Tay et al., ; Child et al., ; Beltagy et al., ; Zaheer et al., ).
However, as pointed out in Dao et al. (), these algorithms can hardly bring wall-clock speed-up due to poor compatibility with accelerators.
Meanwhile, Sun et al. () revives the encoder-decoder architecture to save KV cache, where
all decoder layers consume the same KV vectors from the encoder side with cross attention.
However, it’s unclear how to use the architecture in context extension, as it diverges significantly from default decoder-only models.

### 5.2 Inference Time KV Cache Reduction

Recently, many attempts have been made to evict KV cache at inference time for faster long context serving scenarios (Miao et al., ).
Zhang et al. (); Liu et al. () proposes to discard less important KV vectors based on the accumulated attention score.
Han et al. (); Xiao et al. () find LLMs tend to store important global information into initial tokens, known as attention sink.
By keeping only the attention sink and recent tokens in cache, LLMs can maintain a reasonably well performance.
Ge et al. (); Li et al. () design algorithms to adaptively maintain KV cache with hybrid eviction policies.
PyramidKV (Zhang et al., ) find that the sparsity in attention varies between layers and proposes to dynamically allocate KV cache budget for each layer.
This line of work shows that LLMs can work with partial context at inference time.
(Li et al., ; Zhang et al., )
However, Jiang et al. (); Zhang et al. () shows that many of the methods will have significant degradation in long context tasks.
Also, it’s questionable whether complex eviction methods are compatible with realistic serving systems, as calculating accumulative attention score and releasing arbitrary memories are challenging to implement in PagedAttention and Prefix Caching  (Kwon et al., ; Zheng et al., ).

## 6 Conclusion

In this paper, we propose LongGen to finetune a pretrained LLM into an efficient architecture during length extension.
We incorporate GPU-friendly sparse attention to enable practical efficiency gain.
Additionally, we find using a hybrid architecture with 1/3 full attention layers and
2/3 efficient ones can achieve a balanced trade-off between efficiency and long-context performance.
Through light-weight training, we extend the context length of llama2-7B and 70B from 4K to 128K.
Evaluations on multiple long-context benchmarks suggest that LongGen reaches on-par performance with full attention on long-context retrieval and reasoning.
While at the same time, it reduces wall-clock training time by 36% and KV-cache by 62%, reaching 1.67x acceleration on prefilling and 1.41x on decoding stage.
Our study motivates future work on efficient transformer architectures and low-cost methods for long context extension.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 6
presentation: 3
questions: 1. How did you set up experiment parameters to ensure a fair comparison between LONGGEN and previous KV cache reduction methods, such as Attn Sink, H2O, and PyramidKV?
2. Does the sparse attention method remain consistent between training and inference?
3. When you mention savings in profiling and decoding time, is this in comparison to the full attention method? If so, could you also describe how LONGGEN's inference time compares to previous KV cache methods?
soundness: 3
strengths: Originality
	1. hybrid architecture to long-context expanding. Instead of using a full attention mechanism across all layers, the proposed LONGGEN introduces a creative combination of sparse and full attention. This hybrid method allows model generalize well on long-context tasks without high computational overhead. 
	2. The work also introduces an optimized GPU-friendly KV cache management technique, which makes long-context processing feasible and efficient on hardware.

Quality
The paper write with a good formal quality in both methodological and experimental part. Extensive evaluations are presented on benchmarks, highlighting LONGGEN’s performance benefits over alternative KV cache reduction techniques and full-attention baselines. The experimental part also includes ablations analyses such as the position and number of full attention layers.

Clarity
The clarity of the paper is good, particularly in its well explanations of the  KV cache optimization methods and hybrid attention. There are some graphs support the narrative and help show the model’s efficiency benefits.  Clear section headings and figures guide the reader through the technical details, allowing for a smooth understanding of ideas.

Significance
This paper is significant as it addresses a critical challenge in LLM scalability which is efficiently extending the context length while preserving the model’s quality and saving computational cost on hardware. It suitable for some applications that require processing large contexts, such as document analysis and long-form dialogue.
summary: The paper introduces LONGGEN, an efficient architecture designed to extend context length while minimizing computational and memory overhead during training and inference. The key contribution of LONGGEN is its innovative use of a combination of full attention and KV-reduced attention layers during the post-training phase of context length extension. This approach allows the model to adapt to sparse context scenarios, effectively addressing the poor performance observed with previous KV cache reduction methods on long-context tasks. Additionally, LONGGEN introduces static position access and block-wise context handling to mitigate issues related to position embeddings and idling threads. Empirical experiments demonstrate LONGGEN’s superior performance compared to other KV cache reduction methods, highlighting its effectiveness in managing extended contexts efficiently.
weaknesses: 1. This is not the first paper to propose post-training with a sparse attention mechanism; previous works, such as "Sparser is Faster and Less is More," have also introduced sparse attention methods during both training and inference stages.
2. Comparisons are limited to KV cache reduction methods that allocate KV budgets at the pre-filling stage. However, approaches like "Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inference" use query-level sparsity to dynamically activate the KV cache and provide full context access.
3. Additional experiments are needed to strengthen the robustness of the idea that hybrid sparse attention performs best. For a more comprehensive comparison, LONGGEN could also use H2O, RazorAttention, and PyramidKV methods to extend context length during post-training.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
questions: 1) What is the specific evaluation setting for Figure 2 (Right)? Is it tested on 4 A100 GPUs with TP=4 with vLLM? On a single A100, vLLM (W8A8) requires less than 30ms to decode a token for Llama-2-7B (64K sequence length, batch size = 1). Since the complexity of decoding stage attention grows linearly with regard to the sequence length, decoding a token with 128K context sequence length should take no more than 60ms with a single A100. However, in Figure 2 (Right), it takes 60 seconds to decode 512 token (~117 ms/token) with 4 GPUs (dense baseline with 32 full layers). It would be helpful if the authors can provide more details about this evaluation.

2) What is the kernel-level speedup achieved with the sparsity pattern used in LongGen for the attention kernel? For instance, given the 1/64 sparsity (retain 2K tokens in 128K), what is the speed comparison between dense flash attention and the sparse attention kernel used in LongGen? Is it possible for LongGen to achieve measured speedups when serving (i.e., run inference with) sequences shorter than 128K?

3) How is the proposed method compare with other existing KV cache elimination methods (e.g., [6](https://arxiv.org/abs/2310.01801))?
soundness: 2
strengths: 1) LongGen improves both training and inference efficiency for long-context LLMs, accelerating training by 1.55x and inference by 1.41-1.67x, without prominent accuracy loss.
2) LongGen identifies that keeping middle layers with full attention and applying sparse attention on the beginning and final layers achieves better performance. It also finds that keeping 1/3 layers with full attention achieves a balance between accuracy and efficiency.
3) LongGen integrates a triton-based attention kernel supporting structured sparsity for efficient inference and training.
summary: This paper introduces LongGen, which improves both training and inference efficiency of long-context LLMs. The key insights in this paper are: 1) structured sparse attention patterns, with GPU-friendly memory access patterns, enable practical efficiency gains for long-context LLM; 2) a hybrid architecture with full attention layers and sparse ones provides efficiency merits while keeping the models’ performance; 3) context-extension training only requires a lightweight training dataset. Experiments show that LongGen reduces training wall-clock time by 36%. It also reduces KV cache memory by 62% during inference, achieving 1.67x prefilling speedup and 1.41x decoding speedup.
weaknesses: 1) The key insights proposed in this paper have been introduced in other papers. For instance, attention-sink and block-sparse attention are not new attention patterns (e.g., [1](https://arxiv.org/abs/2309.17453), [2](https://arxiv.org/abs/2407.02490)). Additionally, extending the model's context length can be achieved with light-weight training is also observed in existing literature (e.g., [3](https://arxiv.org/abs/2306.15595), [4](https://arxiv.org/abs/2307.03170), [5](https://arxiv.org/abs/2309.12307)). The novelty of this paper is limited.

2) The evaluation details of the efficiency benchmark should be further elaborated.

3) The models used for main evaluation is old. Results on Llama-3 series would be more persuasive.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: 1) What are the results of LONGGEN  on other popular large language models, such as GPT models or Gemini?

2) How would the results of LONGGEN change if applied to tasks other than retrieval or reasoning ones?

3) Should the architecture, specifically the number and placement of full attention layers, be adjusted for different models or tasks?

4) Is there an algorithmic or automated method for determining the optimal number and placement of full-attention layers?
soundness: 3
strengths: 1) Problem importance: LONGGEN addresses the critical problem of context window extension in LLMs, specifically targeting issues of performance degradation, high computational complexity, and excessive memory consumption as context length increases, which is an important challenge for enabling LLMs to handle long-form content efficiently in real-world applications.

2) Training and inference efficiency: The LONGGEN approach enhances both training and inference efficiency by reducing training FLOPs, KV cache memory size, prefilling speed, and decoding speed. By employing sparse attention in the outer layers and full attention in the middle, LONGGEN effectively manages computational load and memory usage. This design allows for context length extension up to 128K tokens with lightweight fine-tuning on long-context data, making it suitable for long-context applications. 

3) Results: LONGGEN demonstrates almost the same performance as the full-attention model across key benchmarks, including Needle-in-a-Haystack (NIAH), BABILong, RULER, MMLU, Math, and BigBenchHard (BBH). Additionally, it outperforms other KV cache eviction methods and long-context training approaches on the NIAH and BABILong tasks.
summary: This paper presents LONGGEN, a method designed to enhance the efficiency of long-context training and inference in large language models (LLMs). LONGGEN employs a hybrid attention architecture to achieve an optimal balance between computational efficiency and performance in long-context tasks. The model architecture is segmented into three sections: the first and last thirds of the layers utilize sparse attention mechanisms, while the middle third maintains full attention.
weaknesses: 1) Results of the other models: The results of LONGGEN have only been demonstrated on Llama2-7B and Llama2-70B models, which limits understanding of its effectiveness on other model architectures and sizes (such as GPT models, Gemini, or other Llama models).

2) Results on the other tasks: LONGGEN introduces an hourglass architecture that keeps the middle layers in full attention mode, based on previous studies [1, 2, 3] showing that attention heads are crucial for retrieval and reasoning tasks. However, its performance on other long-context benchmarks has not been explored in this work (such as single/multi-document QA or Summarization).

3) Full attention layers:  The specific selection of the full attention layer has not been explored, which can differ from task to task (or model to model). Additionally, although the authors have conducted an ablation study to determine that 1/3 of the layers should use full attention, this proportion may vary across different models or different tasks and would require separate ablation studies for each.

[1] Wenhao Wu, Yizhong Wang, Guangxuan Xiao, Hao Peng, and Yao Fu. Retrieval head mechanistically explains long-context factuality. arXiv preprint arXiv:2404.15574, 2024.

[2] Yao Fu. How do language models put attention weights over long context. Yao FuâA˘Zs Notion ´ , 2024.

[3] Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark W. Barrett, Zhangyang Wang, and Beidi Chen. H2O: heavy-hitter oracle for efficient generative inference of large language models. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10 - 16, 2023, 2023. URL http://papers.nips.cc/paper_files/paper/2023/hash/6ceefa7b15572587b78ecfcebb2827f8-Abstract-Conference.html.

---

confidence: 4
contribution: 2
initial_rating: 8
presentation: 3
questions: In the experiment setup, the author mentions that the tensor parallel size is set to 8 with 256 GPUs. Are the remaining GPUs used for data parallelism or pipeline parallelism? And what is the framework used to measured the speedup? And how many iterations to calculate the average training/inference time?
soundness: 3
strengths: 1. The paper is well written: especially abstract and introduction is well structured and informative on what the paper is going to about. The figures are well made.
2. The performance is very good: e.g. NIAH result is much better than previous methods such as StreamingLLM.
summary: This paper proposes LongGen, which integrates GPU-friendly KV cache reduction architecture to save both length extrapolation and serving cost. It is built on three observations, on sparse attention and the number of tokens needed. It achieves effective cost reduction in both training and serving cost.
weaknesses: There is no noticeable weaknesses that the reviewer hope the authors shall address (only some small clarification questions). Please see the question section.

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 3
questions: See weakness
soundness: 3
strengths: - The work highlights the importance of including full attention layers for models to achieve accurate retrieval.
- The paper is well-written.
- Experimental results indicate that LongGen accelerates both training and inference for long context while preserving model performance.
summary: The authors finetunes a pretrained LLM into an hybrid architecture that consists 1/3 full attention layers and 2/3 sparse attention layers. By incorporating full attention layers, LongGen allows the model to access to certain positions directly, enhancing the performance on accurate retrieval tasks. Experimental results show that LongGen incurs no loss on the needle-in-a-haystack retrieval task and maintains model performance on tasks with short context, such as MMLU, demonstrating its effectiveness.
weaknesses: - Since training on long context constitutes a small portion of pre-training, the training speedup of LongGen is limited during pre-training.
- LongGen with AtteSink and BlockSparse demonstrates similar performance, necessitating a detailed explanation for this observation.

Meta-review: The paper introduces an efficient architecture, LongGen, designed to extend context length while minimizing computational and memory overhead during training and inference. Essentially, LongGen combines full attention and KV-reduced attention layers during the post-training phase of context length extension. This approach allows the model to adapt to sparse context scenarios, effectively addressing the poor performance observed with previous KV cache reduction methods on long-context tasks. Additionally, LONGGEN introduces static position access and block-wise context handling to mitigate issues related to position embeddings and idling threads. Empirical experiments demonstrate LONGGEN’s superior performance compared to other KV cache reduction methods.

Although this paper combines many of the existing techniques, the authors did a solid execution and show decent performance.

Decision: Accept (Poster)