DropKV: Decoupling Residual-Output Perturbation for Near-Optimal KV-Cache Eviction, has been accepted to NeurIPS 2026 as a poster.
DropKV is an attention-output-sensitive KV-cache eviction method for LLM inference. During prefill, it scores KV positions by their estimated effect on the attention output and retains the highest-scoring fraction. Decode then continues with the compressed cache.
dropkv_cache.py: standaloneDropKVCacheand chunked-prefill helper.dropkv_triton.py: Triton kernels and a PyTorch reference implementation.patches/dropkv_transformers.patch: integration for a pinned Transformers revision, including Llama and Qwen2 query-state plumbing.tests/test_dropkv_triton.py: kernel, cache, chunked-prefill, and guard tests.
Experimental outputs, private cluster paths, unrelated eviction methods, model weights, and vendored third-party repositories are intentionally excluded.
- Linux with an NVIDIA GPU for the Triton path
- Python 3.9 or newer
- PyTorch 2.6 or newer
- Triton 3.0 or newer (optional; PyTorch fallback is available)
- A patched checkout of Hugging Face Transformers for
generate()andlm-evalintegration
The integration patch is pinned to Transformers commit fd6bc380c8854a370fbc9f68a157895d84dce7d7 (the repository identified it as v5.3.0.dev0). Do not assume that it applies to another revision.
The setup script creates a Conda environment, clones the pinned Transformers revision, applies the DropKV patch, and installs the local project:
bash setup.sh
conda activate dropkvFor an existing environment, install the local modules first, then apply the patch to an exact checkout of the pinned revision:
python -m pip install -e '.[triton,test]'
git clone https://github.com/huggingface/transformers.git transformers_src
git -C transformers_src checkout fd6bc380c8854a370fbc9f68a157895d84dce7d7
git -C transformers_src apply ../patches/dropkv_transformers.patch
python -m pip install -e ./transformers_srcThe model attention implementation must pass query_states in cache_kwargs; the supplied Transformers patch does this for Llama and Qwen2.
from dropkv_cache import DropKVCache
cache = DropKVCache(
keep_ratio=0.1,
window_size=8,
kernel_size=11,
use_triton=True,
)
outputs = model(
input_ids,
past_key_values=cache,
use_cache=True,
)from dropkv_triton import dropkv_scores_reference, dropkv_scores_triton
scores = dropkv_scores_triton(Q, K, V, window_size=8)
reference = dropkv_scores_reference(Q, K, V, window_size=8)Q has shape [B, Hq, Lq, D]; K and V have shape
[B, Hkv, Lk, D]. Grouped-query attention is supported when Hq is divisible
by Hkv, and chunked prefill is supported when Lq <= Lk and the query is the
tail-aligned suffix.
After running setup.sh:
bash run_eval.sh \
--model meta-llama/Llama-3.1-8B-Instruct \
--task ruler \
--keep-ratio 0.1 \
--output ./resultsPaths may also be configured by copying .env.example to .env.
python -m pytest tests/test_dropkv_triton.py -vMost tests require CUDA. Several integration tests also download
Qwen/Qwen2-0.5B-Instruct and require the patched Transformers checkout.
Before testing a modified patch, verify that it still applies cleanly:
git -C transformers_src reset --hard fd6bc380c8854a370fbc9f68a157895d84dce7d7
git -C transformers_src apply --check ../patches/dropkv_transformers.patchThe reset --hard command above is only for a disposable Transformers clone;
do not run it in a checkout containing work you need to keep.
DropKV is released under the Apache License 2.0. Portions adapted from Hugging Face Transformers remain under Apache-2.0; see THIRD_PARTY_NOTICES.md for details.
If you find DropKV useful in your research, please cite:
@inproceedings{zhang2026dropkv,
title = {DropKV: Decoupling Residual-Output Perturbation for Near-Optimal KV-Cache Eviction},
author = {Zhang, Aozhong and Gurses, Selcuk and Deng, Yanxia and Wang, Naigang and Liu, Chi-Chun and Wertheimer, Davis and Liu, Derrick and Li, Xin and Yang, Zi and Ye, Felix X.-F. and Yin, Penghang},
booktitle = {Advances in Neural Information Processing Systems},
year = {2026}
}