kvpress

While standard presses work during the prefilling phase, DecodingPress is an experimental wrapper that enables compression during the token generation (decoding) phase.

Instead of a compression ratio, DecodingPress uses a target_size to determine how much to compress. It performs compression every compression_interval steps to ensure the cache size after compression matches the target_size.

Important Constraints:

• DecodingPress only supports ScorerPress types as its base_press.
• It can optionally maintain a buffer of recent hidden states via hidden_states_buffer_size.

from transformers import pipeline
from kvpress import KnormPress, DecodingPress

# Initialize the pipeline
device = "cuda:0"
model = "meta-llama/Llama-3.1-8B-Instruct"
model_kwargs = {"attn_implementation": "flash_attention_2"}
pipe = pipeline("kv-press-text-generation", model=model, device=device, model_kwargs=model_kwargs)

# Configure decoding compression
decoding_press = DecodingPress(
    base_press=KnormPress(),
    compression_interval=10,
    target_size=512
)

context = "A very long text you want to compress during generation"
question = "Tell me a long story about this context"
response = pipe(context, question=question, press=decoding_press)["answer"]

A press works by registering a forward hook (press.forward_hook) to each attention layer during the prefilling phase. You can apply a press to a model by using the press instance as a context manager via the __call__ method. This allows you to wrap model calls or model.generate() calls to apply KV cache compression.

Note: While you can use with press(model): with model.generate(), this method does not allow excluding the question from compression (which is important for methods like SnapKV) and does not support batching multiple questions at once.

import torch
from transformers import AutoModelForCausalLM
from kvpress import KnormPress

device = "cuda:0"
ckpt = "meta-llama/Meta-Llama-3.1-8B-Instruct"
model = AutoModelForCausalLM.from_pretrained(ckpt).to(device)
press = KnormPress(compression_ratio=0.4)

inputs = model.dummy_inputs["input_ids"].to(device)

with torch.no_grad():
    print(model(inputs).past_key_values[0][0].shape)
    # torch.Size([3, 8, 5, 128])
    
with torch.no_grad(), press(model):
    print(model(inputs).past_key_values[0][0].shape)
    # torch.Size([3, 8, 3, 128])

KVPress provides "presses" that compress the KV cache during the prefilling phase. The easiest way to use them is via the kv-press-text-generation pipeline, which is automatically registered with Hugging Face transformers upon import. This pipeline handles chat templates and tokenization automatically.

To use it, initialize a pipeline with the name kv-press-text-generation and pass a press object (like ExpectedAttentionPress) to the pipeline call.

from transformers import pipeline
from kvpress import ExpectedAttentionPress

model = "Qwen/Qwen3-8B"
pipe = pipeline("kv-press-text-generation", model=model, device_map="auto", dtype="auto")

context = "A very long text you want to compress once and for all"
question = "\nA question about the compressed context"

# Apply compression with a specific ratio
press = ExpectedAttentionPress(compression_ratio=0.5)
answer = pipe(context, question=question, press=press)["answer"]