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.txt has 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
-
Safetensors
Model size
47.4M params
Tensor type
F16
ยท
Inference Providers NEW
This model isn't deployed by any Inference Provider. ๐Ÿ™‹ Ask for provider support

Model tree for Manalsu/razor-gate-Qwen3-4B

Finetuned
Qwen/Qwen3-4B
Finetuned
(1,262)
this model