Existing inference acceleration methods often reduce memory usage by discarding “unimportant” tokens. A critical limitation is that they judge importance using current attention weights, overlooking tokens that may become important later in generation. This static assessment cannot adapt to changing context requirements and can limit model performance.
Key observation
Our analysis revealed an important pattern: within a single decoding step, consecutive layers identify highly similar sets of important tokens. This observation suggests a different route to more efficient inference.
Method
Based on this observation, we proposed OmniKV, an inference method that requires no additional training and does not discard tokens. Its core ideas are:
Dynamic context selection: during decoding, a “filter” layer uses attention sparsity to select a small set of important tokens. Other layers then reuse this token subset as their attention context, exploiting the similarity across layers.
Retaining all tokens: OmniKV preserves the KV cache for every token so that the model can access the full history when needed. This avoids the performance degradation caused by permanently losing information.
Experimental results
OmniKV has little impact on model quality in benchmarks such as LongBench and InfiniteBench, and performs particularly well in multi-step reasoning scenarios such as chain-of-thought. It extends the maximum context length of Llama-3-8B on a single A100 GPU from 128K to 450K, an increase of more than 3.5×. It achieves a 1.7× speedup in both Hugging Face Transformers and the lightweight inference engine LightLLM. Through LightLLM, it also supports tensor parallelism for multi-GPU execution.
Later developments: IndexCache and GLM-5.2
In our ICLR 2025 work, OmniKV [1], we found that attention distributions are highly similar across layers. A subset of layers can therefore act as selection layers, identifying important tokens, while other layers compute sparse attention using the tokens chosen by the nearest preceding selection layer. Our subsequent work, DeltaKV [2], again confirmed OmniKV’s strong performance. Contemporaneous work such as TidalDecode [3] explored related ideas.
The GLM team’s more recent work, IndexCache, extends this idea to DSA and cites OmniKV. Some Lightning indexers serve as selection layers. GLM-5.2, whose performance was recently reported to be comparable to Claude 4.8, uses this technique. It is exciting to see an idea from our work finding its way into a new foundation model.
[1] OmniKV: Dynamic Context Selection for Efficient Long-Context LLMs
[2] DeltaKV: Residual-Based KV Cache Compression via Long-Range Similarity
[3] TidalDecode: Fast and Accurate LLM Decoding with Position Persistent Sparse Attention
