### Local Installation with uv

Source: https://github.com/nvidia/kvpress/blob/main/README.md

Clone the repository and use uv for a local installation.

```bash
git clone https://github.com/NVIDIA/kvpress.git
cd kvpress
uv sync
```

--------------------------------

### Install Optimum-Quanto

Source: https://github.com/nvidia/kvpress/blob/main/notebooks/speed_and_memory.ipynb

Run this command to install the Optimum-Quanto library, which is required for using quantized caches.

```python
# Run these below to be able to use quantized cache 
#!pip install -U optimum-quanto
```

--------------------------------

### Use KVPressTextGenerationPipeline with QuantizedCache

Source: https://context7.com/nvidia/kvpress/llms.txt

Example of integrating `QuantizedCache` with the `KVPressTextGenerationPipeline`. This requires the `optimum-quanto` library to be installed.

```python
from transformers import pipeline, QuantizedCache
from kvpress import ExpectedAttentionPress

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

context = "The Treaty of Westphalia (1648) ended the Thirty Years' War..."  # long document

press = ExpectedAttentionPress(compression_ratio=0.4)

# With QuantizedCache (requires pip install optimum-quanto)
cache = QuantizedCache(backend="quanto", nbits=4)
result = pipe(context, question="When was the treaty signed?", press=press, cache=cache)
```

--------------------------------

### Install KVPress with Evaluation Extras

Source: https://context7.com/nvidia/kvpress/llms.txt

Installs the KVPress library with the necessary extras for evaluation tasks.

```bash
pip install kvpress[eval]
```

--------------------------------

### Install kvpress using pip

Source: https://github.com/nvidia/kvpress/blob/main/README.md

Install the kvpress library using pip for basic usage.

```bash
pip install kvpress
```

--------------------------------

### KVzap Usage Example

Source: https://github.com/nvidia/kvpress/blob/main/kvzap/README.md

Demonstrates how to use KVzap for both prefilling-only compression and combined prefilling/decoding compression with thinking enabled. Requires importing `KVzapPress` and `DMSPress` from kvpress.

```python
import requests
from transformers import pipeline
from kvpress import KVzapPress, DMSPress

model = "Qwen/Qwen3-8B"
pipe = pipeline("kv-press-text-generation", model=model, device_map="auto", dtype="auto")
press = DMSPress(KVzapPress(model_type="mlp"), threshold=-4)

# Prefilling compression only, thinking disabled
press.decoding = False
context = requests.get("https://arxiv.org/abs/2601.07891").text
question = "\n What is this article about in 2 sentences ?"
answer = pipe(context, question=question, press=press)["answer"]
print(f"Compression ratio: {press.compression_ratio:.2%}\nAnswer: {answer}")

# Prefilling and decoding compression, thinking enabled
press.decoding = True
prompt = "What is the best hardware to run LLMs and why ?"
answer = pipe(prompt, press=press, enable_thinking=True, max_new_tokens=2000)["answer"]
print(f"Compression ratio: {press.compression_ratio:.2%}\nAnswer: {answer}")
```

--------------------------------

### Local Installation with Optional Dependencies

Source: https://github.com/nvidia/kvpress/blob/main/README.md

Install kvpress locally with optional dependencies for evaluation and flash attention.

```bash
git clone https://github.com/NVIDIA/kvpress.git
cd kvpress
uv sync --extra eval --extra flash-attn
```

--------------------------------

### Implement a Custom Press Subclass

Source: https://context7.com/nvidia/kvpress/llms.txt

Shows how to create a custom KV cache compression strategy by subclassing `BasePress` and implementing the `compress` method. This example keeps the most recent tokens.

```python
import torch
from dataclasses import dataclass
from transformers import AutoModelForCausalLM, AutoTokenizer
from kvpress import BasePress, KnormPress

# Custom press subclass
@dataclass
class MyPress(BasePress):
    compression_ratio: float = 0.5

    def compress(self, module, hidden_states, keys, values, attentions, kwargs):
        k_len = keys.shape[2]
        n_kept = max(1, int(k_len * (1 - self.compression_ratio)))
        # Keep the most recent n_kept tokens
        return keys[:, :, -n_kept:, :], values[:, :, -n_kept:, :]


device = "cuda:0"
model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3.1-8B-Instruct").to(device)
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3.1-8B-Instruct")

my_press = MyPress(compression_ratio=0.3)
inputs = tokenizer("A long context passage...", return_tensors="pt").to(device)

with torch.no_grad(), my_press(model):
    out = model(**inputs)
```

--------------------------------

### KVzap Training Command

Source: https://github.com/nvidia/kvpress/blob/main/kvzap/README.md

Command to install necessary libraries and train a KVzap model. Ensure to replace `<model_name>` and `<output_dir>` with your specific values. Use `--help` for additional options.

```bash
pip install skorch scikit-learn
python train.py --model_name <model_name> --output_dir <output_dir>
```

--------------------------------

### Use Quantized KV Cache with Transformers

Source: https://github.com/nvidia/kvpress/blob/main/README.md

Instantiate and use a QuantizedCache with a Hugging Face pipeline. Requires installing additional dependencies like 'optimum-quanto'.

```python
from transformers import QuantizedCache

cache = QuantizedCache(backend="quanto", nbits=4)

pipe(..., cache=cache)
```

--------------------------------

### Implement Head-wise Compression with RandomHeadPress

Source: https://github.com/nvidia/kvpress/blob/main/notebooks/new_press.ipynb

Implement head-wise compression by creating a custom `BasePress` subclass. This example uses `RandomHeadPress` to randomly prune KV cache based on a compression ratio, faking head-wise compression by setting `masked_key_indices`.

```python
@dataclass
class RandomHeadPress(BasePress):

    compression_ratio: float = 0.0

    def compress(self, module, hidden_states, keys, values, attentions, kwargs):
        assert keys.shape[0] == 1, "Only batch size 1 is supported"
        scores = torch.rand(keys.shape[:-1], device=keys.device)
        mask = scores < torch.quantile(scores, self.compression_ratio)
        module.masked_key_indices = torch.nonzero(mask, as_tuple=True)
        
        return keys, values

for compression_ratio in [0, 0.25, 0.9]:
    press = RandomHeadPress(compression_ratio)
    print(f"\ncompression_ratio: {compression_ratio}")
    print(f"Answer: {pipe(context, question=question, press=press)['answer']}")
```

--------------------------------

### Implement a Custom ScorerPress

Source: https://github.com/nvidia/kvpress/blob/main/notebooks/new_press.ipynb

Shows how to create a custom press by inheriting from `ScorerPress` and overriding the `score` method. This example implements a press where the score is the negative norm of the key vector.

```python
class MyKnormPress(ScorerPress):
    def score(
        self,
        module: nn.Module,
        hidden_states: torch.Tensor,
        keys: torch.Tensor,
        values: torch.Tensor,
        attentions: torch.Tensor,
        kwargs,
    ) -> torch.Tensor:

        scores = -keys.norm(dim=-1)

        # For demonstration, we show some details on the shape for the first layer
        if module.layer_idx == 0:
            print(f"module: {module}")
            print(f"Number of key value heads: {module.config.num_key_value_heads}")
            print(f"Sequence length: {hidden_states.shape[1]}")
            print()
            print(f"hidden_states shape: {hidden_states.shape}")
            print(f"keys shape:          {keys.shape}") # shape (bhnd)
            print(f"values shape:        {values.shape}") # shape (bhnd)
            print(f"score shape:         {scores.shape}") # shape (bhn)
            print()
        
        return scores


press = MyKnormPress(compression_ratio)
print(pipe(context, question=question, press=press)["answer"])
```

--------------------------------

### Re-implement StreamingLLMPress Compression

Source: https://github.com/nvidia/kvpress/blob/main/notebooks/new_press.ipynb

Re-implement the `StreamingLLMPress` to create a more compact version by overriding the `compress` method. This example demonstrates how to define custom compression logic based on `n_first` and `n_last` tokens.

```python
from dataclasses import dataclass
import torch
import torch.nn as nn

@dataclass
class MyStreamingLLMPress(BasePress):
    n_first: int = 1
    n_last: int = 8

    def compress(
        self,
        module: nn.Module,
        hidden_states: torch.Tensor,
        keys: torch.Tensor,
        values: torch.Tensor,
        attentions: torch.Tensor,
        kwargs: dict,
    ) -> tuple[torch.Tensor, torch.Tensor]:

        mask = torch.ones(keys.shape[-2], dtype=torch.bool, device=keys.device)
        mask[self.n_first : -self.n_last] = False
        return keys[:, :, mask, :], values[:, :, mask, :]


for n_last in [2, 4, 8]:
    press = MyStreamingLLMPress(n_last=n_last)
    print(f"\nn_last: {n_last}")
    print(f"Last tokens seen by the model: {pipe.tokenizer.decode(tokens.input_ids[0, -n_last:])}")
    print(f"Answer: {pipe(context, question=question, press=press)['answer']}")
```

--------------------------------

### Initialize Pipeline for Decoding Compression

Source: https://github.com/nvidia/kvpress/blob/main/README.md

Set up the text generation pipeline with specific model arguments, such as using flash attention.

```python
from transformers import pipeline
from kvpress import KnormPress
from kvpress import 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)
```

--------------------------------

### Use KVPressTextGenerationPipeline with a Single Question

Source: https://context7.com/nvidia/kvpress/llms.txt

Demonstrates initializing and using the `KVPressTextGenerationPipeline` for a single question with a specified press. Ensure the model and press are correctly configured.

```python
from transformers import pipeline
from kvpress import ExpectedAttentionPress

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

context = "The Treaty of Westphalia (1648) ended the Thirty Years' War..."  # long document
questions = ["When was the treaty signed?", "Which war did it end?"]

press = ExpectedAttentionPress(compression_ratio=0.4)

# Single question
result = pipe(context, question="When was the treaty signed?", press=press, max_new_tokens=64)
print(result["answer"])
# Expected: "The Treaty of Westphalia was signed in 1648."
```

--------------------------------

### Initialize and Run Decoding with Compression

Source: https://github.com/nvidia/kvpress/blob/main/notebooks/kvpress_decoding_aime25.ipynb

Initialize the DecodingPress with compression parameters and then use it to decode a question. This is useful for scenarios where intermediate compression of the KV cache is desired to manage memory or improve efficiency.

```python
compression_interval = 1024  # compress every compression_steps
target_size = 512 # number of tokens to keep after compression. Note that actual cache size lies in [target_size, compression_interval]

press = DecodingPress(base_press=ExpectedAttentionPress(), compression_interval=compression_interval, target_size=target_size)

question = sample["problem"]
true_answer = sample["answer"]
pred_answer = pipe( " ", question=question, press=press, max_new_tokens=16_000)["answer"]

print(f"Question:   {question}")
print(f"Answer:     {true_answer}")
print(f"Prediction: {pred_answer}")
```

--------------------------------

### Initialize and Use FinchPress

Source: https://context7.com/nvidia/kvpress/llms.txt

FinchPress uses a dynamic window size based on a delimiter token and supports key re-rotation. Ensure the input context contains the specified delimiter token.

```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from kvpress import FinchPress

model_name = "meta-llama/Meta-Llama-3.1-8B-Instruct"
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(model_name)

press = FinchPress(compression_ratio=0.5, rerotate_keys=True)
press.update_model_and_tokenizer(model, tokenizer, delimiter_token="\n\nQuestion:")

# Input MUST contain the delimiter
context = "A long research document...\n\nQuestion:"
question = "What is the main finding?"

with torch.no_grad(), press(model):
    inputs = tokenizer(context + question, return_tensors="pt").to(model.device)
    out = model.generate(**inputs, max_new_tokens=80)
    print(tokenizer.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True))
```

--------------------------------

### Basic Decoding Compression with KVPress

Source: https://github.com/nvidia/kvpress/blob/main/README.md

Initialize a pipeline with a decoding press to compress text every N steps during generation. Configure the compression ratio and token buffer size as needed.

```python
from transformers import pipeline
from kvpress import KnormPress
from kvpress import 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)

# Create a decoding press that compresses every 10 steps to 512 tokens
decoding_press = DecodingPress(
    base_press=KnormPress(),
    compression_steps=10,
    token_buffer_size=512
)

# Use with pipeline
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"]
```

--------------------------------

### Get Generation Statistics with Dynamic Cache

Source: https://github.com/nvidia/kvpress/blob/main/notebooks/speed_and_memory.ipynb

Measures total time and peak GPU memory during text generation. Disables EOS token criteria and sampling for consistent timing. Loads model within the function to prevent GPU leakage.

```python
def get_generation_stats(press, n_tokens, max_new_tokens = 100, cache_implementation="dynamic"):
    torch.cuda.reset_peak_memory_stats()
    torch.cuda.empty_cache()
    idle_peak_memory = torch.cuda.max_memory_allocated()
    model = AutoModelForCausalLM.from_pretrained(ckpt, dtype="auto", attn_implementation="flash_attention_2").to(device)
    # disable EosTokenCriteria stopping criteria
    model.generation_config.eos_token_id = None
    model.generation_config.stop_strings = None
    
    model.generation_config.do_sample = False
    model.generation_config.temperature = None
    model.generation_config.top_p = None

    initial_peak_memory = torch.cuda.max_memory_allocated()
        
    inputs =torch.arange(n_tokens).reshape([1, n_tokens]).to(device)

    with press(model):
        kwargs = dict()
        if cache_implementation == "quantized":
            kwargs = dict(cache_implementation="quantized",
                          cache_config={"backend": "quanto", "nbits": 4})
            
        start = time()
        outputs = model.generate(inputs, max_new_tokens=max_new_tokens,
                                 generation_config=model.generation_config,
                                 pad_token_id=-1, **kwargs)
        total_time = time() - start
        assert outputs.shape == (1, n_tokens + max_new_tokens), outputs.shape

    peak_memory = torch.cuda.max_memory_allocated()

    model.cpu()
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()
    return {"Idle Peak memory": idle_peak_memory / 1024**3,
            "Initial Peak memory": initial_peak_memory / 1024**3,
            "Total time": total_time, 
            "Peak memory usage": peak_memory / 1024**3
           }
```

--------------------------------

### Initialize and Use KVzipPress

Source: https://context7.com/nvidia/kvpress/llms.txt

KVzipPress achieves near-lossless compression by scoring token contributions to context reconstruction, incurring a 2-3x computational overhead. Configure layerwise compression and the number of sink tokens.

```python
from transformers import pipeline
from kvpress import KVzipPress

pipe = pipeline("kv-press-text-generation", model="meta-llama/Meta-Llama-3.1-8B-Instruct", device_map="auto")

press = KVzipPress(
    compression_ratio=0.5,
    layerwise=False,   # use non-uniform layer compression
    n_sink=4,
)

result = pipe(
    "A critical technical specification...",
    question="What are the safety requirements?",
    press=press,
    max_new_tokens=200,
)
print(result["answer"])
```

--------------------------------

### Get Prefilling Statistics with Dynamic Cache

Source: https://github.com/nvidia/kvpress/blob/main/notebooks/speed_and_memory.ipynb

Measures GPU memory and prefilling time for a given model and input size. Resets CUDA stats and empties cache before measurement. Loads model within the function to prevent GPU leakage.

```python
def get_prefilling_stats(press, n_tokens, cache_implementation="dynamic"):
    torch.cuda.reset_peak_memory_stats()
    torch.cuda.empty_cache()
    idle_peak_memory = torch.cuda.max_memory_allocated()
    model = AutoModelForCausalLM.from_pretrained(ckpt, dtype="auto", attn_implementation="flash_attention_2").to(device)
    initial_peak_memory = torch.cuda.max_memory_allocated()

    inputs =torch.arange(n_tokens).reshape([1, n_tokens]).to(device)
    # Model warmup (for better prefilling time estimation)
    with torch.no_grad():
        model(inputs[:, :100])
    torch.cuda.synchronize()
    torch.cuda.empty_cache()

    # Compute cache size and prefilling time
    with torch.no_grad(), press(model):
        if cache_implementation == "dynamic":
            cache = DynamicCache()
        elif cache_implementation == "quantized":
            cache = QuantizedCache(backend="quanto", config=model.config, nbits=4)
        else:
            raise NotImplementedError(f"Cache {cache_implementation} not yet implemented")

        start = time()
        model(inputs, num_logits_to_keep=1, past_key_values=cache)
        prefilling_time = time() - start

        cache_size = get_size_of_cache(cache)
        del cache
        
    
    peak_memory = torch.cuda.max_memory_allocated()
    model.cpu()
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()
    return {"Idle Peak memory": idle_peak_memory  / 1024**3,
            "Initial Peak memory": initial_peak_memory / 1024**3,
            "Prefilling time": prefilling_time, 
            "Peak memory usage": peak_memory / 1024**3,
            "Cache Size": cache_size / 1024**3,
            "Peak memory w/o weights and KV cache (GB)": (peak_memory - cache_size - initial_peak_memory) / 1024**3
           }
```

--------------------------------

### Apply Press as Context Manager

Source: https://github.com/nvidia/kvpress/blob/main/README.md

Demonstrates how to apply a press (e.g., KnormPress) to a model using a context manager to observe changes in KV cache shape during inference.

```python
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])
```

--------------------------------

### Use KVPressTextGenerationPipeline with Multiple Questions

Source: https://context7.com/nvidia/kvpress/llms.txt

Shows how to process multiple questions using the same compressed KV cache from the `KVPressTextGenerationPipeline`. This is efficient for multi-turn Q&A.

```python
from transformers import pipeline
from kvpress import ExpectedAttentionPress

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

context = "The Treaty of Westphalia (1648) ended the Thirty Years' War..."  # long document
questions = ["When was the treaty signed?", "Which war did it end?"]

press = ExpectedAttentionPress(compression_ratio=0.4)

# Multiple questions from a single compressed cache
result = pipe(context, questions=questions, press=press, max_new_tokens=64)
for q, a in zip(questions, result["answers"]):
    print(f"Q: {q}
A: {a}
")
```

--------------------------------

### Benchmark Loop for KVPress

Source: https://github.com/nvidia/kvpress/blob/main/notebooks/speed_and_memory.ipynb

Iterates through different token counts and compression ratios to collect prefilling and generation statistics. Uses tqdm for progress visualization.

```python
stats = {}

compression_ratios = [0.0, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8]
for n_tokens in [8_000, 16_000, 32_000, 64_000, 128_000]:
    prefilling_stats = {compression_ratio : get_prefilling_stats(press=KnormPress(compression_ratio), n_tokens=n_tokens)
                        for compression_ratio in tqdm(compression_ratios)}
    generation_stats = {compression_ratio : get_generation_stats(press=KnormPress(compression_ratio), n_tokens=n_tokens)
```

--------------------------------

### Initialize KVpress strategy

Source: https://github.com/nvidia/kvpress/blob/main/notebooks/wikipedia_demo.ipynb

Selects a KVpress strategy for context compression. ExpectedAttentionPress is used here with a specified compression ratio.

```python
# Pick a press with a compression ratio, you can run the following cells with different presses
compression_ratio = 0.5
press = ExpectedAttentionPress(compression_ratio)
# press = KnormPress(compression_ratio)
# press = RandomPress(compression_ratio)
```

--------------------------------

### Configure Transformers Logging

Source: https://github.com/nvidia/kvpress/blob/main/notebooks/speed_and_memory.ipynb

Sets up the environment by ignoring warnings, setting the Transformers logging verbosity to error, and disabling progress bars for a cleaner output.

```python
warnings.filterwarnings("ignore")
transformers.logging.set_verbosity_error()
disable_progress_bar()
```

--------------------------------

### Load AIME25 Dataset Sample

Source: https://github.com/nvidia/kvpress/blob/main/notebooks/kvpress_decoding_aime25.ipynb

Loads a sample from the AIME25 dataset using the `datasets` library. This is used to obtain a math problem for generation.

```python
# Load a sample
dataset = load_dataset("math-ai/aime25")
sample = dataset["test"][0]
sample
```

--------------------------------

### Use KVPress with Pipeline

Source: https://github.com/nvidia/kvpress/blob/main/README.md

Demonstrates how to use KVPress within a generation pipeline. Ensure that the base press is a ScorerPress for compatibility.

```python
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"]
```

--------------------------------

### Collect Prefilling and Generation Stats (bfloat16)

Source: https://github.com/nvidia/kvpress/blob/main/notebooks/speed_and_memory.ipynb

Collects prefilling and generation statistics for different compression ratios using bfloat16 precision. Asserts that idle peak memory is negligible.

```python
stats = {}

compression_ratios = [0.0, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8]
for n_tokens in [8_000, 16_000, 32_000, 64_000, 128_000]:
    prefilling_stats = {compression_ratio : get_prefilling_stats(press=KnormPress(compression_ratio), n_tokens=n_tokens, cache_implementation="bfloat16")
                        for compression_ratio in tqdm(compression_ratios)}
    generation_stats = {compression_ratio : get_generation_stats(press=KnormPress(compression_ratio), n_tokens=n_tokens, cache_implementation="bfloat16")
                        for compression_ratio in tqdm(compression_ratios)}
        # ensure idle Peak memory is neglibible
    for compression_ratio in compression_ratios:
        assert prefilling_stats[compression_ratio]["Idle Peak memory"] < 0.01
        assert generation_stats[compression_ratio]["Idle Peak memory"] < 0.01
    stats[n_tokens] = combine_stats(prefilling_stats, generation_stats)
```

--------------------------------

### Collect Prefilling and Generation Stats (Quantized)

Source: https://github.com/nvidia/kvpress/blob/main/notebooks/speed_and_memory.ipynb

Collects prefilling and generation statistics for different compression ratios using quantized (4-bit) precision. Asserts that idle peak memory is negligible.

```python
stats_quantized = {}

compression_ratios = [0.0, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8]
for n_tokens in [8_000, 16_000, 32_000, 64_000, 128_000]:
    prefilling_stats = {compression_ratio : get_prefilling_stats(press=KnormPress(compression_ratio), n_tokens=n_tokens, cache_implementation="quantized")
                        for compression_ratio in tqdm(compression_ratios)}
    generation_stats = {compression_ratio : get_generation_stats(press=KnormPress(compression_ratio), n_tokens=n_tokens, cache_implementation="quantized")
                        for compression_ratio in tqdm(compression_ratios)}
        # ensure idle Peak memory is neglibible
    for compression_ratio in compression_ratios:
        assert prefilling_stats[compression_ratio]["Idle Peak memory"] < 0.01
        assert generation_stats[compression_ratio]["Idle Peak memory"] < 0.01
    stats_quantized[n_tokens] = combine_stats(prefilling_stats, generation_stats)
```

--------------------------------

### Import KVPress and Transformers Libraries

Source: https://github.com/nvidia/kvpress/blob/main/notebooks/kvpress_decoding_aime25.ipynb

Imports necessary libraries from `transformers`, `datasets`, and `kvpress` for text generation and KV cache compression.

```python
from transformers import pipeline
from datasets import load_dataset

from kvpress import (
    DecodingPress,
    ExpectedAttentionPress,
    KnormPress,
    ObservedAttentionPress,
    RandomPress,
    SnapKVPress,
    StreamingLLMPress,
    TOVAPress,
)
```

--------------------------------

### PyramidKVPress: Layered Budget Allocation

Source: https://context7.com/nvidia/kvpress/llms.txt

PyramidKVPress extends SnapKVPress with a pyramid-shaped per-layer budget, retaining more tokens in lower layers and fewer in higher layers. The `beta` parameter controls the pyramid's steepness.

```python
from transformers import pipeline
from kvpress import PyramidKVPress

pipe = pipeline("kv-press-text-generation", model="meta-llama/Meta-Llama-3.1-8B-Instruct", device_map="auto")

press = PyramidKVPress(
    compression_ratio=0.5,
    window_size=64,
    kernel_size=5,
    beta=20,  # higher = steeper pyramid
)

result = pipe(
    "A long academic paper on transformer architectures...",
    question="What is the proposed attention mechanism?",
    press=press,
    max_new_tokens=150,
)
print(result["answer"])
```

--------------------------------

### Apply KnormPress using Context Manager

Source: https://context7.com/nvidia/kvpress/llms.txt

Demonstrates using `KnormPress` as a context manager with `model.generate` to apply compression during inference. Hooks are automatically removed after the `with` block.

```python
import torch
from dataclasses import dataclass
from transformers import AutoModelForCausalLM, AutoTokenizer
from kvpress import BasePress, KnormPress

device = "cuda:0"
model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3.1-8B-Instruct").to(device)
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3.1-8B-Instruct")

press = KnormPress(compression_ratio=0.4)

inputs = tokenizer("A long context passage...", return_tensors="pt").to(device)

# Without compression
with torch.no_grad():
    out = model(**inputs)
    print(out.past_key_values[0][0].shape)  # (batch, heads, seq_len, head_dim)

# With compression — hooks are removed automatically after the with block
with torch.no_grad(), press(model):
    out = model(**inputs)
    print(out.past_key_values[0][0].shape)  # seq_len reduced by 40%
```

--------------------------------

### Prepare Input Tokens for the Model

Source: https://github.com/nvidia/kvpress/blob/main/notebooks/new_press.ipynb

Tokenizes the input context and question using the pipeline's tokenizer and moves the tokens to the specified device (e.g., GPU).

```python
# Load data

context = "In this step-by-step guide, you will learn how to create a new press in kvpress !"
question = "\nWhat is the purpose of this guide?"
tokens = pipe.tokenizer(context, return_tensors="pt").to(device)
```

--------------------------------

### Initialize and Use PrefillDecodingPress

Source: https://context7.com/nvidia/kvpress/llms.txt

PrefillDecodingPress combines distinct prefilling and decoding compression strategies for maximum memory savings across both generation phases. Configure separate presses for prefilling and decoding stages.

```python
from transformers import pipeline
from kvpress import PrefillDecodingPress, CriticalKVPress, DecodingPress, KnormPress

pipe = pipeline(
    "kv-press-text-generation",
    model="meta-llama/Meta-Llama-3.1-8B-Instruct",
    device_map="auto",
    model_kwargs={"attn_implementation": "flash_attention_2"},
)

prefill_press = CriticalKVPress(
    press=KnormPress(compression_ratio=0.4),
    first_stage_ratio=0.5,
)
decoding_press = DecodingPress(
    base_press=KnormPress(),
    compression_interval=20,
    target_size=256,
)

combined = PrefillDecodingPress(
    prefilling_press=prefill_press,
    decoding_press=decoding_press,
)

result = pipe(
    "A long reference document...",
    question="Generate a comprehensive 500-word essay about this topic.",
    press=combined,
    max_new_tokens=600,
)
print(result["answer"])
```

--------------------------------

### Apply KnormPress to Compress KV Cache

Source: https://github.com/nvidia/kvpress/blob/main/notebooks/new_press.ipynb

Demonstrates applying a KnormPress to compress the KV cache of a model. It compares the cache shapes before and after compression and prints the generated answer using the press.

```python
compression_ratio = 0.25
press = KnormPress(compression_ratio)

with torch.no_grad():
    outputs_without_press = pipe.model(**tokens, output_hidden_states=True)

with torch.no_grad(), press(pipe.model):
    output_with_press = pipe.model(**tokens)

print(f"Cache shape w/o press: {outputs_without_press.past_key_values[0][0].shape}")
print(f"Cache shape w/ press:  {output_with_press.past_key_values[0][0].shape}\n")

# The `KVPressTextGenerationPipeline` simply applies the `press` as above on the context tokens (see `_forward` method for more details).
print(pipe(context, question=question, press=press)["answer"])
```

--------------------------------

### Dataset Directory Structure

Source: https://github.com/nvidia/kvpress/blob/main/evaluation/README.md

Illustrates the standard directory structure for datasets used in evaluation. Each dataset includes scripts for creating Hugging Face datasets and calculating metrics.

```bash
$dataset
├── README.md
├── calculate_metrics.py
├── create_huggingface_dataset.py
```

--------------------------------

### Load pipeline and data

Source: https://github.com/nvidia/kvpress/blob/main/notebooks/expected_attention.ipynb

Initializes a transformer pipeline and loads data from a Wikipedia URL. It then tokenizes the content and prints the number of tokens. Ensure the model is on the correct device (GPU recommended).

```python
# Load pipeline and data

device = "cuda:0"
ckpt = "meta-llama/Meta-Llama-3.1-8B-Instruct"
# ckpt = "mistralai/Mistral-Nemo-Instruct-2407"
# ckpt = "microsoft/Phi-3.5-mini-instruct"
pipe = pipeline("kv-press-text-generation", model=ckpt, device=device, dtype="auto", model_kwargs={"attn_implementation":"flash_attention_2"})

# Load data
url = "https://en.wikipedia.org/wiki/Nvidia"
content = requests.get(url).content
soup = BeautifulSoup(content, "html.parser")
context = "".join([p.text for p in soup.find_all("p")]) + "\n\n"
tokens = pipe.tokenizer.encode(context, return_tensors="pt").to(device)
n_tokens = tokens.size(1)
print(f"Number of tokens: {n_tokens}")

# Get hidden states
with torch.no_grad():
    outputs = pipe.model(tokens, output_hidden_states=True)
```

--------------------------------

### Load and preprocess dataset

Source: https://github.com/nvidia/kvpress/blob/main/notebooks/per_layer_compression_demo.ipynb

Loads a dataset from Hugging Face datasets library and filters it to keep only the 'niah_single_3' task for demonstration purposes.

```python
import datasets 

df = datasets.load_dataset("simonjegou/ruler", "4096")["test"].to_pandas()
df = df.loc[df["task"] == "niah_single_3"].reset_index(drop=True)
```

--------------------------------

### Load Text Generation Pipeline

Source: https://github.com/nvidia/kvpress/blob/main/notebooks/kvpress_decoding_aime25.ipynb

Initializes a text generation pipeline using the `transformers` library with the specified model, device, and attention implementation. Ensure CUDA is available for GPU acceleration.

```python
# Load the pipeline
device = "cuda:0"
ckpt = "nvidia/OpenMath-Nemotron-7B"
attn_implementation = "flash_attention_2"
pipe = pipeline("kv-press-text-generation", model=ckpt, device=device, dtype="auto", model_kwargs={"attn_implementation":attn_implementation})
```

--------------------------------

### Use Press with model.generate

Source: https://github.com/nvidia/kvpress/blob/main/README.md

Shows how to use a press with the model.generate method by applying the press as a context manager. Note that this method has limitations regarding excluding questions from compression and batching.

```python
with press(model):
    outputs = model.generate(inputs)
```

--------------------------------

### Run pipeline with PerLayerCompressionPress

Source: https://github.com/nvidia/kvpress/blob/main/notebooks/per_layer_compression_demo.ipynb

Executes the pipeline using the PerLayerCompressionPress. This demonstrates the effect of applying layer-specific compression ratios on the model's predictions.

```python
pred_answer = pipe(context, question=question, press=press_per_layer)["answer"]

print(f"Question:   {question}")
print(f"Answer:     {true_answer}")
print(f"Prediction: {pred_answer}")
print(f"Correctly predicted: {true_answer in pred_answer}")
```

--------------------------------

### Load KV-Press Text Generation Pipeline

Source: https://github.com/nvidia/kvpress/blob/main/notebooks/new_press.ipynb

Loads a pre-trained language model for text generation using the KV-Press pipeline. Specify the model checkpoint, device (GPU recommended), and attention implementation for performance.

```python
# Load pipeline

device = "cuda:0"
ckpt = "Qwen/Qwen2.5-1.5B-Instruct"
attn_implementation = "flash_attention_2"
pipe = pipeline("kv-press-text-generation", model=ckpt, device=device, dtype="auto", model_kwargs={"attn_implementation":attn_implementation})
```

--------------------------------

### Initialize PerLayerCompressionPress

Source: https://github.com/nvidia/kvpress/blob/main/notebooks/per_layer_compression_demo.ipynb

Initializes the PerLayerCompressionPress by wrapping an ExpectedAttentionPress with a list of per-layer compression ratios. This enables fine-grained compression control across model layers.

```python
press_per_layer = PerLayerCompressionPress(ExpectedAttentionPress(compression_ratio), PHI_35_COMPRESSION_RATIOS)
```

--------------------------------

### Implement SnapKVPress for Attention-Based Compression

Source: https://context7.com/nvidia/kvpress/llms.txt

SnapKVPress scores tokens by the average attention weight they receive from recent query tokens, smoothed with a pooling kernel. Configure compression ratio, window size for query tokens, and kernel size for smoothing.

```python
from transformers import pipeline
from kvpress import SnapKVPress

pipe = pipeline("kv-press-text-generation", model="meta-llama/Meta-Llama-3.1-8B-Instruct", device_map="auto")

press = SnapKVPress(
    compression_ratio=0.5,
    window_size=64,   # use last 64 query tokens to estimate importance
    kernel_size=5,    # smoothing pool size
)

result = pipe(
    "A long scientific paper about protein folding...",
    question="What is AlphaFold?",
    press=press,
    max_new_tokens=100,
)
print(result["answer"])
```

--------------------------------

### Import necessary libraries

Source: https://github.com/nvidia/kvpress/blob/main/notebooks/wikipedia_demo.ipynb

Imports libraries for making HTTP requests, parsing HTML, and using the transformers pipeline with KVpress components.

```python
import requests
from bs4 import BeautifulSoup

from transformers import pipeline

from kvpress import ExpectedAttentionPress, KnormPress, RandomPress
```

--------------------------------

### Create and Configure DecodingPress

Source: https://github.com/nvidia/kvpress/blob/main/README.md

Instantiate a DecodingPress to compress the KV cache periodically during token generation. This press uses a base press, compression interval, and target token buffer size.

```python
# Create a decoding press that compresses every 10 steps to 512 tokens
decoding_press = DecodingPress(
    base_press=KnormPress(),
    compression_steps=10,
    token_buffer_size=512
)
```

--------------------------------

### Combine Prefilling and Generation Statistics

Source: https://github.com/nvidia/kvpress/blob/main/notebooks/speed_and_memory.ipynb

Merges prefilling and generation performance data into a single dictionary for analysis and plotting. Calculates generation time by subtracting prefilling time from total time.

```python
def combine_stats(prefilling_stats, generation_stats):
    """Combines prefilling and generation data, then plots."""
    combined_stats = {}
    for compression_ratio in prefilling_stats:
        combined_stats[compression_ratio] = dict()
        combined_stats[compression_ratio]['Peak memory usage'] = generation_stats[compression_ratio]['Peak memory usage']
        combined_stats[compression_ratio]['Prefilling time'] = prefilling_stats[compression_ratio]['Prefilling time']
        combined_stats[compression_ratio]['Cache Size'] = prefilling_stats[compression_ratio]['Cache Size']
        combined_stats[compression_ratio]['Total time'] = generation_stats[compression_ratio]['Total time']
        combined_stats[compression_ratio]['Generation time'] = generation_stats[compression_ratio]['Total time'] - prefilling_stats[compression_ratio]['Prefilling time']
        
    return combined_stats
```

--------------------------------

### Evaluate KnormPress on RULER Benchmark

Source: https://context7.com/nvidia/kvpress/llms.txt

Evaluates KnormPress on the RULER benchmark with specified model, context length, and compression ratio. Results are saved to the output directory.

```bash
python -m kvpress.evaluation \
    --model meta-llama/Meta-Llama-3.1-8B-Instruct \
    --press KnormPress \
    --compression_ratio 0.5 \
    --benchmark ruler \
    --context_length 4096 \
    --output_dir results/
```

--------------------------------

### Combined Prefill + Decoding Compression with KVPress

Source: https://github.com/nvidia/kvpress/blob/main/README.md

Utilize PrefillDecodingPress to apply different compression strategies for the initial context (prefill) and subsequent generated text (decoding). This allows for optimized compression throughout the generation process.

```python
from transformers import pipeline
from kvpress import CriticalKVPress, KnormPress
from kvpress import DecodingPress, PrefillDecodingPress

# 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)

# Different strategies for prefill vs decoding
prefill_press = CriticalKVPress(KnormPress())
decoding_press = DecodingPress(
    base_press=KnormPress(compression_ratio=0.2),
    compression_steps=5,
    token_buffer_size=256
)

# Combine them
combined_press = PrefillDecodingPress(
    prefilling_press=prefill_press,
    decoding_press=decoding_press
)

context = "A very long context that will be compressed during prefill"
question = "Generate a detailed analysis that will be compressed during decoding"
response = pipe(context, question=question, press=combined_press)["answer"]
```

--------------------------------

### Implement KnormPress for Key Vector Compression

Source: https://context7.com/nvidia/kvpress/llms.txt

KnormPress prunes tokens with the smallest L2 norm in their key vectors. This method is training-free, fast, and compatible with Flash Attention. Ensure the compression ratio is set to control the percentage of KV pairs retained.

```python
from transformers import pipeline
from kvpress import KnormPress

pipe = pipeline("kv-press-text-generation", model="meta-llama/Meta-Llama-3.1-8B-Instruct", device_map="auto")

press = KnormPress(compression_ratio=0.5)  # keep 50% of KV pairs

result = pipe(
    "Einstein published the special theory of relativity in 1905...",  # long context
    question="What year was the special theory of relativity published?",
    press=press,
    max_new_tokens=50,
)
print(result["answer"])  # "1905"
```

--------------------------------

### Sequential Compression with ComposedPress

Source: https://context7.com/nvidia/kvpress/llms.txt

Chain multiple compression strategies sequentially. Each press compresses the output of the previous one, useful for combining orthogonal compression methods.

```python
from transformers import pipeline
from kvpress import ComposedPress, KnormPress, StreamingLLMPress

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

# First keep tokens by k-norm score, then apply streaming window
press = ComposedPress(presses=[
    KnormPress(compression_ratio=0.3),
    StreamingLLMPress(compression_ratio=0.3, n_sink=4),
])

result = pipe(
    "A long transcript...",
    question="What was the conclusion?",
    press=press,
    max_new_tokens=100,
)
print(result["answer"])
```

--------------------------------

### Combine Prefill and Decode Presses

Source: https://github.com/nvidia/kvpress/blob/main/README.md

Illustrates the concept of combining separate presses for the prefilling and decoding phases of inference.

```python
# Parameters:
# - prefilling_press: Press used during prefill phase
# - decoding_press: Press used during decoding phase
```

--------------------------------

### Run Evaluation with Custom Parameters

Source: https://github.com/nvidia/kvpress/blob/main/evaluation/README.md

Execute the evaluation script with specific dataset, model, and press configurations. This allows for fine-grained control over the evaluation process.

```bash
python evaluate.py --dataset loogle --data_dir shortdep_qa --model meta-llama/Meta-Llama-3.1-8B-Instruct --press_name expected_attention --compression_ratio 0.5
```

--------------------------------

### Multi-GPU Inference with Accelerate

Source: https://github.com/nvidia/kvpress/blob/main/README.md

Configure a Hugging Face pipeline for multi-GPU inference using the 'accelerate' library by setting device_map to 'auto'.

```python
pipe = pipeline("kv-press-text-generation", model=model, device_map="auto")
```

--------------------------------

### Import Libraries

Source: https://github.com/nvidia/kvpress/blob/main/notebooks/speed_and_memory.ipynb

Imports all necessary libraries for the KV press pipeline, including tools for timing, progress bars, serialization, plotting, numerical operations, and Hugging Face Transformers components.

```python
import warnings
from time import time
from tqdm import tqdm
import pickle
import os

import matplotlib.pylab as plt
from matplotlib.colors import LinearSegmentedColormap

import numpy as np
import torch
from transformers import AutoModelForCausalLM, pipeline
from transformers import DynamicCache, QuantizedCache
from transformers.utils.logging import disable_progress_bar
import transformers

from kvpress import KnormPress
```

--------------------------------

### Import necessary libraries

Source: https://github.com/nvidia/kvpress/blob/main/notebooks/expected_attention.ipynb

Imports essential libraries for data loading, model interaction, and visualization.

```python
import requests
import numpy as np
from bs4 import BeautifulSoup
import matplotlib.pyplot as plt

import torch
from transformers import pipeline
from transformers.modeling_attn_mask_utils import AttentionMaskConverter


from kvpress import ExpectedAttentionPress
```

--------------------------------

### Load Hugging Face pipeline

Source: https://github.com/nvidia/kvpress/blob/main/notebooks/per_layer_compression_demo.ipynb

Loads a text generation pipeline from Hugging Face, specifying the model, device, data type, and attention implementation. Ensure the model is moved to GPU if using Flash Attention 2.

```python
# Load pipeline

device = "cuda:0"
ckpt = "microsoft/Phi-3.5-mini-instruct"
attn_implementation = "flash_attention_2"
pipe = pipeline("kv-press-text-generation", model=ckpt, device=device, dtype="auto", model_kwargs={"attn_implementation":attn_implementation})
```

--------------------------------

### Key Dimension Compression with ThinKPress

Source: https://context7.com/nvidia/kvpress/llms.txt

Reduce the key head dimension by zeroing out low-importance key channels. Importance is computed from query-key norm products.

```python
from transformers import pipeline
from kvpress import ThinKPress

pipe = pipeline("kv-press-text-generation", model="meta-llama/Meta-Llama-3.1-8B-Instruct", device_map="auto")

press = ThinKPress(
    key_channel_compression_ratio=0.5,  # remove 50% of key channels
    window_size=32,
)
# press.compression_ratio == 0.25 (half of key_channel_compression_ratio)

result = pipe(
    "A long article about renewable energy...",
    question="What is the most efficient solar panel type?",
    press=press,
    max_new_tokens=80,
)
print(result["answer"])
```

--------------------------------

### Import necessary libraries for KVPress

Source: https://github.com/nvidia/kvpress/blob/main/notebooks/new_press.ipynb

Imports required modules from dataclasses, contextlib, torch, and transformers for KVPress functionality.

```python
from dataclasses import dataclass
from contextlib import contextmanager

import torch
from torch import nn
from transformers import pipeline

from kvpress import BasePress, KnormPress, ScorerPress
```

--------------------------------

### Run Evaluation with Custom Config File

Source: https://github.com/nvidia/kvpress/blob/main/evaluation/README.md

Execute the evaluation script using a custom YAML configuration file. This is useful for managing complex or reusable evaluation settings.

```bash
python evaluate.py --config_file <your_config.yaml>
```

--------------------------------

### Load KVpress text generation pipeline

Source: https://github.com/nvidia/kvpress/blob/main/notebooks/wikipedia_demo.ipynb

Loads the specified transformer model for text generation using KVpress. Ensure 'flash_attention_2' is available or use 'eager'/'sdpa' as alternatives.

```python
# Load pipeline

device = "cuda:0"
ckpt = "Qwen/Qwen3-8B"
attn_implementation = "flash_attention_2"  # use "eager" for ObservedAttentionPress and "sdpa" if you can't use "flash_attention_2"
pipe = pipeline("kv-press-text-generation", model=ckpt, device=device, dtype="auto", model_kwargs={"attn_implementation":attn_implementation})
```

--------------------------------

### Initialize and Use DecodingPress

Source: https://context7.com/nvidia/kvpress/llms.txt

DecodingPress compresses the KV cache during token generation at specified intervals, using a base scorer and a buffer for recent hidden states. Adjust compression interval, target size, and buffer size for optimal performance.

```python
from transformers import pipeline
from kvpress import DecodingPress, KnormPress, SnapKVPress

pipe = pipeline(
    "kv-press-text-generation",
    model="meta-llama/Meta-Llama-3.1-8B-Instruct",
    device_map="auto",
    model_kwargs={"attn_implementation": "flash_attention_2"},
)

# Compress every 10 decoding steps, keeping at most 512 tokens
press = DecodingPress(
    base_press=KnormPress(),
    compression_interval=10,
    target_size=512,
    hidden_states_buffer_size=128,
)

result = pipe(
    "A very long context that grows during generation...",
    question="Write a detailed analysis of this document.",
    press=press,
    max_new_tokens=500,  # long generation benefits most from decoding compression
)
print(result["answer"])
```