Research map / Architectures and efficiency / Attention and transformers
A Little Goes a Long Way: Efficient Long Context Training and Inference with Partial Contexts
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes LongGen, which finetunes a pretrained Llama-2 into a hybrid model with 1/3 full attention layers (middle) and 2/3 GPU-friendly sparse attention layers (attention sink or block sparse) during context extension from 4K to 128K.
Problem
Long-context LLMs are expensive to train (context extension) and serve (KV cache memory), and post-hoc inference-time KV cache eviction methods underperform on long-context retrieval and reasoning.
Approach
Integrate context-length extension with a GPU-friendly KV-reduced hybrid architecture: sparse attention (attention sink or block sparse with static, block-wise access) in top and bottom layers and full attention in the middle layers, trained with a custom FlashAttention-2-based Triton kernel and served with vLLM.
Key results
- LongGen-7B matches full attention on NIAH (100%) and reaches 0.27 vs 0.29 BABILong accuracy, with 1.55x training speedup and 36% less wall-clock training time at 128K.
- At inference, the 1/3 full setting reduces KV cache memory by 62% (69.2 GB to 26.5 GB), with 1.67x prefilling and 1.41x decoding speedup.
- On NIAH, LongGen reaches 100% pass rate versus 28-51% for KV eviction methods (Attention Sink, H2O, RazorAttention, PyramidKV) at comparable KV budgets.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use the LongGen hybrid architecture with 1/3 full-attention layers | Llama2-7B · NIAH | accuracy: 100% -> 100% | No clear effect |
| Use the LongGen hybrid architecture with 1/3 full-attention layers | Llama2-7B · BABILong | average accuracy: 0.29 -> 0.27 | Hurt |
| Use the LongGen hybrid architecture with 1/3 full-attention layers | Llama2-7B · RULER | average accuracy: 0.41 -> 0.38 | Hurt |
| Use the LongGen hybrid architecture with 1/3 full-attention layers | Llama2-70B · BABILong | average accuracy: 0.46 -> 0.46 | No clear effect |
| Use the LongGen hybrid architecture with 1/3 full-attention layers | Llama2-70B · RULER | average accuracy: 0.67 -> 0.65 | Hurt |
| Train with LongGen's sparse attention architecture | Llama2-7B · 128K training | wall-clock training speedup: 1.55x | Helped |
| Use LongGen's partial-context architecture | Llama2-7B · 128K inference | KV cache memory: 69.2 GB -> 26.5 GB; 62% reduction | Helped |
| Use LongGen's partial-context architecture | Llama2-7B · 128K inference | prefilling speedup: 1.67x | Helped |
| Use LongGen's partial-context architecture | Llama2-7B · 128K inference | decoding speedup: 1.41x | Helped |
| Use LongGen w/ AttnSink instead of inference-time KV eviction | Llama2-7B · NIAH | pass rate: 51% -> 100% | Helped |
| Use LongGen w/ BlockSparse instead of inference-time KV eviction | Llama2-7B · NIAH | pass rate: 51% -> 100% | Helped |
| Place the full-attention layers in the middle rather than the top, bottom, or interleaved positions (ablation) | LongGen w/ AttnSink · BABILong | accuracy | Helped |
Surprising findings
- Inference-time KV eviction methods (Attention Sink, H2O, PyramidKV) substantially underperform on NIAH at long contexts, even though they work for perplexity.
Limitations
- Slightly underperforms full attention on harder reasoning tasks such as BABILong and RULER, with the gap larger at long contexts.
- Training FLOP savings do not reduce training memory since all tokens must still be stored for the loss.
Experimental setup
Datasets: SlimPajama
Benchmarks: Needle-in-a-Haystack BABILong RULER MMLU GSM8K BBH
Metrics: pass rate accuracy wall-clock speedup KV cache memory prefilling speedup decoding speedup
Baselines: Attention Sink H2O RazorAttention PyramidKV Sliding Window YOCO Full Attention
Builds on: FlashAttention-2 attention sink (StreamingLLM) Mistral sliding window attention YOCO vLLM
Closest papers by meaning
- SCBench: A KV Cache-Centric Analysis of Long-Context Methods (ICLR 2025 · similarity 0.92)
- DuoAttention: Efficient Long-Context LLM Inference with Retrieval and Streaming Heads (ICLR 2025 · similarity 0.92)
- LongRoPE: Extending LLM Context Window Beyond 2 Million Tokens (ICML 2024 · similarity 0.91)
- CateKV: On Sequential Consistency for Long-Context LLM Inference Acceleration (ICML 2025 · similarity 0.91)
- LongLoRA: Efficient Fine-tuning of Long-Context Large Language Models (ICLR 2024 · similarity 0.91)
- RocketKV: Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression (ICML 2025 · similarity 0.91)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.