From c14f7249b103a6d92b0b648e59f02cdff69308dd Mon Sep 17 00:00:00 2001
From: kschwethelm <k.schwethelm@tum.de>
Date: Wed, 1 Apr 2026 12:28:50 +0200
Subject: [PATCH 01/11] feat: integrate KVPress for KV cache compression

Add NVIDIA KVPress as an optional dependency, enabling 31 KV cache
compression strategies for causal language models. Includes algorithm
class, test tester, and compatibility updates across existing LLM
algorithms.
---
 pyproject.toml                                |   4 +
 src/pruna/algorithms/gptq_model.py            |   2 +-
 src/pruna/algorithms/hqq.py                   |   2 +-
 src/pruna/algorithms/huggingface_llm_int8.py  |   2 +-
 src/pruna/algorithms/kvpress.py               | 178 ++++++++++++++++++
 src/pruna/algorithms/llm_compressor.py        |   2 +-
 src/pruna/algorithms/moe_kernel_tuner.py      |   4 +-
 .../algorithms/torch_compile/torch_compile.py |   1 +
 src/pruna/algorithms/torchao.py               |   1 +
 tests/algorithms/testers/kvpress.py           |  19 ++
 10 files changed, 209 insertions(+), 6 deletions(-)
 create mode 100644 src/pruna/algorithms/kvpress.py
 create mode 100644 tests/algorithms/testers/kvpress.py

diff --git a/pyproject.toml b/pyproject.toml
index d70093bf..7c5dcbd1 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -87,6 +87,7 @@ conflicts = [
     [{ extra = "intel" }, { extra = "stable-fast" }, { extra = "stable-fast-extraindex" }],
     [{ extra = "intel" }, { extra = "full" }, { extra = "stable-fast-extraindex" }],
     [{ extra = "intel" }, { extra = "vllm" }],
+    [{ extra = "kvpress" }, { extra = "vbench" }],
 ]
 
 [tool.uv.sources]
@@ -248,6 +249,9 @@ intel = [
     "torch>=2.7.0,<2.9.0",
     "torchvision>=0.22.0,<0.24.0",
 ]
+kvpress = [
+    "kvpress>=0.4.3,<0.5.0",
+]
 
 [build-system]
 requires = ["hatchling"]
diff --git a/src/pruna/algorithms/gptq_model.py b/src/pruna/algorithms/gptq_model.py
index 18daaa5b..11d5573e 100644
--- a/src/pruna/algorithms/gptq_model.py
+++ b/src/pruna/algorithms/gptq_model.py
@@ -46,7 +46,7 @@ class GPTQ(PrunaAlgorithmBase):
     processor_required: bool = False
     runs_on: list[str] = ["cuda"]
     dataset_required: bool = True
-    compatible_after: Iterable[str] = ["torch_compile", "sage_attn"]
+    compatible_after: Iterable[str] = ["torch_compile", "sage_attn", "kvpress"]
     required_install: str = (
         "You must first install the base package with ``pip install pruna`` "
         "before installing the GPTQ extension with ``pip install pruna[gptq] --extra-index-url https://prunaai.pythonanywhere.com/``"
diff --git a/src/pruna/algorithms/hqq.py b/src/pruna/algorithms/hqq.py
index 263b0f0d..843ba143 100644
--- a/src/pruna/algorithms/hqq.py
+++ b/src/pruna/algorithms/hqq.py
@@ -63,7 +63,7 @@ class HQQ(PrunaAlgorithmBase):
     runs_on: list[str] = ["cuda"]
     dataset_required: bool = False
     compatible_before: Iterable[str] = ["torch_structured", "moe_kernel_tuner"]
-    compatible_after: Iterable[str] = ["torch_compile", "sage_attn", "moe_kernel_tuner"]
+    compatible_after: Iterable[str] = ["torch_compile", "sage_attn", "kvpress", "moe_kernel_tuner"]
     disjointly_compatible_before: Iterable[str] = []
     disjointly_compatible_after: Iterable[str] = ["torchao"]
 
diff --git a/src/pruna/algorithms/huggingface_llm_int8.py b/src/pruna/algorithms/huggingface_llm_int8.py
index 49cb1ead..3e6fb289 100644
--- a/src/pruna/algorithms/huggingface_llm_int8.py
+++ b/src/pruna/algorithms/huggingface_llm_int8.py
@@ -58,7 +58,7 @@ class LLMInt8(PrunaAlgorithmBase):
     runs_on: list[str] = ["cuda", "accelerate"]
     save_fn: None = None
     compatible_before: Iterable[str] = ["moe_kernel_tuner"]
-    compatible_after: Iterable[str] = ["torch_compile", "sage_attn", "moe_kernel_tuner"]
+    compatible_after: Iterable[str] = ["torch_compile", "sage_attn", "kvpress", "moe_kernel_tuner"]
 
     def get_hyperparameters(self) -> list:
         """
diff --git a/src/pruna/algorithms/kvpress.py b/src/pruna/algorithms/kvpress.py
new file mode 100644
index 00000000..f1cd682f
--- /dev/null
+++ b/src/pruna/algorithms/kvpress.py
@@ -0,0 +1,178 @@
+# Copyright 2025 - Pruna AI GmbH. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import functools
+from collections.abc import Iterable
+from typing import Any, Dict
+
+from ConfigSpace import CategoricalHyperparameter, UniformFloatHyperparameter
+
+from pruna.algorithms.base.pruna_base import PrunaAlgorithmBase
+from pruna.algorithms.base.tags import AlgorithmTag as tags
+from pruna.config.smash_config import SmashConfigPrefixWrapper
+from pruna.engine.model_checks import is_causal_lm, is_transformers_pipeline_with_causal_lm
+from pruna.engine.save import SAVE_FUNCTIONS
+
+PRESS_TYPES = [
+    "AdaKVPress",
+    "BlockPress",
+    "ChunkKVPress",
+    "ChunkPress",
+    "CompactorPress",
+    "CriticalAdaKVPress",
+    "CriticalKVPress",
+    "CURPress",
+    "DMSPress",
+    "DuoAttentionPress",
+    "ExpectedAttentionPress",
+    "ExpectedAttentionStatsPress",
+    "FinchPress",
+    "KnormPress",
+    "KVzapPress",
+    "KVzipPress",
+    "KeyDiffPress",
+    "KeyRerotationPress",
+    "LagKVPress",
+    "LeverageScorePress",
+    "NonCausalAttnPress",
+    "ObservedAttentionPress",
+    "PerLayerCompressionPress",
+    "PyramidKVPress",
+    "QFilterPress",
+    "RandomPress",
+    "SimLayerKVPress",
+    "SnapKVPress",
+    "StreamingLLMPress",
+    "ThinKPress",
+    "TOVAPress",
+]
+
+
+class KVPress(PrunaAlgorithmBase):
+    """
+    Integrate KVPress for KV cache compression in causal language models.
+
+    KVPress is a library by NVIDIA that compresses the key-value cache during the prefill phase
+    of autoregressive generation, reducing memory usage for long-context inference. It provides
+    over 30 compression strategies (presses) that score and prune KV pairs based on different
+    importance criteria.
+    """
+
+    algorithm_name: str = "kvpress"
+    group_tags: list[tags] = [tags.CACHER]
+    save_fn: SAVE_FUNCTIONS = SAVE_FUNCTIONS.reapply
+    references: dict[str, str] = {
+        "GitHub": "https://github.com/NVIDIA/kvpress",
+        "Article": "https://huggingface.co/blog/nvidia/kvpress",
+    }
+    required_install: str = "pip install pruna[kvpress]"
+    tokenizer_required: bool = False
+    processor_required: bool = False
+    dataset_required: bool = False
+    runs_on: list[str] = ["cuda"]
+    compatible_before: Iterable[str | tags] = [tags.QUANTIZER, "moe_kernel_tuner"]
+    compatible_after: Iterable[str | tags] = ["torch_compile", "sage_attn", "moe_kernel_tuner"]
+
+    def get_hyperparameters(self) -> list:
+        """
+        Configure all algorithm-specific hyperparameters with ConfigSpace.
+
+        Returns
+        -------
+        list
+            The hyperparameters.
+        """
+        return [
+            CategoricalHyperparameter(
+                "press_type",
+                choices=PRESS_TYPES,
+                default_value="ExpectedAttentionPress",
+                meta={"desc": "The KV cache compression strategy to use."},
+            ),
+            UniformFloatHyperparameter(
+                "compression_ratio",
+                lower=0.0,
+                upper=1.0,
+                default_value=0.5,
+                meta={"desc": "Fraction of KV pairs to remove. 0.0 means no compression."},
+            ),
+        ]
+
+    def model_check_fn(self, model: Any) -> bool:
+        """
+        Check if the model is a causal language model.
+
+        Parameters
+        ----------
+        model : Any
+            The model to check.
+
+        Returns
+        -------
+        bool
+            True if the model is a causal language model, False otherwise.
+        """
+        return is_causal_lm(model) or is_transformers_pipeline_with_causal_lm(model)
+
+    def _apply(self, model: Any, smash_config: SmashConfigPrefixWrapper) -> Any:
+        """
+        Apply KV cache compression to the model.
+
+        Parameters
+        ----------
+        model : Any
+            The model to apply KV cache compression to.
+        smash_config : SmashConfigPrefixWrapper
+            The configuration for the compression.
+
+        Returns
+        -------
+        Any
+            The model with KV cache compression applied to its generate method.
+        """
+        imported_modules = self.import_algorithm_packages()
+
+        press_type = smash_config["press_type"]
+        compression_ratio = smash_config["compression_ratio"]
+
+        press_cls = imported_modules[press_type]
+        press = press_cls(compression_ratio=compression_ratio)
+
+        original_generate = model.generate
+
+        @functools.wraps(original_generate)
+        def generate_with_press(*args, **kwargs):
+            with press(model):
+                return original_generate(*args, **kwargs)
+
+        model.generate = generate_with_press
+        model._kvpress_original_generate = original_generate
+        model._kvpress_press = press
+
+        return model
+
+    def import_algorithm_packages(self) -> Dict[str, Any]:
+        """
+        Import the kvpress package lazily.
+
+        Returns
+        -------
+        Dict[str, Any]
+            A dictionary mapping press class names to their classes.
+        """
+        import kvpress
+
+        return {name: getattr(kvpress, name) for name in PRESS_TYPES}
diff --git a/src/pruna/algorithms/llm_compressor.py b/src/pruna/algorithms/llm_compressor.py
index 146d154b..f964ce91 100644
--- a/src/pruna/algorithms/llm_compressor.py
+++ b/src/pruna/algorithms/llm_compressor.py
@@ -52,7 +52,7 @@ class LLMCompressor(PrunaAlgorithmBase):
     dataset_required: bool = True
     runs_on: list[str] = ["cuda"]
     compatible_before: Iterable[str] = ["moe_kernel_tuner"]
-    compatible_after: Iterable[str] = ["sage_attn", "moe_kernel_tuner"]
+    compatible_after: Iterable[str] = ["sage_attn", "kvpress", "moe_kernel_tuner"]
     required_install = "``uv pip install 'pruna[awq]'``"
 
     def get_hyperparameters(self) -> list:
diff --git a/src/pruna/algorithms/moe_kernel_tuner.py b/src/pruna/algorithms/moe_kernel_tuner.py
index ce9411fb..062bf9d8 100644
--- a/src/pruna/algorithms/moe_kernel_tuner.py
+++ b/src/pruna/algorithms/moe_kernel_tuner.py
@@ -50,13 +50,13 @@ class MoeKernelTuner(PrunaAlgorithmBase):
     dataset_required: bool = False
     compatible_before: Iterable[str] = [
         "awq", "deepcache", "diffusers_int8", "fastercache", "flash_attn3",
-        "fora", "hqq", "hqq_diffusers", "llm_int8", "pab", "padding_pruning",
+        "fora", "hqq", "hqq_diffusers", "kvpress", "llm_int8", "pab", "padding_pruning",
         "qkv_diffusers", "quanto", "reduce_noe", "ring_attn", "sage_attn",
         "torch_compile", "torchao",
     ]
     compatible_after: Iterable[str] = [
         "awq", "deepcache", "diffusers_int8", "fastercache", "flash_attn3",
-        "fora", "hqq", "hqq_diffusers", "llm_int8", "pab", "padding_pruning",
+        "fora", "hqq", "hqq_diffusers", "kvpress", "llm_int8", "pab", "padding_pruning",
         "qkv_diffusers", "quanto", "ring_attn", "sage_attn",
         "torch_compile", "torchao",
     ]
diff --git a/src/pruna/algorithms/torch_compile/torch_compile.py b/src/pruna/algorithms/torch_compile/torch_compile.py
index a04e304a..72574914 100644
--- a/src/pruna/algorithms/torch_compile/torch_compile.py
+++ b/src/pruna/algorithms/torch_compile/torch_compile.py
@@ -74,6 +74,7 @@ class TorchCompile(PrunaAlgorithmBase):
         "hyper",
         "padding_pruning",
         "ring_attn",
+        "kvpress",
         "text_to_image_distillation_inplace_perp",
         "text_to_image_distillation_lora",
         "text_to_image_distillation_perp",
diff --git a/src/pruna/algorithms/torchao.py b/src/pruna/algorithms/torchao.py
index aa0e988c..d3411044 100644
--- a/src/pruna/algorithms/torchao.py
+++ b/src/pruna/algorithms/torchao.py
@@ -99,6 +99,7 @@ class Torchao(PrunaAlgorithmBase):
         "fora",
         "torch_compile",
         "sage_attn",
+        "kvpress",
         "img2img_denoise",
         "realesrgan_upscale",
         "moe_kernel_tuner",
diff --git a/tests/algorithms/testers/kvpress.py b/tests/algorithms/testers/kvpress.py
new file mode 100644
index 00000000..d41d272e
--- /dev/null
+++ b/tests/algorithms/testers/kvpress.py
@@ -0,0 +1,19 @@
+from pruna import PrunaModel
+from pruna.algorithms.kvpress import KVPress
+
+from .base_tester import AlgorithmTesterBase
+
+
+class TestKVPress(AlgorithmTesterBase):
+    """Test the KVPress KV cache compression algorithm."""
+
+    models = ["llama_3_tiny_random"]
+    reject_models = ["sd_tiny_random"]
+    allow_pickle_files = False
+    algorithm_class = KVPress
+    metrics = ["perplexity"]
+
+    def post_smash_hook(self, model: PrunaModel) -> None:
+        """Verify that the press was applied to the model."""
+        assert hasattr(model, "_kvpress_press")
+        assert hasattr(model, "_kvpress_original_generate")

From becc89d701567ba04393eb7530032fa0f20713f4 Mon Sep 17 00:00:00 2001
From: kschwethelm <k.schwethelm@tum.de>
Date: Wed, 1 Apr 2026 17:38:25 +0200
Subject: [PATCH 02/11] feat: bump kvpress to >=0.5.2, add FastKVzipPress

kvpress 0.5.2 relaxes the datasets<3 constraint and reverts to
transformers>=4.56, resolving the dependency conflict. uv sync --extra
kvpress now works without workarounds.
---
 pyproject.toml                  | 2 +-
 src/pruna/algorithms/kvpress.py | 1 +
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index 7c5dcbd1..dc42053e 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -250,7 +250,7 @@ intel = [
     "torchvision>=0.22.0,<0.24.0",
 ]
 kvpress = [
-    "kvpress>=0.4.3,<0.5.0",
+    "kvpress>=0.5.2",
 ]
 
 [build-system]
diff --git a/src/pruna/algorithms/kvpress.py b/src/pruna/algorithms/kvpress.py
index f1cd682f..bc21e6e7 100644
--- a/src/pruna/algorithms/kvpress.py
+++ b/src/pruna/algorithms/kvpress.py
@@ -39,6 +39,7 @@
     "DuoAttentionPress",
     "ExpectedAttentionPress",
     "ExpectedAttentionStatsPress",
+    "FastKVzipPress",
     "FinchPress",
     "KnormPress",
     "KVzapPress",

From 1546da90e6d8b44144a05ff13cff9b47eb9c60a7 Mon Sep 17 00:00:00 2001
From: kschwethelm <k.schwethelm@tum.de>
Date: Fri, 3 Apr 2026 09:39:49 +0200
Subject: [PATCH 03/11] feat: add press_kwargs for press-specific parameters

Allow passing additional keyword arguments to the press constructor
via the press_kwargs hyperparameter, enabling fine-grained control
over press-specific settings like window_size, n_sink, etc.
---
 src/pruna/algorithms/kvpress.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/src/pruna/algorithms/kvpress.py b/src/pruna/algorithms/kvpress.py
index bc21e6e7..0f567888 100644
--- a/src/pruna/algorithms/kvpress.py
+++ b/src/pruna/algorithms/kvpress.py
@@ -22,6 +22,7 @@
 
 from pruna.algorithms.base.pruna_base import PrunaAlgorithmBase
 from pruna.algorithms.base.tags import AlgorithmTag as tags
+from pruna.config.hyperparameters import UnconstrainedHyperparameter
 from pruna.config.smash_config import SmashConfigPrefixWrapper
 from pruna.engine.model_checks import is_causal_lm, is_transformers_pipeline_with_causal_lm
 from pruna.engine.save import SAVE_FUNCTIONS
@@ -110,6 +111,11 @@ def get_hyperparameters(self) -> list:
                 default_value=0.5,
                 meta={"desc": "Fraction of KV pairs to remove. 0.0 means no compression."},
             ),
+            UnconstrainedHyperparameter(
+                "press_kwargs",
+                default_value=None,
+                meta={"desc": "Additional keyword arguments passed to the press constructor."},
+            ),
         ]
 
     def model_check_fn(self, model: Any) -> bool:
@@ -148,9 +154,10 @@ def _apply(self, model: Any, smash_config: SmashConfigPrefixWrapper) -> Any:
 
         press_type = smash_config["press_type"]
         compression_ratio = smash_config["compression_ratio"]
+        press_kwargs = smash_config["press_kwargs"] or {}
 
         press_cls = imported_modules[press_type]
-        press = press_cls(compression_ratio=compression_ratio)
+        press = press_cls(compression_ratio=compression_ratio, **press_kwargs)
 
         original_generate = model.generate
 

From 30e12c4fefb5840dbf34a9a04556de99eb75d714 Mon Sep 17 00:00:00 2001
From: kschwethelm <k.schwethelm@tum.de>
Date: Fri, 3 Apr 2026 12:05:53 +0200
Subject: [PATCH 04/11] fix: compatibility, press_kwargs, unit tests, remove
 wrappers

- Replace tags.QUANTIZER with explicit LLM algorithm names to avoid
  false symmetry matches with diffuser algorithms
- Fix SmashConfig.add() dict flattening: only flatten when key is a
  registered algorithm name, not for dict-valued hyperparameters
- Remove wrapper/special presses from PRESS_TYPES (CriticalKVPress
  and others that don't accept compression_ratio directly)
- Add unit tests for press type validation and kwargs forwarding
- Add SnapKV integration test with press_kwargs
---
 src/pruna/algorithms/half.py        |  1 +
 src/pruna/algorithms/kvpress.py     | 21 +++++++------------
 src/pruna/algorithms/quanto.py      |  1 +
 src/pruna/config/smash_config.py    |  5 +++--
 tests/algorithms/test_kvpress.py    | 31 +++++++++++++++++++++++++++++
 tests/algorithms/testers/kvpress.py | 25 ++++++++++++++++++++++-
 6 files changed, 67 insertions(+), 17 deletions(-)
 create mode 100644 tests/algorithms/test_kvpress.py

diff --git a/src/pruna/algorithms/half.py b/src/pruna/algorithms/half.py
index 53bcc029..e227b1c4 100644
--- a/src/pruna/algorithms/half.py
+++ b/src/pruna/algorithms/half.py
@@ -50,6 +50,7 @@ class Half(PrunaAlgorithmBase):
         "stable_fast",
         "torch_compile",
         "ifw",
+        "kvpress",
         "whisper_s2t",
         "sage_attn",
         "hyper",
diff --git a/src/pruna/algorithms/kvpress.py b/src/pruna/algorithms/kvpress.py
index 0f567888..1766e9a1 100644
--- a/src/pruna/algorithms/kvpress.py
+++ b/src/pruna/algorithms/kvpress.py
@@ -28,16 +28,8 @@
 from pruna.engine.save import SAVE_FUNCTIONS
 
 PRESS_TYPES = [
-    "AdaKVPress",
-    "BlockPress",
-    "ChunkKVPress",
-    "ChunkPress",
     "CompactorPress",
-    "CriticalAdaKVPress",
-    "CriticalKVPress",
     "CURPress",
-    "DMSPress",
-    "DuoAttentionPress",
     "ExpectedAttentionPress",
     "ExpectedAttentionStatsPress",
     "FastKVzipPress",
@@ -46,19 +38,15 @@
     "KVzapPress",
     "KVzipPress",
     "KeyDiffPress",
-    "KeyRerotationPress",
     "LagKVPress",
     "LeverageScorePress",
     "NonCausalAttnPress",
     "ObservedAttentionPress",
-    "PerLayerCompressionPress",
     "PyramidKVPress",
     "QFilterPress",
     "RandomPress",
-    "SimLayerKVPress",
     "SnapKVPress",
     "StreamingLLMPress",
-    "ThinKPress",
     "TOVAPress",
 ]
 
@@ -85,8 +73,13 @@ class KVPress(PrunaAlgorithmBase):
     processor_required: bool = False
     dataset_required: bool = False
     runs_on: list[str] = ["cuda"]
-    compatible_before: Iterable[str | tags] = [tags.QUANTIZER, "moe_kernel_tuner"]
-    compatible_after: Iterable[str | tags] = ["torch_compile", "sage_attn", "moe_kernel_tuner"]
+    compatible_before: Iterable[str] = [
+        "awq", "gptq", "half", "hqq", "hyper", "llm_int8",
+        "padding_pruning", "quanto", "sage_attn", "torchao", "moe_kernel_tuner",
+    ]
+    compatible_after: Iterable[str] = [
+        "img2img_denoise", "realesrgan_upscale", "torch_compile", "moe_kernel_tuner",
+    ]
 
     def get_hyperparameters(self) -> list:
         """
diff --git a/src/pruna/algorithms/quanto.py b/src/pruna/algorithms/quanto.py
index 6c80fe73..7c044cad 100644
--- a/src/pruna/algorithms/quanto.py
+++ b/src/pruna/algorithms/quanto.py
@@ -52,6 +52,7 @@ class Quanto(PrunaAlgorithmBase):
     compatible_before: Iterable[str] = ["qkv_diffusers", "moe_kernel_tuner"]
     compatible_after: Iterable[str] = [
         "deepcache",
+        "kvpress",
         "sage_attn",
         "text_to_image_distillation_inplace_perp",
         "text_to_image_distillation_lora",
diff --git a/src/pruna/config/smash_config.py b/src/pruna/config/smash_config.py
index 0acc1e12..baaba7f7 100644
--- a/src/pruna/config/smash_config.py
+++ b/src/pruna/config/smash_config.py
@@ -623,8 +623,9 @@ def add(self, request: str | list[str] | dict[str, Any]) -> None:
         # request wants to activate a dictionary of algorithms and their hyperparameters
         elif isinstance(request, dict):
             for key, value in request.items():
-                # target modules are a special case, as they are a hyperparameter but their value is a dict
-                if isinstance(value, dict) and "target_module" not in key:
+                # if the key is an algorithm name and the value is a dict, treat it as
+                # algorithm activation + hyperparameter setting (e.g. {"hqq": {"weight_bits": 4}})
+                if isinstance(value, dict) and key in SMASH_SPACE.get_all_algorithms():
                     self._configuration[key] = True
                     for k, v in value.items():
                         if not k.startswith(key):
diff --git a/tests/algorithms/test_kvpress.py b/tests/algorithms/test_kvpress.py
new file mode 100644
index 00000000..a7453aac
--- /dev/null
+++ b/tests/algorithms/test_kvpress.py
@@ -0,0 +1,31 @@
+import pytest
+
+from pruna.algorithms.kvpress import PRESS_TYPES
+
+
+@pytest.fixture()
+def kvpress_modules():
+    """Import kvpress and return all press classes."""
+    import kvpress
+
+    return {name: getattr(kvpress, name) for name in PRESS_TYPES}
+
+
+def test_press_types_exist(kvpress_modules):
+    """Verify all PRESS_TYPES are valid classes in the kvpress module."""
+    for name, cls in kvpress_modules.items():
+        assert isinstance(cls, type), f"{name} is not a class"
+
+
+def test_press_types_accept_compression_ratio(kvpress_modules):
+    """Verify all PRESS_TYPES can be constructed with compression_ratio."""
+    for name, cls in kvpress_modules.items():
+        press = cls(compression_ratio=0.5)
+        assert press.compression_ratio == 0.5, f"{name} did not set compression_ratio"
+
+
+def test_press_kwargs_forwarded(kvpress_modules):
+    """Verify press_kwargs are forwarded to the press constructor."""
+    snap = kvpress_modules["SnapKVPress"](compression_ratio=0.4, window_size=32, kernel_size=3)
+    assert snap.window_size == 32
+    assert snap.kernel_size == 3
diff --git a/tests/algorithms/testers/kvpress.py b/tests/algorithms/testers/kvpress.py
index d41d272e..faa02068 100644
--- a/tests/algorithms/testers/kvpress.py
+++ b/tests/algorithms/testers/kvpress.py
@@ -5,7 +5,7 @@
 
 
 class TestKVPress(AlgorithmTesterBase):
-    """Test the KVPress KV cache compression algorithm."""
+    """Test the KVPress KV cache compression algorithm with default settings."""
 
     models = ["llama_3_tiny_random"]
     reject_models = ["sd_tiny_random"]
@@ -17,3 +17,26 @@ def post_smash_hook(self, model: PrunaModel) -> None:
         """Verify that the press was applied to the model."""
         assert hasattr(model, "_kvpress_press")
         assert hasattr(model, "_kvpress_original_generate")
+
+
+class TestKVPressSnapKV(AlgorithmTesterBase):
+    """Test the KVPress algorithm with SnapKV and custom press_kwargs."""
+
+    models = ["llama_3_tiny_random"]
+    reject_models = ["sd_tiny_random"]
+    allow_pickle_files = False
+    algorithm_class = KVPress
+    metrics = ["perplexity"]
+    hyperparameters = {
+        "kvpress_press_type": "SnapKVPress",
+        "kvpress_compression_ratio": 0.3,
+        "kvpress_press_kwargs": {"window_size": 32, "kernel_size": 3},
+    }
+
+    def post_smash_hook(self, model: PrunaModel) -> None:
+        """Verify that SnapKV press was applied with correct parameters."""
+        assert hasattr(model, "_kvpress_press")
+        press = model._kvpress_press
+        assert type(press).__name__ == "SnapKVPress"
+        assert press.window_size == 32
+        assert press.kernel_size == 3

From 469e5b34cf3aa705d11ca9470021aeeab874e79c Mon Sep 17 00:00:00 2001
From: kschwethelm <k.schwethelm@tum.de>
Date: Fri, 3 Apr 2026 17:47:27 +0200
Subject: [PATCH 05/11] feat: add KV_CACHER tag, replace explicit kvpress
 references

Add a new KV_CACHER algorithm tag for KV cache compression algorithms,
separate from CACHER (used by diffuser cachers). Use the tag in all
LLM algorithm compatibility lists instead of explicit "kvpress" strings.
---
 src/pruna/algorithms/base/tags.py                   |  4 ++++
 src/pruna/algorithms/gptq_model.py                  |  2 +-
 src/pruna/algorithms/half.py                        |  4 ++--
 src/pruna/algorithms/hqq.py                         |  2 +-
 src/pruna/algorithms/huggingface_llm_int8.py        |  2 +-
 src/pruna/algorithms/kvpress.py                     | 10 ++++------
 src/pruna/algorithms/llm_compressor.py              |  2 +-
 src/pruna/algorithms/moe_kernel_tuner.py            |  8 ++++----
 src/pruna/algorithms/quanto.py                      |  4 ++--
 src/pruna/algorithms/sage_attn.py                   |  2 +-
 src/pruna/algorithms/torch_compile/torch_compile.py |  4 ++--
 src/pruna/algorithms/torchao.py                     |  4 ++--
 12 files changed, 25 insertions(+), 23 deletions(-)

diff --git a/src/pruna/algorithms/base/tags.py b/src/pruna/algorithms/base/tags.py
index b648fc78..d672c982 100644
--- a/src/pruna/algorithms/base/tags.py
+++ b/src/pruna/algorithms/base/tags.py
@@ -62,6 +62,10 @@ class AlgorithmTag(Enum):
         "cacher",
         "Caching is a technique used to store intermediate results of computations to speed up subsequent operations, particularly useful in reducing inference time for machine learning models by reusing previously computed results.",
     )
+    KV_CACHER = (
+        "kv_cacher",
+        "KV cache compression reduces memory usage during autoregressive generation by scoring and pruning key-value pairs from the attention cache.",
+    )
     COMPILER = ("compiler", "Compilation optimizes the model for specific hardware.")
     BATCHER = (
         "batcher",
diff --git a/src/pruna/algorithms/gptq_model.py b/src/pruna/algorithms/gptq_model.py
index 11d5573e..94c35749 100644
--- a/src/pruna/algorithms/gptq_model.py
+++ b/src/pruna/algorithms/gptq_model.py
@@ -46,7 +46,7 @@ class GPTQ(PrunaAlgorithmBase):
     processor_required: bool = False
     runs_on: list[str] = ["cuda"]
     dataset_required: bool = True
-    compatible_after: Iterable[str] = ["torch_compile", "sage_attn", "kvpress"]
+    compatible_after: Iterable[str | tags] = ["torch_compile", "sage_attn", tags.KV_CACHER]
     required_install: str = (
         "You must first install the base package with ``pip install pruna`` "
         "before installing the GPTQ extension with ``pip install pruna[gptq] --extra-index-url https://prunaai.pythonanywhere.com/``"
diff --git a/src/pruna/algorithms/half.py b/src/pruna/algorithms/half.py
index e227b1c4..404194ee 100644
--- a/src/pruna/algorithms/half.py
+++ b/src/pruna/algorithms/half.py
@@ -42,7 +42,7 @@ class Half(PrunaAlgorithmBase):
     runs_on: list[str] = ["cpu", "cuda", "accelerate"]
     dataset_required: bool = False
     compatible_before: Iterable[str] = ["torch_structured", "torch_unstructured"]
-    compatible_after: Iterable[str] = [
+    compatible_after: Iterable[str | tags] = [
         "deepcache",
         "c_translate",
         "c_generate",
@@ -50,7 +50,7 @@ class Half(PrunaAlgorithmBase):
         "stable_fast",
         "torch_compile",
         "ifw",
-        "kvpress",
+        tags.KV_CACHER,
         "whisper_s2t",
         "sage_attn",
         "hyper",
diff --git a/src/pruna/algorithms/hqq.py b/src/pruna/algorithms/hqq.py
index 843ba143..6e4bee68 100644
--- a/src/pruna/algorithms/hqq.py
+++ b/src/pruna/algorithms/hqq.py
@@ -63,7 +63,7 @@ class HQQ(PrunaAlgorithmBase):
     runs_on: list[str] = ["cuda"]
     dataset_required: bool = False
     compatible_before: Iterable[str] = ["torch_structured", "moe_kernel_tuner"]
-    compatible_after: Iterable[str] = ["torch_compile", "sage_attn", "kvpress", "moe_kernel_tuner"]
+    compatible_after: Iterable[str | tags] = ["torch_compile", "sage_attn", tags.KV_CACHER, "moe_kernel_tuner"]
     disjointly_compatible_before: Iterable[str] = []
     disjointly_compatible_after: Iterable[str] = ["torchao"]
 
diff --git a/src/pruna/algorithms/huggingface_llm_int8.py b/src/pruna/algorithms/huggingface_llm_int8.py
index 3e6fb289..afb19045 100644
--- a/src/pruna/algorithms/huggingface_llm_int8.py
+++ b/src/pruna/algorithms/huggingface_llm_int8.py
@@ -58,7 +58,7 @@ class LLMInt8(PrunaAlgorithmBase):
     runs_on: list[str] = ["cuda", "accelerate"]
     save_fn: None = None
     compatible_before: Iterable[str] = ["moe_kernel_tuner"]
-    compatible_after: Iterable[str] = ["torch_compile", "sage_attn", "kvpress", "moe_kernel_tuner"]
+    compatible_after: Iterable[str | tags] = ["torch_compile", "sage_attn", tags.KV_CACHER, "moe_kernel_tuner"]
 
     def get_hyperparameters(self) -> list:
         """
diff --git a/src/pruna/algorithms/kvpress.py b/src/pruna/algorithms/kvpress.py
index 1766e9a1..e3a934a2 100644
--- a/src/pruna/algorithms/kvpress.py
+++ b/src/pruna/algorithms/kvpress.py
@@ -62,7 +62,7 @@ class KVPress(PrunaAlgorithmBase):
     """
 
     algorithm_name: str = "kvpress"
-    group_tags: list[tags] = [tags.CACHER]
+    group_tags: list[tags] = [tags.KV_CACHER]
     save_fn: SAVE_FUNCTIONS = SAVE_FUNCTIONS.reapply
     references: dict[str, str] = {
         "GitHub": "https://github.com/NVIDIA/kvpress",
@@ -74,12 +74,10 @@ class KVPress(PrunaAlgorithmBase):
     dataset_required: bool = False
     runs_on: list[str] = ["cuda"]
     compatible_before: Iterable[str] = [
-        "awq", "gptq", "half", "hqq", "hyper", "llm_int8",
-        "padding_pruning", "quanto", "sage_attn", "torchao", "moe_kernel_tuner",
-    ]
-    compatible_after: Iterable[str] = [
-        "img2img_denoise", "realesrgan_upscale", "torch_compile", "moe_kernel_tuner",
+        "awq", "gptq", "half", "hqq", "llm_int8",
+        "quanto", "sage_attn", "torchao", "moe_kernel_tuner",
     ]
+    compatible_after: Iterable[str] = ["torch_compile", "moe_kernel_tuner"]
 
     def get_hyperparameters(self) -> list:
         """
diff --git a/src/pruna/algorithms/llm_compressor.py b/src/pruna/algorithms/llm_compressor.py
index f964ce91..456e7132 100644
--- a/src/pruna/algorithms/llm_compressor.py
+++ b/src/pruna/algorithms/llm_compressor.py
@@ -52,7 +52,7 @@ class LLMCompressor(PrunaAlgorithmBase):
     dataset_required: bool = True
     runs_on: list[str] = ["cuda"]
     compatible_before: Iterable[str] = ["moe_kernel_tuner"]
-    compatible_after: Iterable[str] = ["sage_attn", "kvpress", "moe_kernel_tuner"]
+    compatible_after: Iterable[str | tags] = ["sage_attn", tags.KV_CACHER, "moe_kernel_tuner"]
     required_install = "``uv pip install 'pruna[awq]'``"
 
     def get_hyperparameters(self) -> list:
diff --git a/src/pruna/algorithms/moe_kernel_tuner.py b/src/pruna/algorithms/moe_kernel_tuner.py
index 062bf9d8..49ecbc83 100644
--- a/src/pruna/algorithms/moe_kernel_tuner.py
+++ b/src/pruna/algorithms/moe_kernel_tuner.py
@@ -48,15 +48,15 @@ class MoeKernelTuner(PrunaAlgorithmBase):
     processor_required: bool = False
     runs_on: list[str] = ["cuda", "accelerate"]
     dataset_required: bool = False
-    compatible_before: Iterable[str] = [
+    compatible_before: Iterable[str | tags] = [
         "awq", "deepcache", "diffusers_int8", "fastercache", "flash_attn3",
-        "fora", "hqq", "hqq_diffusers", "kvpress", "llm_int8", "pab", "padding_pruning",
+        "fora", "hqq", "hqq_diffusers", tags.KV_CACHER, "llm_int8", "pab", "padding_pruning",
         "qkv_diffusers", "quanto", "reduce_noe", "ring_attn", "sage_attn",
         "torch_compile", "torchao",
     ]
-    compatible_after: Iterable[str] = [
+    compatible_after: Iterable[str | tags] = [
         "awq", "deepcache", "diffusers_int8", "fastercache", "flash_attn3",
-        "fora", "hqq", "hqq_diffusers", "kvpress", "llm_int8", "pab", "padding_pruning",
+        "fora", "hqq", "hqq_diffusers", tags.KV_CACHER, "llm_int8", "pab", "padding_pruning",
         "qkv_diffusers", "quanto", "ring_attn", "sage_attn",
         "torch_compile", "torchao",
     ]
diff --git a/src/pruna/algorithms/quanto.py b/src/pruna/algorithms/quanto.py
index 7c044cad..1d613325 100644
--- a/src/pruna/algorithms/quanto.py
+++ b/src/pruna/algorithms/quanto.py
@@ -50,9 +50,9 @@ class Quanto(PrunaAlgorithmBase):
     dataset_required: bool = False
     runs_on: list[str] = ["cuda"]
     compatible_before: Iterable[str] = ["qkv_diffusers", "moe_kernel_tuner"]
-    compatible_after: Iterable[str] = [
+    compatible_after: Iterable[str | tags] = [
         "deepcache",
-        "kvpress",
+        tags.KV_CACHER,
         "sage_attn",
         "text_to_image_distillation_inplace_perp",
         "text_to_image_distillation_lora",
diff --git a/src/pruna/algorithms/sage_attn.py b/src/pruna/algorithms/sage_attn.py
index 45d8d042..24069775 100644
--- a/src/pruna/algorithms/sage_attn.py
+++ b/src/pruna/algorithms/sage_attn.py
@@ -50,7 +50,7 @@ class SageAttn(PrunaAlgorithmBase):
     runs_on: list[str] = ["cuda", "accelerate"]
     dataset_required: bool = False
     compatible_before: Iterable[str | tags] = [tags.QUANTIZER, "moe_kernel_tuner"]
-    compatible_after: Iterable[str | tags] = ["torch_compile", tags.CACHER, "moe_kernel_tuner"]
+    compatible_after: Iterable[str | tags] = ["torch_compile", tags.CACHER, tags.KV_CACHER, "moe_kernel_tuner"]
 
     def model_check_fn(self, model: Any) -> bool:
         """
diff --git a/src/pruna/algorithms/torch_compile/torch_compile.py b/src/pruna/algorithms/torch_compile/torch_compile.py
index 72574914..b3b2107e 100644
--- a/src/pruna/algorithms/torch_compile/torch_compile.py
+++ b/src/pruna/algorithms/torch_compile/torch_compile.py
@@ -57,7 +57,7 @@ class TorchCompile(PrunaAlgorithmBase):
     processor_required: bool = False
     runs_on: list[str] = ["cpu", "cuda"]
     dataset_required: bool = False
-    compatible_before: Iterable[str] = [
+    compatible_before: Iterable[str | tags] = [
         "qkv_diffusers",
         "torch_structured",
         "half",
@@ -74,7 +74,7 @@ class TorchCompile(PrunaAlgorithmBase):
         "hyper",
         "padding_pruning",
         "ring_attn",
-        "kvpress",
+        tags.KV_CACHER,
         "text_to_image_distillation_inplace_perp",
         "text_to_image_distillation_lora",
         "text_to_image_distillation_perp",
diff --git a/src/pruna/algorithms/torchao.py b/src/pruna/algorithms/torchao.py
index d3411044..c81477ed 100644
--- a/src/pruna/algorithms/torchao.py
+++ b/src/pruna/algorithms/torchao.py
@@ -94,12 +94,12 @@ class Torchao(PrunaAlgorithmBase):
     runs_on: list[str] = ["cpu", "cuda", "accelerate"]
     dataset_required: bool = False
     compatible_before: Iterable[str] = ["qkv_diffusers", "torch_structured", "padding_pruning", "moe_kernel_tuner"]
-    compatible_after: Iterable[str] = [
+    compatible_after: Iterable[str | tags] = [
         "flash_attn3",
         "fora",
         "torch_compile",
         "sage_attn",
-        "kvpress",
+        tags.KV_CACHER,
         "img2img_denoise",
         "realesrgan_upscale",
         "moe_kernel_tuner",

From 1ab2e3b880f46c9a86bfd63db637055fa092bcdc Mon Sep 17 00:00:00 2001
From: kschwethelm <k.schwethelm@tum.de>
Date: Fri, 10 Apr 2026 14:39:14 +0200
Subject: [PATCH 06/11] refactor: rename KV_CACHER tag to KV_COMPRESSOR,
 improve docstrings

---
 src/pruna/algorithms/base/tags.py             |  6 ++---
 src/pruna/algorithms/gptq_model.py            |  2 +-
 src/pruna/algorithms/half.py                  |  2 +-
 src/pruna/algorithms/hqq.py                   |  2 +-
 src/pruna/algorithms/huggingface_llm_int8.py  |  2 +-
 src/pruna/algorithms/kvpress.py               | 25 +++++++++----------
 src/pruna/algorithms/llm_compressor.py        |  2 +-
 src/pruna/algorithms/moe_kernel_tuner.py      |  4 +--
 src/pruna/algorithms/quanto.py                |  2 +-
 src/pruna/algorithms/sage_attn.py             |  2 +-
 .../algorithms/torch_compile/torch_compile.py |  2 +-
 src/pruna/algorithms/torchao.py               |  2 +-
 12 files changed, 26 insertions(+), 27 deletions(-)

diff --git a/src/pruna/algorithms/base/tags.py b/src/pruna/algorithms/base/tags.py
index d672c982..d0807940 100644
--- a/src/pruna/algorithms/base/tags.py
+++ b/src/pruna/algorithms/base/tags.py
@@ -62,9 +62,9 @@ class AlgorithmTag(Enum):
         "cacher",
         "Caching is a technique used to store intermediate results of computations to speed up subsequent operations, particularly useful in reducing inference time for machine learning models by reusing previously computed results.",
     )
-    KV_CACHER = (
-        "kv_cacher",
-        "KV cache compression reduces memory usage during autoregressive generation by scoring and pruning key-value pairs from the attention cache.",
+    KV_COMPRESSOR = (
+        "kv_compressor",
+        "KV cache compression reduces the memory footprint of the key-value cache during autoregressive generation.",
     )
     COMPILER = ("compiler", "Compilation optimizes the model for specific hardware.")
     BATCHER = (
diff --git a/src/pruna/algorithms/gptq_model.py b/src/pruna/algorithms/gptq_model.py
index 94c35749..808610a0 100644
--- a/src/pruna/algorithms/gptq_model.py
+++ b/src/pruna/algorithms/gptq_model.py
@@ -46,7 +46,7 @@ class GPTQ(PrunaAlgorithmBase):
     processor_required: bool = False
     runs_on: list[str] = ["cuda"]
     dataset_required: bool = True
-    compatible_after: Iterable[str | tags] = ["torch_compile", "sage_attn", tags.KV_CACHER]
+    compatible_after: Iterable[str | tags] = ["torch_compile", "sage_attn", tags.KV_COMPRESSOR]
     required_install: str = (
         "You must first install the base package with ``pip install pruna`` "
         "before installing the GPTQ extension with ``pip install pruna[gptq] --extra-index-url https://prunaai.pythonanywhere.com/``"
diff --git a/src/pruna/algorithms/half.py b/src/pruna/algorithms/half.py
index 404194ee..7c9bb944 100644
--- a/src/pruna/algorithms/half.py
+++ b/src/pruna/algorithms/half.py
@@ -50,7 +50,7 @@ class Half(PrunaAlgorithmBase):
         "stable_fast",
         "torch_compile",
         "ifw",
-        tags.KV_CACHER,
+        tags.KV_COMPRESSOR,
         "whisper_s2t",
         "sage_attn",
         "hyper",
diff --git a/src/pruna/algorithms/hqq.py b/src/pruna/algorithms/hqq.py
index 6e4bee68..4330bf0d 100644
--- a/src/pruna/algorithms/hqq.py
+++ b/src/pruna/algorithms/hqq.py
@@ -63,7 +63,7 @@ class HQQ(PrunaAlgorithmBase):
     runs_on: list[str] = ["cuda"]
     dataset_required: bool = False
     compatible_before: Iterable[str] = ["torch_structured", "moe_kernel_tuner"]
-    compatible_after: Iterable[str | tags] = ["torch_compile", "sage_attn", tags.KV_CACHER, "moe_kernel_tuner"]
+    compatible_after: Iterable[str | tags] = ["torch_compile", "sage_attn", tags.KV_COMPRESSOR, "moe_kernel_tuner"]
     disjointly_compatible_before: Iterable[str] = []
     disjointly_compatible_after: Iterable[str] = ["torchao"]
 
diff --git a/src/pruna/algorithms/huggingface_llm_int8.py b/src/pruna/algorithms/huggingface_llm_int8.py
index afb19045..38dc7604 100644
--- a/src/pruna/algorithms/huggingface_llm_int8.py
+++ b/src/pruna/algorithms/huggingface_llm_int8.py
@@ -58,7 +58,7 @@ class LLMInt8(PrunaAlgorithmBase):
     runs_on: list[str] = ["cuda", "accelerate"]
     save_fn: None = None
     compatible_before: Iterable[str] = ["moe_kernel_tuner"]
-    compatible_after: Iterable[str | tags] = ["torch_compile", "sage_attn", tags.KV_CACHER, "moe_kernel_tuner"]
+    compatible_after: Iterable[str | tags] = ["torch_compile", "sage_attn", tags.KV_COMPRESSOR, "moe_kernel_tuner"]
 
     def get_hyperparameters(self) -> list:
         """
diff --git a/src/pruna/algorithms/kvpress.py b/src/pruna/algorithms/kvpress.py
index e3a934a2..f4e3fc37 100644
--- a/src/pruna/algorithms/kvpress.py
+++ b/src/pruna/algorithms/kvpress.py
@@ -53,16 +53,15 @@
 
 class KVPress(PrunaAlgorithmBase):
     """
-    Integrate KVPress for KV cache compression in causal language models.
+    Compress the KV cache of causal language models using KVPress.
 
-    KVPress is a library by NVIDIA that compresses the key-value cache during the prefill phase
-    of autoregressive generation, reducing memory usage for long-context inference. It provides
-    over 30 compression strategies (presses) that score and prune KV pairs based on different
-    importance criteria.
+    KVPress is a library by NVIDIA that provides over 20 compression strategies (presses) for
+    reducing the memory footprint of the key-value cache during long-context inference. Each press
+    scores and prunes KV pairs after the prefill phase according to a chosen importance criterion.
     """
 
     algorithm_name: str = "kvpress"
-    group_tags: list[tags] = [tags.KV_CACHER]
+    group_tags: list[tags] = [tags.KV_COMPRESSOR]
     save_fn: SAVE_FUNCTIONS = SAVE_FUNCTIONS.reapply
     references: dict[str, str] = {
         "GitHub": "https://github.com/NVIDIA/kvpress",
@@ -111,7 +110,7 @@ def get_hyperparameters(self) -> list:
 
     def model_check_fn(self, model: Any) -> bool:
         """
-        Check if the model is a causal language model.
+        Check if the model is a causal language model or a pipeline wrapping one.
 
         Parameters
         ----------
@@ -121,25 +120,25 @@ def model_check_fn(self, model: Any) -> bool:
         Returns
         -------
         bool
-            True if the model is a causal language model, False otherwise.
+            True if the model is compatible with KV cache compression, False otherwise.
         """
         return is_causal_lm(model) or is_transformers_pipeline_with_causal_lm(model)
 
     def _apply(self, model: Any, smash_config: SmashConfigPrefixWrapper) -> Any:
         """
-        Apply KV cache compression to the model.
+        Wrap the model's generate method to apply KV cache compression via a press context manager.
 
         Parameters
         ----------
         model : Any
-            The model to apply KV cache compression to.
+            The causal language model to compress.
         smash_config : SmashConfigPrefixWrapper
-            The configuration for the compression.
+            The algorithm-prefixed configuration containing press_type, compression_ratio, and press_kwargs.
 
         Returns
         -------
         Any
-            The model with KV cache compression applied to its generate method.
+            The model with its generate method wrapped to compress the KV cache on each call.
         """
         imported_modules = self.import_algorithm_packages()
 
@@ -165,7 +164,7 @@ def generate_with_press(*args, **kwargs):
 
     def import_algorithm_packages(self) -> Dict[str, Any]:
         """
-        Import the kvpress package lazily.
+        Lazily import kvpress and collect all supported press classes.
 
         Returns
         -------
diff --git a/src/pruna/algorithms/llm_compressor.py b/src/pruna/algorithms/llm_compressor.py
index 456e7132..6d0da9dd 100644
--- a/src/pruna/algorithms/llm_compressor.py
+++ b/src/pruna/algorithms/llm_compressor.py
@@ -52,7 +52,7 @@ class LLMCompressor(PrunaAlgorithmBase):
     dataset_required: bool = True
     runs_on: list[str] = ["cuda"]
     compatible_before: Iterable[str] = ["moe_kernel_tuner"]
-    compatible_after: Iterable[str | tags] = ["sage_attn", tags.KV_CACHER, "moe_kernel_tuner"]
+    compatible_after: Iterable[str | tags] = ["sage_attn", tags.KV_COMPRESSOR, "moe_kernel_tuner"]
     required_install = "``uv pip install 'pruna[awq]'``"
 
     def get_hyperparameters(self) -> list:
diff --git a/src/pruna/algorithms/moe_kernel_tuner.py b/src/pruna/algorithms/moe_kernel_tuner.py
index 49ecbc83..25de7f01 100644
--- a/src/pruna/algorithms/moe_kernel_tuner.py
+++ b/src/pruna/algorithms/moe_kernel_tuner.py
@@ -50,13 +50,13 @@ class MoeKernelTuner(PrunaAlgorithmBase):
     dataset_required: bool = False
     compatible_before: Iterable[str | tags] = [
         "awq", "deepcache", "diffusers_int8", "fastercache", "flash_attn3",
-        "fora", "hqq", "hqq_diffusers", tags.KV_CACHER, "llm_int8", "pab", "padding_pruning",
+        "fora", "hqq", "hqq_diffusers", tags.KV_COMPRESSOR, "llm_int8", "pab", "padding_pruning",
         "qkv_diffusers", "quanto", "reduce_noe", "ring_attn", "sage_attn",
         "torch_compile", "torchao",
     ]
     compatible_after: Iterable[str | tags] = [
         "awq", "deepcache", "diffusers_int8", "fastercache", "flash_attn3",
-        "fora", "hqq", "hqq_diffusers", tags.KV_CACHER, "llm_int8", "pab", "padding_pruning",
+        "fora", "hqq", "hqq_diffusers", tags.KV_COMPRESSOR, "llm_int8", "pab", "padding_pruning",
         "qkv_diffusers", "quanto", "ring_attn", "sage_attn",
         "torch_compile", "torchao",
     ]
diff --git a/src/pruna/algorithms/quanto.py b/src/pruna/algorithms/quanto.py
index 1d613325..bd517c40 100644
--- a/src/pruna/algorithms/quanto.py
+++ b/src/pruna/algorithms/quanto.py
@@ -52,7 +52,7 @@ class Quanto(PrunaAlgorithmBase):
     compatible_before: Iterable[str] = ["qkv_diffusers", "moe_kernel_tuner"]
     compatible_after: Iterable[str | tags] = [
         "deepcache",
-        tags.KV_CACHER,
+        tags.KV_COMPRESSOR,
         "sage_attn",
         "text_to_image_distillation_inplace_perp",
         "text_to_image_distillation_lora",
diff --git a/src/pruna/algorithms/sage_attn.py b/src/pruna/algorithms/sage_attn.py
index 24069775..49cd7569 100644
--- a/src/pruna/algorithms/sage_attn.py
+++ b/src/pruna/algorithms/sage_attn.py
@@ -50,7 +50,7 @@ class SageAttn(PrunaAlgorithmBase):
     runs_on: list[str] = ["cuda", "accelerate"]
     dataset_required: bool = False
     compatible_before: Iterable[str | tags] = [tags.QUANTIZER, "moe_kernel_tuner"]
-    compatible_after: Iterable[str | tags] = ["torch_compile", tags.CACHER, tags.KV_CACHER, "moe_kernel_tuner"]
+    compatible_after: Iterable[str | tags] = ["torch_compile", tags.CACHER, tags.KV_COMPRESSOR, "moe_kernel_tuner"]
 
     def model_check_fn(self, model: Any) -> bool:
         """
diff --git a/src/pruna/algorithms/torch_compile/torch_compile.py b/src/pruna/algorithms/torch_compile/torch_compile.py
index b3b2107e..d8904152 100644
--- a/src/pruna/algorithms/torch_compile/torch_compile.py
+++ b/src/pruna/algorithms/torch_compile/torch_compile.py
@@ -74,7 +74,7 @@ class TorchCompile(PrunaAlgorithmBase):
         "hyper",
         "padding_pruning",
         "ring_attn",
-        tags.KV_CACHER,
+        tags.KV_COMPRESSOR,
         "text_to_image_distillation_inplace_perp",
         "text_to_image_distillation_lora",
         "text_to_image_distillation_perp",
diff --git a/src/pruna/algorithms/torchao.py b/src/pruna/algorithms/torchao.py
index c81477ed..f68446f9 100644
--- a/src/pruna/algorithms/torchao.py
+++ b/src/pruna/algorithms/torchao.py
@@ -99,7 +99,7 @@ class Torchao(PrunaAlgorithmBase):
         "fora",
         "torch_compile",
         "sage_attn",
-        tags.KV_CACHER,
+        tags.KV_COMPRESSOR,
         "img2img_denoise",
         "realesrgan_upscale",
         "moe_kernel_tuner",

From 63a1315af2bbc2daeef9940e1c797a4a4ab8880b Mon Sep 17 00:00:00 2001
From: kschwethelm <k.schwethelm@tum.de>
Date: Fri, 10 Apr 2026 14:51:12 +0200
Subject: [PATCH 07/11] docs: document excluded wrapper presses in kvpress
 docstring

---
 src/pruna/algorithms/kvpress.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/src/pruna/algorithms/kvpress.py b/src/pruna/algorithms/kvpress.py
index f4e3fc37..02749890 100644
--- a/src/pruna/algorithms/kvpress.py
+++ b/src/pruna/algorithms/kvpress.py
@@ -58,6 +58,11 @@ class KVPress(PrunaAlgorithmBase):
     KVPress is a library by NVIDIA that provides over 20 compression strategies (presses) for
     reducing the memory footprint of the key-value cache during long-context inference. Each press
     scores and prunes KV pairs after the prefill phase according to a chosen importance criterion.
+
+    This integration covers all scorer and standalone presses. Wrapper presses (e.g., ChunkPress,
+    AdaKVPress, PerLayerCompressionPress) that require a nested scorer press as input are not
+    included, as well as ThinKPress which compresses along the channel dimension with a different
+    parameter interface.
     """
 
     algorithm_name: str = "kvpress"

From bfa72ca20816b94a0b59f8df20ab9385a15c8ced Mon Sep 17 00:00:00 2001
From: kschwethelm <k.schwethelm@tum.de>
Date: Wed, 15 Apr 2026 18:43:09 +0200
Subject: [PATCH 08/11] refactor: remove KV_COMPRESSOR tag, reference kvpress
 by name

Drop the dedicated KV_COMPRESSOR tag and use tags.PRUNER as kvpress's
group tag, matching how other pruners are categorized. Replace all
tags.KV_COMPRESSOR references in compatible_before/after lists with the
string "kvpress" to align with the repo convention of naming specific
algorithms in compatibility lists.
---
 src/pruna/algorithms/base/tags.py                   | 4 ----
 src/pruna/algorithms/gptq_model.py                  | 2 +-
 src/pruna/algorithms/half.py                        | 4 ++--
 src/pruna/algorithms/hqq.py                         | 2 +-
 src/pruna/algorithms/huggingface_llm_int8.py        | 2 +-
 src/pruna/algorithms/kvpress.py                     | 2 +-
 src/pruna/algorithms/llm_compressor.py              | 2 +-
 src/pruna/algorithms/moe_kernel_tuner.py            | 8 ++++----
 src/pruna/algorithms/quanto.py                      | 4 ++--
 src/pruna/algorithms/sage_attn.py                   | 2 +-
 src/pruna/algorithms/torch_compile/torch_compile.py | 4 ++--
 src/pruna/algorithms/torchao.py                     | 4 ++--
 12 files changed, 18 insertions(+), 22 deletions(-)

diff --git a/src/pruna/algorithms/base/tags.py b/src/pruna/algorithms/base/tags.py
index d0807940..b648fc78 100644
--- a/src/pruna/algorithms/base/tags.py
+++ b/src/pruna/algorithms/base/tags.py
@@ -62,10 +62,6 @@ class AlgorithmTag(Enum):
         "cacher",
         "Caching is a technique used to store intermediate results of computations to speed up subsequent operations, particularly useful in reducing inference time for machine learning models by reusing previously computed results.",
     )
-    KV_COMPRESSOR = (
-        "kv_compressor",
-        "KV cache compression reduces the memory footprint of the key-value cache during autoregressive generation.",
-    )
     COMPILER = ("compiler", "Compilation optimizes the model for specific hardware.")
     BATCHER = (
         "batcher",
diff --git a/src/pruna/algorithms/gptq_model.py b/src/pruna/algorithms/gptq_model.py
index 808610a0..11d5573e 100644
--- a/src/pruna/algorithms/gptq_model.py
+++ b/src/pruna/algorithms/gptq_model.py
@@ -46,7 +46,7 @@ class GPTQ(PrunaAlgorithmBase):
     processor_required: bool = False
     runs_on: list[str] = ["cuda"]
     dataset_required: bool = True
-    compatible_after: Iterable[str | tags] = ["torch_compile", "sage_attn", tags.KV_COMPRESSOR]
+    compatible_after: Iterable[str] = ["torch_compile", "sage_attn", "kvpress"]
     required_install: str = (
         "You must first install the base package with ``pip install pruna`` "
         "before installing the GPTQ extension with ``pip install pruna[gptq] --extra-index-url https://prunaai.pythonanywhere.com/``"
diff --git a/src/pruna/algorithms/half.py b/src/pruna/algorithms/half.py
index 7c9bb944..e227b1c4 100644
--- a/src/pruna/algorithms/half.py
+++ b/src/pruna/algorithms/half.py
@@ -42,7 +42,7 @@ class Half(PrunaAlgorithmBase):
     runs_on: list[str] = ["cpu", "cuda", "accelerate"]
     dataset_required: bool = False
     compatible_before: Iterable[str] = ["torch_structured", "torch_unstructured"]
-    compatible_after: Iterable[str | tags] = [
+    compatible_after: Iterable[str] = [
         "deepcache",
         "c_translate",
         "c_generate",
@@ -50,7 +50,7 @@ class Half(PrunaAlgorithmBase):
         "stable_fast",
         "torch_compile",
         "ifw",
-        tags.KV_COMPRESSOR,
+        "kvpress",
         "whisper_s2t",
         "sage_attn",
         "hyper",
diff --git a/src/pruna/algorithms/hqq.py b/src/pruna/algorithms/hqq.py
index 4330bf0d..843ba143 100644
--- a/src/pruna/algorithms/hqq.py
+++ b/src/pruna/algorithms/hqq.py
@@ -63,7 +63,7 @@ class HQQ(PrunaAlgorithmBase):
     runs_on: list[str] = ["cuda"]
     dataset_required: bool = False
     compatible_before: Iterable[str] = ["torch_structured", "moe_kernel_tuner"]
-    compatible_after: Iterable[str | tags] = ["torch_compile", "sage_attn", tags.KV_COMPRESSOR, "moe_kernel_tuner"]
+    compatible_after: Iterable[str] = ["torch_compile", "sage_attn", "kvpress", "moe_kernel_tuner"]
     disjointly_compatible_before: Iterable[str] = []
     disjointly_compatible_after: Iterable[str] = ["torchao"]
 
diff --git a/src/pruna/algorithms/huggingface_llm_int8.py b/src/pruna/algorithms/huggingface_llm_int8.py
index 38dc7604..3e6fb289 100644
--- a/src/pruna/algorithms/huggingface_llm_int8.py
+++ b/src/pruna/algorithms/huggingface_llm_int8.py
@@ -58,7 +58,7 @@ class LLMInt8(PrunaAlgorithmBase):
     runs_on: list[str] = ["cuda", "accelerate"]
     save_fn: None = None
     compatible_before: Iterable[str] = ["moe_kernel_tuner"]
-    compatible_after: Iterable[str | tags] = ["torch_compile", "sage_attn", tags.KV_COMPRESSOR, "moe_kernel_tuner"]
+    compatible_after: Iterable[str] = ["torch_compile", "sage_attn", "kvpress", "moe_kernel_tuner"]
 
     def get_hyperparameters(self) -> list:
         """
diff --git a/src/pruna/algorithms/kvpress.py b/src/pruna/algorithms/kvpress.py
index 02749890..696ca7a7 100644
--- a/src/pruna/algorithms/kvpress.py
+++ b/src/pruna/algorithms/kvpress.py
@@ -66,7 +66,7 @@ class KVPress(PrunaAlgorithmBase):
     """
 
     algorithm_name: str = "kvpress"
-    group_tags: list[tags] = [tags.KV_COMPRESSOR]
+    group_tags: list[tags] = [tags.PRUNER]
     save_fn: SAVE_FUNCTIONS = SAVE_FUNCTIONS.reapply
     references: dict[str, str] = {
         "GitHub": "https://github.com/NVIDIA/kvpress",
diff --git a/src/pruna/algorithms/llm_compressor.py b/src/pruna/algorithms/llm_compressor.py
index 6d0da9dd..f964ce91 100644
--- a/src/pruna/algorithms/llm_compressor.py
+++ b/src/pruna/algorithms/llm_compressor.py
@@ -52,7 +52,7 @@ class LLMCompressor(PrunaAlgorithmBase):
     dataset_required: bool = True
     runs_on: list[str] = ["cuda"]
     compatible_before: Iterable[str] = ["moe_kernel_tuner"]
-    compatible_after: Iterable[str | tags] = ["sage_attn", tags.KV_COMPRESSOR, "moe_kernel_tuner"]
+    compatible_after: Iterable[str] = ["sage_attn", "kvpress", "moe_kernel_tuner"]
     required_install = "``uv pip install 'pruna[awq]'``"
 
     def get_hyperparameters(self) -> list:
diff --git a/src/pruna/algorithms/moe_kernel_tuner.py b/src/pruna/algorithms/moe_kernel_tuner.py
index 25de7f01..062bf9d8 100644
--- a/src/pruna/algorithms/moe_kernel_tuner.py
+++ b/src/pruna/algorithms/moe_kernel_tuner.py
@@ -48,15 +48,15 @@ class MoeKernelTuner(PrunaAlgorithmBase):
     processor_required: bool = False
     runs_on: list[str] = ["cuda", "accelerate"]
     dataset_required: bool = False
-    compatible_before: Iterable[str | tags] = [
+    compatible_before: Iterable[str] = [
         "awq", "deepcache", "diffusers_int8", "fastercache", "flash_attn3",
-        "fora", "hqq", "hqq_diffusers", tags.KV_COMPRESSOR, "llm_int8", "pab", "padding_pruning",
+        "fora", "hqq", "hqq_diffusers", "kvpress", "llm_int8", "pab", "padding_pruning",
         "qkv_diffusers", "quanto", "reduce_noe", "ring_attn", "sage_attn",
         "torch_compile", "torchao",
     ]
-    compatible_after: Iterable[str | tags] = [
+    compatible_after: Iterable[str] = [
         "awq", "deepcache", "diffusers_int8", "fastercache", "flash_attn3",
-        "fora", "hqq", "hqq_diffusers", tags.KV_COMPRESSOR, "llm_int8", "pab", "padding_pruning",
+        "fora", "hqq", "hqq_diffusers", "kvpress", "llm_int8", "pab", "padding_pruning",
         "qkv_diffusers", "quanto", "ring_attn", "sage_attn",
         "torch_compile", "torchao",
     ]
diff --git a/src/pruna/algorithms/quanto.py b/src/pruna/algorithms/quanto.py
index bd517c40..7c044cad 100644
--- a/src/pruna/algorithms/quanto.py
+++ b/src/pruna/algorithms/quanto.py
@@ -50,9 +50,9 @@ class Quanto(PrunaAlgorithmBase):
     dataset_required: bool = False
     runs_on: list[str] = ["cuda"]
     compatible_before: Iterable[str] = ["qkv_diffusers", "moe_kernel_tuner"]
-    compatible_after: Iterable[str | tags] = [
+    compatible_after: Iterable[str] = [
         "deepcache",
-        tags.KV_COMPRESSOR,
+        "kvpress",
         "sage_attn",
         "text_to_image_distillation_inplace_perp",
         "text_to_image_distillation_lora",
diff --git a/src/pruna/algorithms/sage_attn.py b/src/pruna/algorithms/sage_attn.py
index 49cd7569..8738375c 100644
--- a/src/pruna/algorithms/sage_attn.py
+++ b/src/pruna/algorithms/sage_attn.py
@@ -50,7 +50,7 @@ class SageAttn(PrunaAlgorithmBase):
     runs_on: list[str] = ["cuda", "accelerate"]
     dataset_required: bool = False
     compatible_before: Iterable[str | tags] = [tags.QUANTIZER, "moe_kernel_tuner"]
-    compatible_after: Iterable[str | tags] = ["torch_compile", tags.CACHER, tags.KV_COMPRESSOR, "moe_kernel_tuner"]
+    compatible_after: Iterable[str | tags] = ["torch_compile", tags.CACHER, "kvpress", "moe_kernel_tuner"]
 
     def model_check_fn(self, model: Any) -> bool:
         """
diff --git a/src/pruna/algorithms/torch_compile/torch_compile.py b/src/pruna/algorithms/torch_compile/torch_compile.py
index d8904152..72574914 100644
--- a/src/pruna/algorithms/torch_compile/torch_compile.py
+++ b/src/pruna/algorithms/torch_compile/torch_compile.py
@@ -57,7 +57,7 @@ class TorchCompile(PrunaAlgorithmBase):
     processor_required: bool = False
     runs_on: list[str] = ["cpu", "cuda"]
     dataset_required: bool = False
-    compatible_before: Iterable[str | tags] = [
+    compatible_before: Iterable[str] = [
         "qkv_diffusers",
         "torch_structured",
         "half",
@@ -74,7 +74,7 @@ class TorchCompile(PrunaAlgorithmBase):
         "hyper",
         "padding_pruning",
         "ring_attn",
-        tags.KV_COMPRESSOR,
+        "kvpress",
         "text_to_image_distillation_inplace_perp",
         "text_to_image_distillation_lora",
         "text_to_image_distillation_perp",
diff --git a/src/pruna/algorithms/torchao.py b/src/pruna/algorithms/torchao.py
index f68446f9..d3411044 100644
--- a/src/pruna/algorithms/torchao.py
+++ b/src/pruna/algorithms/torchao.py
@@ -94,12 +94,12 @@ class Torchao(PrunaAlgorithmBase):
     runs_on: list[str] = ["cpu", "cuda", "accelerate"]
     dataset_required: bool = False
     compatible_before: Iterable[str] = ["qkv_diffusers", "torch_structured", "padding_pruning", "moe_kernel_tuner"]
-    compatible_after: Iterable[str | tags] = [
+    compatible_after: Iterable[str] = [
         "flash_attn3",
         "fora",
         "torch_compile",
         "sage_attn",
-        tags.KV_COMPRESSOR,
+        "kvpress",
         "img2img_denoise",
         "realesrgan_upscale",
         "moe_kernel_tuner",

From e046d4a9ade2a5e311f4a003b29f4d1a13ab090a Mon Sep 17 00:00:00 2001
From: kschwethelm <k.schwethelm@tum.de>
Date: Fri, 17 Apr 2026 09:09:18 +0200
Subject: [PATCH 09/11] fix: handle transformers pipeline in kvpress _apply

Add pipeline guard at the top of _apply to delegate to
_apply_to_model_within_transformers_pipeline when the model is a
TextGenerationPipeline, matching the pattern used by gptq, torch_compile,
and other algorithms.
---
 src/pruna/algorithms/kvpress.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/pruna/algorithms/kvpress.py b/src/pruna/algorithms/kvpress.py
index 696ca7a7..60afddaa 100644
--- a/src/pruna/algorithms/kvpress.py
+++ b/src/pruna/algorithms/kvpress.py
@@ -145,6 +145,9 @@ def _apply(self, model: Any, smash_config: SmashConfigPrefixWrapper) -> Any:
         Any
             The model with its generate method wrapped to compress the KV cache on each call.
         """
+        if is_transformers_pipeline_with_causal_lm(model):
+            return self._apply_to_model_within_transformers_pipeline(model, smash_config)
+
         imported_modules = self.import_algorithm_packages()
 
         press_type = smash_config["press_type"]

From c0b1fa86e873ea51b39717f45a8f1709560ed2a0 Mon Sep 17 00:00:00 2001
From: kschwethelm <k.schwethelm@tum.de>
Date: Sun, 3 May 2026 09:36:14 +0200
Subject: [PATCH 10/11] ci: register requires_kvpress marker for optional extra

---
 tests/algorithms/test_kvpress.py | 31 -------------------------------
 tests/conftest.py                |  1 +
 2 files changed, 1 insertion(+), 31 deletions(-)
 delete mode 100644 tests/algorithms/test_kvpress.py

diff --git a/tests/algorithms/test_kvpress.py b/tests/algorithms/test_kvpress.py
deleted file mode 100644
index a7453aac..00000000
--- a/tests/algorithms/test_kvpress.py
+++ /dev/null
@@ -1,31 +0,0 @@
-import pytest
-
-from pruna.algorithms.kvpress import PRESS_TYPES
-
-
-@pytest.fixture()
-def kvpress_modules():
-    """Import kvpress and return all press classes."""
-    import kvpress
-
-    return {name: getattr(kvpress, name) for name in PRESS_TYPES}
-
-
-def test_press_types_exist(kvpress_modules):
-    """Verify all PRESS_TYPES are valid classes in the kvpress module."""
-    for name, cls in kvpress_modules.items():
-        assert isinstance(cls, type), f"{name} is not a class"
-
-
-def test_press_types_accept_compression_ratio(kvpress_modules):
-    """Verify all PRESS_TYPES can be constructed with compression_ratio."""
-    for name, cls in kvpress_modules.items():
-        press = cls(compression_ratio=0.5)
-        assert press.compression_ratio == 0.5, f"{name} did not set compression_ratio"
-
-
-def test_press_kwargs_forwarded(kvpress_modules):
-    """Verify press_kwargs are forwarded to the press constructor."""
-    snap = kvpress_modules["SnapKVPress"](compression_ratio=0.4, window_size=32, kernel_size=3)
-    assert snap.window_size == 32
-    assert snap.kernel_size == 3
diff --git a/tests/conftest.py b/tests/conftest.py
index c7bea91e..80d54825 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -20,6 +20,7 @@
     "requires_whisper": "mark test that needs pruna[whisper]",
     "requires_upscale": "mark test that needs pruna[upscale]",
     "requires_rapidata": "mark test that needs pruna[rapidata]",
+    "requires_kvpress": "mark test that needs pruna[kvpress]",
 }
 
 

From 9302ef23af9dae1cf851f2c772dfc77bce6841b2 Mon Sep 17 00:00:00 2001
From: Gaspar Rochette <gaspar.rochette@pruna.ai>
Date: Wed, 6 May 2026 13:50:37 +0000
Subject: [PATCH 11/11] test: mark kvpress as require_kvpress

---
 tests/algorithms/testers/kvpress.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tests/algorithms/testers/kvpress.py b/tests/algorithms/testers/kvpress.py
index faa02068..e2bf60ae 100644
--- a/tests/algorithms/testers/kvpress.py
+++ b/tests/algorithms/testers/kvpress.py
@@ -1,9 +1,12 @@
+import pytest
+
 from pruna import PrunaModel
 from pruna.algorithms.kvpress import KVPress
 
 from .base_tester import AlgorithmTesterBase
 
 
+@pytest.mark.requires_kvpress
 class TestKVPress(AlgorithmTesterBase):
     """Test the KVPress KV cache compression algorithm with default settings."""
 
@@ -19,6 +22,7 @@ def post_smash_hook(self, model: PrunaModel) -> None:
         assert hasattr(model, "_kvpress_original_generate")
 
 
+@pytest.mark.requires_kvpress
 class TestKVPressSnapKV(AlgorithmTesterBase):
     """Test the KVPress algorithm with SnapKV and custom press_kwargs."""
 
