Loading…
Loading…
A learned KV-cache forget gate for Qwen/Qwen3-4B. For every token, as it is written to the cache, it gives one keep score per KV head. Tokens scored below the threshold are evicted. The base model is not modified.
KVzapModel.train_log.txt has the run.Code, training and evaluation: https://github.com/Manaslu-AI/razor. Paper: razor.pdf in that repo (arXiv version
pending).
All on 16 GB T4s.
RULER 4k (n=195): 93.9 at 74% evicted, vs 94.3 uncompressed. All 8 needle tasks stay at 100. The losses are common-word extraction (95 to 88) and frequent-word extraction (83 to 79). No compressed baselines were run at 4B.
Longest context on one 16 GB T4 (passkey retrieval, fp16, YaRN x4, chunked prefill of 4096):
| Tokens | Full cache peak GB | RAZOR peak GB / KV GB | RAZOR passkey |
|---|---|---|---|
| 31k | 15.0 (limit) | 9.6 / 0.72 | correct |
| 61k | OOM | 11.1 / 1.41 | correct |
| 96k | OOM | 12.6 / 2.22 | correct |
| 123k | OOM | 13.9 / 2.84 | correct |
Decoding at 29k tokens is 2.9x faster than with the full cache (7.5 vs 2.6 tokens/s). At 8k it is 7% slower.
InfiniteBench at 100k tokens (n=12 per task): passkey 0.67, number_string 0.67, kv_retrieval 0.08, all equal to the full cache, and longbook_choice 0.67 vs 0.58.
pip install torch 'transformers>=4.56,<5.3' kvpress
git clone https://github.com/Manaslu-AI/razor && cd razor
from kvpress import KVPressTextGenerationPipeline # noqa: F401, registers the pipeline
from transformers import pipeline
from razor_kv.press import LocalGatePress
from razor_kv.ragged import RaggedCache, RaggedDMSPress, enable
pipe = pipeline("kv-press-text-generation", model="Qwen/Qwen3-4B", dtype="float16", device_map="cuda")
enable(pipe.model) # per-head ragged attention
press = RaggedDMSPress(LocalGatePress(path="Manalsu/razor-gate-Qwen3-4B"), threshold=0.0, prefill_chunk=4096)
print(pipe(long_document, question="...", press=press, cache=RaggedCache())["answer"]) # evicted KV is freed, not masked
Beyond Qwen3's native 32k, load the model with YaRN factor 4 (see bench_memory.py in the repo). RaggedDMSPress
supports batch size 1.
@misc{rawal2026razor,
title = {RAZOR: Learning What to Forget in the KV Cache with a Description-Length Price},
author = {Rawal, Himansu},
year = {2026},
note = {Manaslu Labs. arXiv version pending},
url = {https://github.com/Manaslu-AI/razor}
}
Qwen/Qwen3-4B