A Little Goes a Long Way: Efficient Long Context Training and Inference with Partial Contexts
Authors: Suyu Ge
Xihui Lin
Yunan Zhang
Jiawei Han
Hao Peng
Venue: ICLR 2025 Poster (2025)
Keywords: Long-Context LLM
Efficient LLM
Context Extension
KV Cache Reduction
Primary area: foundation or frontier models, including LLMs
TL;DR: This paper argues that integrating length extension with a GPU-friendly KV cache reduction architecture not only reduces training and inference overhead, but also achieve better long-context performance.

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.

Published as a conference paper at ICLR 2025
A LITTLE GOES A LONG WAY: EFFICIENT LONG CONTEXT TRAINING AND INFERENCE WITH PARTIAL CONTEXTS
Suyu Gihui Lin Yunan Zhang" Jiawei Han', Hao Peng' l University of Illinois Urbana-Chane Ohaana-Champaign Microsoft (suyuge2, haopeng) @ iniiiiis s. edu
ABSTRACT
Training and serving long-context large language models (LLMs) incurs substantial overhead. To address this, two critical steps are often required: a pretrained L.M typically undergoes 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- long-context perfor- mance. 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 have direct access to all tokens. A hybrid architecture with 1/3 full attention layers and 2/3 efficient ones achieves 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 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 1.55x training speedup and reduces wall-clock time by 36%, compared to full-attention baseline. During inference, LONGGEN reduces KV cache memory by 62%, achieving 1.67x prefilling speedup and 1.41x decoding speedup. Compared
baselines that apply KV-cache reduction techniques full-attention long-come oong-contett LLMs, LONGGEN achieves substantially stronger per- formance not only on the Weedle-in-a-Hasstre retrieval task, but also on more challenging long-context reasoning tasks, including BABILong and RULER.
INTRODUCTION
Transformer-hased large language models (LLLMs) capable processing long contexts have unlocked new opportunities across wide range of applications, such 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. [024bb Sun et al. 2024a). To address this,
typical approaches involve two-stage framework (Xiong et al. 2023 Dubey al. 2024
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 smaller amount of long-context data (e.g., 5B of 128K tokens).
Authors contrihuted equally to this research.

---

Published as a conference paper at ICLR 2025
full
context in
middle layers
KV-Reduced Attention
Full Attention
KV-Reduced Attention
LONGGEN AttnSink
LONGGEN w BlockSparse
context
middle
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 V-reduced attention variants are explored.
Wall-Clock Training Speedup
Memory Reduction
End-to-End Inference on VLLM
Tokens Second
092
278354
LongSen
2
60
70
FlachAttentionn22
LongGen
69.2
Latency (secondg)
80
70
Prefiling Latency (scc/128K tekens)
Decoding (sec 12
Memory (GB)
250K
252715
80
40
50
30
26.5
200K
15.9
40
163201
10:
11.7
30
150K
20
FUll Full Full
Full Layers
5.7
12
Layers
Layers
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- 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.
Serving long-context L.M is also challenging due to high memory requirements for caching key and value vectors (KV cache). For example, a 7B Ilama-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-con contect 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.
Our proposed approach LONGGEN, as illustrated Figure 1 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 window attention (Jiang et al. 2023), attention sink (Xiao et al.) 2024), and blockwise strided attention (Zaheer et al.) (020b). 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 2023). but achieves faster speed in sparse settings, and the inference kernel is well-suited to vLLM (won et al. (2023) for this sparse attention in 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 long-con context information (4.4)
serving. (2) LONGGEN uses
Experiments are conducted on Llama-2-77 base model Touvron et al. 2023) and its 70B counterpart with group query attention. We extend their context length from 4K to 128K and perform evaluations
2

---

Published as a conference paper at ICLR 2025
on multiple long-context benchmarks, including the needle-in-a-hayse (NIAH) retrieval, and two more distinguishing benchmarks, BABILong (Kuratov et (0224) and RULER (Hsich et al. (024). where long-contens reasoning is also needed. Through ablations the position and number of sparse attention layers, we study how sparsity level affects long-conters 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. 2023) tokens (less than 0.1% GPU hours of pretraining) and maintains its short context capability, as indicated by the MMLU (Hendrycks) al., 2020) score. With our customized kernel, it achieves 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.4 xx 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-traini stage. State-of-the-art models, such as Gemini (Team et al. 2023). Llama series TTuvron et al. 2023 Dubey et al. 2024) and Qwen (Bai et al., 2023) are typically pretrained on large short-sequence corpora, then undergo length extension relatively smaller amounts of longer sequences. For example, Llama-3 is pre-trained on 15T tokens within an 8K context length, then post- tostttrained on an additional 800B tokens to extend 128K length (Dubey et al. 2024 Xiong al., 2023). A related data engineering work (Fu et al., (024) significantly reduce the tokens needed for context extension to only 5B by carefully balancing the data source 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. 2024a). Therefore, techniques are required to reduce the inference- -ime memory overhead of these models.
LLoMN>2-78-120K
LLaMMh-278--2206 of Atrršink LLeMM-2-78-1288 PyremidKV
L.7 LongGen
Depth Parcert
20
0.8
40
0,6
60
D4
80
02
100
32k
64k
128k
32%
64k
128k
64k
128k
64k
128k.
0.0
Context Length
Context Length
Context Length
Length
Figure 3: Inference time KV cache reduction methods fail on long context.
Inference-time KV reduction fails to generalize to long contexts. 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 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.) 2024] Liu et al. 2023) Ge et al.) 2024). However, as we show here, they often underperform in long-contess scenarios, where KV cache reduction is needed the most.
We evaluate established KV cache reduction methods, Attention Sink (Xiao al. 2024). and PyramidKV V (Zhang et al. 2024a) on a context-extended Llama2-7B- 18K model Attention Sink (AttnSink) only retains initial and local tokens, while PyramidKV identifies the unique attention dis- tribution at each layer and allocates a dynamic KV budget. Similarly to Heavy-Hitter (H20) (Zhang al. 2023). PyramidKV preserves tokens that have the largest cumulative attention scores. Ad- ditionally, we include the original Llama2-78 1-ama2-7B-1n 2---2----1288 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, V, and LONGGEN, i.e., 60% of the original full transformer: Figure shows the results on the needle-in-a-haystade (NIAH) retrieval task. Both the original Llama2-7B. - base model and
Details about this Llama2-BB-128K model can be found in 4.1
3

---

Published as a conference paper at ICLR 2025
LONGGEN can pass the needle retrieval task with 100% accuracy. which will analyze later in the experiment section $4.2 However, both KV eviction methods underperform on NIAH, especially with > 32K 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, has been observed that the widely used RoPE-style position embedding (Su et al., [021) struggles to generalize to unseen position ranges (Wang et al. 2024). After KV reduction, the position range becomes non-contiguous, creating 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 's key design choices: (1) retaining full context at some layers, and (2) training the model adapt sparse attention 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 D We will explore various strategies to initialize the sparse attention, followed by a comprehensive analysis of FLOPS and memory savings demonstrate the theoretical efficiency benefits of LONGGEN.
3.1 CONTEXT EXTENSION WITH EFFICIENT ARCHITECTURES
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 al. (2024): Fu (2024); Zhang et al.) (2024b). Inspired by their findings, we design LONGGEN to be in an hourglass shape sparse attention in both ends 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 kemel and memory IO. During inference, it is also feasible enforce the fixed cache management strategy, thus avoiding extra computation and tensor replacement. From both sides, static attention patterns lead 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 [ Our first variant is Attention Sink, which combines initial tokens with local blocks of context window (Xiao et al.) 2024). This technique builds upon the proven effectiveness of sliding window attention in long context training, as demonstrated by Mistral (Jiang et al. 2023). 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. 2020al Qiu et al. 2019).
4

---

Published as a conference paper at ICLR 2025
Training FLOPs
KV Cache Mem Inference Prefilling Time
Full Attention
(L D + D))
(L
(LLual (N ND9))
Sparse Attention in LONGGEN
O Lsparee (SD2 + D))
OLLapassS
O (52 SD?))
Table 1: Efficiency comparison between conventional full attention and the sparse attention in LONGGEN. Training FLOPs count both attention forward and backward operations. Luall and Lpuriall are the number 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 < N.
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 improve transformer efficiency, such as flash attention. In fact, we build our customized kernel upon flash attention. Similarly to FlashAttention (Dao et al. 2022). LONGGEN takes an attention mask matrix as input. The attention mask is represented in compressed sparse format, which consumes ((N) memory instead of the original ( 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, skip computing position by avoiding loading it into HBM its mask value zero. This implementation enables FL.OPs 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 Fltentntention, probably because of reduced SRAM usage per loading cycle. For the attention sink initialization. follow Mistral to use the rolling buffer cache for local sliding window tokens (Jiang et al. 2023
3.3 THEORETICAL TRAINING AND INFERENCE ADVANTAGE
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 substantial reduction in training FLOOPs. Table [ demonstttttt that when the context length N significantly exceeds the block size S, the FLOPs required for sparse attention layers become negligible (S < N). Consequently, LONGGEN effectively reduces the total FLOPs to fraction of Lia (Laill + Laparse) compared to the original transformer model. Empirical studies indicate that ratio of Liull Lsparse = 1 2 provides 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 [ 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 < N, the cost induced by sparse attention layers becomes negligible, reducing both metrics to Lian/ (Lruil + Lsparse) 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 will further illustrate in 4.3
EXPERIMENTS
4.1 EXPERIMENTAL SETTINGS
In order to evaluate LONGGEN across different model scales and architecture, experiment with Llama- Hama-2--B- 2-7B base (regular multihead attention) and Llama-2-70B 2--00------ base (grouped-query attention). Following Fu et al. (2024) continue pretraining them the same SlimPajama data blend, open-
5

---

Published as a conference paper at ICLR 2025
KV Cache Eviction Methods
Attn Sink
H2O
RazorAttention
PyramidKV
NIAH Pass Rate
28%
32%
46%
51%
Long-contedd
LONGGEN w/ AttnSink
LONGGEN w/ BlockSparse
Sliding Window
YOCO
Training Methods
NIAH Pass Rate
48%
88%
100%
100%
Table 2: Averaged pass rate of different methods on the needle-in-a -aaee-ii---hystackk retrieval task. All
methods are based on a Llama2-7B model. Evaluation context length ranges from 512 to 128K.
Full Attention
512
1k
2k
4k
8k
16k
32k
64k
70k
80k
90k 100k
100
AVG 0.29
Full
0.39
0.37
0.36
0.37
0.34
0.36
0.26
0.24
0.2
0.21
0.22
0.22
0.19
KV Cache Eviction Methods
Attn Sink
0.36 0.34 0.34
0.32
0.28
0.22
0.15
0.1 0.08 0.06 0. 0.06
0.07 0.05
0.05
0.1
H2O
0.37 0.36
0.35
0.36
0.28
0.24
0.15 0.13 0 09 0. 0.07 0
0.06
0.07
0.06
0.05
0.18
RazorAttention
0.37 0.35 0.35
0.35
0.29
0.25
0.2 0.17 0.16 0.1 0.12
0.1 0.09
0.06
0.07
0.2
PyramidKV
0.38 0.36 0.35
0.36
0.28
0.26 0.21 0.19 0.160.13 0.1 0.12 0.07
0.08
0.21
Long Context Training Methods
Sliding Window
0.36 0.34 0.33
0.34
0.31
0.26 0.19 0.13 0.13 0.11 0.1 0.11 0.09 0.
0 00
0.2
YOCO
0.37 0.36 0.37 0.33
0.28
0.24
0.22 0.17 0.18 0.19 0.16 0.16 0.0.180.1144 0.14
0.23
LONGGEN AttnSink
0.41 0.38 0.36
0.37
0.3
0.3 0.25
0.23
0.22
0.2 0.22
0.2 0.2
0.18 0.17
0.27
LONGGEN BlockSpurse
0.39 0,36 0.36 0.34
0.31 0.3 0.25 0.24 0.23 0.21 0.21 0.21 0.22 0.2 0.18 0.27
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 tasks at given sequence length, calculated over 1000 samples.
source reproduction of Llama-2's pretraining corpus (Shen et al. 2023 training data by concatenating short texts 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 32, and we modify the RoPE base to 5M. Each kernel block handles 64 tokens. For LONGGEN with attention sink, 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 only retain 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 128K sequence length. We run inference on the vLLM benchmark and compare efficiency accordingly.
We curate long sequence
We evaluate on three long-context Long (Kuratov et al. 2024), and RULER (Hsieh et al.) 2024). 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 long sequence. Following BABILong paper, we used 3 different random seeds to generate 2250 test samples for each task, and average model perfor- mance on the first five tasks. For RULER, 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
(NIAH) retrieval, BABI-
fluctuate up 40%. (2) Compared with other synthetic counting tasks, the two multi-hop hop QA tasks are more practical and close to real- world application. For short-context evaluation,
followed Llama2 test MMLU (Hendrycks et al.)
2020). Math (Hendrycks et al. 202T), and BigBenchHard((BB1)) (Suzgun et al. 2022).

---

Published as a conference paper at ICLR 2025
NIAH
BABiLong
RULER
MMLU
Math
BBH
Full Attention-7B
100%
0.29
0.41
0.419
0.12
0.32
LONGGEN-7B
100%
0.27
0.38
0.415
0.13
0.31
Full Attention-70B
100%
0.46
0.67
0.661
0.32
0.49
LONGGEN-70B 100%
0.46
0.65
0.658
0.32
0.49
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 the soadl--in-a-haystace 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.
4.2 OVERALL PERFORMANCE
Baselines We evaluate LONGGEN against various baselines, reporting performance on NIAH and
BABILong datasets in Tables 2 and 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., 2024) and Heavy Hitter (H20) (Zhang et al. (022). Additionally, we include recently proposed methods such as RazorAttentin (Tang et al., 2024) and PyramidKV (Zhang et al., (024a)). 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.) (022bb), and our proposed approach with both patterns. Sliding window attention at all layers is employed by Mistral for long context training, while YOCO utilizes cross-decoder design reuse the KV cache from earlier layers. To ensure fair comparison, we maintain consistent KV cache budgets across all methods by adjusting the corresponding hyperparameters. Due variations in implementation, K. 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 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 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 challenging reasoning tasks.
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/ AutSiink and LONGGEN w/ BlockSparse since their accuracies exhibited minor fluctuations across different training iterations. For simplicity, present the evaluation results of LONGGEN w/ AttnSink on both short-context and long-context benchmarks Table 4 is worth noting that we do not express 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
7

---

Published as a conference paper at ICLR 2025
Full Layers Position
1k
2k
32k
64k
70k
80k
90k
100k
110k
120k
128k
AVG
Top
0.39
0.38
0.35
0.31
0.26
0.22
0.16
0.08
0.12
0.09
0.09
0.09
0.08
0.08
08
0.07
0.18
(LONGGRN)
0.41
0.38
0.36
0.37
0.3
0.25
0.23
0.22
0.2
0.22
0.2
0.2
0.18
0.17
0.27
Bottom
0.38
0.35
0.32
0.27
0.25
0.22
0.2
0.14
0.14
0.13
11
0.12
0.11
0.13
0.11
0.2
Interleaving
0.35 0.33 0.33 0.31 0.29 0.25 0.2 0.21 0.18 0. 0.11 0.18
0.
0.15
0.24
Table 5: Ablation 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 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.
Full Layers Num
70k
80k
AVG
All Full
0.39
0.37
0.36
0.37
0.34
0.36
0.3
0.26
0.24
0.25
0.2
0.21
0.22
0.22
0.19
0.29
1/3 Full
0.41 0.38
0.36
0.37
0.23
0.22
0.2 0.22 0.2
0.17
0.27
1/5 Full
0.39 0.37 0.36
0.36
0.31
0.29
0.24
0.22
0.21
0.21
0.2
0.18
0.17
0.17
0.15
0.26
1/16 Full
0.38 0.37 0.34 0.33 0.32
0.27 0.21 0.15 0.14 0.12 0.11 0.13 0. 12 0.11
0.09 0.21
Table 6: Ablation on the number of full layers in LONGGEN w/ AttnSink. The base model is Llama2-7B-122 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.
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 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 more expansive parameter spaces.
4.3 EFFICIENCY IMPROVEMENT
Training Similar to FlashAttention (Da et al. 2022), we use wall-clock training speedup to evaluate training efficiency. We measured the latency of training a Llama2-7b model 128K sequence length with 256 A 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 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 $4.4
Inference To verify that LONGGEN will bring practical system-level improvement during serving, pair inference kernel to be compatible with vLLM. Using vLLM's official benchmarking tools, we report KV cache memory saving and inference latency reduction Figure 2 Both memory and latency measured figure, our chosen "1/3 Full (12 Full layers)" setting reduces memory consumption from 69.2 GB to 26.5 GB 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 performance consideration.
single 128K sequence with tensor parallel size of 4. From the
4.4 IDENTIFYING ESSENTIAL FACTORS FOR CONG-COTE
In this section, we investigate the impact of varying the number and location of full attention layers on long-context performance.
8

---

Published as a conference paper at ICLR 2025
Position of Full Layers We first examine the effect of full layer positioning by maintaining constant number of 12 full layers and altering their placement. We explore four configurations: (I) 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 2023) to enhance training efficiency, state-space models emerged as they only require linear computation complexity. [uu et al. (2020 2021)
first mitigate the training scalability issue of sequential models with hippo matrix and paralle iizblle
operands. Mamba (Gu & Dao 2023 Dao & Gu 2024) introduces input-dependent state parametriza- tion to SSMs and optimized aardware-aware imp emeetttioos (Park et al. 2024 Zuo et al. 2022 Ma
et al. 2023 Ren et al.) 2023 Glorioso et al. 2024 Ren et al.) 2024) 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. 2024
Arora et al., 2024). Earlier, several works attempts inject sparsity into transformers improve efficiency in long context settings (Tay et al. 2023 Child et al. 2019)] [ellagy et al.) 2020 Zaheer in Dao et al. (2022). these algorithms can hardly bring
al., 2020a). However, as pointed wall-clock speed-up due to poor compatibility with accelerators. Meanwhile, [Sun et al. (2024b) revives the -ncoder-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 context extension, as it diverges significantly from default decoder-only models.
5.2 INFERENCE TIME KV CACHE REDUCTION
Recently, many attempts have been made evict KV cache at inference time for faster long context serving scenarios (Miao et al., 2023). Zhang et al. (023) Liu et al.] (2023) proposes to discard less important KV vectors based the accumulated attention score. Han et al. (2023); Xiao al. (2024) 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. (202)); Li et al. (2024) design algorithms to adaptively maintain KV cache with hybrid eviction policies. PyramidKy V (Zhang al. 2024b) 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. 20224] Zhang et al. 2024b) However, Jiang et al. (2024a): Zhang et al. (2024b) shows that many of the methods will have significant degradation in long context tasks. Also, it's questionable whether complex eviction methods compatible with realistic serving systems, as calculating accumulative attention score and releasing arbitrary memories are challenging to implement in PagedAttention and Prefix Caching
(K et al.
2023 Zheng al. 2023

---

Published as a conference paper at ICLR 2025
6 CONCLUSION
In this paper, we propose LONGGEN to finetune a pretrained LLM into an efficient architecture during length extension. We incorporate GPU- UPU-friendly sparse attention enable practical efficiency gain. Additionally, we find using 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 Ilama2-7B and 70B from 4K to 128K. Evaluations on multiple long-contexs benchmarks suggest that LONGGEN reaches on par performance with full attention 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. 1.67x acceleration on prefilling and 1.41x on decoding stage. Our study motivates future work efficient transformer architectures and low-cost methods for long context extension.
REFERENCES
Simran Arora, Sabri Eyuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, James Zou, Atri Rudra, and Christopher Ré. Simple linear attention language models balance the recall-throughput tradeoff. In Forty-first International Conference on Machine Learning, ICML 2024. Vienna, Austria. July 21-27, 2024. OpenReview.net, 2024. URL e e93f FDcpH3
net forum
Jinze Bai, Shuai Bai, Yunfei Chu, Zeyu Cui, Kai Dang, Xiaodong Deng, Yang Fan, Wenbin Ge, Yu Han, Fei Huang, et al. Qwen technical report. arXiv preprint arXiv:2309 16609, 2023.
Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long-document transformer. CoRR, abs/2004.05150, 2020. URL
Rewon Child, Scott Gray, Alec Radford, and llya Sutskever. lonerating sequence with sparse transformers. CoRR, abs/1904 10509, 2019. URL http://////a.....////////00.
Tri Dao. Flashattention Faster attention with better parallelism and work partitioning. arXiv preprint arXivv2307.08661, 2023.
Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. In Forty-first International Conference Machime Learning. ICML 2024, Vienna, Austria, July 27, 2024. OpenReview.net, 2024. URL net htt/:/oormmnd22m2mmmm
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. In Sanmi Koyejo, S. Mo- hamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Process- ing Systems 2022, NeurlPS 2022, New Orleans, LA, USA, November 28 December 9, 2022, 2022. URL http::/ppeeerrs......ppppppeeeep httpappper..llleeppprr///2 31d57722222200lla733220......... hat.raaeeeeeeeettttttttooo-orrrr Conference.ecccce. html
nips cc wwtp:ippppppppppeeaeeeeeppppeeee
/hash/
Abs
Abhimanyu Dubey, Abhinay Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, al. The llama 3 herd of models. arXiv preprint arXiv:2407.2178. 2024.
Yao Fu. How do language models put attention weights over long context. Yao FuâÄžs Notion, 2024.
Yao Fu, Rameswar Panda, Xinyao Niu, Xiang Yue, Hannaneh Hajishirzi, Yoon Kim, and Hao Peng. Data engineering for scaling language models 128k context. arXiv preprint 100v:2442.10171,, 2024.
Suyu Ge, Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, and Jianfeng Gao. Model tells you what discard: Adaptive KV cache compression for llms. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna. Austria, May 7-11, 2024. OpenReview.net, 2024.
URL httpt::t////t..... http:hts:///pppeeiiiiiitoo / forum?
F o
10

---

Published as a conference paper at ICLR 2025
Paolo Glorioso, Quentin Anthony. Yury Tokpanov, James Whittington, Jonathan Pilault, Adam Ibrahim, and Beren Millidge. Zambu: A compact 7b SSM hybrid model. CoRR, abs/2405.16712, 2024. doi: 1048500ARXIV.2455.6600
URL ottps:://o...rg//1... 48550//r0085555511111
2405. 6712
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. CoRR, abs/2312.007.52, 2023. doi: 10448550/ARXIIV2312.00777 URL /ttpp:://oi..nnn org/100
Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. Hippo: Recurrent memory with optimal polynomial projections. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information P Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020. URL http:/ttttpptteeeeeeeeee www.riiiinnnnnnnnnneoeeemm nwwrips. www.nnuwww...ppppeeeeeee 2020 hta0//aah//aa 102 www.000beeee3e622222333333
wwweeebeebbbtttttteeeee html
Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré. Combining recurrent, convolutional, and continuous-timg models with linear state space layers. In Marc' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.), Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pp. 572 585, 2021. URL http:tts///rcccccciiiooo neurips CC/ ct/ htcp/apper/222//aaa www.--bsssttt.a.
Chi Han, Qifan Wang, Wenhan Xiong. Yu Chen, Heng Ji, and Sinong Wang. Lm-infinite: Simple on-the-fly length generalization for large language models. CoRR, 2t0///300.1637,, 2023. doi: 10. 48550/ARXIV.2308.666 16137. /RLLltttps///di.orgg 088.:/....0.0...........
20t.044455////...383303 2308. 16137
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. arXiv preprint arXiv:2009.03300 2020.
Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. arXiv preprint arXiv:2103.03882 2021.
Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, and Boris Ginsburg. Ruler: What's the real context size of your long-context language models? arXiv preprint ar 24V::240000066544 2024.
Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, et al. Mistral 7b. arXiv preprint 1:Xiv:2310.06825, 2023.
Huiqiang Jiang, Yucheng Li, Chengruidong Zhang, Qianhui Wu, Xufang Luo, Surin Ahn, Zhenhua Han, Amir H. Abdi, Dongsheng Li, Chin- Yew Lin, Yuqing Yang, and Lili Qiu. Minference 1.0: Accelerating pre-filling for long-context Ilms via dynamic sparse attention. CoRR, abs/2407.02490. 2024a. doi: 00-48550/ARXIV.2400 200000000000 URL 2407. 02490
Huiqiang Jiang, Yucheng Li, Chengruidong Zhang, Qianhui Wu, Xufang Luo, Surin Ahn, Zhenhua Han, Amir H Abdi, Dongsheng Li, Chin- Yew Lin, et al. Minference 1.0: Accelerating pre-filling for long-context Ilms via dynamic sparse attention. arXiv preprint arXiv:2407.024990 2024b.
Yuri Kuratov, Aydar Bulatov, Petr Anokhin, Ivan Rodkin, Dmitry Sorokin, Artyom Sorokin, and Mikhail Burtsev. Babilong: Testing the limits of Ilms with long context reasoning-in----aastacck arXiv preprint 2rXiv:2406.10144 2024.
Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention. In Jason Flinn, Margo I. Seltzer, Peter Druschel, Antoine Kaufmann, and Jonathan Mace (eds.), Proceedings of the 29th Symposium on Operating Systems Principles, SOSP 2023, Koblenz, Germany, October 23 26-26, 2023, pp. 61 626. ACM, 2023, doi: 10. 145/ 3613165
3600006.3611165. URL. URR.https::0 org 10. 110.

---

Published as a conference paper at ICLR 2025
Yuhong Li, Yingbing Huang, Bowen Yang. Bharat Venkitesh, Acyr Locatelli, Hanchen Ye, Tianle Cai, Patrick Lewis, and Deming Chen. Snapkv: LLM knows what you are looking for before generation. CoRR, abs/2404 2024. doi: 10448500ARXIV
URL https:
10lolloogg/.....5///aaaaa.. 40854885550iii 240.://2xxii........... 2404. 14469)
Zichang Liu, Aditya Desai, Fangshuo Liao, Weitao Wang, Victor Xie, Zhaozhuo Xu, Anas- tasios Kyrillidis, and Anshumali Shrivastava. Scissorhands: Exploiting the persistence of importance hypothesis for LLM KV cache compression at test time. In Alice Oh, Tris- tan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Ad- vances in Neural Information Processing Systems 36: Annual Conference on Neural Infor- mation Processing Systems 2023, NeurlPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023. URL /papstsssppppprr.... nips hwtpp/ppee..llleeeppppeeeeeeee 3 htw.ee/eeeeeeee--ooooaoooooo--ee
html
Xuezhe Ma, Chunting Zhou, Xiang Kong, Junxian He, Liangke Gui, Graham Neubig, Jonathan May, and Luke Zettlemoyer. Mega: Moving average equipped gated attention. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023. URL https://openzeview. net www.wooooooieeeeeeeeeieeoe 3iq2E1
Xupeng Miao, Gabricle Oliaro, Zhihao Zhang. Xinhao Cheng, Hongyi Jin, Tianqi Chen, and Zhihao Jia. Towards efficient generative large language model serving: A survey from algorithms to systems. CORR, http//2312.15222. 2023. doi: 10:48550/ARXIV.232..5555
URL https
//doi. 48550/ 2312. 15234
Jongho Park, Jaeseung Park, Zheyang Xiong, Nayoung Lee, Jaewoong Cho, Samet Oymak, Kangwook Lee, and Dimitris Papailiopoulos. Can mamba learn how to learn? A com- parative study on in-context learning tasks. In Forty-first International Conference on Ma- chine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net, 2024. URL net www.ooiiddooooo. E
https://openrett
Jiezhong Qiu, Hao Ma, Omer Levy, Scott Wen-tau Yih, Sinong Wang, and Jie Tang. Blockwise self-attention for long document understanding. arXiv preprint arXiv:1911 191.00972,, 2019.
Liliang Ren, Yang Liu, Shuohang Wang, Yichong Xu, Chenguang Zhu, and ChengXiang Zhai. Sparse modular activation for efficient sequence modeling. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Ad- vances in Neural Information Processing Systems 36: Annual Conference on Neural Infor- mation Processing Systems 2023, NeurIPS 2023, New Orleans. LA, USA, December 10 16, 2023, 2023. http:/tt:///ppeeee nips http:/ccepppppeeeeeeeeeee fila htt.ffll/leappppp22.233. hashh oag.3330h0..ccccccaaaaaaaaaaaaa.
Liliang Ren, Yang Liu, Yadong Lu, Yelong Shen, Chen Liang, and Weizhu Chen. Samba: Simple hy- brid state space models for efficient unlimited context language modeling. CoRR, abs/2406.07522. 2024. doi: 00448550/AXXIV.24000777752 1406607522
URL https: org/10, 48550/arXiv
Zhiqiang Shen, Tianhua Tao, Liqun Ma. Willie Neiswanger, Zhengzhong Liu, Hongyi Wang, Bowen Tan, Joel Hestness, Natalia Vassilieva, Daria Soboleva, et al. Slimpajama-de: Understanding data combinations for Ilm training. arXiv preprint arX9v: 10818 2023.
Jianlin Su, Yu Lu, Shengfeng Pan, Bo Wen, and Yunfeng Liu. Roformer: enhanced transformer with rotary position embedding. corr (0bs/2104.00l 09864 (2021). arXiv preprint arXiv:2104 09864, 2021.
Yutao Sun, Li Dong, Yi Zhu, Shaohan Huang, Wenhui Wang, Shuming Ma. Quanlu Zhang, Jianyong Wang, and Furu Wei. You only cache once: Decoder-decoder architectures for language models. arXiv preprint ar arXiv:240s 245505544 2024a.
Yutao Sun, Li Dong, Yi Zhu, Shaohan Huang, Wenhui Wang, Shuming Ma, Quanlu Zhang, Jianyong Wang, and Furu Wei. You only cache once: Decoder-decoder architectures for language models. CoRR, abs/2405.05 055544 2024b. doi: 1044555/ARXXI224055..55 URL https: w/tp://ooi.o... / 10. 00554855//220 20wpaaawww..220....00 05254]
12

---

Published as a conference paper at ICLR 2025
Mirac Suzgun, Nathan Scales, Nathanael Schärli, Sebastian Gehrmann, Yi Tay, Hyung Won Chung, Aakanksha Chowdhery, Quoe Le. Ed H Chi, Denny Zhou, et al. Challenging big-bench tasks and whether chain-of-thought can solve them. arXiv preprint arXiv:2210.09200 2022.
Hanlin Tang, Yang Lin, Jing Lin, Qingsen Han, Shikuan Hong, Yiwu Yao, and Gongyi Wang. Razorat- tention: Efficient kv cache compression through retrieval heads. arXiv preprint arXiv:2407. 15891, 2024.
Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: survey, ACM Comput, Surv., 15:6): 2023. doi: 145/3333333111 URL htLpptss:///bi.oog 10.1145/353 3530811
Gemini Team, Rohan Anil, Sehastian Borgeaud, Yonghui Wu, Jean-Baptiste Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew M Dai, Anja Hauth, et al. Gemini: family of highly capable multimodal models. arXiv preprint arXiv:2312.11880 2023.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288. 2023.
Suyuchen Wang, Ivan Kobyzev, Peng Lu, Mehdi Rezagholizadeh, and Bang Liu. Resonance rope: Im- proving context length generalization of large language models. arXiv preprint arXiv:2403.00071, 2024.
Kaiyue Wen, Xingyu Dang, and Kaifeng Lyu. Rnns are not transformers (yet): The key bottleneck on in-context retrieval. CoRR, abs/2402. 182 18510, 2024, doi: 10.48500/ARXIV2400.55111 URL https://doi wtl.g//.....5//aaa..//
/t0.g//0.000 48550 ww./0//arxvvvvm 2402 18510
Wenhao Wu, Yizhong Wang, Guangxuan Xiao, Hao Peng, and Yao Fu. Retrieval head mechanistically explains long-context factuality. arXiv preprint arXiv:2404. 15574, 2024.
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024. URL https:://ppenree ew net www.roorrmduonno s55 ZVF
Wenhan Xiong, Jingyu Liu, Igor Molybog. Hejia Zhang, Prajjwal Bhargava, Rui Hou, Louis Martin, Rashi Rungta, Karthik Abinav Sankararaman, Barlas Oguz, et al. Effective long-context scaling of foundation models. arXiv preprint ar arXiv:2309.160.9, 2023.
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santi- ago Ontañón, Philip Pham. Anirudh Ravula, Qifan Wang, Li Yang, and Amr Ahmed. Big bird: Transformers for longer sequences. In Hugo Larochelle, Mare' Aurelio Ranzato. Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tiee Lin (eds.), Advances in Neural Information Process- ing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6- 2020, virtual, 2020a. URL htttp:://pooceeeiiiiiiee newriggs..e.i..... cc/ Aww.:ett..bbtttttt....oo .tml
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, al. Big bird: Transformers for longer sequences. Advances in neural information processing systems, 33: 1283-727777 2020b.
Yichi Zhang, Bofei Gao, Tianyu Liu, Keming Lu, Wayne Xiong, Yue Dong, Baobao Chang, Junjie Hu, Wen Xiao, al. Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling, arXiv preprint arXiv:2406.00 020609 2024a.
Yichi Zhang, Bofei Gao, Tianyu Liu, Keming Lu, Wayne Xiong, Yue Dong, Baobao Chang, Junjie Hu, Wen Xiao, et al. Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling, arXiv preprint 4rX.v.2406.02069 2024b.
13

---

Published as a conference paper at ICLR 2025
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 Neural In- formation Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023. wtt.ppppp/pppper........c.. 0etee3721557721112eccc22222660-3
/
html
Lianmin Zheng, Liangsheng Yin, Zhiqiang Xie, Jeff Huang, Chuyue Sun, Cody Hao Yu, Shiyi Cao, Christos Kozyrakis, Ion Stoica, Joseph E. Gonzalez, Clark W. Barrett, and Ying Sheng. Efficiently programming large language models using sglang. CoRR, abs/2312.07 2023. doi: 10.48550/ ARXIV.2312.07100 URL htt.hhtts::///m
patttps://oi.oooo /oii http://b...oog1..4455500000 http://0.4455/arx/1... wwt.as///axiii..o 2ww..iii................ 07104
Simiao Zuo, Xiaodong Liu, Jian Jiao, Denis Charles, Eren Manavoglu, Tuo Zhao, and Jianfeng Gao. Efficient long sequence modeling via state space augmented transformer. CoRR, abs/2212.08136 2022. doi: 10445550/RRIV.2212..00 2212. 08136
URL
48550 /rrviv
14


Reviews:
[{"review_id": "4LjLteWAxz", "reviewer": "Reviewer_GZk8", "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.", "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?\n2. Does the sparse attention method remain consistent between training and inference?\n3. 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?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "Originality\n\t1. 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. \n\t2. The work also introduces an optimized GPU-friendly KV cache management technique, which makes long-context processing feasible and efficient on hardware.\n\nQuality\nThe 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.\n\nClarity\nThe 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.\n\nSignificance\nThis 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.", "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.\n2. 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.\n3. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "l5Y80Uu6n8", "reviewer": "Reviewer_Gkd8", "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.", "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.\n\n2) 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?\n\n3) How is the proposed method compare with other existing KV cache elimination methods (e.g., [6](https://arxiv.org/abs/2310.01801))?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "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.\n2) 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.\n3) LongGen integrates a triton-based attention kernel supporting structured sparsity for efficient inference and training.", "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.\n\n2) The evaluation details of the efficiency benchmark should be further elaborated.\n\n3) The models used for main evaluation is old. Results on Llama-3 series would be more persuasive.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "x0MQGeJFPh", "reviewer": "Reviewer_E6wM", "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.", "questions": "1) What are the results of LONGGEN  on other popular large language models, such as GPT models or Gemini?\n\n2) How would the results of LONGGEN change if applied to tasks other than retrieval or reasoning ones?\n\n3) Should the architecture, specifically the number and placement of full attention layers, be adjusted for different models or tasks?\n\n4) Is there an algorithmic or automated method for determining the optimal number and placement of full-attention layers?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 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.\n\n2) 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. \n\n3) 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.", "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).\n\n2) 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).\n\n3) 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.\n\n[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.\n\n[2] Yao Fu. How do language models put attention weights over long context. Yao FuâA˘Zs Notion ´ , 2024.\n\n[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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2VSIrDdRp8", "reviewer": "Reviewer_eef1", "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.", "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?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "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.\n2. The performance is very good: e.g. NIAH result is much better than previous methods such as StreamingLLM.", "weaknesses": "There is no noticeable weaknesses that the reviewer hope the authors shall address (only some small clarification questions). Please see the question section.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0mJD7OJmps", "reviewer": "Reviewer_83A1", "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.", "questions": "See weakness", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The work highlights the importance of including full attention layers for models to achieve accurate retrieval.\n- The paper is well-written.\n- Experimental results indicate that LongGen accelerates both training and inference for long context while preserving model performance.", "weaknesses": "- Since training on long context constitutes a small portion of pre-training, the training speedup of LongGen is limited during pre-training.\n- LongGen with AtteSink and BlockSparse demonstrates similar performance, necessitating a detailed explanation for this observation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)