Add model card and metadata
#1
by
nielsr
HF Staff
- opened
README.md
CHANGED
|
@@ -1,3 +1,45 @@
|
|
| 1 |
-
---
|
| 2 |
-
license: mit
|
| 3 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
license: mit
|
| 3 |
+
pipeline_tag: text-generation
|
| 4 |
+
---
|
| 5 |
+
|
| 6 |
+
# KV Admission: Learning What to Write for Efficient Long-Context Inference
|
| 7 |
+
|
| 8 |
+
This repository contains the official weights for the **Write-Gate MLP** introduced in the paper [KV Admission: Learning What to Write for Efficient Long-Context Inference](https://huggingface.co/papers/2512.17452).
|
| 9 |
+
|
| 10 |
+
## Abstract
|
| 11 |
+
|
| 12 |
+
Long-context LLM inference is bottlenecked by the quadratic attention complexity and linear KV cache growth. Prior approaches (KV Selection or Eviction) mitigate this post-hoc, but overlook the root inefficiency: indiscriminate writing to memory. We propose **Write-Gated KV (WG-KV)** to introduce a missing primitive: **KV Admission**. Instead of blindly persisting every token, WG-KV employs a lightweight, learnable mechanism to predict token utility before cache entry. By filtering out low-utility states early to maintain a compact global cache alongside a sliding local cache, WG-KV reduces memory usage by 46-68% and delivers significant prefill and decode speedups, while maintaining compatibility with FlashAttention and Paged-KV systems.
|
| 13 |
+
|
| 14 |
+
## Resources
|
| 15 |
+
|
| 16 |
+
- **Paper:** [KV Admission: Learning What to Write for Efficient Long-Context Inference](https://huggingface.co/papers/2512.17452)
|
| 17 |
+
- **GitHub Repository:** [EMCLab-Sinica/WG-KV](https://github.com/EMCLab-Sinica/WG-KV)
|
| 18 |
+
|
| 19 |
+
## Usage
|
| 20 |
+
|
| 21 |
+
Using these checkpoints requires the environment and custom implementations provided in the official repository. This includes a modified version of the Transformers library and a custom vLLM fork for sparse prefill kernels. Please refer to the [official installation guide](https://github.com/EMCLab-Sinica/WG-KV#%EF%B8%8F-installation).
|
| 22 |
+
|
| 23 |
+
After setting up the environment, you can run inference using the trained gate by specifying the checkpoint path:
|
| 24 |
+
|
| 25 |
+
```bash
|
| 26 |
+
python scripts/inference.py \
|
| 27 |
+
--model_name meta-llama/Llama-3.1-8B-Instruct \
|
| 28 |
+
--filtering_path weights/llama-3.1-8b-instruct-0.04.pt
|
| 29 |
+
```
|
| 30 |
+
|
| 31 |
+
The checkpoints follow the naming convention `{model_name}-{lambda}.pt`, where `lambda` (λ) controls the trade-off between sparsity and accuracy.
|
| 32 |
+
|
| 33 |
+
## Citation
|
| 34 |
+
|
| 35 |
+
```bibtex
|
| 36 |
+
@misc{wgkv,
|
| 37 |
+
title={KV Admission: Learning What to Write for Efficient Long-Context Inference},
|
| 38 |
+
author={Yen-Chieh Huang and Pi-Cheng Hsiu and Rui Fang and Ming-Syan Chen},
|
| 39 |
+
year={2025},
|
| 40 |
+
eprint={2512.17452},
|
| 41 |
+
archivePrefix={arXiv},
|
| 42 |
+
primaryClass={cs.LG},
|
| 43 |
+
url={https://arxiv.org/abs/2512.17452},
|
| 44 |
+
}
|
| 45 |
+
```
|