RAZOR gate for Qwen3-4B
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.
- Architecture: one MLP per layer on the layer's hidden state (2560 -> 512 -> 8 KV heads), 36 layers,
47.4M parameters. Same architecture as kvpress's
KVzapModel. - Training: KL divergence to the uncompressed model's next-token distribution, hard keep/drop decisions with
straight-through gradients, and a dual-ascent price that holds the keep rate at 30%. 1.7 hours on one T4, documents
of at most 1k tokens (Nemotron data).
train_log.txthas the run. - Inference: keep a token if its score is above 0 (DMS-style threshold with a 128-token sliding window).
Code, training and evaluation: https://github.com/Manaslu-AI/razor. Paper: razor.pdf in that repo (arXiv version
pending).
Results
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.
Use
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.
Citation
@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}
}
- Downloads last month
- -