Views
No views yet
DBTrimKV is the dynamic-budget variant of TrimKV: a single global KV budget is shared across layers and heads and reallocated on the fly, with the retention-gate's final projection tied across layers.
Qwen/Qwen3-4B-Instruct-2507 (131072-token training context, M = 512). The base-model weights are not included — they are loaded from Qwen/Qwen3-4B-Instruct-2507 at runtime and the retention-gate weights from trimkv_weights.pth are overlaid on top.trimkv library from the official repository.1import torch
2from trimkv.models.qwen3 import TrimKVQwen3ForCausalLM
3from trimkv.cache_utils import PagedTrimKVCache
4from transformers import AutoTokenizer
5
6model = TrimKVQwen3ForCausalLM.from_pretrained(
7 "ngocbh/DBTrimKV-Qwen3-4B-Instruct-2507",
8 torch_dtype=torch.bfloat16,
9 load_trimkv_weights=True,
10 download_from="huggingface",
11 use_cache=True,
12 device_map="cuda",
13)
14model.config._attn_implementation = "flash_attention_2"
15
16tokenizer = AutoTokenizer.from_pretrained(
17 model.config.base_model, use_fast=True, padding_side="left"
18)
19
20past_key_values = PagedTrimKVCache(
21 num_layers=model.config.num_hidden_layers,
22 num_heads=model.config.num_key_value_heads,
23 max_seq_len=32768,
24 memory_size=512,
25 num_blocks_ratio=1.0,
26 buffer_size=32,
27 strategy="fixed_budget",
28 device="cuda",
29)
30
31# Use as a normal HF model — pass `past_key_values=past_key_values` to .generateexamples/test_qwen3.py in the GitHub repo for a full runnable example.Qwen/Qwen3-4B-Instruct-2507retention_gate=rg10)fwkl_ntprg_attn_flex1@article{bui2025make,
2 title={Make Each Token Count: Towards Improving Long-Context Performance with KV Cache Eviction},
3 author={Bui, Ngoc and Nguyen, Hieu Trung and Cohan, Arman and Ying, Rex},
4 journal={arXiv preprint arXiv:2512.03324},
5 year={2025}
6}