Add model card and metadata

#1
by nielsr HF Staff - opened
Files changed (1) hide show
  1. README.md +45 -3
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
+ ```