From 3bad2ca290377e30ee65adf9c7d9de1a00376344 Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 01:59:45 +0200 Subject: [PATCH 01/57] =?UTF-8?q?Convert=20InfiniteBench=20Hub=20labels=20?= =?UTF-8?q?to=20the=20upstream=20format=20before=20scoring=20=F0=9F=A4=96?= =?UTF-8?q?=F0=9F=A4=96=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The MaxJeblick/InfiniteBench answers arrive as numpy arrays of strings and multiple-choice answers only keep the letter, so with perfect predictions code_run scored 0, code_debug raised an IndexError, math_find a TypeError, math_calc an AssertionError, and passkey/number_string relied on a deprecated ndarray-to-float conversion. calculate_metrics now converts each label to a list, casts code_run, math_find and math_calc labels to numbers, and rebuilds the upstream [option_text, letter] labels of code_debug and longbook_choice_eng from the options at the end of the question (or of the context with query_aware). The upstream scoring functions are unchanged. Signed-off-by: Maximilian Jeblick --- .../infinite_bench/calculate_metrics.py | 29 ++++++++- tests/test_evaluation_utils.py | 63 +++++++++++++++++++ 2 files changed, 91 insertions(+), 1 deletion(-) diff --git a/evaluation/benchmarks/infinite_bench/calculate_metrics.py b/evaluation/benchmarks/infinite_bench/calculate_metrics.py index e99378967..88c421594 100644 --- a/evaluation/benchmarks/infinite_bench/calculate_metrics.py +++ b/evaluation/benchmarks/infinite_bench/calculate_metrics.py @@ -14,12 +14,39 @@ def calculate_metrics(df): preds = df["predicted_answer"].tolist() - labels = df["answer"].tolist() task = df["task"].tolist()[0] + labels = [to_upstream_label(row, task) for _, row in df.iterrows()] model = "dummy_model" # not really used return get_score(labels, preds, task, model) +# The Hub dataset stores every answer as a list of strings and keeps only the letter of multiple-choice answers, +# while the upstream scorers below expect numbers for code_run, math_find and math_calc, and [option_text, letter] +# for code_debug and longbook_choice_eng (see get_answer in InfiniteBench/src/eval_utils.py). +OPTIONS_PATTERN = re.compile(r"\nA\. (.*)\nB\. (.*)\nC\. (.*)\nD\. (.*)\Z") + + +def to_upstream_label(row, task: str) -> list: + label = list(row["answer"]) + if task in ["code_run", "math_find", "math_calc"]: + return [to_number(value) for value in label] + if task in ["code_debug", "longbook_choice_eng"]: + # The options end the question, which query-aware evaluation moves to the end of the context + match = OPTIONS_PATTERN.search(row["question"] or row["context"]) + if match is None: + raise ValueError(f"Could not find the A-D options of the {task} question") + letter = label[0] + return [match.group("ABCD".index(letter) + 1), letter] + return label + + +def to_number(value: str) -> int | float: + try: + return int(value) + except ValueError: + return float(value) + + def normalize_answer(s: str) -> str: """Lower text and remove punctuation, articles and extra whitespace.""" diff --git a/tests/test_evaluation_utils.py b/tests/test_evaluation_utils.py index 934cbdb4a..862d6b06e 100644 --- a/tests/test_evaluation_utils.py +++ b/tests/test_evaluation_utils.py @@ -1,12 +1,37 @@ # SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import numpy as np +import pandas as pd import pytest from evaluation.benchmarks.aime25.calculate_metrics import score_aime as score_aime25 +from evaluation.benchmarks.infinite_bench.calculate_metrics import calculate_metrics as infinite_bench_scorer from evaluation.benchmarks.math500.calculate_metrics import score_aime as score_math500 from evaluation.benchmarks.utils import extract_boxed +CODE_DEBUG_QUESTION = ( + "Which funtion has deliberate error?\n\nOptions:\nA. Resource.pkgname\nB. repack_carchive\nC. cmd_gen\nD. _init" +) +LONGBOOK_CHOICE_QUESTION = ( + "Which of the following is NOT one of Alain's chores at Hall Farm?\n\nOnly one of the following options is " + "correct, tell me the answer using one single letter (A, B, C, or D). Don't say anything else.\n" + "A. Walking Georgie\nB. Taking care of Totty\nC. Working in the dairy\nD. Light housework" +) + + +def infinite_bench_df(task, question, answer, prediction, context="Some long context."): + # Same layout as load_dataset("MaxJeblick/InfiniteBench", data_dir=task, split="test").to_pandas() + return pd.DataFrame( + { + "context": [context], + "question": [question], + "answer": [np.array(answer, dtype=object)], + "task": [task], + "predicted_answer": [prediction], + } + ) + @pytest.mark.parametrize( ("text", "expected"), @@ -43,3 +68,41 @@ def test_aime25_scores_last_boxed_answer(): def test_math500_scores_first_boxed_answer(): assert score_math500(r"First: \boxed{\frac{1}{2}}. Later: \boxed{1}.", r"\frac{1}{2}") + + +# task: (question, answer, perfect prediction) +INFINITE_BENCH_ROWS = { + "passkey": ("What is the pass key?", ["71432"], "71432"), + "number_string": ("What is the sequence of digits?", ["2200012222"], "2200012222"), + "kv_retrieval": ('\nKey: "798c2306-5ad1-42a9"\n ', ["5e6b7b90-710d-4953"], "5e6b7b90-710d-4953"), + "code_run": ("Please give me the exact number of the return value of func_6577(-7).", ["-14"], "-14"), + "code_debug": (CODE_DEBUG_QUESTION, ["B"], "B"), + "math_find": ("The largest number of the list is: ", ["88"], "88"), + "math_calc": ("What's the intermediate results of the given numerical expression?", ["79", "5", "73"], "79, 5, 73"), + "longbook_qa_eng": ("Which among Annalisa, Seb and Peyton is not Mrs. Bronwyn's child?", ['"Peyton"'], "Peyton"), + "longbook_qa_chn": ("云景, 素素, 云林哪一个人物是第二个登场？", ["素素"], "素素"), + "longbook_choice_eng": (LONGBOOK_CHOICE_QUESTION, ["A"], "A"), + "longdialogue_qa_eng": ("Which character is $$MASK$$ ?", ["ACE", "ACE ROTHSTEIN"], "ACE"), +} + + +@pytest.mark.filterwarnings("error::DeprecationWarning") +@pytest.mark.parametrize("task", list(INFINITE_BENCH_ROWS)) +def test_infinite_bench_scores_perfect_prediction(task): + question, answer, prediction = INFINITE_BENCH_ROWS[task] + assert infinite_bench_scorer(infinite_bench_df(task, question, answer, prediction)) == 1.0 + + +@pytest.mark.parametrize("query_aware", [False, True]) +@pytest.mark.parametrize( + ("task", "correct_option", "wrong_option"), + [("code_debug", "repack_carchive", "cmd_gen"), ("longbook_choice_eng", "Walking Georgie", "Light housework")], +) +def test_infinite_bench_choice_label_includes_option_text(task, correct_option, wrong_option, query_aware): + question, answer, _ = INFINITE_BENCH_ROWS[task] + context = "Some long context." + if query_aware: + context, question = context + question, "" + + assert infinite_bench_scorer(infinite_bench_df(task, question, answer, correct_option, context)) == 1.0 + assert infinite_bench_scorer(infinite_bench_df(task, question, answer, wrong_option, context)) == 0.0 From 3ca7aa4c9f5446d58e6d814bc637a9287b343c4d Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 02:02:24 +0200 Subject: [PATCH 02/57] =?UTF-8?q?Fall=20back=20to=20an=20empty=20BOS=20tok?= =?UTF-8?q?en=20in=20the=20pipeline=20=F0=9F=A4=96=F0=9F=A4=96=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Tokenizers without a BOS token (e.g. Qwen) return None for bos_token, so models without a chat template failed with a TypeError when building the context. Signed-off-by: Maximilian Jeblick --- kvpress/pipeline.py | 2 +- tests/test_pipeline_fixes.py | 20 ++++++++++++++++++++ 2 files changed, 21 insertions(+), 1 deletion(-) create mode 100644 tests/test_pipeline_fixes.py diff --git a/kvpress/pipeline.py b/kvpress/pipeline.py index eb366a554..4d4e5b626 100644 --- a/kvpress/pipeline.py +++ b/kvpress/pipeline.py @@ -138,7 +138,7 @@ def preprocess( # Apply chat template if available if self.tokenizer.chat_template is None: - bos_token = getattr(self.tokenizer, "bos_token", "") + bos_token = self.tokenizer.bos_token or "" context = bos_token + context question_suffix = "\n" # to separate the question from the answer else: diff --git a/tests/test_pipeline_fixes.py b/tests/test_pipeline_fixes.py new file mode 100644 index 000000000..050f5c64e --- /dev/null +++ b/tests/test_pipeline_fixes.py @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import torch +from transformers import AutoTokenizer + +from kvpress.pipeline import KVPressTextGenerationPipeline +from tests.fixtures import unit_test_model # noqa: F401 + + +def test_preprocess_without_chat_template_and_bos_token(unit_test_model): # noqa: F811 + tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct") + tokenizer.chat_template = None + assert tokenizer.bos_token is None + + pipe = KVPressTextGenerationPipeline(model=unit_test_model, tokenizer=tokenizer, device=unit_test_model.device) + inputs = pipe.preprocess("A short context.", ["A question?"], answer_prefix="", max_context_length=100) + + expected_ids = tokenizer.encode("A short context.", return_tensors="pt", add_special_tokens=False) + assert torch.equal(inputs["context_ids"], expected_ids) From 4c86a2df27162ec974dae2b2a99ad03f218562fc Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 02:03:42 +0200 Subject: [PATCH 03/57] =?UTF-8?q?Format=20kvcompose=5Fpress.py=20with=20bl?= =?UTF-8?q?ack=20and=20isort=20=F0=9F=A4=96=F0=9F=A4=96=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Maximilian Jeblick --- kvpress/presses/kvcompose_press.py | 70 ++++++++++++++---------------- 1 file changed, 33 insertions(+), 37 deletions(-) diff --git a/kvpress/presses/kvcompose_press.py b/kvpress/presses/kvcompose_press.py index ed34a8b11..7499858c0 100644 --- a/kvpress/presses/kvcompose_press.py +++ b/kvpress/presses/kvcompose_press.py @@ -3,20 +3,19 @@ import logging +import types from abc import ABC, abstractmethod from contextlib import contextmanager from dataclasses import dataclass from functools import partial from typing import Generator, Union - import numpy as np import torch -import types from torch import nn -from transformers.models.llama import LlamaForCausalLM -from transformers.modeling_utils import PreTrainedModel from transformers.cache_utils import DynamicCache +from transformers.modeling_utils import PreTrainedModel +from transformers.models.llama import LlamaForCausalLM from transformers.models.qwen2.modeling_qwen2 import Qwen2ForCausalLM from transformers.models.qwen3.modeling_qwen3 import Qwen3ForCausalLM @@ -72,7 +71,7 @@ def __init__(self, n, device): super().__init__(n, device) def _partial_fit(self, nd_data: torch.Tensor): - new_max_data = nd_data.amax(dim=tuple(range(len(nd_data.shape)-1))) + new_max_data = nd_data.amax(dim=tuple(range(len(nd_data.shape) - 1))) self.data = torch.maximum(self.data, new_max_data) @@ -81,14 +80,14 @@ class MeanAggregator(Aggregator): count_data: torch.Tensor def __init__(self, n, device): - self.neutral = 0. + self.neutral = 0.0 super().__init__(n, device) - self.sum_data = torch.full((n, ), self.neutral, device=self.device) - self.count_data = torch.full((n, ), self.neutral, device=self.device) + self.sum_data = torch.full((n,), self.neutral, device=self.device) + self.count_data = torch.full((n,), self.neutral, device=self.device) def _partial_fit(self, nd_data: torch.Tensor): - new_sum_data = nd_data.sum(dim=tuple(range(len(nd_data.shape)-1))) - new_count_data = torch.ones_like(nd_data, device=self.device).sum(dim=tuple(range(len(nd_data.shape)-1))) + new_sum_data = nd_data.sum(dim=tuple(range(len(nd_data.shape) - 1))) + new_count_data = torch.ones_like(nd_data, device=self.device).sum(dim=tuple(range(len(nd_data.shape) - 1))) self.sum_data += new_sum_data self.count_data += new_count_data self.data = self.sum_data / self.count_data @@ -209,7 +208,7 @@ def forward_hook(self, module: nn.Module, input: list[torch.Tensor], kwargs: dic if layer_attentions.shape[3] == layer_attentions.shape[2]: # Skip self-to-self attention (prefill step), only record context-to-query attentions continue - layer_att_head_attention = layer_attentions[:, att_head, :, :self.context_len] + layer_att_head_attention = layer_attentions[:, att_head, :, : self.context_len] self.task_aggregators[layer][att_head].partial_fit(layer_att_head_attention) # Clean up attention to save memory. @@ -253,17 +252,17 @@ def compute_composite_scores(self): - self.composite_scores_per_layer of shape (num_layers, context_len): structured compression. """ self.composite_scores_per_head = self.scores.sort(dim=-1, descending=True)[0] - self.composite_scores_per_head[..., :self.keep_token_lower_bound] += 1e9 + self.composite_scores_per_head[..., : self.keep_token_lower_bound] += 1e9 - self.composite_scores_per_layer = torch.full((self.num_layers, self.context_len), 0., device=self.device) + self.composite_scores_per_layer = torch.full((self.num_layers, self.context_len), 0.0, device=self.device) for layer in range(self.num_layers): layer_aggregator = aggregator_by_name[self.agg_head](self.context_len, self.device) for kv_head in range(self.num_kv_heads): layer_aggregator.partial_fit(self.scores[layer, kv_head].sort(descending=True)[0]) self.composite_scores_per_layer[layer] = layer_aggregator.transform() - self.composite_scores_per_layer[..., :self.keep_token_lower_bound] += 1e9 - self.composite_scores_per_layer[0] = \ - self.composite_scores_per_layer.max(dim=0).values # Ensures first layer is the largest. + self.composite_scores_per_layer[..., : self.keep_token_lower_bound] += 1e9 + # Ensures first layer is the largest. + self.composite_scores_per_layer[0] = self.composite_scores_per_layer.max(dim=0).values def compute_important_per_layer(self): """ @@ -303,7 +302,7 @@ def prepare_important_masks(self): self.important_mask_per_kv_head = [ [ - torch.zeros(size=(self.context_len, ), device=self.device, dtype=torch.bool) + torch.zeros(size=(self.context_len,), device=self.device, dtype=torch.bool) for _ in range(self.num_kv_heads) ] for _ in range(self.num_layers) @@ -312,9 +311,7 @@ def prepare_important_masks(self): for layer in range(self.num_layers): for kv_head in range(self.num_kv_heads): count_of_important = ( - self.important_per_layer[layer] - if self.structured - else self.important_per_head[layer, kv_head] + self.important_per_layer[layer] if self.structured else self.important_per_head[layer, kv_head] ) important_indices = torch.argsort(self.scores[layer, kv_head], descending=True)[:count_of_important] self.important_mask_per_kv_head[layer][kv_head][important_indices] = True @@ -329,8 +326,8 @@ def compress_structured(self) -> None: for kv_head in range(self.num_kv_heads): important_mask = self.important_mask_per_kv_head[layer][kv_head] - keys = self.cache.layers[layer].keys[0, kv_head][:self.context_len] - values = self.cache.layers[layer].values[0, kv_head][:self.context_len] + keys = self.cache.layers[layer].keys[0, kv_head][: self.context_len] + values = self.cache.layers[layer].values[0, kv_head][: self.context_len] keys = keys[important_mask] values = values[important_mask] kv_over_layer[0].append(keys) @@ -355,8 +352,8 @@ def compress_unstructured(self, model: PreTrainedModel) -> None: for kv_head in range(self.num_kv_heads): non_important_mask = ~self.important_mask_per_kv_head[layer_idx][kv_head] num_non_important_tokens = int(non_important_mask.sum().item()) - batch_indices = torch.full((num_non_important_tokens, ), 0, device=self.device) - head_indices = torch.full((num_non_important_tokens, ), kv_head, device=self.device) + batch_indices = torch.full((num_non_important_tokens,), 0, device=self.device) + head_indices = torch.full((num_non_important_tokens,), kv_head, device=self.device) seq_indices = non_important_mask.nonzero(as_tuple=True)[0] masked_over_layer[0].append(batch_indices) masked_over_layer[1].append(head_indices) @@ -381,21 +378,20 @@ def __call__(self, model: PreTrainedModel) -> Generator: Model to apply the compression method to """ - logger.warning( - "KVComposePress temporarily creates a KV cache of ~2x the context length during prefill; " - ) + logger.warning("KVComposePress temporarily creates a KV cache of ~2x the context length during prefill; ") if not isinstance(model, (LlamaForCausalLM, Qwen2ForCausalLM, Qwen3ForCausalLM)): logger.warning(f"Model {type(model)} not tested") self._register_model(model) - def new_forward(self, - input_ids, - past_key_values, - *args, - press: KVComposePress, - **kwargs, - ): + def new_forward( + self, + input_ids, + past_key_values, + *args, + press: KVComposePress, + **kwargs, + ): press.register_context_ids(input_ids) original_attn_implementation = self.model.config._attn_implementation @@ -405,17 +401,17 @@ def new_forward(self, past_key_values=past_key_values, *args, **kwargs, - ) + ) press._register_cache(past_key_values) - for prompt_ids in (press.prompt_ids or [press.context_ids]): + for prompt_ids in press.prompt_ids or [press.context_ids]: cache = past_key_values self.original_forward_KVComposePress( input_ids=prompt_ids.to(self.model.device), past_key_values=cache, *args, **kwargs, - ) + ) self.model.config._attn_implementation = original_attn_implementation return outputs From 8f0b4fc33bd1939e94c75a350ed67ee8452632e7 Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 02:03:49 +0200 Subject: [PATCH 04/57] =?UTF-8?q?Parse=20boolean=20CLI=20strings=20in=20ev?= =?UTF-8?q?aluate.py=20=F0=9F=A4=96=F0=9F=A4=96=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fire passes "--trust_remote_code false" (likewise --query_aware and --fp8) as the truthy string "false", which silently enabled remote code. Coerce "true"/"false" strings (case-insensitive) for every bool field of EvaluationConfig and raise a ValueError for any other non-bool value. Signed-off-by: Maximilian Jeblick --- evaluation/evaluate.py | 12 +++++- tests/test_evaluate_config.py | 72 +++++++++++++++++++++++++++++++++++ 2 files changed, 83 insertions(+), 1 deletion(-) create mode 100644 tests/test_evaluate_config.py diff --git a/evaluation/evaluate.py b/evaluation/evaluate.py index 02d50c7cf..0bb061290 100644 --- a/evaluation/evaluate.py +++ b/evaluation/evaluate.py @@ -5,7 +5,7 @@ import logging import random import sys -from dataclasses import asdict, dataclass +from dataclasses import asdict, dataclass, fields from pathlib import Path from typing import Any, Dict, Optional @@ -80,6 +80,16 @@ class EvaluationConfig: def __post_init__(self): """Validate configuration after initialization.""" + # Fire parses "--trust_remote_code false" as the (truthy) string "false" + for config_field in fields(self): + if config_field.type is bool: + value = getattr(self, config_field.name) + if isinstance(value, str) and value.lower() in ["true", "false"]: + value = value.lower() == "true" + setattr(self, config_field.name, value) + if not isinstance(value, bool): + raise ValueError(f"{config_field.name} must be a boolean, got {value!r}") + # Validate dataset assert self.dataset in DATASET_REGISTRY, f"No dataset found for {self.dataset}" assert self.dataset in SCORER_REGISTRY, f"No scorer found for {self.dataset}" diff --git a/tests/test_evaluate_config.py b/tests/test_evaluate_config.py new file mode 100644 index 000000000..d089c71fe --- /dev/null +++ b/tests/test_evaluate_config.py @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import importlib +import importlib.util +import sys +from pathlib import Path +from unittest.mock import MagicMock + +import fire +import pytest + +EVALUATION_DIR = Path(__file__).resolve().parents[1] / "evaluation" + +# Imported by the scorers in evaluate_registry.py, but only needed to compute metrics ("eval" extra) +SCORER_DEPENDENCIES = [ + "bert_score", + "fuzzywuzzy", + "jieba", + "nltk", + "nltk.translate", + "nltk.translate.bleu_score", + "nltk.translate.meteor_score", + "rouge", +] + + +@pytest.fixture(scope="module") +def evaluate(): + missing = [name for name in SCORER_DEPENDENCIES if importlib.util.find_spec(name.split(".")[0]) is None] + with pytest.MonkeyPatch.context() as mp: + # evaluate.py imports evaluate_registry and benchmarks as top-level modules + mp.syspath_prepend(str(EVALUATION_DIR)) + for name in missing: + mp.setitem(sys.modules, name, MagicMock()) + yield importlib.import_module("evaluation.evaluate") + + +@pytest.fixture +def run_cli(evaluate, monkeypatch): + """Parses command line arguments like `python evaluate.py` and returns the config it would evaluate.""" + configs = [] + + class RecordingRunner: + def __init__(self, config): + configs.append(config) + + def run_evaluation(self): + pass + + monkeypatch.setattr(evaluate, "EvaluationRunner", RecordingRunner) + + def run(*args): + fire.Fire(evaluate.CliEntryPoint, command=list(args)) + return configs[-1] + + return run + + +@pytest.mark.parametrize(("value", "expected"), [("false", False), ("FALSE", False), ("true", True), ("True", True)]) +def test_cli_parses_boolean_strings(run_cli, value, expected): + config = run_cli("--trust_remote_code", value, "--query_aware", value, "--fp8", value) + + assert config.trust_remote_code is expected + assert config.query_aware is expected + assert config.fp8 is expected + + +@pytest.mark.parametrize("value", ["yes", 1, None]) +def test_config_rejects_non_boolean_values(evaluate, value): + with pytest.raises(ValueError, match="trust_remote_code"): + evaluate.EvaluationConfig(trust_remote_code=value) From a84c9f95ceffc56d8fc51d6ac939740d8a84901c Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 02:04:01 +0200 Subject: [PATCH 05/57] =?UTF-8?q?Reject=20multiple=20questions=20only=20fo?= =?UTF-8?q?r=20presses=20that=20compress=20during=20decoding=20?= =?UTF-8?q?=F0=9F=A4=96=F0=9F=A4=96=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit DMSPress(decoding=True) keeps masks from the first answer that point past the trimmed cache, so the next question indexes out of bounds. PrefillDecodingPress without a decoding press is safe for multiple questions. Signed-off-by: Maximilian Jeblick --- kvpress/pipeline.py | 10 +++++++-- tests/test_pipeline_fixes.py | 40 +++++++++++++++++++++++++++++++++++- 2 files changed, 47 insertions(+), 3 deletions(-) diff --git a/kvpress/pipeline.py b/kvpress/pipeline.py index 4d4e5b626..de1149d40 100644 --- a/kvpress/pipeline.py +++ b/kvpress/pipeline.py @@ -199,9 +199,15 @@ def _forward( list[str] Generated answers for each input question. """ - if isinstance(press, (DecodingPress, PrefillDecodingPress)) and len(input_tensors["questions_ids"]) > 1: + compresses_during_decoding = ( + isinstance(press, DecodingPress) + or (isinstance(press, PrefillDecodingPress) and press.decoding_press is not None) + or (isinstance(press, DMSPress) and press.decoding) + ) + if compresses_during_decoding and len(input_tensors["questions_ids"]) > 1: raise ValueError( - "DecodingPress is not compatible with multiple questions. Please specify a single question." + f"{type(press).__name__} compresses the cache during decoding and is not compatible with multiple " + "questions. Please specify a single question." ) context_ids = input_tensors["context_ids"].to(self.model.device) diff --git a/tests/test_pipeline_fixes.py b/tests/test_pipeline_fixes.py index 050f5c64e..97518b62b 100644 --- a/tests/test_pipeline_fixes.py +++ b/tests/test_pipeline_fixes.py @@ -1,11 +1,16 @@ # SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import pytest import torch from transformers import AutoTokenizer +from kvpress import DecodingPress, DMSPress, KnormPress, PrefillDecodingPress, RandomPress from kvpress.pipeline import KVPressTextGenerationPipeline -from tests.fixtures import unit_test_model # noqa: F401 +from tests.fixtures import kv_press_unit_test_pipeline, unit_test_model # noqa: F401 + +CONTEXT = "This is a test article. It was written on 2022-01-01. " * 4 +QUESTIONS = ["When was this article written?", "What is this?"] def test_preprocess_without_chat_template_and_bos_token(unit_test_model): # noqa: F811 @@ -18,3 +23,36 @@ def test_preprocess_without_chat_template_and_bos_token(unit_test_model): # noq expected_ids = tokenizer.encode("A short context.", return_tensors="pt", add_special_tokens=False) assert torch.equal(inputs["context_ids"], expected_ids) + + +@pytest.mark.parametrize( + "make_press", + [ + lambda: DecodingPress(base_press=KnormPress(), compression_interval=2, target_size=16), + lambda: PrefillDecodingPress( + decoding_press=DecodingPress(base_press=KnormPress(), compression_interval=2, target_size=16) + ), + lambda: DMSPress(press=RandomPress(), threshold=0.5, sliding_window_size=0, decoding=True), + ], + ids=["decoding", "prefill_decoding", "dms_decoding"], +) +def test_pipeline_rejects_multiple_questions_with_decoding_compression( + kv_press_unit_test_pipeline, make_press # noqa: F811 +): + with pytest.raises(ValueError, match="not compatible with multiple questions"): + kv_press_unit_test_pipeline(CONTEXT, questions=QUESTIONS, press=make_press(), max_new_tokens=3) + + +@pytest.mark.parametrize( + "make_press", + [ + lambda: PrefillDecodingPress(prefilling_press=KnormPress(compression_ratio=0.5)), + lambda: DMSPress(press=RandomPress(), threshold=0.5, sliding_window_size=0), + ], + ids=["prefill_decoding", "dms_prefill"], +) +def test_pipeline_accepts_multiple_questions_with_prefill_compression( + kv_press_unit_test_pipeline, make_press # noqa: F811 +): + answers = kv_press_unit_test_pipeline(CONTEXT, questions=QUESTIONS, press=make_press(), max_new_tokens=3) + assert len(answers["answers"]) == 2 From 679fdea2ff2b7622fcc6aafcaaf293d59e6c9dac Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 02:06:21 +0200 Subject: [PATCH 06/57] =?UTF-8?q?Also=20stop=20generation=20at=20the=20tok?= =?UTF-8?q?enizer=20EOS=20token=20=F0=9F=A4=96=F0=9F=A4=96=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit generation_config.eos_token_id may be None or miss the tokenizer EOS token, in which case the pipeline always generated max_new_tokens tokens. Signed-off-by: Maximilian Jeblick --- kvpress/pipeline.py | 9 ++++++--- tests/test_pipeline_fixes.py | 17 +++++++++++++++++ 2 files changed, 23 insertions(+), 3 deletions(-) diff --git a/kvpress/pipeline.py b/kvpress/pipeline.py index de1149d40..f1a70bdad 100644 --- a/kvpress/pipeline.py +++ b/kvpress/pipeline.py @@ -311,9 +311,12 @@ def generate_answer( position_ids = position_ids[:, -1:] + 1 generated_ids = [outputs.logits[0, -1].argmax()] - should_stop_token_ids = self.model.generation_config.eos_token_id - if not isinstance(should_stop_token_ids, list): - should_stop_token_ids = [should_stop_token_ids] + eos_token_ids = self.model.generation_config.eos_token_id + if not isinstance(eos_token_ids, list): + eos_token_ids = [eos_token_ids] + should_stop_token_ids = { + token_id for token_id in [*eos_token_ids, self.tokenizer.eos_token_id] if token_id is not None + } for i in range(max_new_tokens - 1): outputs = self.model( diff --git a/tests/test_pipeline_fixes.py b/tests/test_pipeline_fixes.py index 97518b62b..ac7135346 100644 --- a/tests/test_pipeline_fixes.py +++ b/tests/test_pipeline_fixes.py @@ -56,3 +56,20 @@ def test_pipeline_accepts_multiple_questions_with_prefill_compression( ): answers = kv_press_unit_test_pipeline(CONTEXT, questions=QUESTIONS, press=make_press(), max_new_tokens=3) assert len(answers["answers"]) == 2 + + +def test_generate_answer_stops_at_tokenizer_eos_token(unit_test_model, monkeypatch): # noqa: F811 + tokenizer = AutoTokenizer.from_pretrained(unit_test_model.config.name_or_path) + pipe = KVPressTextGenerationPipeline(model=unit_test_model, tokenizer=tokenizer, device=unit_test_model.device) + monkeypatch.setattr(unit_test_model.generation_config, "eos_token_id", None) + generated_ids = [] + monkeypatch.setattr(tokenizer, "decode", lambda ids, **kwargs: generated_ids.append(ids.tolist()) or "") + + pipe(CONTEXT, question=QUESTIONS[0], max_new_tokens=8) + reference = generated_ids.pop() + assert len(reference) == 8 + + # The first generated token is never checked, so the second one is the earliest possible stop + tokenizer.eos_token = tokenizer.convert_ids_to_tokens(reference[1]) + pipe(CONTEXT, question=QUESTIONS[0], max_new_tokens=8) + assert generated_ids.pop() == reference[:2] From 181fabe44786c00f80cad29f61919355400733f0 Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 02:06:41 +0200 Subject: [PATCH 07/57] =?UTF-8?q?Respect=20a=20configured=20attn=5Fimpleme?= =?UTF-8?q?ntation=20in=20evaluate.py=20=F0=9F=A4=96=F0=9F=A4=96?= =?UTF-8?q?=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _setup_model_pipeline overwrote model_kwargs["attn_implementation"] with flash_attention_2 whenever flash_attn was importable, even when the user set it (evaluate_config.yaml advertises model_kwargs). Only select flash attention when attn_implementation is unset. Signed-off-by: Maximilian Jeblick --- evaluation/evaluate.py | 2 ++ tests/test_evaluate_config.py | 29 +++++++++++++++++++++++++++++ 2 files changed, 31 insertions(+) diff --git a/evaluation/evaluate.py b/evaluation/evaluate.py index 0bb061290..c8dd0cc5a 100644 --- a/evaluation/evaluate.py +++ b/evaluation/evaluate.py @@ -385,6 +385,8 @@ def _setup_model_pipeline(self): if isinstance(self.press, ObservedAttentionPress): model_kwargs["attn_implementation"] = "eager" logger.info("ObservedAttentionPress detected, setting attn_implementation to 'eager'.") + elif model_kwargs.get("attn_implementation") is not None: + logger.info(f"Using attn_implementation='{model_kwargs['attn_implementation']}' from model_kwargs.") else: try: import flash_attn # noqa: F401 diff --git a/tests/test_evaluate_config.py b/tests/test_evaluate_config.py index d089c71fe..aac7c939b 100644 --- a/tests/test_evaluate_config.py +++ b/tests/test_evaluate_config.py @@ -10,6 +10,8 @@ import fire import pytest +from kvpress import KnormPress + EVALUATION_DIR = Path(__file__).resolve().parents[1] / "evaluation" # Imported by the scorers in evaluate_registry.py, but only needed to compute metrics ("eval" extra) @@ -57,6 +59,19 @@ def run(*args): return run +@pytest.fixture +def pipeline_calls(evaluate, monkeypatch): + """Records the keyword arguments of the transformers pipeline() calls instead of loading a model.""" + calls = [] + + def fake_pipeline(task, **kwargs): + calls.append(kwargs) + return MagicMock() + + monkeypatch.setattr(evaluate, "pipeline", fake_pipeline) + return calls + + @pytest.mark.parametrize(("value", "expected"), [("false", False), ("FALSE", False), ("true", True), ("True", True)]) def test_cli_parses_boolean_strings(run_cli, value, expected): config = run_cli("--trust_remote_code", value, "--query_aware", value, "--fp8", value) @@ -70,3 +85,17 @@ def test_cli_parses_boolean_strings(run_cli, value, expected): def test_config_rejects_non_boolean_values(evaluate, value): with pytest.raises(ValueError, match="trust_remote_code"): evaluate.EvaluationConfig(trust_remote_code=value) + + +@pytest.mark.parametrize("flash_attn_available", [True, False]) +def test_model_pipeline_keeps_configured_attn_implementation( + evaluate, pipeline_calls, monkeypatch, flash_attn_available +): + monkeypatch.setitem(sys.modules, "flash_attn", MagicMock() if flash_attn_available else None) + runner = evaluate.EvaluationRunner( + evaluate.EvaluationConfig(device="cpu", model_kwargs={"attn_implementation": "sdpa"}) + ) + runner.press = KnormPress() + runner._setup_model_pipeline() + + assert pipeline_calls[-1]["model_kwargs"]["attn_implementation"] == "sdpa" From ef40360f760705a3d73b55de3012f1472889399f Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 02:07:57 +0200 Subject: [PATCH 08/57] =?UTF-8?q?Keep=20runtime=20model=20kwargs=20out=20o?= =?UTF-8?q?f=20the=20saved=20evaluation=20config=20=F0=9F=A4=96?= =?UTF-8?q?=F0=9F=A4=96=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit model_kwargs in _setup_model_pipeline aliased config.model_kwargs, so with --fp8 a FineGrainedFP8Config object (and the automatic attn_implementation) ended up in config.yaml, which yaml.dump wrote as a python/object tag that yaml.safe_load cannot read. Copy the dict before adding entries, and save the config with yaml.safe_dump (which also writes tuples, e.g. from --needle_depth 10,50, as plain lists). Signed-off-by: Maximilian Jeblick --- evaluation/evaluate.py | 4 ++-- tests/test_evaluate_config.py | 29 ++++++++++++++++++++++++----- 2 files changed, 26 insertions(+), 7 deletions(-) diff --git a/evaluation/evaluate.py b/evaluation/evaluate.py index c8dd0cc5a..3cb2de6c6 100644 --- a/evaluation/evaluate.py +++ b/evaluation/evaluate.py @@ -183,7 +183,7 @@ def save_config(self, config_filename: Path): if self.head_compression_ratio is None: config_dict.pop("head_compression_ratio", None) with open(str(config_filename), "w") as f: - yaml.dump(config_dict, f, default_flow_style=False, indent=2, sort_keys=False) + yaml.safe_dump(config_dict, f, default_flow_style=False, indent=2, sort_keys=False) def _load_yaml_config(path: str | Path) -> dict: @@ -376,7 +376,7 @@ def _setup_model_pipeline(self): device = "auto" if torch.cuda.is_available() else "cpu" logger.info(f"No device specified, auto-detected device: {device}") - model_kwargs = self.config.model_kwargs or {} + model_kwargs = dict(self.config.model_kwargs or {}) if self.config.fp8: model_kwargs["quantization_config"] = FineGrainedFP8Config() diff --git a/tests/test_evaluate_config.py b/tests/test_evaluate_config.py index aac7c939b..52cbd9b9a 100644 --- a/tests/test_evaluate_config.py +++ b/tests/test_evaluate_config.py @@ -9,6 +9,8 @@ import fire import pytest +import yaml +from transformers import FineGrainedFP8Config from kvpress import KnormPress @@ -87,15 +89,32 @@ def test_config_rejects_non_boolean_values(evaluate, value): evaluate.EvaluationConfig(trust_remote_code=value) +def setup_model_pipeline(evaluate, press, **config_kwargs): + runner = evaluate.EvaluationRunner(evaluate.EvaluationConfig(device="cpu", **config_kwargs)) + runner.press = press + runner._setup_model_pipeline() + return runner + + @pytest.mark.parametrize("flash_attn_available", [True, False]) def test_model_pipeline_keeps_configured_attn_implementation( evaluate, pipeline_calls, monkeypatch, flash_attn_available ): monkeypatch.setitem(sys.modules, "flash_attn", MagicMock() if flash_attn_available else None) - runner = evaluate.EvaluationRunner( - evaluate.EvaluationConfig(device="cpu", model_kwargs={"attn_implementation": "sdpa"}) - ) - runner.press = KnormPress() - runner._setup_model_pipeline() + setup_model_pipeline(evaluate, KnormPress(), model_kwargs={"attn_implementation": "sdpa"}) assert pipeline_calls[-1]["model_kwargs"]["attn_implementation"] == "sdpa" + + +def test_saved_config_can_be_read_with_safe_load(evaluate, pipeline_calls, tmp_path): + runner = setup_model_pipeline( + evaluate, KnormPress(), fp8=True, model_kwargs={"dtype": "auto"}, needle_depth=(10, 50) + ) + runner.config.save_config(tmp_path / "config.yaml") + with open(tmp_path / "config.yaml") as f: + saved_config = yaml.safe_load(f) + + assert isinstance(pipeline_calls[-1]["model_kwargs"]["quantization_config"], FineGrainedFP8Config) + assert runner.config.model_kwargs == {"dtype": "auto"} + assert saved_config["model_kwargs"] == {"dtype": "auto"} + assert saved_config["needle_depth"] == [10, 50] From 3aeaa85355772263fae614d7c1b182da115a5d59 Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 02:09:10 +0200 Subject: [PATCH 09/57] =?UTF-8?q?Restrict=20KVComposePress=20value=20norms?= =?UTF-8?q?=20to=20the=20context=20=F0=9F=A4=96=F0=9F=A4=96=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit With add_v_norm=True the token scores (context length) were multiplied by the norms of all cached values, which also cover the re-encoded copy of the context, so the press always crashed with a size mismatch. Signed-off-by: Maximilian Jeblick --- kvpress/presses/kvcompose_press.py | 2 +- tests/presses/test_kvcompose_fixes.py | 26 ++++++++++++++++++++++++++ 2 files changed, 27 insertions(+), 1 deletion(-) create mode 100644 tests/presses/test_kvcompose_fixes.py diff --git a/kvpress/presses/kvcompose_press.py b/kvpress/presses/kvcompose_press.py index 7499858c0..d31cf43e6 100644 --- a/kvpress/presses/kvcompose_press.py +++ b/kvpress/presses/kvcompose_press.py @@ -239,7 +239,7 @@ def enhance_scores(self): for layer in range(self.num_layers): if self.add_v_norm: for kv_head in range(self.num_kv_heads): - v = self.cache.layers[layer].values[0, kv_head].detach() + v = self.cache.layers[layer].values[0, kv_head, : self.context_len].detach() self.scores[layer, kv_head] = self.scores[layer, kv_head] * v.norm(dim=1) if self.add_mean_across_heads: self.scores[layer] += self.scores[layer].mean(dim=0, keepdim=True) diff --git a/tests/presses/test_kvcompose_fixes.py b/tests/presses/test_kvcompose_fixes.py new file mode 100644 index 000000000..663b5d66d --- /dev/null +++ b/tests/presses/test_kvcompose_fixes.py @@ -0,0 +1,26 @@ +# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch +from transformers import DynamicCache + +from kvpress import KVComposePress +from tests.fixtures import unit_test_model # noqa: F401 + + +@pytest.mark.parametrize("structured", [True, False]) +def test_kvcompose_add_v_norm(unit_test_model, structured): # noqa: F811 + context_len = 64 + for layer in unit_test_model.model.layers: + layer.self_attn.masked_key_indices = None + press = KVComposePress(structured=structured, compression_ratio=0.5, add_v_norm=True) + cache = DynamicCache() + with press(unit_test_model): + input_ids = torch.randint(0, 1024, (1, context_len), device=unit_test_model.device) + unit_test_model(input_ids, past_key_values=cache) + + if structured: + assert 0 < cache.get_seq_length() < context_len + else: + assert unit_test_model.model.layers[0].self_attn.masked_key_indices is not None From 348f6932838bab490bf56b14d6ad26faf2869027 Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 02:09:59 +0200 Subject: [PATCH 10/57] =?UTF-8?q?Fix=20needle=5Fin=5Fhaystack=20with=20fra?= =?UTF-8?q?ction=20<=201=20and=20too=20small=20contexts=20=F0=9F=A4=96?= =?UTF-8?q?=F0=9F=A4=96=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The haystack dataset has a single row, so fraction <= 0.5 sampled it away and insert_needle_in_haystack then failed with KeyError: 0 (it read the first row by label). Skip fraction sampling for this dataset with a warning, read the first row by position, and raise a ValueError when max_context_length leaves no room for the haystack once the needle and the 150 reserved prompt tokens are subtracted (a negative budget used to keep almost the whole essay). Signed-off-by: Maximilian Jeblick --- .../benchmarks/needle_in_haystack/utils.py | 15 +++++--- evaluation/evaluate.py | 4 ++- tests/test_evaluate_config.py | 14 ++++++++ tests/test_evaluation_utils.py | 35 +++++++++++++++++++ 4 files changed, 62 insertions(+), 6 deletions(-) diff --git a/evaluation/benchmarks/needle_in_haystack/utils.py b/evaluation/benchmarks/needle_in_haystack/utils.py index edec1766d..5f049ff88 100644 --- a/evaluation/benchmarks/needle_in_haystack/utils.py +++ b/evaluation/benchmarks/needle_in_haystack/utils.py @@ -51,16 +51,21 @@ def insert_needle_in_haystack( """ # Store the original context and needle to be reused for each depth - original_context = df["context"][0] - needle_text = needle_text or df["needle"][0] - question_text = question_text or df["question"][0] - answer_prefix = answer_prefix or df["answer_prefix"][0] - max_new_tokens = df["max_new_tokens"][0] + original_context = df["context"].iloc[0] + needle_text = needle_text or df["needle"].iloc[0] + question_text = question_text or df["question"].iloc[0] + answer_prefix = answer_prefix or df["answer_prefix"].iloc[0] + max_new_tokens = df["max_new_tokens"].iloc[0] logger.info(f"Preparing dataset for inference. Needle: {needle_text}") tokenized_needle = tokenizer.encode(needle_text, add_special_tokens=False) # Account for system prompts and other overhead context_length_limit = max_context_length - len(tokenized_needle) - 150 + if context_length_limit <= 0: + raise ValueError( + f"max_context_length ({max_context_length}) must exceed the needle length ({len(tokenized_needle)} " + "tokens) plus 150 tokens reserved for the prompt" + ) # Tokenize the original context once tokenized_context = tokenizer.encode(original_context, add_special_tokens=False)[:context_length_limit] # Initialize a list to hold the new rows diff --git a/evaluation/evaluate.py b/evaluation/evaluate.py index 3cb2de6c6..0e473d596 100644 --- a/evaluation/evaluate.py +++ b/evaluation/evaluate.py @@ -337,7 +337,9 @@ def _load_and_prepare_dataset(self): logger.info(f"Loading dataset: {DATASET_REGISTRY[dataset_name]} (data_dir: {data_dir})") df = load_dataset(DATASET_REGISTRY[dataset_name], data_dir=data_dir, split="test").to_pandas() - if fraction < 1.0: + if fraction < 1.0 and dataset_name == "needle_in_haystack": + logger.warning("Ignoring fraction for needle_in_haystack: the haystack dataset has a single row.") + elif fraction < 1.0: original_len = len(df) df = df.sample(frac=fraction, random_state=self.config.seed) logger.info(f"Sampled {len(df)} samples ({fraction:.2f}) from original {original_len} samples.") diff --git a/tests/test_evaluate_config.py b/tests/test_evaluate_config.py index 52cbd9b9a..9fca3bdc2 100644 --- a/tests/test_evaluate_config.py +++ b/tests/test_evaluate_config.py @@ -13,6 +13,7 @@ from transformers import FineGrainedFP8Config from kvpress import KnormPress +from tests.test_evaluation_utils import WordTokenizer, haystack_df EVALUATION_DIR = Path(__file__).resolve().parents[1] / "evaluation" @@ -118,3 +119,16 @@ def test_saved_config_can_be_read_with_safe_load(evaluate, pipeline_calls, tmp_p assert runner.config.model_kwargs == {"dtype": "auto"} assert saved_config["model_kwargs"] == {"dtype": "auto"} assert saved_config["needle_depth"] == [10, 50] + + +def test_needle_in_haystack_ignores_fraction(evaluate, monkeypatch): + monkeypatch.setattr(evaluate, "load_dataset", lambda *args, **kwargs: MagicMock(to_pandas=lambda: haystack_df(0))) + config = evaluate.EvaluationConfig( + dataset="needle_in_haystack", needle_depth=[10, 90], max_context_length=1000, fraction=0.1 + ) + runner = evaluate.EvaluationRunner(config) + runner.pipeline = MagicMock(tokenizer=WordTokenizer()) + + runner._load_and_prepare_dataset() + + assert runner.df["needle_depth"].tolist() == [10, 90] diff --git a/tests/test_evaluation_utils.py b/tests/test_evaluation_utils.py index 862d6b06e..ea52f7b64 100644 --- a/tests/test_evaluation_utils.py +++ b/tests/test_evaluation_utils.py @@ -8,6 +8,7 @@ from evaluation.benchmarks.aime25.calculate_metrics import score_aime as score_aime25 from evaluation.benchmarks.infinite_bench.calculate_metrics import calculate_metrics as infinite_bench_scorer from evaluation.benchmarks.math500.calculate_metrics import score_aime as score_math500 +from evaluation.benchmarks.needle_in_haystack.utils import insert_needle_in_haystack from evaluation.benchmarks.utils import extract_boxed CODE_DEBUG_QUESTION = ( @@ -20,6 +21,27 @@ ) +class WordTokenizer: + def encode(self, text, add_special_tokens=False): + return text.split() + + def decode(self, tokens, skip_special_tokens=True): + return " ".join(tokens) + + +def haystack_df(index): + return pd.DataFrame( + { + "context": [" ".join(f"word{i}" for i in range(1000))], + "needle": ["The secret is 42."], + "question": ["What is the secret?"], + "answer_prefix": ["Answer:"], + "max_new_tokens": [10], + }, + index=[index], + ) + + def infinite_bench_df(task, question, answer, prediction, context="Some long context."): # Same layout as load_dataset("MaxJeblick/InfiniteBench", data_dir=task, split="test").to_pandas() return pd.DataFrame( @@ -86,6 +108,19 @@ def test_math500_scores_first_boxed_answer(): } +def test_insert_needle_in_haystack_reads_first_row_by_position(): + df = insert_needle_in_haystack(haystack_df(index=7), WordTokenizer(), max_context_length=400, needle_depth=[0, 50]) + + assert df["needle_depth"].tolist() == [0, 50] + assert all("The secret is 42." in context for context in df["context"]) + assert df["question"].tolist() == ["What is the secret?"] * 2 + + +def test_insert_needle_in_haystack_rejects_too_small_max_context_length(): + with pytest.raises(ValueError, match="max_context_length"): + insert_needle_in_haystack(haystack_df(index=0), WordTokenizer(), max_context_length=150, needle_depth=50) + + @pytest.mark.filterwarnings("error::DeprecationWarning") @pytest.mark.parametrize("task", list(INFINITE_BENCH_ROWS)) def test_infinite_bench_scores_perfect_prediction(task): From af8daa3621d47df01a1ef765569a0d7e9e82202b Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 02:15:55 +0200 Subject: [PATCH 11/57] =?UTF-8?q?Score=20only=20the=20first=20prefill=20in?= =?UTF-8?q?=20KVComposePress=20and=20drop=20its=20replayed=20copy=20?= =?UTF-8?q?=F0=9F=A4=96=F0=9F=A4=96=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The patched forward re-encoded the context after every call, decoding steps included, so under model.generate the answer changed even at compression_ratio=0 and the cache ended up holding the last token only. The re-encoded copy of the context also stayed in the cache: structured mode cropped it when compressing, but in unstructured mode a 64-token context left 128 entries and the question positions collided with the copy. Only the first prefill (empty cache) is now scored, with the forward hooks registered for its two passes only, and the copy is cropped right after the scoring pass. Later calls run unchanged, and the replay no longer reuses the positions and masks of the prefill call. The scoring uses the cache returned by the model, so calling the model without a cache also works. Structured results are unchanged. Unstructured results change: the kvcompose_unstructured registry entry needs to be re-evaluated. Signed-off-by: Maximilian Jeblick --- kvpress/presses/kvcompose_press.py | 50 +++++++++++++++++---------- tests/presses/test_kvcompose_fixes.py | 38 ++++++++++++++++++++ 2 files changed, 70 insertions(+), 18 deletions(-) diff --git a/kvpress/presses/kvcompose_press.py b/kvpress/presses/kvcompose_press.py index d31cf43e6..76a750820 100644 --- a/kvpress/presses/kvcompose_press.py +++ b/kvpress/presses/kvcompose_press.py @@ -383,6 +383,7 @@ def __call__(self, model: PreTrainedModel) -> Generator: logger.warning(f"Model {type(model)} not tested") self._register_model(model) + self._reset_state() def new_forward( self, @@ -392,35 +393,50 @@ def new_forward( press: KVComposePress, **kwargs, ): + if press.context_ids is not None or (past_key_values is not None and past_key_values.get_seq_length() > 0): + # Only the first prefill is scored, later calls (e.g. decoding steps) run unchanged + return self.original_forward_KVComposePress( + input_ids=input_ids, + past_key_values=past_key_values, + *args, + **kwargs, + ) + press.register_context_ids(input_ids) + hooks = [ + layer.self_attn.register_forward_hook(press.forward_hook, with_kwargs=True) + for layer in self.model.layers + ] original_attn_implementation = self.model.config._attn_implementation self.model.config._attn_implementation = "eager" - outputs = self.original_forward_KVComposePress( - input_ids=input_ids, - past_key_values=past_key_values, - *args, - **kwargs, - ) - - press._register_cache(past_key_values) - for prompt_ids in press.prompt_ids or [press.context_ids]: - cache = past_key_values - self.original_forward_KVComposePress( - input_ids=prompt_ids.to(self.model.device), - past_key_values=cache, + try: + outputs = self.original_forward_KVComposePress( + input_ids=input_ids, + past_key_values=past_key_values, *args, **kwargs, ) - self.model.config._attn_implementation = original_attn_implementation + cache = outputs.past_key_values + press._register_cache(cache) + for prompt_ids in press.prompt_ids or [press.context_ids]: + # The positions and masks of the prefill (e.g. passed by model.generate) do not apply here + self.original_forward_KVComposePress( + input_ids=prompt_ids.to(self.model.device), + past_key_values=cache, + ) + # Later forward passes must only attend to the context + cache.crop(press.context_len) + finally: + self.model.config._attn_implementation = original_attn_implementation + for hook in hooks: + hook.remove() return outputs - hooks = [] try: for layer in model.model.layers: layer.self_attn.rotary_emb = model.model.rotary_emb - hooks.append(layer.self_attn.register_forward_hook(self.forward_hook, with_kwargs=True)) setattr(model, "original_forward_KVComposePress", model.model.forward) new_forward_with_press = partial(new_forward, press=self) @@ -430,8 +446,6 @@ def new_forward( finally: model.model.forward = getattr(model, "original_forward_KVComposePress") delattr(model, "original_forward_KVComposePress") - for forward_hook in hooks: - forward_hook.remove() self.prepare_important_masks() self.compress_cache(model) self._reset_state() diff --git a/tests/presses/test_kvcompose_fixes.py b/tests/presses/test_kvcompose_fixes.py index 663b5d66d..7e374d286 100644 --- a/tests/presses/test_kvcompose_fixes.py +++ b/tests/presses/test_kvcompose_fixes.py @@ -24,3 +24,41 @@ def test_kvcompose_add_v_norm(unit_test_model, structured): # noqa: F811 assert 0 < cache.get_seq_length() < context_len else: assert unit_test_model.model.layers[0].self_attn.masked_key_indices is not None + + +def test_kvcompose_unstructured_cache_only_holds_the_context(unit_test_model): # noqa: F811 + context_len = 64 + press = KVComposePress(structured=False, compression_ratio=0.5) + cache = DynamicCache() + with press(unit_test_model): + input_ids = torch.randint(0, 1024, (1, context_len), device=unit_test_model.device) + unit_test_model(input_ids, past_key_values=cache) + + assert cache.get_seq_length() == context_len + for layer in unit_test_model.model.layers: + assert layer.self_attn.masked_key_indices[2].max() < context_len + + +@torch.no_grad() +@pytest.mark.parametrize("structured", [True, False]) +def test_kvcompose_generate_decodes_with_the_uncompressed_context(unit_test_model, structured): # noqa: F811 + """Under model.generate, only the prefill is scored and the compression happens when leaving the + context manager, so at compression_ratio=0 the decoding steps must match the model without press.""" + context_len, max_new_tokens = 32, 8 + input_ids = torch.randint(0, 1024, (1, context_len), device=unit_test_model.device) + generate_kwargs = dict( + max_new_tokens=max_new_tokens, do_sample=False, output_scores=True, return_dict_in_generate=True + ) + reference = unit_test_model.generate(input_ids, **generate_kwargs) + + cache = DynamicCache() + with KVComposePress(structured=structured, compression_ratio=0.0)(unit_test_model): + outputs = unit_test_model.generate(input_ids, past_key_values=cache, **generate_kwargs) + + assert len(outputs.scores) == len(reference.scores) == max_new_tokens + for scores, reference_scores in zip(outputs.scores, reference.scores): + # The press prefills with eager attention, hence the tolerance + torch.testing.assert_close(scores, reference_scores, rtol=1e-3, atol=1e-3) + # The structured compression rebuilds the cache from the context only + expected_length = context_len if structured else context_len + max_new_tokens - 1 + assert cache.get_seq_length() == expected_length From e5d3be0e4c9f3aacbc980842d247ba9d1df9ea4f Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 02:17:00 +0200 Subject: [PATCH 12/57] =?UTF-8?q?Discard=20the=20masked=20keys=20of=20a=20?= =?UTF-8?q?previous=20press=20in=20KVComposePress=20prefill=20=F0=9F=A4=96?= =?UTF-8?q?=F0=9F=A4=96=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit kvpress' attention patch resets masked_key_indices on prefill, but KVComposePress prefills with eager attention, which is not patched. The masks left by a previous head-wise press (e.g. KVzipPress) then applied to the structured cache during decoding. Signed-off-by: Maximilian Jeblick --- kvpress/presses/kvcompose_press.py | 3 +++ tests/presses/test_kvcompose_fixes.py | 12 ++++++++++++ 2 files changed, 15 insertions(+) diff --git a/kvpress/presses/kvcompose_press.py b/kvpress/presses/kvcompose_press.py index 76a750820..868d8875f 100644 --- a/kvpress/presses/kvcompose_press.py +++ b/kvpress/presses/kvcompose_press.py @@ -403,6 +403,9 @@ def new_forward( ) press.register_context_ids(input_ids) + # Eager attention bypasses kvpress' attention patch, which discards the masked keys on prefill + for layer in self.model.layers: + layer.self_attn.masked_key_indices = None hooks = [ layer.self_attn.register_forward_hook(press.forward_hook, with_kwargs=True) diff --git a/tests/presses/test_kvcompose_fixes.py b/tests/presses/test_kvcompose_fixes.py index 7e374d286..eac58ef56 100644 --- a/tests/presses/test_kvcompose_fixes.py +++ b/tests/presses/test_kvcompose_fixes.py @@ -39,6 +39,18 @@ def test_kvcompose_unstructured_cache_only_holds_the_context(unit_test_model): assert layer.self_attn.masked_key_indices[2].max() < context_len +def test_kvcompose_prefill_discards_masks_of_a_previous_press(unit_test_model): # noqa: F811 + stale_masks = (torch.tensor([0]), torch.tensor([0]), torch.tensor([10_000])) + for layer in unit_test_model.model.layers: + layer.self_attn.masked_key_indices = stale_masks + with KVComposePress(compression_ratio=0.5)(unit_test_model): + input_ids = torch.randint(0, 1024, (1, 32), device=unit_test_model.device) + unit_test_model(input_ids, past_key_values=DynamicCache()) + + for layer in unit_test_model.model.layers: + assert layer.self_attn.masked_key_indices is None + + @torch.no_grad() @pytest.mark.parametrize("structured", [True, False]) def test_kvcompose_generate_decodes_with_the_uncompressed_context(unit_test_model, structured): # noqa: F811 From e1fa4e1ae288e7d882913884604976d2cae6afd8 Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 02:17:51 +0200 Subject: [PATCH 13/57] =?UTF-8?q?Keep=20the=20full=20precision=20residual?= =?UTF-8?q?=20of=20QuantizedCache=20during=20compression=20=F0=9F=A4=96?= =?UTF-8?q?=F0=9F=A4=96=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The newest tokens of a transformers QuantizedLayer are stored in full precision in layer.keys and layer.values. Decoding compression (DecodingPress, CAMPress) ignored them and silently dropped the newest tokens. Extraction now appends the residual, and all write-backs go through set_keys_and_values, which re-quantizes every token and empties the residual. Signed-off-by: Maximilian Jeblick --- kvpress/presses/base_press.py | 14 +----- kvpress/presses/cam_press.py | 14 +----- kvpress/presses/decoding_press.py | 14 +----- kvpress/utils.py | 30 ++++++++++++ tests/test_quantized_cache.py | 80 +++++++++++++++++++++++++++++++ 5 files changed, 116 insertions(+), 36 deletions(-) create mode 100644 tests/test_quantized_cache.py diff --git a/kvpress/presses/base_press.py b/kvpress/presses/base_press.py index 5c73ee0fa..b81c2adef 100644 --- a/kvpress/presses/base_press.py +++ b/kvpress/presses/base_press.py @@ -15,12 +15,11 @@ MistralForCausalLM, Phi3ForCausalLM, PreTrainedModel, - QuantizedCache, Qwen2ForCausalLM, Qwen3ForCausalLM, ) -from kvpress.utils import extract_keys_and_values +from kvpress.utils import extract_keys_and_values, set_keys_and_values logger = logging.getLogger(__name__) @@ -138,7 +137,6 @@ def forward_hook(self, module: nn.Module, input: list[torch.Tensor], kwargs: dic hidden_states = kwargs["hidden_states"] cache = kwargs["past_key_values"] - cache_layer = cache.layers[module.layer_idx] q_len = hidden_states.shape[1] # Don't compress after pre-filling @@ -149,15 +147,7 @@ def forward_hook(self, module: nn.Module, input: list[torch.Tensor], kwargs: dic keys, values = self.compress(module, hidden_states, keys, values, output[1], kwargs) - if isinstance(cache, QuantizedCache): - cache_layer._quantized_keys = cache_layer._quantize(keys, axis=cache_layer.axis_key) - cache_layer._quantized_values = cache_layer._quantize(values, axis=cache_layer.axis_value) - cache_layer.keys = torch.zeros(0, dtype=keys.dtype, device=keys.device) # type: ignore[index] - cache_layer.values = torch.zeros(0, dtype=keys.dtype, device=keys.device) # type: ignore[index] - cache_layer.cumulative_length = keys.shape[2] - else: - cache_layer.keys = keys - cache_layer.values = values + set_keys_and_values(cache, module.layer_idx, keys, values) return output diff --git a/kvpress/presses/cam_press.py b/kvpress/presses/cam_press.py index 8b521235a..60e977001 100644 --- a/kvpress/presses/cam_press.py +++ b/kvpress/presses/cam_press.py @@ -9,14 +9,13 @@ import torch import torch.nn as nn -from transformers import QuantizedCache from transformers.models.llama.modeling_llama import repeat_kv from kvpress.presses.adakv_press import AdaKVPress from kvpress.presses.base_press import is_prefilling from kvpress.presses.decoding_press import DecodingPress from kvpress.presses.scorer_press import ScorerPress -from kvpress.utils import extract_keys_and_values, get_query_states +from kvpress.utils import extract_keys_and_values, get_query_states, set_keys_and_values logger = logging.getLogger(__name__) @@ -251,7 +250,6 @@ def forward_hook( # All hidden_states_buffer code is borrowed from DecodingPress self.hidden_states_buffer[layer_idx].append(hidden_states.detach().clone()) - cache_layer = cache.layers[module.layer_idx] keys, values = extract_keys_and_values(cache, layer_idx) bsz, num_key_value_heads, seq_len, _ = keys.shape @@ -295,15 +293,7 @@ def forward_hook( keys, values = self.compress(module, buffered_hidden_states, keys, values, attn_squeezed, kwargs) # Update cache with compressed keys and values - if isinstance(cache, QuantizedCache): - cache_layer._quantized_keys = cache_layer._quantize(keys, axis=cache_layer.axis_key) - cache_layer._quantized_values = cache_layer._quantize(values, axis=cache_layer.axis_value) - cache_layer.keys = torch.zeros(0, dtype=keys.dtype, device=keys.device) # type: ignore[index] - cache_layer.values = torch.zeros(0, dtype=keys.dtype, device=keys.device) # type: ignore[index] - cache_layer.cumulative_length = keys.shape[2] - else: - cache_layer.keys = keys - cache_layer.values = values + set_keys_and_values(cache, layer_idx, keys, values) self.layer_step_counts[layer_idx] = 0 # Always clear the buffer after compression - otherwise there's a mismatch between diff --git a/kvpress/presses/decoding_press.py b/kvpress/presses/decoding_press.py index b492ac5ba..019d1905f 100644 --- a/kvpress/presses/decoding_press.py +++ b/kvpress/presses/decoding_press.py @@ -9,12 +9,11 @@ import torch import torch.nn as nn from transformers import PreTrainedModel -from transformers.cache_utils import QuantizedCache from kvpress.presses.adakv_press import AdaKVPress from kvpress.presses.base_press import BasePress, is_prefilling from kvpress.presses.scorer_press import ScorerPress -from kvpress.utils import extract_keys_and_values +from kvpress.utils import extract_keys_and_values, set_keys_and_values logger = logging.getLogger(__name__) @@ -143,7 +142,6 @@ def forward_hook(self, module: nn.Module, input: list[torch.Tensor], kwargs: dic f"Applying decoding compression: layer_step_count ({self.layer_step_counts[layer_idx]}) >= compression_steps ({self.compression_interval})" # noqa: E501 ) - cache_layer = cache.layers[module.layer_idx] keys, values = extract_keys_and_values(cache, module.layer_idx) # Get attention weights from output @@ -155,15 +153,7 @@ def forward_hook(self, module: nn.Module, input: list[torch.Tensor], kwargs: dic logger.debug(f"Applied decoding compression: " f"keys.shape: {keys.shape}, values.shape: {values.shape}") # Update cache with compressed keys and values - if isinstance(cache, QuantizedCache): - cache_layer._quantized_keys = cache_layer._quantize(keys, axis=cache_layer.axis_key) - cache_layer._quantized_values = cache_layer._quantize(values, axis=cache_layer.axis_value) - cache_layer.keys = torch.zeros(0, dtype=keys.dtype, device=keys.device) # type: ignore[index] - cache_layer.values = torch.zeros(0, dtype=keys.dtype, device=keys.device) # type: ignore[index] - cache_layer.cumulative_length = keys.shape[2] - else: - cache_layer.keys = keys - cache_layer.values = values + set_keys_and_values(cache, module.layer_idx, keys, values) # Reset step count and clear buffer for this layer self.layer_step_counts[layer_idx] = 0 diff --git a/kvpress/utils.py b/kvpress/utils.py index 460d39fd0..805e0746d 100644 --- a/kvpress/utils.py +++ b/kvpress/utils.py @@ -129,11 +129,29 @@ def get_prerope_key_states(module: nn.Module, hidden_states: torch.Tensor) -> to def dequantize_layer(cache_layer) -> tuple[torch.Tensor, torch.Tensor]: + """ + Dequantizes a quantized cache layer. The most recent tokens are stored in full precision in + ``cache_layer.keys`` and ``cache_layer.values`` until the residual length is reached. + """ keys = cache_layer._dequantize(cache_layer._quantized_keys) values = cache_layer._dequantize(cache_layer._quantized_values) + if cache_layer.keys.dim() == 4 and cache_layer.keys.shape[-2] > 0: + keys = torch.cat([keys, cache_layer.keys], dim=-2) + values = torch.cat([values, cache_layer.values], dim=-2) return keys, values +def quantize_layer(cache_layer, keys: torch.Tensor, values: torch.Tensor): + """ + Quantizes all keys and values of a quantized cache layer and empties its full precision residual. + """ + cache_layer._quantized_keys = cache_layer._quantize(keys.contiguous(), axis=cache_layer.axis_key) + cache_layer._quantized_values = cache_layer._quantize(values.contiguous(), axis=cache_layer.axis_value) + cache_layer.keys = torch.tensor([], dtype=keys.dtype, device=keys.device) + cache_layer.values = torch.tensor([], dtype=values.dtype, device=values.device) + cache_layer.cumulative_length = keys.shape[2] + + def compute_n_kept(k_len: int, compression_ratio: float) -> int: """Number of KV pairs to keep after applying ``compression_ratio``. @@ -158,3 +176,15 @@ def extract_keys_and_values(cache: Cache, layer_idx: int) -> tuple[torch.Tensor, keys = cache.layers[layer_idx].keys values = cache.layers[layer_idx].values return keys, values + + +def set_keys_and_values(cache: Cache, layer_idx: int, keys: torch.Tensor, values: torch.Tensor): + """ + Writes the keys and values of a given cache layer, + handling both quantized and unquantized caches. + """ + if isinstance(cache, QuantizedCache): + quantize_layer(cache.layers[layer_idx], keys, values) + else: + cache.layers[layer_idx].keys = keys + cache.layers[layer_idx].values = values diff --git a/tests/test_quantized_cache.py b/tests/test_quantized_cache.py new file mode 100644 index 000000000..3bf92326c --- /dev/null +++ b/tests/test_quantized_cache.py @@ -0,0 +1,80 @@ +# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch +from transformers import Cache, DynamicCache, QuantizedCache +from transformers.cache_utils import QuantizedLayer + +from kvpress import CAMPress, DecodingPress, KnormPress +from kvpress.utils import extract_keys_and_values +from tests.fixtures import unit_test_model # noqa: F401 + + +class IdentityQuantizedLayer(QuantizedLayer): + """Quantized layer without quantization error, so that no quantization backend is needed.""" + + def _quantize(self, tensor, axis): + return tensor.clone() + + def _dequantize(self, q_tensor): + return q_tensor + + +def make_quantized_cache(num_layers: int, residual_length: int = 128) -> QuantizedCache: + cache = QuantizedCache.__new__(QuantizedCache) + Cache.__init__(cache, layers=[IdentityQuantizedLayer(residual_length=residual_length) for _ in range(num_layers)]) + return cache + + +def test_extract_keys_and_values_includes_the_full_precision_residual(): + keys, values = torch.randn(2, 1, 2, 10, 4).unbind(0) + cache = make_quantized_cache(num_layers=1, residual_length=8) + cache.update(keys[:, :, :6], values[:, :, :6], layer_idx=0) + for i in range(6, 10): + cache.update(keys[:, :, i : i + 1], values[:, :, i : i + 1], layer_idx=0) + assert cache.layers[0].keys.shape[2] == 4 + + extracted_keys, extracted_values = extract_keys_and_values(cache, 0) + assert torch.equal(extracted_keys, keys) + assert torch.equal(extracted_values, values) + + +@torch.no_grad() +def test_prefill_compression_with_quantized_cache(unit_test_model): # noqa: F811 + input_ids = torch.randint(0, 1000, (1, 32), device=unit_test_model.device) + dynamic_cache = DynamicCache() + quantized_cache = make_quantized_cache(unit_test_model.config.num_hidden_layers) + for cache in [dynamic_cache, quantized_cache]: + with KnormPress(compression_ratio=0.5)(unit_test_model): + unit_test_model(input_ids, past_key_values=cache) + + for layer_idx, layer in enumerate(quantized_cache.layers): + keys, values = extract_keys_and_values(quantized_cache, layer_idx) + assert torch.equal(keys, dynamic_cache.layers[layer_idx].keys) + assert torch.equal(values, dynamic_cache.layers[layer_idx].values) + assert layer.keys.numel() == layer.values.numel() == 0 + assert layer.get_seq_length() == 16 + + +@pytest.mark.parametrize("press_cls", [DecodingPress, CAMPress]) +@torch.no_grad() +def test_decoding_compression_with_quantized_cache_matches_dynamic_cache(unit_test_model, press_cls): # noqa: F811 + input_ids = torch.randint(0, 1000, (1, 32), device=unit_test_model.device) + dynamic_cache = DynamicCache() + quantized_cache = make_quantized_cache(unit_test_model.config.num_hidden_layers) + outputs = [] + for cache in [dynamic_cache, quantized_cache]: + torch.manual_seed(0) + press = press_cls(base_press=KnormPress(), compression_interval=4, target_size=24) + with press(unit_test_model): + outputs.append( + unit_test_model.generate(input_ids, past_key_values=cache, max_new_tokens=10, do_sample=False) + ) + + assert torch.equal(outputs[0], outputs[1]) + for layer_idx, layer in enumerate(quantized_cache.layers): + keys, values = extract_keys_and_values(quantized_cache, layer_idx) + assert torch.equal(keys, dynamic_cache.layers[layer_idx].keys) + assert torch.equal(values, dynamic_cache.layers[layer_idx].values) + assert layer.get_seq_length() == keys.shape[2] From 0f0797665ec0370a83846ff3708f7cc85012fa6c Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 12:27:39 +0200 Subject: [PATCH 14/57] =?UTF-8?q?Compress=20with=20KVComposePress=20only?= =?UTF-8?q?=20after=20a=20completed=20prefill=20=F0=9F=A4=96=F0=9F=A4=96?= =?UTF-8?q?=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The compression ran in the context manager's finally block, so a block without a forward pass raised AttributeError and an exception raised in the block was replaced by the one raised while compressing. Compress only when the block completed and a context was registered. Signed-off-by: Maximilian Jeblick --- kvpress/presses/kvcompose_press.py | 24 ++++++++++++++---------- tests/presses/test_kvcompose_fixes.py | 12 ++++++++++++ 2 files changed, 26 insertions(+), 10 deletions(-) diff --git a/kvpress/presses/kvcompose_press.py b/kvpress/presses/kvcompose_press.py index 868d8875f..657a5c4f1 100644 --- a/kvpress/presses/kvcompose_press.py +++ b/kvpress/presses/kvcompose_press.py @@ -438,17 +438,21 @@ def new_forward( return outputs try: - for layer in model.model.layers: - layer.self_attn.rotary_emb = model.model.rotary_emb + try: + for layer in model.model.layers: + layer.self_attn.rotary_emb = model.model.rotary_emb + + setattr(model, "original_forward_KVComposePress", model.model.forward) + new_forward_with_press = partial(new_forward, press=self) + model.model.forward = types.MethodType(new_forward_with_press, model) - setattr(model, "original_forward_KVComposePress", model.model.forward) - new_forward_with_press = partial(new_forward, press=self) - model.model.forward = types.MethodType(new_forward_with_press, model) + yield + finally: + model.model.forward = getattr(model, "original_forward_KVComposePress") + delattr(model, "original_forward_KVComposePress") - yield + if self.context_ids is not None: + self.prepare_important_masks() + self.compress_cache(model) finally: - model.model.forward = getattr(model, "original_forward_KVComposePress") - delattr(model, "original_forward_KVComposePress") - self.prepare_important_masks() - self.compress_cache(model) self._reset_state() diff --git a/tests/presses/test_kvcompose_fixes.py b/tests/presses/test_kvcompose_fixes.py index eac58ef56..cb196ce01 100644 --- a/tests/presses/test_kvcompose_fixes.py +++ b/tests/presses/test_kvcompose_fixes.py @@ -9,6 +9,18 @@ from tests.fixtures import unit_test_model # noqa: F401 +@pytest.mark.parametrize("structured", [True, False]) +def test_kvcompose_without_forward_pass_is_a_no_op(unit_test_model, structured): # noqa: F811 + with KVComposePress(structured=structured, compression_ratio=0.5)(unit_test_model): + pass + + +def test_kvcompose_does_not_mask_errors_raised_in_the_context(unit_test_model): # noqa: F811 + with pytest.raises(RuntimeError, match="boom"): + with KVComposePress(compression_ratio=0.5)(unit_test_model): + raise RuntimeError("boom") + + @pytest.mark.parametrize("structured", [True, False]) def test_kvcompose_add_v_norm(unit_test_model, structured): # noqa: F811 context_len = 64 From 15087f29f5798add3ee39b98e4e15e04c0badee0 Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 02:18:24 +0200 Subject: [PATCH 15/57] =?UTF-8?q?Reject=20batch=20sizes=20above=201=20in?= =?UTF-8?q?=20KVComposePress=20=F0=9F=A4=96=F0=9F=A4=96=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Structured compression rebuilt the cache from the first batch element only, silently dropping the others, and unstructured compression raised only after both prefill passes. Raise a ValueError when the context is registered instead. test_kvcompose_structured_batch_size_2_does_not_collapse_on_short_contexts relied on batch size 2 being accepted and is removed. Signed-off-by: Maximilian Jeblick --- kvpress/presses/kvcompose_press.py | 2 ++ tests/presses/test_budget_floor_sweep.py | 10 ---------- tests/presses/test_kvcompose_fixes.py | 8 ++++++++ 3 files changed, 10 insertions(+), 10 deletions(-) diff --git a/kvpress/presses/kvcompose_press.py b/kvpress/presses/kvcompose_press.py index 657a5c4f1..ddc84dfe2 100644 --- a/kvpress/presses/kvcompose_press.py +++ b/kvpress/presses/kvcompose_press.py @@ -169,6 +169,8 @@ def _register_model(self, model: PreTrainedModel): self.device = next(model.parameters()).device def register_context_ids(self, context_ids: torch.Tensor): + if context_ids.shape[0] != 1: + raise ValueError(f"KVComposePress supports only batch size 1, got {context_ids.shape[0]}") self.context_ids = context_ids self.context_len = self.context_ids.shape[-1] self.prompt_ids: list[torch.Tensor] = [] diff --git a/tests/presses/test_budget_floor_sweep.py b/tests/presses/test_budget_floor_sweep.py index 429caaa0c..f9375722f 100644 --- a/tests/presses/test_budget_floor_sweep.py +++ b/tests/presses/test_budget_floor_sweep.py @@ -92,16 +92,6 @@ def test_kvcompose_keep_token_lower_bound_survives_on_short_contexts(unit_test_m assert layer.keys.shape[2] >= press.keep_token_lower_bound -def test_kvcompose_structured_batch_size_2_does_not_collapse_on_short_contexts(unit_test_model): # noqa: F811 - """The same zero budget also collapses a batch of sequences to an empty cache.""" - press = KVComposePress(compression_ratio=0.95) - cache = compress_short_context(press, unit_test_model, batch_size=2) - - assert cache.get_seq_length() >= 1 - for layer in cache.layers: - assert layer.keys.shape[2] >= 1 - - def test_kvcompose_unstructured_keeps_at_least_one_token_on_short_contexts(unit_test_model): # noqa: F811 """Unstructured compression flags evicted tokens via ``masked_key_indices``; with a zero budget every token of every head is flagged for eviction. Per-head zero counts remain diff --git a/tests/presses/test_kvcompose_fixes.py b/tests/presses/test_kvcompose_fixes.py index cb196ce01..885eaa4db 100644 --- a/tests/presses/test_kvcompose_fixes.py +++ b/tests/presses/test_kvcompose_fixes.py @@ -15,6 +15,14 @@ def test_kvcompose_without_forward_pass_is_a_no_op(unit_test_model, structured): pass +@pytest.mark.parametrize("structured", [True, False]) +def test_kvcompose_rejects_batch_size_above_1(unit_test_model, structured): # noqa: F811 + with pytest.raises(ValueError, match="batch size 1"): + with KVComposePress(structured=structured, compression_ratio=0.5)(unit_test_model): + input_ids = torch.randint(0, 1024, (2, 16), device=unit_test_model.device) + unit_test_model(input_ids, past_key_values=DynamicCache()) + + def test_kvcompose_does_not_mask_errors_raised_in_the_context(unit_test_model): # noqa: F811 with pytest.raises(RuntimeError, match="boom"): with KVComposePress(compression_ratio=0.5)(unit_test_model): From 26807626ed97d9cd0088ab01e68e597c6ab66d84 Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 02:19:25 +0200 Subject: [PATCH 16/57] =?UTF-8?q?Reject=20eager=20attention=20in=20unstruc?= =?UTF-8?q?tured=20KVComposePress=20=F0=9F=A4=96=F0=9F=A4=96=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Unstructured compression masks the evicted keys through kvpress' attention patch, which does not wrap eager attention, so with an eager model the masks were silently ignored. Raise a ValueError when entering the context manager. Signed-off-by: Maximilian Jeblick --- kvpress/presses/kvcompose_press.py | 5 +++++ tests/presses/test_kvcompose_fixes.py | 18 +++++++++++++++++- 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/kvpress/presses/kvcompose_press.py b/kvpress/presses/kvcompose_press.py index ddc84dfe2..e6d203237 100644 --- a/kvpress/presses/kvcompose_press.py +++ b/kvpress/presses/kvcompose_press.py @@ -383,6 +383,11 @@ def __call__(self, model: PreTrainedModel) -> Generator: logger.warning("KVComposePress temporarily creates a KV cache of ~2x the context length during prefill; ") if not isinstance(model, (LlamaForCausalLM, Qwen2ForCausalLM, Qwen3ForCausalLM)): logger.warning(f"Model {type(model)} not tested") + if not self.structured and model.config._attn_implementation == "eager": + raise ValueError( + "Unstructured KVComposePress does not support eager attention: the evicted keys are masked by " + "kvpress' attention patch, which does not apply to eager attention" + ) self._register_model(model) self._reset_state() diff --git a/tests/presses/test_kvcompose_fixes.py b/tests/presses/test_kvcompose_fixes.py index 885eaa4db..7ba4ce600 100644 --- a/tests/presses/test_kvcompose_fixes.py +++ b/tests/presses/test_kvcompose_fixes.py @@ -6,7 +6,7 @@ from transformers import DynamicCache from kvpress import KVComposePress -from tests.fixtures import unit_test_model # noqa: F401 +from tests.fixtures import unit_test_model, unit_test_model_output_attention # noqa: F401 @pytest.mark.parametrize("structured", [True, False]) @@ -23,6 +23,22 @@ def test_kvcompose_rejects_batch_size_above_1(unit_test_model, structured): # n unit_test_model(input_ids, past_key_values=DynamicCache()) +def test_kvcompose_unstructured_rejects_eager_attention(unit_test_model_output_attention): # noqa: F811 + with pytest.raises(ValueError, match="eager"): + with KVComposePress(structured=False, compression_ratio=0.5)(unit_test_model_output_attention): + pytest.fail("the press should raise before running the context") + + +def test_kvcompose_structured_supports_eager_attention(unit_test_model_output_attention): # noqa: F811 + model = unit_test_model_output_attention + cache = DynamicCache() + with KVComposePress(compression_ratio=0.5)(model): + model(torch.randint(0, 1024, (1, 32), device=model.device), past_key_values=cache) + + assert 0 < cache.get_seq_length() < 32 + assert model.config._attn_implementation == "eager" + + def test_kvcompose_does_not_mask_errors_raised_in_the_context(unit_test_model): # noqa: F811 with pytest.raises(RuntimeError, match="boom"): with KVComposePress(compression_ratio=0.5)(unit_test_model): From 36385e21c56ac0454708d8b83707064946f041e3 Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 02:20:14 +0200 Subject: [PATCH 17/57] =?UTF-8?q?Raise=20a=20clear=20error=20for=20head-wi?= =?UTF-8?q?se=20compression=20under=20eager=20attention=20=F0=9F=A4=96?= =?UTF-8?q?=F0=9F=A4=96=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit In transformers 5, eager attention is not part of ALL_ATTENTION_FUNCTIONS, so it is never patched and masked_key_indices are silently ignored. AdaKVPress, CriticalAdaKVPress, LUKVPress and DMSPress now share check_masked_key_indices_support, which raises a ValueError for eager. BasePress also initializes masked_key_indices on hooked layers (without overwriting existing masks), so DMSPress no longer fails with an AttributeError. Signed-off-by: Maximilian Jeblick --- kvpress/attention_patch.py | 12 +++++ kvpress/presses/adakv_press.py | 3 +- kvpress/presses/base_press.py | 3 ++ kvpress/presses/criticalkv_press.py | 3 +- kvpress/presses/dms_press.py | 2 + kvpress/presses/lukv_press.py | 3 +- tests/presses/test_headwise_masking.py | 62 ++++++++++++++++++++++++++ 7 files changed, 85 insertions(+), 3 deletions(-) create mode 100644 tests/presses/test_headwise_masking.py diff --git a/kvpress/attention_patch.py b/kvpress/attention_patch.py index 270ad9e6f..90e472feb 100644 --- a/kvpress/attention_patch.py +++ b/kvpress/attention_patch.py @@ -48,6 +48,18 @@ def search_hyperplane(X, max_iter: int = 1000): raise ValueError("Could not find fake keys such that for every query q, exp() = 0") +def check_masked_key_indices_support(module): + """ + Raise a ValueError if the attention implementation of the module ignores module.masked_key_indices. + Eager attention is not registered in ALL_ATTENTION_FUNCTIONS, so it is not patched by patch_attention_functions. + """ + if module.config._attn_implementation == "eager": + raise ValueError( + "Head-wise compression is not supported with attn_implementation='eager' as masked keys would be ignored. " + "Please use another attention implementation such as 'sdpa' or 'flash_attention_2'." + ) + + def attention_patch(func): """ Decorator to update the keys before the attention computation at the indices provided in module.masked_key_indices diff --git a/kvpress/presses/adakv_press.py b/kvpress/presses/adakv_press.py index 4dfcc86a2..bfe92f65d 100644 --- a/kvpress/presses/adakv_press.py +++ b/kvpress/presses/adakv_press.py @@ -6,6 +6,7 @@ import torch +from kvpress.attention_patch import check_masked_key_indices_support from kvpress.presses.base_press import BasePress from kvpress.presses.scorer_press import ScorerPress from kvpress.utils import compute_n_kept @@ -55,7 +56,7 @@ def compress(self, module, hidden_states, keys, values, attentions, kwargs): if self.compression_ratio == 0: return keys, values - assert module.config._attn_implementation != "eager", "eager mode not supported" + check_masked_key_indices_support(module) # Compute scores scores = self.press.score(module, hidden_states, keys, values, attentions, kwargs) diff --git a/kvpress/presses/base_press.py b/kvpress/presses/base_press.py index b81c2adef..aa7d6f8df 100644 --- a/kvpress/presses/base_press.py +++ b/kvpress/presses/base_press.py @@ -190,6 +190,9 @@ def __call__(self, model: PreTrainedModel) -> Generator: # Skip layers with sliding window attention, only for Gemma3 continue layer.self_attn.rotary_emb = language_model.rotary_emb + # Masks set by head-wise presses during prefilling must survive later press contexts (e.g. decoding) + if not hasattr(layer.self_attn, "masked_key_indices"): + layer.self_attn.masked_key_indices = None hooks.append(layer.self_attn.register_forward_hook(self.forward_hook, with_kwargs=True)) yield finally: diff --git a/kvpress/presses/criticalkv_press.py b/kvpress/presses/criticalkv_press.py index 2db080439..3b526fbb8 100644 --- a/kvpress/presses/criticalkv_press.py +++ b/kvpress/presses/criticalkv_press.py @@ -7,6 +7,7 @@ import torch from transformers.models.llama.modeling_llama import repeat_kv +from kvpress.attention_patch import check_masked_key_indices_support from kvpress.presses.base_press import BasePress from kvpress.presses.expected_attention_press import ExpectedAttentionPress from kvpress.presses.scorer_press import ScorerPress @@ -144,7 +145,7 @@ def compress(self, module, hidden_states, keys, values, attentions, kwargs): if self.compression_ratio == 0: return keys, values - assert module.config._attn_implementation != "eager", "eager mode not supported" + check_masked_key_indices_support(module) # Compute scores scores = self.press.score(module, hidden_states, keys, values, attentions, kwargs) diff --git a/kvpress/presses/dms_press.py b/kvpress/presses/dms_press.py index 29634a30e..2a2e1c636 100644 --- a/kvpress/presses/dms_press.py +++ b/kvpress/presses/dms_press.py @@ -7,6 +7,7 @@ import torch import torch.nn as nn +from kvpress.attention_patch import check_masked_key_indices_support from kvpress.presses.base_press import BasePress, is_prefilling from kvpress.presses.scorer_press import ScorerPress from kvpress.utils import extract_keys_and_values @@ -67,6 +68,7 @@ def compression_ratio(self, value): raise AttributeError(f"compression ratio cannot be set for {type(self).__name__}") def forward_hook(self, module: nn.Module, input: list[torch.Tensor], kwargs: dict, output: list): + check_masked_key_indices_support(module) hidden_states = kwargs["hidden_states"] cache = kwargs["past_key_values"] q_len = hidden_states.shape[1] diff --git a/kvpress/presses/lukv_press.py b/kvpress/presses/lukv_press.py index 5dec73aef..0edf28781 100644 --- a/kvpress/presses/lukv_press.py +++ b/kvpress/presses/lukv_press.py @@ -12,6 +12,7 @@ from torch import nn from transformers import PreTrainedModel +from kvpress.attention_patch import check_masked_key_indices_support from kvpress.presses.base_press import BasePress from kvpress.presses.expected_attention_press import ExpectedAttentionPress from kvpress.presses.scorer_press import ScorerPress @@ -132,7 +133,7 @@ def compress( return keys, values if self._budget_curves is None: raise ValueError("LU-KV budget curves are not loaded. Use LUKVPress as a model context manager first.") - assert module.config._attn_implementation != "eager", "eager mode not supported" + check_masked_key_indices_support(module) bsz, num_key_value_heads, seq_len, _ = keys.shape scores = self.press.score(module, hidden_states, keys, values, attentions, kwargs) diff --git a/tests/presses/test_headwise_masking.py b/tests/presses/test_headwise_masking.py new file mode 100644 index 000000000..c3a1e87e1 --- /dev/null +++ b/tests/presses/test_headwise_masking.py @@ -0,0 +1,62 @@ +# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch +from transformers import DynamicCache, LlamaConfig, LlamaForCausalLM + +from kvpress import AdaKVPress, CriticalAdaKVPress, DMSPress, KnormPress, RandomPress +from tests import default_presses +from tests.fixtures import unit_test_model_output_attention # noqa: F401 + + +def make_tiny_llama() -> LlamaForCausalLM: + config = LlamaConfig( + hidden_size=64, + intermediate_size=128, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + vocab_size=128, + ) + return LlamaForCausalLM(config).eval() + + +@pytest.mark.parametrize( + "make_press", + [ + lambda: AdaKVPress(KnormPress(compression_ratio=0.5)), + lambda: CriticalAdaKVPress(KnormPress(compression_ratio=0.5)), + lambda: default_presses.TestLUKVPress(KnormPress(), compression_ratio=0.5), + lambda: DMSPress(press=RandomPress(), threshold=0.5, sliding_window_size=0), + ], + ids=["adakv", "critical_adakv", "lukv", "dms"], +) +def test_head_wise_presses_reject_eager_attention(unit_test_model_output_attention, make_press): # noqa: F811 + model = unit_test_model_output_attention + input_ids = torch.randint(0, 1024, (1, 64), device=model.device) + with pytest.raises(ValueError, match="eager"): + with make_press()(model): + model(input_ids, past_key_values=DynamicCache()) + + +def test_press_initializes_masked_key_indices(): + model = make_tiny_llama() + assert not hasattr(model.model.layers[0].self_attn, "masked_key_indices") + with KnormPress(compression_ratio=0.5)(model): + assert all(layer.self_attn.masked_key_indices is None for layer in model.model.layers) + + +@torch.no_grad() +def test_press_keeps_existing_masked_key_indices(): + model = make_tiny_llama() + with AdaKVPress(KnormPress(compression_ratio=0.5))(model): + model(torch.randint(0, 128, (1, 32)), past_key_values=DynamicCache()) + masked_key_indices = [layer.self_attn.masked_key_indices for layer in model.model.layers] + assert all(indices is not None for indices in masked_key_indices) + + with KnormPress()(model): + assert all( + layer.self_attn.masked_key_indices is indices + for layer, indices in zip(model.model.layers, masked_key_indices) + ) From 3da14612fcd619f75304d5676fd10b7efec9fe80 Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 02:24:01 +0200 Subject: [PATCH 18/57] =?UTF-8?q?Compress=20the=20cache=20the=20model=20ac?= =?UTF-8?q?tually=20used=20in=20KVzipPress=20=F0=9F=A4=96=F0=9F=A4=96?= =?UTF-8?q?=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit KVzipPress stored kwargs["past_key_values"], which is None when the model is called without a cache (e.g. model(inputs) as in the README). The replay then scored a throwaway cache while the masks applied to the cache returned by the model, and KVgradPress crashed. The cache is now taken from the forward output, input_ids may be passed positionally, and a clear ValueError is raised without input_ids or without a cache. Signed-off-by: Maximilian Jeblick --- kvpress/presses/kvzip_press.py | 13 +++++-- tests/presses/test_kvzip_fixes.py | 59 +++++++++++++++++++++++++++++++ 2 files changed, 69 insertions(+), 3 deletions(-) create mode 100644 tests/presses/test_kvzip_fixes.py diff --git a/kvpress/presses/kvzip_press.py b/kvpress/presses/kvzip_press.py index 3e6bd4d39..600e79c81 100644 --- a/kvpress/presses/kvzip_press.py +++ b/kvpress/presses/kvzip_press.py @@ -123,9 +123,16 @@ def __call__(self, model: PreTrainedModel) -> Generator: original_forward = model.model.forward def wrapped_forward(model_self, *args, **kwargs): - self._context_ids = kwargs["input_ids"] - self._cache = kwargs["past_key_values"] - return original_forward(*args, **kwargs) + input_ids = kwargs["input_ids"] if "input_ids" in kwargs else (args[0] if args else None) + if input_ids is None: + raise ValueError(f"{type(self).__name__} replays the context, call the model with input_ids") + outputs = original_forward(*args, **kwargs) + # The model creates the cache if none is passed + if outputs.past_key_values is None: + raise ValueError(f"{type(self).__name__} compresses the KV cache, call the model with use_cache=True") + self._context_ids = input_ids + self._cache = outputs.past_key_values + return outputs model.model.forward = MethodType(wrapped_forward, model.model) diff --git a/tests/presses/test_kvzip_fixes.py b/tests/presses/test_kvzip_fixes.py new file mode 100644 index 000000000..8da28f08f --- /dev/null +++ b/tests/presses/test_kvzip_fixes.py @@ -0,0 +1,59 @@ +# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch +from transformers import DynamicCache + +from kvpress import KVgradPress, KVzipPress +from tests.fixtures import unit_test_model # noqa: F401 + +PRESSES = [(KVzipPress, {"chunk_size": 32}), (KVgradPress, {"chunk_size": 32})] + + +def masked_keys(model, press, run): + """Run `run(model)` within `press` and return the masked key indices of every layer.""" + with press(model): + run(model) + return [tuple(indices.clone() for indices in layer.self_attn.masked_key_indices) for layer in model.model.layers] + + +def assert_same_masks(masks, reference_masks): + for layer_masks, reference_layer_masks in zip(masks, reference_masks): + for indices, reference_indices in zip(layer_masks, reference_layer_masks): + assert torch.equal(indices, reference_indices) + + +@pytest.mark.parametrize("press_cls, kwargs", PRESSES) +def test_kvzip_compresses_the_cache_created_by_the_model(unit_test_model, press_cls, kwargs): # noqa: F811 + input_ids = torch.randint(0, 1024, (1, 64), device=unit_test_model.device) + reference = masked_keys( + unit_test_model, + press_cls(compression_ratio=0.5, **kwargs), + lambda m: m(input_ids, past_key_values=DynamicCache()), + ) + + assert_same_masks( + masked_keys(unit_test_model, press_cls(compression_ratio=0.5, **kwargs), lambda m: m(input_ids)), reference + ) + assert_same_masks( + masked_keys( + unit_test_model, + press_cls(compression_ratio=0.5, **kwargs), + lambda m: m.model(input_ids, past_key_values=DynamicCache()), + ), + reference, + ) + + +def test_kvzip_requires_input_ids(unit_test_model): # noqa: F811 + inputs_embeds = unit_test_model.model.embed_tokens(torch.randint(0, 1024, (1, 16), device=unit_test_model.device)) + with pytest.raises(ValueError, match="input_ids"): + with KVzipPress(compression_ratio=0.5)(unit_test_model): + unit_test_model(inputs_embeds=inputs_embeds, past_key_values=DynamicCache()) + + +def test_kvzip_requires_a_cache(unit_test_model): # noqa: F811 + with pytest.raises(ValueError, match="use_cache=True"): + with KVzipPress(compression_ratio=0.5)(unit_test_model): + unit_test_model(torch.randint(0, 1024, (1, 16), device=unit_test_model.device), use_cache=False) From e3aa8aa98fe8a3ba6e65fd1f4309e40b3e896a88 Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 02:24:33 +0200 Subject: [PATCH 19/57] =?UTF-8?q?Scale=20attention=20logits=20by=20module.?= =?UTF-8?q?scaling=20=F0=9F=A4=96=F0=9F=A4=96=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit SnapKVPress, ExpectedAttentionPress and CAMPress divided by sqrt(head_dim), which differs from the model attention when query_pre_attn_scalar != head_dim (e.g. Gemma3-27B). Signed-off-by: Maximilian Jeblick --- kvpress/presses/cam_press.py | 5 +- kvpress/presses/expected_attention_press.py | 11 ++- kvpress/presses/snapkv_press.py | 4 +- tests/presses/test_attention_scaling.py | 85 +++++++++++++++++++++ 4 files changed, 93 insertions(+), 12 deletions(-) create mode 100644 tests/presses/test_attention_scaling.py diff --git a/kvpress/presses/cam_press.py b/kvpress/presses/cam_press.py index 60e977001..062ae3cbc 100644 --- a/kvpress/presses/cam_press.py +++ b/kvpress/presses/cam_press.py @@ -4,7 +4,6 @@ from __future__ import annotations import logging -import math from dataclasses import dataclass import torch @@ -321,7 +320,7 @@ def _compute_current_token_attention( kwargs: dict, ) -> torch.Tensor: """Compute softmax attention from the last query token to all cached keys.""" - _, num_key_value_heads, cache_len, head_dim = keys.shape + num_key_value_heads = keys.shape[1] num_query_heads = module.config.num_attention_heads num_key_value_groups = num_query_heads // num_key_value_heads @@ -331,7 +330,7 @@ def _compute_current_token_attention( query_states = query_states[:, :, -1:, :] keys_repeated = repeat_kv(keys, num_key_value_groups) - scores = torch.matmul(query_states, keys_repeated.transpose(-2, -1)) / math.sqrt(head_dim) + scores = torch.matmul(query_states, keys_repeated.transpose(-2, -1)) * module.scaling return torch.nn.functional.softmax(scores, dim=-1, dtype=torch.float32).to(query_states.dtype) @staticmethod diff --git a/kvpress/presses/expected_attention_press.py b/kvpress/presses/expected_attention_press.py index 88ffe021a..6f73e0998 100644 --- a/kvpress/presses/expected_attention_press.py +++ b/kvpress/presses/expected_attention_press.py @@ -2,7 +2,6 @@ # SPDX-License-Identifier: Apache-2.0 -import math from dataclasses import dataclass import torch @@ -26,8 +25,8 @@ class ExpectedAttentionPress(ScorerPress): 1. Compute the mean and covariance matrix of the queries before RoPE. 2. Compute the RoPE rotation matrix R on next n_future_positions and average it 3. Apply R to the mean and covariance matrice of the queries. - 4. As attention A = exp(Q @ K / sqrt(d)), we compute the expected attention - E(A) = exp(K @ mean.T / sqrt(d) + 1/2 K @ cov @ K.T / d) + 4. As attention A = exp(Q @ K * s) with s = module.scaling (usually 1 / sqrt(d)), we compute the expected + attention E(A) = exp(K @ mean.T * s + 1/2 K @ cov @ K.T * s^2) 5. Rescale the scores using (scores + epsilon) * ||V||_2 Parameters @@ -142,13 +141,13 @@ def score( mean_query, cov_query = self.get_query_statistics(module, hidden_states) # Compute scores - bsz, num_key_value_heads, q_len, d = keys.shape + bsz, num_key_value_heads, q_len, _ = keys.shape num_key_value_groups = module.config.num_attention_heads // num_key_value_heads keys = repeat_kv(keys, num_key_value_groups).transpose(2, 3) - scores = torch.matmul(mean_query.unsqueeze(2), keys).squeeze(2) / math.sqrt(d) + scores = torch.matmul(mean_query.unsqueeze(2), keys).squeeze(2) * module.scaling if self.use_covariance: - scores += torch.einsum("bhin, bhij, bhjn->bhn", keys, cov_query, keys) / d / 2 + scores += torch.einsum("bhin, bhij, bhjn->bhn", keys, cov_query, keys) * module.scaling**2 / 2 scores = F.softmax(scores, dim=-1) # Average scores across groups diff --git a/kvpress/presses/snapkv_press.py b/kvpress/presses/snapkv_press.py index 61cd85055..3c15bbd50 100644 --- a/kvpress/presses/snapkv_press.py +++ b/kvpress/presses/snapkv_press.py @@ -2,7 +2,6 @@ # SPDX-License-Identifier: Apache-2.0 -import math from dataclasses import dataclass import torch @@ -46,7 +45,6 @@ def compute_window_attention(module, hidden_states, keys, window_size, position_ bsz, _, k_len, _ = keys.shape num_heads = module.config.num_attention_heads - head_dim = module.head_dim num_key_value_groups = num_heads // module.config.num_key_value_heads # Get last window_size queries and apply RoPE @@ -57,7 +55,7 @@ def compute_window_attention(module, hidden_states, keys, window_size, position_ # Compute attention for first q_len - window_size tokens key_states = repeat_kv(keys, num_key_value_groups) - attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(head_dim) + attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) * module.scaling attention_mask = torch.ones_like(attn_weights) * float("-inf") attention_mask = torch.triu(attention_mask, diagonal=k_len - window_size + 1) attn_weights += attention_mask diff --git a/tests/presses/test_attention_scaling.py b/tests/presses/test_attention_scaling.py new file mode 100644 index 000000000..521d43297 --- /dev/null +++ b/tests/presses/test_attention_scaling.py @@ -0,0 +1,85 @@ +# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch +from torch.nn import functional as F +from transformers import DynamicCache, Gemma3ForCausalLM, Gemma3TextConfig, LlamaConfig, LlamaForCausalLM +from transformers.models.llama.modeling_llama import repeat_kv + +from kvpress import CAMPress, ExpectedAttentionPress, KnormPress, SnapKVPress + + +@pytest.fixture(scope="module") +def gemma3_attention_inputs(): + """Inputs of an eager Gemma3 attention layer whose scaling differs from 1 / sqrt(head_dim).""" + config = Gemma3TextConfig( + hidden_size=64, + intermediate_size=128, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=16, + query_pre_attn_scalar=64, + vocab_size=128, + layer_types=["full_attention"], + attn_implementation="eager", + ) + torch.manual_seed(0) + model = Gemma3ForCausalLM(config).eval() + module = model.model.layers[0].self_attn + assert module.scaling != module.head_dim**-0.5 + + captured = {} + + def hook(module, args, kwargs, output): + captured.update(kwargs=kwargs, attentions=output[1]) + + handle = module.register_forward_hook(hook, with_kwargs=True) + with torch.no_grad(): + cache = DynamicCache() + model(torch.randint(0, 128, (1, 24)), past_key_values=cache) + handle.remove() + module.rotary_emb = model.model.rotary_emb + return module, captured["kwargs"], cache.layers[0].keys, cache.layers[0].values, captured["attentions"] + + +def test_snapkv_window_attention_matches_eager_attention(gemma3_attention_inputs): + module, kwargs, keys, _, attentions = gemma3_attention_inputs + window_size = 4 + attn_weights = SnapKVPress.compute_window_attention( + module, kwargs["hidden_states"], keys, window_size, kwargs["position_embeddings"] + ) + torch.testing.assert_close(attn_weights, attentions[..., -window_size:, :-window_size]) + + +def test_cam_current_token_attention_matches_eager_attention(gemma3_attention_inputs): + module, kwargs, keys, _, attentions = gemma3_attention_inputs + attn_weights = CAMPress(base_press=KnormPress())._compute_current_token_attention( + module, kwargs["hidden_states"], keys, kwargs + ) + torch.testing.assert_close(attn_weights, attentions[..., -1:, :]) + + +def test_expected_attention_uses_module_scaling(): + config = LlamaConfig( + hidden_size=64, intermediate_size=128, num_hidden_layers=1, num_attention_heads=4, num_key_value_heads=2 + ) + model = LlamaForCausalLM(config).eval() + module = model.model.layers[0].self_attn + module.rotary_emb = model.model.rotary_emb + module.scaling = 0.1 + hidden_states = torch.randn(1, 24, config.hidden_size) + keys, values = torch.randn(2, 1, config.num_key_value_heads, 24, module.head_dim).unbind(0) + + press = ExpectedAttentionPress(n_sink=0, use_vnorm=False) + kwargs = {"cache_position": torch.arange(24)} + with torch.no_grad(): + scores = press.score(module, hidden_states, keys, values, None, kwargs) + mean_query, cov_query = press.get_query_statistics(module, hidden_states) + num_key_value_groups = module.config.num_attention_heads // keys.shape[1] + repeated_keys = repeat_kv(keys, num_key_value_groups) + logits = torch.einsum("bhd,bhnd->bhn", mean_query, repeated_keys) * module.scaling + logits += torch.einsum("bhnd,bhde,bhne->bhn", repeated_keys, cov_query, repeated_keys) * module.scaling**2 / 2 + expected = F.softmax(logits, dim=-1).view(1, keys.shape[1], num_key_value_groups, -1).mean(dim=2) + torch.testing.assert_close(scores, expected) From 7930d3ac957c3ea861cd0e8b05d4f0b519493157 Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 02:26:36 +0200 Subject: [PATCH 20/57] =?UTF-8?q?Replay=20the=20first=20prefill=20in=20KVz?= =?UTF-8?q?ipPress,=20not=20the=20last=20forward=20pass=20=F0=9F=A4=96?= =?UTF-8?q?=F0=9F=A4=96=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The patched forward stored the input_ids and cache of every call, so after model.generate the replay reconstructed the last decoded token only and the cache was truncated to 1 token (KVgradPress scored that token only). Only the first prefill (empty cache) is now captured and later calls run unchanged. Tokens decoded after the prefill are cropped before the replay, so the scores match the ones of a plain prefill. This applies to KVgradPress and RestoreKVPress too. Signed-off-by: Maximilian Jeblick --- kvpress/presses/kvzip_press.py | 6 ++++++ tests/presses/test_kvzip_fixes.py | 21 +++++++++++++++++++++ 2 files changed, 27 insertions(+) diff --git a/kvpress/presses/kvzip_press.py b/kvpress/presses/kvzip_press.py index 600e79c81..26a6c7136 100644 --- a/kvpress/presses/kvzip_press.py +++ b/kvpress/presses/kvzip_press.py @@ -123,6 +123,10 @@ def __call__(self, model: PreTrainedModel) -> Generator: original_forward = model.model.forward def wrapped_forward(model_self, *args, **kwargs): + cache = kwargs.get("past_key_values") + if self._context_ids is not None or (cache is not None and cache.get_seq_length() > 0): + # Only the first prefill is replayed, later calls (e.g. decoding steps) run unchanged + return original_forward(*args, **kwargs) input_ids = kwargs["input_ids"] if "input_ids" in kwargs else (args[0] if args else None) if input_ids is None: raise ValueError(f"{type(self).__name__} replays the context, call the model with input_ids") @@ -147,6 +151,8 @@ def wrapped_forward(model_self, *args, **kwargs): for layer in model.model.layers: layer.self_attn.rotary_emb = model.model.rotary_emb + # Tokens decoded after the prefill (e.g. by model.generate) are not part of the context + self._cache.crop(self._context_ids.shape[1]) self._perform_kvzip_compression(model, tokenizer) finally: self._reset_internal_parameters() diff --git a/tests/presses/test_kvzip_fixes.py b/tests/presses/test_kvzip_fixes.py index 8da28f08f..08f234f20 100644 --- a/tests/presses/test_kvzip_fixes.py +++ b/tests/presses/test_kvzip_fixes.py @@ -46,6 +46,27 @@ def test_kvzip_compresses_the_cache_created_by_the_model(unit_test_model, press_ ) +@pytest.mark.parametrize("press_cls, kwargs", PRESSES) +def test_kvzip_generate_replays_the_prefilled_context(unit_test_model, press_cls, kwargs): # noqa: F811 + context_len = 64 + input_ids = torch.randint(0, 1024, (1, context_len), device=unit_test_model.device) + reference = masked_keys( + unit_test_model, + press_cls(compression_ratio=0.5, **kwargs), + lambda m: m(input_ids, past_key_values=DynamicCache()), + ) + + cache = DynamicCache() + masks = masked_keys( + unit_test_model, + press_cls(compression_ratio=0.5, **kwargs), + lambda m: m.generate(input_ids, past_key_values=cache, max_new_tokens=5, do_sample=False), + ) + + assert_same_masks(masks, reference) + assert cache.get_seq_length() == context_len + + def test_kvzip_requires_input_ids(unit_test_model): # noqa: F811 inputs_embeds = unit_test_model.model.embed_tokens(torch.randint(0, 1024, (1, 16), device=unit_test_model.device)) with pytest.raises(ValueError, match="input_ids"): From 4961a49276bcba43a5d3cc2b776dc2fbbb350cc4 Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 02:30:29 +0200 Subject: [PATCH 21/57] =?UTF-8?q?Place=20ExpectedAttentionPress=20future?= =?UTF-8?q?=20queries=20after=20the=20absolute=20position=20=F0=9F=A4=96?= =?UTF-8?q?=F0=9F=A4=96=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Future positions started at hidden_states.shape[1], which is only the next position during a plain prefill. Under DecodingPress (buffered hidden states), ChunkPress or BlockPress they were placed too early. The keys are now rotated back by the difference to the next position (from position_ids, falling back to cache_position), which is equivalent for RoPE and leaves the get_query_statistics/apply_avg_rope signatures unchanged for ExpectedAttentionStatsPress. Prefill scores are unchanged. Signed-off-by: Maximilian Jeblick --- kvpress/presses/expected_attention_press.py | 31 ++++++++- .../test_expected_attention_positions.py | 67 +++++++++++++++++++ 2 files changed, 97 insertions(+), 1 deletion(-) create mode 100644 tests/presses/test_expected_attention_positions.py diff --git a/kvpress/presses/expected_attention_press.py b/kvpress/presses/expected_attention_press.py index 6f73e0998..77e8f6191 100644 --- a/kvpress/presses/expected_attention_press.py +++ b/kvpress/presses/expected_attention_press.py @@ -7,7 +7,7 @@ import torch from torch import nn from torch.nn import functional as F -from transformers.models.llama.modeling_llama import repeat_kv +from transformers.models.llama.modeling_llama import repeat_kv, rotate_half from kvpress.presses.scorer_press import ScorerPress from kvpress.utils import get_prerope_query_states @@ -122,6 +122,28 @@ def apply_avg_rope(self, module: nn.Module, mu: torch.Tensor, cov: torch.Tensor, cov = torch.matmul(R, torch.matmul(cov, R.T)) return mu, cov + @staticmethod + def get_next_position(kwargs: dict, hidden_states: torch.Tensor) -> int: + """ + Position of the token following the current forward pass. RoPE uses position_ids, which may differ from + cache_position once the cache has been compressed. Defaults to the number of hidden states. + """ + positions = kwargs.get("position_ids") + if positions is None: + positions = kwargs.get("cache_position") + if positions is None: + return hidden_states.shape[1] + return int(positions.max()) + 1 + + @staticmethod + def rotate_keys(module, keys: torch.Tensor, n_positions: int) -> torch.Tensor: + """ + Rotate post-RoPE keys by n_positions positions. + """ + inv_freq = module.rotary_emb.inv_freq + angles = torch.cat([inv_freq, inv_freq]).to(keys.device, torch.float32) * n_positions + return keys * angles.cos().to(keys.dtype) + rotate_half(keys) * angles.sin().to(keys.dtype) + def score( self, module: nn.Module, @@ -140,6 +162,13 @@ def score( # Compute query statistics mean_query, cov_query = self.get_query_statistics(module, hidden_states) + # The statistics place the future queries right after the hidden_states.shape[1] first positions. If the + # next position differs (e.g. for chunks or buffered hidden states), rotating the keys back by the difference + # is equivalent to moving the future queries after the next position. + shift = self.get_next_position(kwargs, hidden_states) - hidden_states.shape[1] + if shift != 0: + keys = self.rotate_keys(module, keys, -shift) + # Compute scores bsz, num_key_value_heads, q_len, _ = keys.shape num_key_value_groups = module.config.num_attention_heads // num_key_value_heads diff --git a/tests/presses/test_expected_attention_positions.py b/tests/presses/test_expected_attention_positions.py new file mode 100644 index 000000000..d2c2655e2 --- /dev/null +++ b/tests/presses/test_expected_attention_positions.py @@ -0,0 +1,67 @@ +# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import dataclasses + +import pytest +import torch +from transformers import LlamaConfig, LlamaForCausalLM + +from kvpress import ExpectedAttentionPress + + +@pytest.fixture(scope="module") +def attention_module(): + config = LlamaConfig( + hidden_size=64, + intermediate_size=128, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + vocab_size=128, + ) + torch.manual_seed(0) + model = LlamaForCausalLM(config).eval() + module = model.model.layers[0].self_attn + module.rotary_emb = model.model.rotary_emb + return module + + +def make_inputs(module, n_hidden_states, n_keys): + torch.manual_seed(1) + hidden_states = torch.randn(1, n_hidden_states, module.config.hidden_size) + keys, values = torch.randn(2, 1, module.config.num_key_value_heads, n_keys, module.head_dim).unbind(0) + return hidden_states, keys, values + + +def reference_scores(press, module, hidden_states, keys, values, next_position): + """Scores computed with the future queries placed after next_position by construction.""" + reference = dataclasses.replace(press) + apply_avg_rope = reference.apply_avg_rope + reference.apply_avg_rope = lambda module, mu, cov, q_len: apply_avg_rope(module, mu, cov, next_position) + kwargs = {"position_ids": torch.arange(hidden_states.shape[1]).unsqueeze(0)} + return reference.score(module, hidden_states, keys, values, None, kwargs) + + +@torch.no_grad() +def test_prefill_scores_are_unchanged(attention_module): + press = ExpectedAttentionPress(compression_ratio=0.5) + hidden_states, keys, values = make_inputs(attention_module, 32, 32) + kwargs = {"position_ids": torch.arange(32).unsqueeze(0), "cache_position": torch.arange(32)} + + scores = press.score(attention_module, hidden_states, keys, values, None, kwargs) + assert torch.equal(scores, reference_scores(press, attention_module, hidden_states, keys, values, 32)) + + +@pytest.mark.parametrize("use_covariance", [True, False]) +@torch.no_grad() +def test_future_positions_follow_the_absolute_position(attention_module, use_covariance): + press = ExpectedAttentionPress(compression_ratio=0.5, use_covariance=use_covariance) + hidden_states, keys, values = make_inputs(attention_module, 16, 140) + # e.g. 16 buffered hidden states during decoding, with a compressed cache shorter than the positions + kwargs = {"position_ids": torch.arange(184, 200).unsqueeze(0), "cache_position": torch.arange(124, 140)} + + scores = press.score(attention_module, hidden_states, keys, values, None, kwargs) + expected = reference_scores(press, attention_module, hidden_states, keys, values, 200) + torch.testing.assert_close(scores, expected) + assert not torch.allclose(scores, reference_scores(press, attention_module, hidden_states, keys, values, 16)) From 39964b7e2abb9b715b6add9abf9568be653b089f Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 02:30:33 +0200 Subject: [PATCH 22/57] =?UTF-8?q?Reject=20QuantizedCache=20in=20KVzipPress?= =?UTF-8?q?=20and=20KVComposePress=20=F0=9F=A4=96=F0=9F=A4=96=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The replayed chunks land in the residual buffer of the quantized layers, so KVzip scored the wrong keys (up to 0.93 off with an identity quantizer), and KVComposePress cannot crop or rebuild quantized layers. Raise a ValueError on the prefill instead, for KVzipPress and its subclasses (KVgradPress, RestoreKVPress) and for KVComposePress, and drop the now unreachable QuantizedCache branch of KVzipPress's hook. Signed-off-by: Maximilian Jeblick --- kvpress/presses/kvcompose_press.py | 4 +++- kvpress/presses/kvzip_press.py | 14 ++++---------- tests/presses/test_kvzip_fixes.py | 31 ++++++++++++++++++++++++++++-- 3 files changed, 36 insertions(+), 13 deletions(-) diff --git a/kvpress/presses/kvcompose_press.py b/kvpress/presses/kvcompose_press.py index e6d203237..23ad7665d 100644 --- a/kvpress/presses/kvcompose_press.py +++ b/kvpress/presses/kvcompose_press.py @@ -13,7 +13,7 @@ import numpy as np import torch from torch import nn -from transformers.cache_utils import DynamicCache +from transformers.cache_utils import DynamicCache, QuantizedCache from transformers.modeling_utils import PreTrainedModel from transformers.models.llama import LlamaForCausalLM from transformers.models.qwen2.modeling_qwen2 import Qwen2ForCausalLM @@ -408,6 +408,8 @@ def new_forward( *args, **kwargs, ) + if isinstance(past_key_values, QuantizedCache): + raise ValueError("KVComposePress does not support QuantizedCache") press.register_context_ids(input_ids) # Eager attention bypasses kvpress' attention patch, which discards the masked keys on prefill diff --git a/kvpress/presses/kvzip_press.py b/kvpress/presses/kvzip_press.py index 26a6c7136..122d8c896 100644 --- a/kvpress/presses/kvzip_press.py +++ b/kvpress/presses/kvzip_press.py @@ -127,6 +127,8 @@ def wrapped_forward(model_self, *args, **kwargs): if self._context_ids is not None or (cache is not None and cache.get_seq_length() > 0): # Only the first prefill is replayed, later calls (e.g. decoding steps) run unchanged return original_forward(*args, **kwargs) + if isinstance(cache, QuantizedCache): + raise ValueError(f"{type(self).__name__} does not support QuantizedCache") input_ids = kwargs["input_ids"] if "input_ids" in kwargs else (args[0] if args else None) if input_ids is None: raise ValueError(f"{type(self).__name__} replays the context, call the model with input_ids") @@ -174,16 +176,8 @@ def forward_hook(self, module: nn.Module, input: list[torch.Tensor], kwargs: dic # retaining only the originally prefilled KV pairs. keys, values = self.score_kvzip(module, hidden_states, keys, values, output[1], kwargs) - if isinstance(cache, QuantizedCache): - # Update cache with compressed keys and values - cache_layer._quantized_keys = cache_layer._quantize(keys, axis=cache_layer.axis_key) - cache_layer._quantized_values = cache_layer._quantize(values, axis=cache_layer.axis_value) - cache_layer.keys = torch.zeros(0, dtype=keys.dtype, device=keys.device) # type: ignore[index] - cache_layer.values = torch.zeros(0, dtype=keys.dtype, device=keys.device) # type: ignore[index] - cache_layer.cumulative_length = keys.shape[2] - else: - cache_layer.keys = keys - cache_layer.values = values + cache_layer.keys = keys + cache_layer.values = values return output diff --git a/tests/presses/test_kvzip_fixes.py b/tests/presses/test_kvzip_fixes.py index 08f234f20..8b46a1569 100644 --- a/tests/presses/test_kvzip_fixes.py +++ b/tests/presses/test_kvzip_fixes.py @@ -3,14 +3,30 @@ import pytest import torch -from transformers import DynamicCache +from transformers import Cache, DynamicCache, QuantizedCache +from transformers.cache_utils import QuantizedLayer -from kvpress import KVgradPress, KVzipPress +from kvpress import KVComposePress, KVgradPress, KVzipPress from tests.fixtures import unit_test_model # noqa: F401 PRESSES = [(KVzipPress, {"chunk_size": 32}), (KVgradPress, {"chunk_size": 32})] +class IdentityQuantizedLayer(QuantizedLayer): + def _quantize(self, tensor, axis): + return tensor + + def _dequantize(self, q_tensor): + return q_tensor + + +class IdentityQuantizedCache(QuantizedCache): + """QuantizedCache without quantization backend.""" + + def __init__(self, num_layers: int): + Cache.__init__(self, layers=[IdentityQuantizedLayer() for _ in range(num_layers)]) + + def masked_keys(model, press, run): """Run `run(model)` within `press` and return the masked key indices of every layer.""" with press(model): @@ -67,6 +83,17 @@ def test_kvzip_generate_replays_the_prefilled_context(unit_test_model, press_cls assert cache.get_seq_length() == context_len +@pytest.mark.parametrize( + "press", + [KVzipPress(compression_ratio=0.5), KVgradPress(compression_ratio=0.5), KVComposePress(compression_ratio=0.5)], +) +def test_multipass_presses_reject_quantized_cache(unit_test_model, press): # noqa: F811 + cache = IdentityQuantizedCache(unit_test_model.config.num_hidden_layers) + with pytest.raises(ValueError, match="QuantizedCache"): + with press(unit_test_model): + unit_test_model(torch.randint(0, 1024, (1, 16), device=unit_test_model.device), past_key_values=cache) + + def test_kvzip_requires_input_ids(unit_test_model): # noqa: F811 inputs_embeds = unit_test_model.model.embed_tokens(torch.randint(0, 1024, (1, 16), device=unit_test_model.device)) with pytest.raises(ValueError, match="input_ids"): From 15a267db1c4c367b8626e6ee5cfdd965acdfc50e Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 02:31:54 +0200 Subject: [PATCH 23/57] =?UTF-8?q?Mask=20padded=20tokens=20in=20NonCausalAt?= =?UTF-8?q?tnPress=20chunked=20attention=20=F0=9F=A4=96=F0=9F=A4=96?= =?UTF-8?q?=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Padded keys were masked with -1e-9 instead of a large negative value, and the zero-padded queries of the last chunk spread a uniform softmax over all keys of that chunk. Padded keys are now masked with the float32 minimum and padded query rows are zeroed after the softmax, so the column sums total the number of real queries (#295). Signed-off-by: Maximilian Jeblick --- kvpress/presses/non_causal_attention_press.py | 12 +++---- tests/presses/test_non_causal_attention.py | 31 +++++++++++++++++++ 2 files changed, 37 insertions(+), 6 deletions(-) create mode 100644 tests/presses/test_non_causal_attention.py diff --git a/kvpress/presses/non_causal_attention_press.py b/kvpress/presses/non_causal_attention_press.py index 81c5415a9..a12491a93 100644 --- a/kvpress/presses/non_causal_attention_press.py +++ b/kvpress/presses/non_causal_attention_press.py @@ -85,11 +85,11 @@ def non_causal_chunked_attn(q: torch.Tensor, k: torch.Tensor, chunk_size: int) - q_chunks = q_padded.view(B, H, num_chunks, chunk_size, d) k_chunks = k_padded.view(B, H, num_chunks, chunk_size, d) - # (B, H, num_chunks, chunk_size, chunk_size) - dots = torch.matmul(q_chunks, k_chunks.transpose(-2, -1)) - dots[:, :, -1].masked_fill_(query_mask.unsqueeze(-1), 0) - dots[:, :, -1].masked_fill_(key_mask.unsqueeze(-2), -1e-9) - attn = torch.softmax(dots.to(torch.float32), dim=-1) + # (B, H, num_chunks, chunk_size, chunk_size), unscaled as in the official implementation (sm_scale=1.0) + dots = torch.matmul(q_chunks, k_chunks.transpose(-2, -1)).to(torch.float32) + dots[:, :, -1].masked_fill_(key_mask.unsqueeze(-2), torch.finfo(dots.dtype).min) + attn = torch.softmax(dots, dim=-1) + attn[:, :, -1].masked_fill_(query_mask.unsqueeze(-1), 0) # sum over query and trim padding return attn.sum(dim=-2).view(B, H, S_pad)[..., :S] @@ -116,5 +116,5 @@ def score( scores = A * values.norm(dim=-1) # (B, H_kv, S) scores = F.avg_pool1d(scores, kernel_size=3, padding=1, stride=1) - z_scores = (scores - scores.mean()) / scores.std().clamp_min(1e-6) # head-wise z-norm + z_scores = (scores - scores.mean()) / scores.std().clamp_min(1e-6) # z-norm over all heads and tokens return z_scores diff --git a/tests/presses/test_non_causal_attention.py b/tests/presses/test_non_causal_attention.py new file mode 100644 index 000000000..e3463e44b --- /dev/null +++ b/tests/presses/test_non_causal_attention.py @@ -0,0 +1,31 @@ +# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch + +from kvpress import NonCausalAttnPress + + +def reference_chunked_attn(q, k, chunk_size): + """Column sums of the softmax attention computed independently on each chunk, without padding.""" + column_sums = [] + for start in range(0, q.shape[-2], chunk_size): + q_chunk, k_chunk = q[..., start : start + chunk_size, :], k[..., start : start + chunk_size, :] + column_sums.append(torch.softmax(q_chunk @ k_chunk.transpose(-2, -1), dim=-1).sum(dim=-2)) + return torch.cat(column_sums, dim=-1) + + +def test_padding_adds_no_attention_mass(): + torch.manual_seed(0) + q, k = torch.randn(2, 1, 3, 2, 8).unbind(0) + scores = NonCausalAttnPress.non_causal_chunked_attn(q, k, chunk_size=4) + torch.testing.assert_close(scores.sum(dim=-1), torch.full((1, 3), 2.0)) + + +@pytest.mark.parametrize("seq_len", [7, 8, 10, 13]) +def test_matches_unpadded_chunked_attention(seq_len): + torch.manual_seed(0) + q, k = torch.randn(2, 2, 3, seq_len, 8).unbind(0) + scores = NonCausalAttnPress.non_causal_chunked_attn(q, k, chunk_size=4) + torch.testing.assert_close(scores, reference_chunked_attn(q, k, chunk_size=4)) From 6a2453018d3d4991a4c3ab47786f716040aaf10c Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 12:30:59 +0200 Subject: [PATCH 24/57] =?UTF-8?q?Forward=20post=5Finit=5Ffrom=5Fmodel=20in?= =?UTF-8?q?=20PerLayerCompressionPress=20=F0=9F=A4=96=F0=9F=A4=96?= =?UTF-8?q?=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Inner presses that load data from the model (QFilterPress, KVzapPress, DuoAttentionPress, ...) failed because post_init_from_model was not delegated to them. Signed-off-by: Maximilian Jeblick --- .../presses/per_layer_compression_press.py | 3 +++ .../test_per_layer_compression_fixes.py | 22 +++++++++++++++++++ 2 files changed, 25 insertions(+) create mode 100644 tests/presses/test_per_layer_compression_fixes.py diff --git a/kvpress/presses/per_layer_compression_press.py b/kvpress/presses/per_layer_compression_press.py index b0222c723..184e4ff02 100644 --- a/kvpress/presses/per_layer_compression_press.py +++ b/kvpress/presses/per_layer_compression_press.py @@ -53,6 +53,9 @@ def __post_init__(self): ), f"compression_ratio can't be set in the provided press: {self.press.__class__}" assert isinstance(self.press, ScorerPress), "PerLayerCompressionPress requires a ScorerPress as input" + def post_init_from_model(self, model): + self.press.post_init_from_model(model) + def forward_hook(self, module: nn.Module, input: list[torch.Tensor], kwargs: dict, output: list): original_compression_ratio = self.press.compression_ratio # type:ignore[index] self.press.compression_ratio = self.compression_ratios[module.layer_idx] # type:ignore[index] diff --git a/tests/presses/test_per_layer_compression_fixes.py b/tests/presses/test_per_layer_compression_fixes.py new file mode 100644 index 000000000..76b2cfd46 --- /dev/null +++ b/tests/presses/test_per_layer_compression_fixes.py @@ -0,0 +1,22 @@ +# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from dataclasses import dataclass, field + +from kvpress import KnormPress, PerLayerCompressionPress +from tests.fixtures import unit_test_model # noqa: F401 + + +@dataclass +class RecordingKnormPress(KnormPress): + models: list = field(default_factory=list) + + def post_init_from_model(self, model): + self.models.append(model) + + +def test_post_init_from_model_is_forwarded(unit_test_model): # noqa: F811 + inner_press = RecordingKnormPress() + with PerLayerCompressionPress(press=inner_press, compression_ratios=[0.5, 0.25])(unit_test_model): + pass + assert inner_press.models == [unit_test_model] From 14c645fdb1c84ffb5001bef8aef8359b59bd6ced Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 12:32:21 +0200 Subject: [PATCH 25/57] =?UTF-8?q?Protect=20CompactorPress=20sinks=20with?= =?UTF-8?q?=20max=20score=20+=201=20=F0=9F=A4=96=F0=9F=A4=96=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Sinks were padded with the maximum score, so they tied with the top token and could be pruned. Signed-off-by: Maximilian Jeblick --- kvpress/presses/compactor_press.py | 2 +- tests/presses/test_compactor_fixes.py | 35 +++++++++++++++++++++++++++ 2 files changed, 36 insertions(+), 1 deletion(-) create mode 100644 tests/presses/test_compactor_fixes.py diff --git a/kvpress/presses/compactor_press.py b/kvpress/presses/compactor_press.py index 338ad7342..a72e67182 100644 --- a/kvpress/presses/compactor_press.py +++ b/kvpress/presses/compactor_press.py @@ -118,5 +118,5 @@ def score( blending = 0.35 if blending is None else blending scores = blending * l_scores + attn_scores # protect sinks by padding - scores = F.pad(scores, (left_keep, right_keep), value=scores.detach().max()) + scores = F.pad(scores, (left_keep, right_keep), value=scores.max().item() + 1) return scores diff --git a/tests/presses/test_compactor_fixes.py b/tests/presses/test_compactor_fixes.py new file mode 100644 index 000000000..d682f45fe --- /dev/null +++ b/tests/presses/test_compactor_fixes.py @@ -0,0 +1,35 @@ +# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch +from transformers import LlamaConfig, LlamaForCausalLM + +from kvpress import CompactorPress + + +@pytest.fixture(scope="module") +def attention_inputs(): + config = LlamaConfig( + hidden_size=64, + intermediate_size=128, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + vocab_size=128, + ) + torch.manual_seed(0) + module = LlamaForCausalLM(config).eval().model.layers[0].self_attn + hidden_states = torch.randn(1, 64, config.hidden_size) + keys, values = torch.randn(2, 1, config.num_key_value_heads, 64, module.head_dim).unbind(0) + cos, sin = torch.randn(2, 1, 64, module.head_dim).unbind(0) + return module, hidden_states, keys, values, {"position_embeddings": (cos, sin)} + + +@torch.no_grad() +def test_compactor_sinks_score_above_every_other_token(attention_inputs): + module, hidden_states, keys, values, kwargs = attention_inputs + press = CompactorPress(compression_ratio=0.5, sink_size_start=8, sink_size_end=4) + scores = press.score(module, hidden_states, keys, values, None, kwargs) + sinks = torch.cat([scores[..., :8], scores[..., -4:]], dim=-1) + assert sinks.min() > scores[..., 8:-4].max() From d0fa84023d4f8cb420f073ff547a6976253212a1 Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 02:45:11 +0200 Subject: [PATCH 26/57] =?UTF-8?q?Use=20module.head=5Fdim=20in=20CriticalKV?= =?UTF-8?q?Press=20=F0=9F=A4=96=F0=9F=A4=96=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Qwen2, Qwen2.5 and Phi3 configs have no head_dim attribute, so CriticalKVPress and CriticalAdaKVPress failed on these models (#213). Signed-off-by: Maximilian Jeblick --- kvpress/presses/criticalkv_press.py | 2 +- tests/presses/test_criticalkv_fixes.py | 32 ++++++++++++++++++++++++++ 2 files changed, 33 insertions(+), 1 deletion(-) create mode 100644 tests/presses/test_criticalkv_fixes.py diff --git a/kvpress/presses/criticalkv_press.py b/kvpress/presses/criticalkv_press.py index 3b526fbb8..8aaede31d 100644 --- a/kvpress/presses/criticalkv_press.py +++ b/kvpress/presses/criticalkv_press.py @@ -61,7 +61,7 @@ def vwl1norm(values, module): bsz, num_key_value_heads, k_len, _ = values.shape num_key_value_groups = module.config.num_attention_heads // num_key_value_heads Wo = module.o_proj.weight.transpose(0, 1) - Wo = Wo.view(module.config.num_attention_heads, module.config.head_dim, module.config.hidden_size) + Wo = Wo.view(module.config.num_attention_heads, module.head_dim, module.config.hidden_size) V = repeat_kv(values, num_key_value_groups) # We use head-wise computation instead of direct matmul to reduce the memory usage of WoV. diff --git a/tests/presses/test_criticalkv_fixes.py b/tests/presses/test_criticalkv_fixes.py new file mode 100644 index 000000000..340364d51 --- /dev/null +++ b/tests/presses/test_criticalkv_fixes.py @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch +from transformers import DynamicCache, Qwen2Config, Qwen2ForCausalLM + +from kvpress import CriticalAdaKVPress, CriticalKVPress, KnormPress + + +@pytest.fixture(scope="module") +def qwen2_model(): + config = Qwen2Config( + hidden_size=64, + intermediate_size=128, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + vocab_size=128, + ) + assert getattr(config, "head_dim", None) is None + return Qwen2ForCausalLM(config).eval() + + +@pytest.mark.parametrize("press_cls", [CriticalKVPress, CriticalAdaKVPress]) +@torch.no_grad() +def test_critical_presses_run_on_qwen2(qwen2_model, press_cls): + cache = DynamicCache() + with press_cls(KnormPress(compression_ratio=0.5))(qwen2_model): + qwen2_model(torch.randint(0, 128, (1, 32)), past_key_values=cache) + expected_length = 16 if press_cls is CriticalKVPress else 32 + assert cache.get_seq_length() == expected_length From 81c2456960450b82e4db39ac75d09c5049c6cbb5 Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 12:33:06 +0200 Subject: [PATCH 27/57] =?UTF-8?q?Forward=20post=5Finit=5Ffrom=5Fmodel=20in?= =?UTF-8?q?=20CriticalAdaKVPress=20=F0=9F=A4=96=F0=9F=A4=96=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Inner presses that load data from the model were never initialized. Signed-off-by: Maximilian Jeblick --- kvpress/presses/criticalkv_press.py | 3 +++ tests/presses/test_criticalkv_fixes.py | 18 ++++++++++++++++++ 2 files changed, 21 insertions(+) diff --git a/kvpress/presses/criticalkv_press.py b/kvpress/presses/criticalkv_press.py index 8aaede31d..214f7e6c1 100644 --- a/kvpress/presses/criticalkv_press.py +++ b/kvpress/presses/criticalkv_press.py @@ -132,6 +132,9 @@ def __post_init__(self): if isinstance(self.press, ExpectedAttentionPress) and self.press.use_vnorm: logger.warning("use_vnorm should be disabled for CriticalAdaKVPress") + def post_init_from_model(self, model): + self.press.post_init_from_model(model) + @property def compression_ratio(self): return self.press.compression_ratio diff --git a/tests/presses/test_criticalkv_fixes.py b/tests/presses/test_criticalkv_fixes.py index 340364d51..c3fea8662 100644 --- a/tests/presses/test_criticalkv_fixes.py +++ b/tests/presses/test_criticalkv_fixes.py @@ -1,6 +1,8 @@ # SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + import pytest import torch from transformers import DynamicCache, Qwen2Config, Qwen2ForCausalLM @@ -8,6 +10,14 @@ from kvpress import CriticalAdaKVPress, CriticalKVPress, KnormPress +@dataclass +class RecordingKnormPress(KnormPress): + models: list = field(default_factory=list) + + def post_init_from_model(self, model): + self.models.append(model) + + @pytest.fixture(scope="module") def qwen2_model(): config = Qwen2Config( @@ -30,3 +40,11 @@ def test_critical_presses_run_on_qwen2(qwen2_model, press_cls): qwen2_model(torch.randint(0, 128, (1, 32)), past_key_values=cache) expected_length = 16 if press_cls is CriticalKVPress else 32 assert cache.get_seq_length() == expected_length + + +@pytest.mark.parametrize("press_cls", [CriticalKVPress, CriticalAdaKVPress]) +def test_post_init_from_model_is_forwarded(qwen2_model, press_cls): + inner_press = RecordingKnormPress(compression_ratio=0.5) + with press_cls(inner_press)(qwen2_model): + pass + assert inner_press.models == [qwen2_model] From 7755e2d09aae7e9344b24fa87f9ca7064317dfb9 Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 02:52:15 +0200 Subject: [PATCH 28/57] =?UTF-8?q?Cast=20the=20FastKVzip=20gate=20inputs=20?= =?UTF-8?q?to=20the=20gate=20dtype=20=F0=9F=A4=96=F0=9F=A4=96=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The published gates store their weights in bfloat16 and are only moved to the model device, so float16 and float32 models crashed with "mat1 and mat2 must have the same dtype". The test mock builds its gates in the model dtype, which hid the issue. bfloat16 models are unaffected. Signed-off-by: Maximilian Jeblick --- kvpress/presses/fastkvzip_press.py | 2 +- tests/presses/test_fastkvzip_fixes.py | 53 +++++++++++++++++++++++++++ 2 files changed, 54 insertions(+), 1 deletion(-) create mode 100644 tests/presses/test_fastkvzip_fixes.py diff --git a/kvpress/presses/fastkvzip_press.py b/kvpress/presses/fastkvzip_press.py index c6897be49..2c482e884 100644 --- a/kvpress/presses/fastkvzip_press.py +++ b/kvpress/presses/fastkvzip_press.py @@ -53,7 +53,7 @@ def __init__( self.d = math.sqrt(self.output_dim) def forward(self, hidden_states: torch.Tensor): - hidden_states = hidden_states.squeeze(0) # bsz = 1 + hidden_states = hidden_states.squeeze(0).to(self.q_proj.weight.dtype) # bsz = 1 nseq = hidden_states.shape[0] # sequence x dim hidden_shape = (nseq, self.nhead, -1, self.output_dim) diff --git a/tests/presses/test_fastkvzip_fixes.py b/tests/presses/test_fastkvzip_fixes.py new file mode 100644 index 000000000..d7679d255 --- /dev/null +++ b/tests/presses/test_fastkvzip_fixes.py @@ -0,0 +1,53 @@ +# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch +from transformers import AutoConfig, AutoModelForCausalLM, DynamicCache + +from kvpress import FastKVzipPress +from kvpress.presses import fastkvzip_press +from kvpress.presses.fastkvzip_press import FastKVzipGate +from tests.fixtures import get_device + +MODEL_NAME = "MaxJeblick/llama2-0b-unit-test" + + +def use_random_gates(monkeypatch, gate_dtype=torch.bfloat16): + """Replace the gate weights downloaded from the Hub by random ones, stored in `gate_dtype` as published.""" + loaded = [] + + def get_gate_weight(model_name): + loaded.append(model_name) + config = AutoConfig.from_pretrained(model_name) + ngroup = config.num_attention_heads // config.num_key_value_heads + gates = [ + FastKVzipGate(idx, config.hidden_size, config.num_key_value_heads, ngroup, gate_dtype, sink=16) + for idx in range(config.num_hidden_layers) + ] + return [gate.state_dict() for gate in gates], f"{model_name.split('/')[-1]}/q{ngroup}_dim16_sink16.pt" + + monkeypatch.setattr(fastkvzip_press, "get_gate_weight", get_gate_weight) + return loaded + + +def run_press(model, press, context_len=64): + with press(model): + model(torch.randint(0, 1024, (1, context_len), device=model.device), past_key_values=DynamicCache()) + + +def n_masked(model): + return sum(len(layer.self_attn.masked_key_indices[0]) for layer in model.model.layers) + + +@pytest.mark.parametrize("model_dtype", [torch.float32, torch.float16]) +def test_fastkvzip_bfloat16_gates_on_other_model_dtypes(monkeypatch, model_dtype): + use_random_gates(monkeypatch, gate_dtype=torch.bfloat16) + model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, dtype=model_dtype).eval().to(get_device()) + context_len, press = 64, FastKVzipPress(compression_ratio=0.5) + run_press(model, press, context_len) + + assert press.gates[0].q_proj.weight.dtype == torch.bfloat16 + config = model.config + assert n_masked(model) == int(config.num_hidden_layers * config.num_key_value_heads * context_len * 0.5) + From aa6a4dbb07a678293dd31e3a0ae5aa556a864996 Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 02:54:00 +0200 Subject: [PATCH 29/57] =?UTF-8?q?Support=20Gemma3=20in=20FastKVzipPress=20?= =?UTF-8?q?=F0=9F=A4=96=F0=9F=A4=96=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Gemma3Model has no layers attribute (they live in its language_model), and the skipped sliding window layers left None scores that could not be stacked, so FastKVzipPress failed on Gemma3 although it is listed as supported. Look the layers up as BasePress.__call__ does and only stack and compress the scored layers. A block without forward pass no longer raises either. Signed-off-by: Maximilian Jeblick --- kvpress/presses/fastkvzip_press.py | 20 ++++++---- tests/presses/test_fastkvzip_fixes.py | 57 ++++++++++++++++++++++++++- 2 files changed, 67 insertions(+), 10 deletions(-) diff --git a/kvpress/presses/fastkvzip_press.py b/kvpress/presses/fastkvzip_press.py index 2c482e884..ce3dab9c4 100644 --- a/kvpress/presses/fastkvzip_press.py +++ b/kvpress/presses/fastkvzip_press.py @@ -208,8 +208,8 @@ def __call__(self, model: PreTrainedModel) -> Generator: self.post_init_from_model(model) hooks = [] try: - self.score_val = [None for _ in range(len(model.model.layers))] # reset every prefilling language_model = model.model.language_model if hasattr(model.model, "language_model") else model.model + self.score_val = [None for _ in range(len(language_model.layers))] # reset every prefilling for layer in language_model.layers: if isinstance(model, Gemma3ForConditionalGeneration) and layer.self_attn.is_sliding: # Skip layers with sliding window attention, only for Gemma3 @@ -267,7 +267,14 @@ def compress_post(self, model: PreTrainedModel): Obtain the indices of KV pairs to be evicted. Adopted from adakv_press.compress (fake compression). KVzip does not rely on safeguards. """ - self.score_val = torch.stack(self.score_val, dim=0) + language_model = model.model.language_model if hasattr(model.model, "language_model") else model.model + # Layers without scores (e.g. sliding window layers of Gemma3) are not compressed + modules = [ + layer.self_attn for layer in language_model.layers if self.score_val[layer.self_attn.layer_idx] is not None + ] + if not modules: + return + self.score_val = torch.stack([self.score_val[module.layer_idx] for module in modules], dim=0) if self.compression_ratio > 0: n_layer, bsz, num_key_value_heads, ctx_len = self.score_val.shape @@ -282,16 +289,13 @@ def compress_post(self, model: PreTrainedModel): n_tokens_per_layer = bsz * num_key_value_heads * ctx_len n_pruned_layers = torch.bincount(pruned_indices // n_tokens_per_layer, minlength=n_layer).int() - for layer in model.model.layers: - module = layer.self_attn - layer_idx = int(module.layer_idx) - + for i, module in enumerate(modules): assert module.config._attn_implementation != "eager", "eager mode not supported" - scores = self.score_val[layer_idx] + scores = self.score_val[i] # Compute bottom-k across heads - n_pruned = n_pruned_layers[layer_idx].cpu() + n_pruned = n_pruned_layers[i].cpu() indices = torch.topk(-scores.reshape(bsz, -1), n_pruned, dim=1).indices.flatten().cpu() # Save indices to mask during the attention mechanism. Please refer to attention_patch.py for details diff --git a/tests/presses/test_fastkvzip_fixes.py b/tests/presses/test_fastkvzip_fixes.py index d7679d255..0b716f272 100644 --- a/tests/presses/test_fastkvzip_fixes.py +++ b/tests/presses/test_fastkvzip_fixes.py @@ -3,12 +3,12 @@ import pytest import torch -from transformers import AutoConfig, AutoModelForCausalLM, DynamicCache +from transformers import AutoConfig, AutoModelForCausalLM, DynamicCache, Gemma3Config, Gemma3ForConditionalGeneration from kvpress import FastKVzipPress from kvpress.presses import fastkvzip_press from kvpress.presses.fastkvzip_press import FastKVzipGate -from tests.fixtures import get_device +from tests.fixtures import get_device, unit_test_model # noqa: F401 MODEL_NAME = "MaxJeblick/llama2-0b-unit-test" @@ -20,6 +20,7 @@ def use_random_gates(monkeypatch, gate_dtype=torch.bfloat16): def get_gate_weight(model_name): loaded.append(model_name) config = AutoConfig.from_pretrained(model_name) + config = getattr(config, "text_config", config) ngroup = config.num_attention_heads // config.num_key_value_heads gates = [ FastKVzipGate(idx, config.hidden_size, config.num_key_value_heads, ngroup, gate_dtype, sink=16) @@ -51,3 +52,55 @@ def test_fastkvzip_bfloat16_gates_on_other_model_dtypes(monkeypatch, model_dtype config = model.config assert n_masked(model) == int(config.num_hidden_layers * config.num_key_value_heads * context_len * 0.5) + +def tiny_gemma3(path): + """Random Gemma3 with alternating sliding window and full attention layers.""" + config = Gemma3Config( + text_config=dict( + vocab_size=1024, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=4, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=16, + sliding_window=8, + layer_types=["sliding_attention", "full_attention"] * 2, + ), + vision_config=dict( + hidden_size=32, + intermediate_size=64, + num_hidden_layers=1, + num_attention_heads=2, + image_size=28, + patch_size=14, + ), + mm_tokens_per_image=4, + ) + torch.manual_seed(0) + Gemma3ForConditionalGeneration(config).save_pretrained(path) + return Gemma3ForConditionalGeneration.from_pretrained(path).eval().to(get_device()) + + +def test_fastkvzip_gemma3_compresses_full_attention_layers(monkeypatch, tmp_path): + use_random_gates(monkeypatch) + model = tiny_gemma3(tmp_path) + context_len, press = 64, FastKVzipPress(compression_ratio=0.5) + with press(model): + input_ids = torch.randint(0, 1024, (1, context_len), device=model.device) + model(input_ids=input_ids, past_key_values=DynamicCache(config=model.config)) + + layers = model.model.language_model.layers + full_attention_modules = [layer.self_attn for layer in layers if not layer.self_attn.is_sliding] + assert press.score_val.shape[0] == len(full_attention_modules) == 2 + n_pruned = sum(len(module.masked_key_indices[0]) for module in full_attention_modules) + assert n_pruned == int( + len(full_attention_modules) * model.config.text_config.num_key_value_heads * context_len * 0.5 + ) + assert all(layer.self_attn.masked_key_indices is None for layer in layers if layer.self_attn.is_sliding) + + +def test_fastkvzip_without_forward_pass_is_a_no_op(unit_test_model, monkeypatch): # noqa: F811 + use_random_gates(monkeypatch) + with FastKVzipPress(compression_ratio=0.5)(unit_test_model): + pass From ad3db055283dc88ffbbec17c7823cf4191a6ae7d Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 02:55:32 +0200 Subject: [PATCH 30/57] =?UTF-8?q?Gather=20the=20FastKVzip=20scores=20on=20?= =?UTF-8?q?the=20model=20device=20=F0=9F=A4=96=F0=9F=A4=96=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Each layer scores its KV pairs on its own device, so torch.stack failed when the layers were spread over several devices (device_map). Signed-off-by: Maximilian Jeblick --- kvpress/presses/fastkvzip_press.py | 2 +- tests/presses/test_fastkvzip_fixes.py | 15 +++++++++++++++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/kvpress/presses/fastkvzip_press.py b/kvpress/presses/fastkvzip_press.py index ce3dab9c4..2ebe4e03f 100644 --- a/kvpress/presses/fastkvzip_press.py +++ b/kvpress/presses/fastkvzip_press.py @@ -274,7 +274,7 @@ def compress_post(self, model: PreTrainedModel): ] if not modules: return - self.score_val = torch.stack([self.score_val[module.layer_idx] for module in modules], dim=0) + self.score_val = torch.stack([self.score_val[module.layer_idx].to(model.device) for module in modules], dim=0) if self.compression_ratio > 0: n_layer, bsz, num_key_value_heads, ctx_len = self.score_val.shape diff --git a/tests/presses/test_fastkvzip_fixes.py b/tests/presses/test_fastkvzip_fixes.py index 0b716f272..7774173cc 100644 --- a/tests/presses/test_fastkvzip_fixes.py +++ b/tests/presses/test_fastkvzip_fixes.py @@ -100,6 +100,21 @@ def test_fastkvzip_gemma3_compresses_full_attention_layers(monkeypatch, tmp_path assert all(layer.self_attn.masked_key_indices is None for layer in layers if layer.self_attn.is_sliding) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Needs a second device") +def test_fastkvzip_gathers_the_scores_of_layers_on_other_devices(monkeypatch): + use_random_gates(monkeypatch) + model = AutoModelForCausalLM.from_pretrained(MODEL_NAME).eval() + context_len, press = 64, FastKVzipPress(compression_ratio=0.5) + with press(model): + model(torch.randint(0, 1024, (1, context_len)), past_key_values=DynamicCache()) + # Scores are computed on the device of each layer + press.score_val[-1] = press.score_val[-1].cuda() + + assert press.score_val.device == model.device + config = model.config + assert n_masked(model) == int(config.num_hidden_layers * config.num_key_value_heads * context_len * 0.5) + + def test_fastkvzip_without_forward_pass_is_a_no_op(unit_test_model, monkeypatch): # noqa: F811 use_random_gates(monkeypatch) with FastKVzipPress(compression_ratio=0.5)(unit_test_model): From 18923ed76ad5f7c48b1be0b53889a48e4c3bff93 Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 03:01:59 +0200 Subject: [PATCH 31/57] =?UTF-8?q?Match=20the=20DuoAttention=20patterns=20o?= =?UTF-8?q?f=20gradientai=20and=20Llama-3.1=20checkpoints=20=F0=9F=A4=96?= =?UTF-8?q?=F0=9F=A4=96=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The two gradientai entries of PATTERNS_DICT had a double slash, so their checkpoints never matched, and only the deprecated meta-llama/Meta-Llama-3.1-8B-Instruct name was listed for Llama 3.1. meta-llama/Llama-3.1-8B-Instruct now uses the same pattern. Signed-off-by: Maximilian Jeblick --- kvpress/presses/duo_attention_press.py | 5 +++-- tests/presses/test_duo_attention_fixes.py | 10 ++++++++++ 2 files changed, 13 insertions(+), 2 deletions(-) create mode 100644 tests/presses/test_duo_attention_fixes.py diff --git a/kvpress/presses/duo_attention_press.py b/kvpress/presses/duo_attention_press.py index 646e53528..228ab6035 100644 --- a/kvpress/presses/duo_attention_press.py +++ b/kvpress/presses/duo_attention_press.py @@ -18,9 +18,10 @@ PATTERNS_DICT = { "togethercomputer/Llama-2-7B-32K-Instruct": "Llama-2-7B-32K-Instruct/lr%3D0.02-reg%3D0.05-ctx%3D1000_32000-multi_passkey10", # noqa: E501 - "gradientai//Llama-3-8B-Instruct-Gradient-1048k": "Llama-3-8B-Instruct-Gradient-1048k/lr%3D0.02-reg%3D0.05-ctx%3D1000_32000-multi_passkey10", # noqa: E501 - "gradientai//Llama-3-8B-Instruct-Gradient-4194k": "Llama-3-8B-Instruct-Gradient-4194k/lr%3D0.02-reg%3D0.05-ctx%3D1000_32000-multi_passkey10", # noqa: E501 + "gradientai/Llama-3-8B-Instruct-Gradient-1048k": "Llama-3-8B-Instruct-Gradient-1048k/lr%3D0.02-reg%3D0.05-ctx%3D1000_32000-multi_passkey10", # noqa: E501 + "gradientai/Llama-3-8B-Instruct-Gradient-4194k": "Llama-3-8B-Instruct-Gradient-4194k/lr%3D0.02-reg%3D0.05-ctx%3D1000_32000-multi_passkey10", # noqa: E501 "meta-llama/Meta-Llama-3.1-8B-Instruct": "Meta-Llama-3.1-8B-Instruct/lr=0.02-reg=0.05-ctx=1000_128000-multi_passkey10", # noqa: E501 + "meta-llama/Llama-3.1-8B-Instruct": "Meta-Llama-3.1-8B-Instruct/lr=0.02-reg=0.05-ctx=1000_128000-multi_passkey10", # noqa: E501 "mistralai/Mistral-7B-Instruct-v0.2": "Mistral-7B-Instruct-v0.2/lr%3D0.02-reg%3D0.05-ctx%3D1000_32000-multi_passkey10", # noqa: E501 "mistralai/Mistral-7B-Instruct-v0.3": "Mistral-7B-Instruct-v0.3/lr%3D0.02-reg%3D0.05-ctx%3D1000_32000-multi_passkey10", # noqa: E501 } diff --git a/tests/presses/test_duo_attention_fixes.py b/tests/presses/test_duo_attention_fixes.py new file mode 100644 index 000000000..06527923d --- /dev/null +++ b/tests/presses/test_duo_attention_fixes.py @@ -0,0 +1,10 @@ +# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from kvpress.presses.duo_attention_press import PATTERNS_DICT + + +def test_duo_attention_patterns_match_hub_checkpoints(): + assert all(name.count("/") == 1 for name in PATTERNS_DICT) + assert "gradientai/Llama-3-8B-Instruct-Gradient-1048k" in PATTERNS_DICT + assert PATTERNS_DICT["meta-llama/Llama-3.1-8B-Instruct"] == PATTERNS_DICT["meta-llama/Meta-Llama-3.1-8B-Instruct"] From 6272d25b038dab154aa0b57cbb0edf947bd3050a Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 03:21:58 +0200 Subject: [PATCH 32/57] =?UTF-8?q?Support=20configs=20without=20head=5Fdim?= =?UTF-8?q?=20in=20ExpectedAttentionStatsPress=20=F0=9F=A4=96=F0=9F=A4=96?= =?UTF-8?q?=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Qwen2 and Qwen2.5 configs have no head_dim attribute, so both the stats lookup and the statistics script raised an AttributeError. Fall back to hidden_size // num_attention_heads. Signed-off-by: Maximilian Jeblick --- .../presses/expected_attention_with_stats.py | 8 ++++-- tests/presses/test_ea_with_stats_fixes.py | 27 +++++++++++++++++++ 2 files changed, 33 insertions(+), 2 deletions(-) create mode 100644 tests/presses/test_ea_with_stats_fixes.py diff --git a/kvpress/presses/expected_attention_with_stats.py b/kvpress/presses/expected_attention_with_stats.py index bd776ca37..bb5f5cfee 100644 --- a/kvpress/presses/expected_attention_with_stats.py +++ b/kvpress/presses/expected_attention_with_stats.py @@ -18,6 +18,10 @@ from kvpress.presses.expected_attention_press import ExpectedAttentionPress +def get_head_dim(config) -> int: + return getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads + + @dataclass class ExpectedAttentionStatsPress(ExpectedAttentionPress): """ @@ -87,7 +91,7 @@ def _maybe_load_stats_from_hub(self, model: PreTrainedModel): model_name=model.config.name_or_path, num_layers=model.config.num_hidden_layers, num_heads=model.config.num_attention_heads, - head_dim=model.config.head_dim, + head_dim=get_head_dim(model.config), dataset_name=self.dataset_name, num_samples=self.num_samples, sample_seq_len=self.sample_seq_len, @@ -270,7 +274,7 @@ def main( stats = ExpectedAttentionStats( num_layers=model.config.num_hidden_layers, num_heads=model.config.num_attention_heads, - head_dim=model.config.head_dim, + head_dim=get_head_dim(model.config), dataset_name=dataset_name, model_name=model_name, num_samples=num_samples, diff --git a/tests/presses/test_ea_with_stats_fixes.py b/tests/presses/test_ea_with_stats_fixes.py new file mode 100644 index 000000000..40018e773 --- /dev/null +++ b/tests/presses/test_ea_with_stats_fixes.py @@ -0,0 +1,27 @@ +# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace + +from transformers import Qwen2Config + +from kvpress import ExpectedAttentionStatsPress +from kvpress.presses import expected_attention_with_stats +from kvpress.presses.expected_attention_with_stats import ExpectedAttentionStats + + +def fake_model(config, name_or_path): + config.name_or_path = name_or_path + return SimpleNamespace(config=config) + + +def test_ea_stats_of_models_without_head_dim_in_their_config(monkeypatch): + stats_ids = [] + monkeypatch.setattr(ExpectedAttentionStats, "from_pretrained", lambda stats_id: stats_ids.append(stats_id)) + config = Qwen2Config(hidden_size=64, num_attention_heads=4, num_key_value_heads=2, num_hidden_layers=2) + assert getattr(config, "head_dim", None) is None + model = fake_model(config, "Qwen/Qwen2.5-0.5B-Instruct") + + ExpectedAttentionStatsPress()._maybe_load_stats_from_hub(model) + assert stats_ids == ["alessiodevoto/exp_att_stats_Qwen_Qwen2.5-0.5B-Instruct_kmfoda_booksum_100_1000_4"] + assert expected_attention_with_stats.get_head_dim(config) == 16 From a6f4f17fcf8a5d4e88b036c4bce51021f17868e0 Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 03:23:23 +0200 Subject: [PATCH 33/57] =?UTF-8?q?Move=20the=20ExpectedAttentionStatsPress?= =?UTF-8?q?=20statistics=20to=20each=20layer's=20device=20=F0=9F=A4=96?= =?UTF-8?q?=F0=9F=A4=96=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The query statistics are loaded on model.device, so scoring failed for the layers placed on other devices (device_map). Move the statistics of a layer to the device of its hidden states when they are used. Signed-off-by: Maximilian Jeblick --- .../presses/expected_attention_with_stats.py | 4 +- tests/presses/test_ea_with_stats_fixes.py | 38 ++++++++++++++++++- 2 files changed, 40 insertions(+), 2 deletions(-) diff --git a/kvpress/presses/expected_attention_with_stats.py b/kvpress/presses/expected_attention_with_stats.py index bb5f5cfee..915fdbd21 100644 --- a/kvpress/presses/expected_attention_with_stats.py +++ b/kvpress/presses/expected_attention_with_stats.py @@ -65,7 +65,9 @@ def get_query_statistics(self, module: nn.Module, hidden_states: torch.Tensor): """ q_len = hidden_states.shape[1] layer_idx = module.layer_idx - mu, cov = self.apply_avg_rope(module, self.mu[layer_idx], self.cov[layer_idx], q_len) # type: ignore + mu = self.mu[layer_idx].to(hidden_states.device) # type: ignore[index] + cov = self.cov[layer_idx].to(hidden_states.device) # type: ignore[index] + mu, cov = self.apply_avg_rope(module, mu, cov, q_len) return mu.unsqueeze(0), cov.unsqueeze(0) @staticmethod diff --git a/tests/presses/test_ea_with_stats_fixes.py b/tests/presses/test_ea_with_stats_fixes.py index 40018e773..6b6a1945e 100644 --- a/tests/presses/test_ea_with_stats_fixes.py +++ b/tests/presses/test_ea_with_stats_fixes.py @@ -3,18 +3,54 @@ from types import SimpleNamespace -from transformers import Qwen2Config +import pytest +import torch +from transformers import AutoModelForCausalLM, DynamicCache, Qwen2Config from kvpress import ExpectedAttentionStatsPress from kvpress.presses import expected_attention_with_stats from kvpress.presses.expected_attention_with_stats import ExpectedAttentionStats +MODEL_NAME = "MaxJeblick/llama2-0b-unit-test" + def fake_model(config, name_or_path): config.name_or_path = name_or_path return SimpleNamespace(config=config) +def random_stats(model): + config = model.config + stats = ExpectedAttentionStats( + num_layers=config.num_hidden_layers, + num_heads=config.num_attention_heads, + head_dim=config.head_dim, + dataset_name="kmfoda/booksum", + model_name=config.name_or_path, + num_samples=100, + sample_seq_len=1000, + n_sink=4, + ) + stats.query_mean.data.normal_() + stats.query_cov.data = torch.eye(config.head_dim).expand_as(stats.query_cov).clone() + return stats + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Needs a second device") +def test_ea_stats_are_moved_to_the_device_of_each_layer(tmp_path): + model = AutoModelForCausalLM.from_pretrained(MODEL_NAME).eval() + random_stats(model).save_pretrained(tmp_path) + press = ExpectedAttentionStatsPress(compression_ratio=0.5, stats_folder=str(tmp_path)) + press.post_init_from_model(model) + # The statistics are loaded on the model device, which is not the one of every layer with a device_map + press.mu, press.cov = press.mu.cuda(), press.cov.cuda() + + cache = DynamicCache() + with press(model): + model(torch.randint(0, 1024, (1, 32)), past_key_values=cache) + assert cache.get_seq_length() == 16 + + def test_ea_stats_of_models_without_head_dim_in_their_config(monkeypatch): stats_ids = [] monkeypatch.setattr(ExpectedAttentionStats, "from_pretrained", lambda stats_id: stats_ids.append(stats_id)) From c9e8bf7bb0e96967be1642ba4e6cc615d02305ad Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 03:26:49 +0200 Subject: [PATCH 34/57] =?UTF-8?q?Move=20the=20CapPress=20RoPE=20embeddings?= =?UTF-8?q?=20to=20the=20device=20of=20the=20queries=20=F0=9F=A4=96?= =?UTF-8?q?=F0=9F=A4=96=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The rotary embedding is shared by all layers, and with a device_map it can return cos and sin on another device than the layer being scored, so building the averaged RoPE matrix failed. Signed-off-by: Maximilian Jeblick --- kvpress/presses/cap_press.py | 3 ++- tests/presses/test_cap_press_fixes.py | 29 +++++++++++++++++++++++++++ 2 files changed, 31 insertions(+), 1 deletion(-) create mode 100644 tests/presses/test_cap_press_fixes.py diff --git a/kvpress/presses/cap_press.py b/kvpress/presses/cap_press.py index 446b7a91b..9655cb6e4 100644 --- a/kvpress/presses/cap_press.py +++ b/kvpress/presses/cap_press.py @@ -115,8 +115,9 @@ def _avg_rope_matrix( head_dim = module.head_dim dummy = torch.zeros((1, 1, head_dim), device=device, dtype=dtype) + # The rotary embedding is shared by all layers and may be placed on another device cos, sin = module.rotary_emb(dummy, position_ids) - cos, sin = cos[0], sin[0] # [n_future_positions, D_head] + cos, sin = cos[0].to(device), sin[0].to(device) # [n_future_positions, D_head] identity = torch.eye(head_dim, device=device, dtype=dtype) rotate_half = torch.zeros((head_dim, head_dim), device=device, dtype=dtype) diff --git a/tests/presses/test_cap_press_fixes.py b/tests/presses/test_cap_press_fixes.py new file mode 100644 index 000000000..177354e74 --- /dev/null +++ b/tests/presses/test_cap_press_fixes.py @@ -0,0 +1,29 @@ +# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import torch +from torch import nn + +from kvpress import CapPress +from tests.fixtures import unit_test_model # noqa: F401 + + +class CPURotaryEmbedding(nn.Module): + """Rotary embedding returning its outputs on the CPU, like one placed on another device by a device_map.""" + + def forward(self, x, position_ids): + shape = (position_ids.shape[0], position_ids.shape[1], x.shape[-1]) + return torch.ones(shape, dtype=x.dtype), torch.zeros(shape, dtype=x.dtype) + + +def test_cap_rope_matrix_is_built_on_the_device_of_the_queries(unit_test_model): # noqa: F811 + module = unit_test_model.model.layers[0].self_attn + module.rotary_emb = CPURotaryEmbedding() + try: + device = torch.device("meta") + rotation = CapPress()._avg_rope_matrix(module, q_len=8, device=device, dtype=torch.float32) + finally: + module.rotary_emb = unit_test_model.model.rotary_emb + + assert rotation.device == device + assert rotation.shape == (module.head_dim, module.head_dim) From 3b55c4c5962985ca7cb8045206fe6840e969ef6d Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 03:28:11 +0200 Subject: [PATCH 35/57] =?UTF-8?q?Start=20the=20CapPress=20future=20positio?= =?UTF-8?q?ns=20after=20the=20last=20token=20=F0=9F=A4=96=F0=9F=A4=96?= =?UTF-8?q?=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The averaged RoPE rotation used the positions following hidden_states.shape[1], the number of tokens of the forward pass, rather than the position following the last token (cache_position[-1] + 1). Both match for a plain prefill, but not when the hidden states follow cached tokens, e.g. with chunked prefilling or decoding presses. Signed-off-by: Maximilian Jeblick --- kvpress/presses/cap_press.py | 10 ++++++---- tests/presses/test_cap_press_fixes.py | 24 ++++++++++++++++++++++++ 2 files changed, 30 insertions(+), 4 deletions(-) diff --git a/kvpress/presses/cap_press.py b/kvpress/presses/cap_press.py index 9655cb6e4..301862490 100644 --- a/kvpress/presses/cap_press.py +++ b/kvpress/presses/cap_press.py @@ -166,6 +166,7 @@ def _query_anchor( self, module: nn.Module, hidden_states: torch.Tensor, + cache_position: torch.Tensor, ) -> torch.Tensor: """ Compute the historical-query anchor. @@ -173,11 +174,12 @@ def _query_anchor( Args: module: Attention module. hidden_states: Hidden states with shape [B, T, D_model]. + cache_position: Positions of the hidden states in the sequence, with shape [T]. Returns: Query anchor with shape [B, H, D_head]. """ - q_len = hidden_states.shape[1] + q_len = int(cache_position[-1]) + 1 query_states = self._query_states_pre_rope(module, hidden_states) query_states = self._apply_avg_rope(module, query_states, q_len) return query_states.mean(dim=2) @@ -239,13 +241,13 @@ def score( keys: Cached keys with shape [B, H_kv, T_cache, D_head]. values: Cached values with shape [B, H_kv, T_cache, D_head]. attentions: Unused; included for KVPress scorer compatibility. - kwargs: Unused; included for KVPress scorer compatibility. + kwargs: Keyword arguments of the attention layer, including `cache_position`. Returns: Retention scores with shape [B, H_kv, T_cache]. Larger values indicate higher retention priority. """ - del attentions, kwargs + del attentions if keys.size(2) <= self.n_sink: raise ValueError(f"Input cache length ({keys.size(2)}) must be larger than " f"n_sink={self.n_sink}.") @@ -261,7 +263,7 @@ def score( values_rep = repeat_kv(values_no_sink, num_key_value_groups) # Query relevance. - query_anchor = self._query_anchor(module, hidden_states) + query_anchor = self._query_anchor(module, hidden_states, kwargs["cache_position"]) statistic = self._compute_query_key_statistic(query_anchor, keys_rep) weights = self._compute_query_relevance_weights(statistic) diff --git a/tests/presses/test_cap_press_fixes.py b/tests/presses/test_cap_press_fixes.py index 177354e74..66b697b19 100644 --- a/tests/presses/test_cap_press_fixes.py +++ b/tests/presses/test_cap_press_fixes.py @@ -1,6 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import pytest import torch from torch import nn @@ -27,3 +28,26 @@ def test_cap_rope_matrix_is_built_on_the_device_of_the_queries(unit_test_model): assert rotation.device == device assert rotation.shape == (module.head_dim, module.head_dim) + + +@torch.no_grad() +@pytest.mark.parametrize("n_cached", [0, 100]) +def test_cap_future_positions_start_after_the_last_token(unit_test_model, monkeypatch, n_cached): # noqa: F811 + start_positions = [] + avg_rope_matrix = CapPress._avg_rope_matrix + + def recording_avg_rope_matrix(self, module, q_len, device, dtype): + start_positions.append(q_len) + return avg_rope_matrix(self, module, q_len, device, dtype) + + monkeypatch.setattr(CapPress, "_avg_rope_matrix", recording_avg_rope_matrix) + model, q_len = unit_test_model, 16 + module = model.model.layers[0].self_attn + module.rotary_emb = model.model.rotary_emb + hidden_states = torch.randn(1, q_len, model.config.hidden_size, dtype=model.dtype, device=model.device) + keys = torch.randn(1, model.config.num_key_value_heads, n_cached + q_len, module.head_dim, dtype=model.dtype) + keys = keys.to(model.device) + cache_position = torch.arange(n_cached, n_cached + q_len, device=model.device) + + CapPress(compression_ratio=0.5).score(module, hidden_states, keys, keys, None, {"cache_position": cache_position}) + assert start_positions == [n_cached + q_len] From 642eb7fe3cf61fa5749b4f226a1809ab41fa7e18 Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 03:29:38 +0200 Subject: [PATCH 36/57] =?UTF-8?q?Scale=20the=20DropKVPress=20recomputed=20?= =?UTF-8?q?attention=20with=20module.scaling=20=F0=9F=A4=96=F0=9F=A4=96?= =?UTF-8?q?=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Outside eager attention, DropKVPress recomputes the window attention with 1/sqrt(head_dim) instead of the attention scaling of the model. Both match for all supported models except Gemma3-27B, whose scaling is query_pre_attn_scalar**-0.5. Signed-off-by: Maximilian Jeblick --- kvpress/presses/dropkv_press.py | 11 ++++---- tests/presses/test_dropkv_fixes.py | 40 ++++++++++++++++++++++++++++++ tests/presses/test_dropkv_press.py | 2 +- 3 files changed, 47 insertions(+), 6 deletions(-) create mode 100644 tests/presses/test_dropkv_fixes.py diff --git a/kvpress/presses/dropkv_press.py b/kvpress/presses/dropkv_press.py index 1784afd90..a9bcee184 100644 --- a/kvpress/presses/dropkv_press.py +++ b/kvpress/presses/dropkv_press.py @@ -1,7 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -import math from dataclasses import dataclass import torch @@ -102,19 +101,21 @@ def _get_window_queries( (cos[:, -window_size:], sin[:, -window_size:]), ) - def _compute_window_probabilities(self, query_states: torch.Tensor, keys: torch.Tensor) -> torch.Tensor: + def _compute_window_probabilities( + self, query_states: torch.Tensor, keys: torch.Tensor, scaling: float + ) -> torch.Tensor: """Attention probabilities of the window queries over the whole cache. Only used when the attention weights are not already available from the forward pass, i.e. outside of the eager attention implementation. """ num_query_heads, query_len = query_states.shape[1], query_states.shape[2] - num_kv_heads, seq_len, head_dim = keys.shape[1], keys.shape[2], keys.shape[3] + num_kv_heads, seq_len = keys.shape[1], keys.shape[2] if num_query_heads % num_kv_heads != 0: raise ValueError(f"Query heads {num_query_heads} must be divisible by KV heads {num_kv_heads}") repeated_keys = repeat_kv(keys, num_query_heads // num_kv_heads) - attention_weights = torch.matmul(query_states, repeated_keys.transpose(2, 3)) / math.sqrt(head_dim) + attention_weights = torch.matmul(query_states, repeated_keys.transpose(2, 3)) * scaling causal_mask = torch.full( (query_len, seq_len), @@ -185,7 +186,7 @@ def score( probabilities = attentions[..., -window_size:, :] else: query_states = self._get_window_queries(module, hidden_states, kwargs["position_embeddings"], window_size) - probabilities = self._compute_window_probabilities(query_states, keys) + probabilities = self._compute_window_probabilities(query_states, keys, module.scaling) scores = self._compute_scores(probabilities, keys, values) diff --git a/tests/presses/test_dropkv_fixes.py b/tests/presses/test_dropkv_fixes.py new file mode 100644 index 000000000..ce5fb25e6 --- /dev/null +++ b/tests/presses/test_dropkv_fixes.py @@ -0,0 +1,40 @@ +# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import torch +from transformers import DynamicCache, Gemma3ForCausalLM, Gemma3TextConfig + +from kvpress import DropKVPress + + +@torch.no_grad() +def test_dropkv_recomputed_attention_uses_the_attention_scaling(): + """Gemma3 scales the attention logits by query_pre_attn_scalar**-0.5 rather than head_dim**-0.5.""" + config = Gemma3TextConfig( + vocab_size=1024, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=16, + query_pre_attn_scalar=64, + sliding_window=64, + ) + config._attn_implementation = "eager" + torch.manual_seed(0) + model = Gemma3ForCausalLM(config).eval() + recorded: dict[str, list] = {"reused": [], "recomputed": []} + + class RecordingDropKVPress(DropKVPress): + def score(self, module, hidden_states, keys, values, attentions, kwargs): + recorded["reused"].append(super().score(module, hidden_states, keys, values, attentions, kwargs)) + recorded["recomputed"].append(super().score(module, hidden_states, keys, values, None, kwargs)) + return recorded["reused"][-1] + + with RecordingDropKVPress(compression_ratio=0.5, window_size=8, kernel_size=3)(model): + model(torch.randint(0, 1024, (1, 40)), past_key_values=DynamicCache(), output_attentions=True) + + assert len(recorded["reused"]) == config.num_hidden_layers + for reused, recomputed in zip(recorded["reused"], recorded["recomputed"]): + torch.testing.assert_close(reused, recomputed, rtol=1e-4, atol=1e-4) diff --git a/tests/presses/test_dropkv_press.py b/tests/presses/test_dropkv_press.py index 13ade8080..6407f6534 100644 --- a/tests/presses/test_dropkv_press.py +++ b/tests/presses/test_dropkv_press.py @@ -78,7 +78,7 @@ def test_dropkv_scores_match_reference(num_query_heads, num_kv_heads): press = DropKVPress(window_size=window_size) expected = _reference_scores(query_states, keys, values, press.epsilon) - probabilities = press._compute_window_probabilities(query_states, keys) + probabilities = press._compute_window_probabilities(query_states, keys, 1 / math.sqrt(head_dim)) actual = press._compute_scores(probabilities, keys, values) torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5) From 7d65abfe3ae79d2f7a69c8028c26084cd6451d8c Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 03:42:07 +0200 Subject: [PATCH 37/57] =?UTF-8?q?Map=20the=20FastKVzip=20gates=20of=20Gemm?= =?UTF-8?q?a3=20to=20its=20full=20attention=20layers=20=F0=9F=A4=96?= =?UTF-8?q?=F0=9F=A4=96=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The published Gemma3 gates only cover the full attention layers (8 gates for the 48 layers of gemma-3-12b-it), in order, as in the original implementation (layer_id_to_static_id). Indexing them by layer index gave layer 5 the gate of the sixth full attention layer and raised an IndexError from layer 11. Gates are now looked up by the position of the layer among the scored layers, which is the layer index for the other models. Signed-off-by: Maximilian Jeblick --- kvpress/presses/fastkvzip_press.py | 9 +++++++-- tests/presses/test_fastkvzip_fixes.py | 5 ++++- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/kvpress/presses/fastkvzip_press.py b/kvpress/presses/fastkvzip_press.py index 2ebe4e03f..168bcd5d1 100644 --- a/kvpress/presses/fastkvzip_press.py +++ b/kvpress/presses/fastkvzip_press.py @@ -178,6 +178,7 @@ class FastKVzipPress(BasePress): gates: list[nn.Module] | None = field(init=False, default=None) score_val: list[torch.Tensor] | torch.Tensor | None = field(init=False, default=None) + gate_indices: dict[int, int] = field(init=False, default_factory=dict) def post_init_from_model(self, model): """ @@ -210,10 +211,13 @@ def __call__(self, model: PreTrainedModel) -> Generator: try: language_model = model.model.language_model if hasattr(model.model, "language_model") else model.model self.score_val = [None for _ in range(len(language_model.layers))] # reset every prefilling + # The gates are only trained for the scored layers, in order + self.gate_indices = {} for layer in language_model.layers: if isinstance(model, Gemma3ForConditionalGeneration) and layer.self_attn.is_sliding: # Skip layers with sliding window attention, only for Gemma3 continue + self.gate_indices[layer.self_attn.layer_idx] = len(self.gate_indices) layer.self_attn.rotary_emb = language_model.rotary_emb hooks.append(layer.self_attn.register_forward_hook(self.forward_hook, with_kwargs=True)) yield @@ -246,9 +250,10 @@ def _score_fast(self, module: nn.Module, hidden_states: torch.Tensor): Calculate the KV importance scores. """ layer_idx = int(module.layer_idx) + gate_idx = self.gate_indices[layer_idx] - self.gates[layer_idx] = self.gates[layer_idx].to(hidden_states.device) - scores = self.gates[layer_idx](hidden_states) + self.gates[gate_idx] = self.gates[gate_idx].to(hidden_states.device) + scores = self.gates[gate_idx](hidden_states) scores[:, :, : self.n_sink] = 1.0 ctx_len = scores.size(-1) diff --git a/tests/presses/test_fastkvzip_fixes.py b/tests/presses/test_fastkvzip_fixes.py index 7774173cc..0772ff04a 100644 --- a/tests/presses/test_fastkvzip_fixes.py +++ b/tests/presses/test_fastkvzip_fixes.py @@ -22,9 +22,12 @@ def get_gate_weight(model_name): config = AutoConfig.from_pretrained(model_name) config = getattr(config, "text_config", config) ngroup = config.num_attention_heads // config.num_key_value_heads + # As for the published gates, there is no gate for the sliding window layers of Gemma3 + layer_types = getattr(config, "layer_types", None) + n_gates = config.num_hidden_layers if layer_types is None else layer_types.count("full_attention") gates = [ FastKVzipGate(idx, config.hidden_size, config.num_key_value_heads, ngroup, gate_dtype, sink=16) - for idx in range(config.num_hidden_layers) + for idx in range(n_gates) ] return [gate.state_dict() for gate in gates], f"{model_name.split('/')[-1]}/q{ngroup}_dim16_sink16.pt" From 0dc3358740a7bf619acdd511a92109d5c665065f Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 02:55:10 +0200 Subject: [PATCH 38/57] =?UTF-8?q?Pass=20RoPE=20embeddings=20matching=20the?= =?UTF-8?q?=20buffered=20hidden=20states=20in=20DecodingPress=20?= =?UTF-8?q?=F0=9F=A4=96=F0=9F=A4=96=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Window-based scorers received the buffered hidden states with the RoPE embeddings of the current step only, so every buffered query was rotated at the current position. The embeddings are now buffered with the hidden states (also in CAMPress) and passed in a kwargs copy. Attention weights are no longer passed to the base press as they only cover the current query. Signed-off-by: Maximilian Jeblick --- kvpress/presses/cam_press.py | 6 +-- kvpress/presses/decoding_press.py | 33 +++++++++++------ tests/presses/test_decoding_press_fixes.py | 43 ++++++++++++++++++++++ 3 files changed, 67 insertions(+), 15 deletions(-) create mode 100644 tests/presses/test_decoding_press_fixes.py diff --git a/kvpress/presses/cam_press.py b/kvpress/presses/cam_press.py index 062ae3cbc..14df10311 100644 --- a/kvpress/presses/cam_press.py +++ b/kvpress/presses/cam_press.py @@ -247,7 +247,7 @@ def forward_hook( return output # All hidden_states_buffer code is borrowed from DecodingPress - self.hidden_states_buffer[layer_idx].append(hidden_states.detach().clone()) + self._append_to_buffer(layer_idx, hidden_states, kwargs) keys, values = extract_keys_and_values(cache, layer_idx) bsz, num_key_value_heads, seq_len, _ = keys.shape @@ -288,8 +288,8 @@ def forward_hook( # Apply compression using cumulative attention scores and buffered hidden states attn_squeezed = self._running_attn_sum[layer_idx] - buffered_hidden_states = torch.cat(self.hidden_states_buffer[layer_idx], dim=1) - keys, values = self.compress(module, buffered_hidden_states, keys, values, attn_squeezed, kwargs) + buffered_hidden_states, buffered_kwargs = self._get_buffered_inputs(layer_idx, kwargs) + keys, values = self.compress(module, buffered_hidden_states, keys, values, attn_squeezed, buffered_kwargs) # Update cache with compressed keys and values set_keys_and_values(cache, layer_idx, keys, values) diff --git a/kvpress/presses/decoding_press.py b/kvpress/presses/decoding_press.py index 019d1905f..5cdd8b14c 100644 --- a/kvpress/presses/decoding_press.py +++ b/kvpress/presses/decoding_press.py @@ -47,9 +47,9 @@ class DecodingPress(BasePress): hidden_states_buffer_size: int = 256 def __post_init__(self): - # Buffer to store hidden states during decoding (per layer) assert isinstance(self.base_press, (ScorerPress, AdaKVPress)), "DecodingPress requires a ScorerPress as input" - self.hidden_states_buffer = defaultdict(list) # Per-layer buffer + # Buffer of (hidden_states, cos, sin) during decoding (per layer) + self.hidden_states_buffer = defaultdict(list) self.layer_step_counts = defaultdict(int) # Track step count per layer assert self.compression_interval > 0, "compression_interval must be greater than 0" @@ -82,7 +82,7 @@ def compress( keys: Key cache from all previous steps including current (shape: [batch, n_heads, seq_len, head_dim]) values: Value cache from all previous steps including current (shape: [batch, n_heads, seq_len, head_dim]) attentions: Attention weights (shape varies by implementation) - kwargs: Additional keyword arguments + kwargs: Additional keyword arguments, with position_embeddings matching the buffered hidden states Returns: tuple[torch.Tensor, torch.Tensor]: Compressed (keys, values) tensors @@ -115,7 +115,7 @@ def forward_hook(self, module: nn.Module, input: list[torch.Tensor], kwargs: dic This hook: 1. Detects when we're in decoding phase (not prefilling) - 2. Accumulates hidden states in a buffer + 2. Accumulates hidden states and their RoPE embeddings in a buffer 3. Applies compression every N steps 4. Clears the buffer after compression """ @@ -129,8 +129,8 @@ def forward_hook(self, module: nn.Module, input: list[torch.Tensor], kwargs: dic # We're still in prefilling phase, don't do anything return output # print(f"Adding hidden states to buffer: {hidden_states.shape}") - # Add current hidden states to buffer for this layer - self.hidden_states_buffer[layer_idx].append(hidden_states.detach().clone()) + # Add current hidden states and their RoPE embeddings to buffer for this layer + self._append_to_buffer(layer_idx, hidden_states, kwargs) # print(f"Layer step counts: {self.layer_step_counts[layer_idx]}") self.layer_step_counts[layer_idx] += 1 @@ -144,12 +144,10 @@ def forward_hook(self, module: nn.Module, input: list[torch.Tensor], kwargs: dic keys, values = extract_keys_and_values(cache, module.layer_idx) - # Get attention weights from output - attentions = output[1] if len(output) > 1 and output[1] is not None else None - - # Apply compression using buffered hidden states for this layer - buffered_hidden_states = torch.cat(self.hidden_states_buffer[layer_idx], dim=1) - keys, values = self.compress(module, buffered_hidden_states, keys, values, attentions, kwargs) + # Apply compression using buffered hidden states for this layer. Attention weights are not used + # as they only cover the current queries. + buffered_hidden_states, buffered_kwargs = self._get_buffered_inputs(layer_idx, kwargs) + keys, values = self.compress(module, buffered_hidden_states, keys, values, None, buffered_kwargs) logger.debug(f"Applied decoding compression: " f"keys.shape: {keys.shape}, values.shape: {values.shape}") # Update cache with compressed keys and values @@ -168,6 +166,17 @@ def forward_hook(self, module: nn.Module, input: list[torch.Tensor], kwargs: dic ) return output + def _append_to_buffer(self, layer_idx: int, hidden_states: torch.Tensor, kwargs: dict): + cos, sin = kwargs["position_embeddings"] + self.hidden_states_buffer[layer_idx].append((hidden_states.detach().clone(), cos, sin)) + + def _get_buffered_inputs(self, layer_idx: int, kwargs: dict) -> tuple[torch.Tensor, dict]: + """ + Return the buffered hidden states of a layer and a copy of kwargs with the matching RoPE embeddings. + """ + hidden_states, cos, sin = (torch.cat(tensors, dim=1) for tensors in zip(*self.hidden_states_buffer[layer_idx])) + return hidden_states, {**kwargs, "position_embeddings": (cos, sin)} + def reset(self): """Reset the decoding press state.""" self.hidden_states_buffer = defaultdict(list) diff --git a/tests/presses/test_decoding_press_fixes.py b/tests/presses/test_decoding_press_fixes.py new file mode 100644 index 000000000..dafba69ed --- /dev/null +++ b/tests/presses/test_decoding_press_fixes.py @@ -0,0 +1,43 @@ +# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from dataclasses import dataclass, field + +import pytest +import torch + +from kvpress import CAMPress, DecodingPress, ScorerPress +from tests.fixtures import unit_test_model_output_attention # noqa: F401 + + +@dataclass +class RecordingKnormPress(ScorerPress): + calls: list = field(default_factory=list) + fail: bool = False + + def score(self, module, hidden_states, keys, values, attentions, kwargs): + if self.fail: + raise RuntimeError("scoring failed") + self.calls.append((hidden_states, kwargs["position_embeddings"], attentions)) + return -keys.norm(dim=-1) + + +@pytest.mark.parametrize("press_cls", [DecodingPress, CAMPress]) +@torch.no_grad() +def test_buffered_hidden_states_come_with_their_rope_embeddings( + unit_test_model_output_attention, press_cls # noqa: F811 +): + model = unit_test_model_output_attention + base_press = RecordingKnormPress() + with press_cls(base_press=base_press, compression_interval=4, target_size=12)(model): + model.generate(torch.randint(0, 1000, (1, 16), device=model.device), max_new_tokens=6, do_sample=False) + + # The first compression happens after the decoding steps at positions 16, 17, 18 and 19 + hidden_states, (cos, sin), _ = base_press.calls[0] + assert hidden_states.shape[1] == 4 + position_ids = torch.arange(16, 20, device=model.device).unsqueeze(0) + expected_cos, expected_sin = model.model.rotary_emb(hidden_states, position_ids) + torch.testing.assert_close(cos, expected_cos) + torch.testing.assert_close(sin, expected_sin) + # Attention weights of the eager model only cover the current query + assert all(attentions is None for _, _, attentions in base_press.calls) From 42b1c2d35efac842c92bcbd821244ed618cd10c1 Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 02:55:38 +0200 Subject: [PATCH 39/57] =?UTF-8?q?Reset=20DecodingPress=20state=20on=20pref?= =?UTF-8?q?ill=20=F0=9F=A4=96=F0=9F=A4=96=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The buffer and step counts of a layer survived a new prefill in the same press context (e.g. two generate calls), so the next compression mixed hidden states of the previous sequence. They are now reset per layer on prefill, as in CAMPress. Signed-off-by: Maximilian Jeblick --- kvpress/presses/decoding_press.py | 8 +++++--- tests/presses/test_decoding_press_fixes.py | 18 ++++++++++++++++-- 2 files changed, 21 insertions(+), 5 deletions(-) diff --git a/kvpress/presses/decoding_press.py b/kvpress/presses/decoding_press.py index 5cdd8b14c..ac3e79e06 100644 --- a/kvpress/presses/decoding_press.py +++ b/kvpress/presses/decoding_press.py @@ -114,8 +114,8 @@ def forward_hook(self, module: nn.Module, input: list[torch.Tensor], kwargs: dic Forward hook that manages decoding-specific compression logic. This hook: - 1. Detects when we're in decoding phase (not prefilling) - 2. Accumulates hidden states and their RoPE embeddings in a buffer + 1. Resets the state of the layer during prefilling + 2. Accumulates hidden states and their RoPE embeddings in a buffer during decoding 3. Applies compression every N steps 4. Clears the buffer after compression """ @@ -126,7 +126,9 @@ def forward_hook(self, module: nn.Module, input: list[torch.Tensor], kwargs: dic # Only operate during decoding phase (after prefilling) if is_prefilling(kwargs["cache_position"], q_len): - # We're still in prefilling phase, don't do anything + # A new sequence starts, drop the state left over from a previous one + self.hidden_states_buffer[layer_idx] = [] + self.layer_step_counts[layer_idx] = 0 return output # print(f"Adding hidden states to buffer: {hidden_states.shape}") # Add current hidden states and their RoPE embeddings to buffer for this layer diff --git a/tests/presses/test_decoding_press_fixes.py b/tests/presses/test_decoding_press_fixes.py index dafba69ed..5d400b0d0 100644 --- a/tests/presses/test_decoding_press_fixes.py +++ b/tests/presses/test_decoding_press_fixes.py @@ -5,9 +5,10 @@ import pytest import torch +from transformers import DynamicCache -from kvpress import CAMPress, DecodingPress, ScorerPress -from tests.fixtures import unit_test_model_output_attention # noqa: F401 +from kvpress import CAMPress, DecodingPress, KnormPress, ScorerPress +from tests.fixtures import unit_test_model, unit_test_model_output_attention # noqa: F401 @dataclass @@ -41,3 +42,16 @@ def test_buffered_hidden_states_come_with_their_rope_embeddings( torch.testing.assert_close(sin, expected_sin) # Attention weights of the eager model only cover the current query assert all(attentions is None for _, _, attentions in base_press.calls) + + +@torch.no_grad() +def test_state_is_reset_on_prefill(unit_test_model): # noqa: F811 + press = DecodingPress(base_press=KnormPress(), compression_interval=10, target_size=64) + input_ids = torch.randint(0, 1000, (1, 16), device=unit_test_model.device) + with press(unit_test_model): + unit_test_model.generate(input_ids, max_new_tokens=4, do_sample=False) + assert set(press.layer_step_counts.values()) == {3} + + unit_test_model(input_ids, past_key_values=DynamicCache()) + assert set(press.layer_step_counts.values()) == {0} + assert all(len(buffer) == 0 for buffer in press.hidden_states_buffer.values()) From 0ba9162180a76eff85cdbd3c8b67f9a864bd2b32 Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 03:00:15 +0200 Subject: [PATCH 40/57] =?UTF-8?q?Only=20accept=20ScorerPress=20base=20pres?= =?UTF-8?q?ses=20in=20CAMPress=20=F0=9F=A4=96=F0=9F=A4=96=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CAMPress calls base_press.score, which AdaKVPress does not implement, so the cam_adakv_snapkv registry entry could not run. It is removed from the registry. Signed-off-by: Maximilian Jeblick --- evaluation/evaluate_registry.py | 1 - kvpress/presses/cam_press.py | 5 +++-- tests/presses/test_cam_press_fixes.py | 11 +++++++++++ 3 files changed, 14 insertions(+), 3 deletions(-) create mode 100644 tests/presses/test_cam_press_fixes.py diff --git a/evaluation/evaluate_registry.py b/evaluation/evaluate_registry.py index 9552e9b07..4c0bbd767 100644 --- a/evaluation/evaluate_registry.py +++ b/evaluation/evaluate_registry.py @@ -125,7 +125,6 @@ "no_press": None, "cam_streaming_llm": CAMPress(base_press=StreamingLLMPress()), "cam_knorm": CAMPress(base_press=KnormPress()), - "cam_adakv_snapkv": CAMPress(base_press=AdaKVPress(SnapKVPress())), "cam_tova": CAMPress(base_press=TOVAPress()), "decoding_knorm": DecodingPress(base_press=KnormPress()), "decoding_streaming_llm": DecodingPress(base_press=StreamingLLMPress()), diff --git a/kvpress/presses/cam_press.py b/kvpress/presses/cam_press.py index 14df10311..2e4766a58 100644 --- a/kvpress/presses/cam_press.py +++ b/kvpress/presses/cam_press.py @@ -10,7 +10,6 @@ import torch.nn as nn from transformers.models.llama.modeling_llama import repeat_kv -from kvpress.presses.adakv_press import AdaKVPress from kvpress.presses.base_press import is_prefilling from kvpress.presses.decoding_press import DecodingPress from kvpress.presses.scorer_press import ScorerPress @@ -56,13 +55,15 @@ class CAMPress(DecodingPress): more evenly. """ - base_press: ScorerPress | AdaKVPress + base_press: ScorerPress compression_interval: int = 512 target_size: int = 2048 hidden_states_buffer_size: int = 256 merge_budget: int = 32 def __post_init__(self): + if not isinstance(self.base_press, ScorerPress): + raise ValueError(f"CAMPress requires a ScorerPress as base_press, got {type(self.base_press).__name__}") super().__post_init__() assert self.merge_budget > 0, "merge_budget must be positive " diff --git a/tests/presses/test_cam_press_fixes.py b/tests/presses/test_cam_press_fixes.py new file mode 100644 index 000000000..fbf3380b1 --- /dev/null +++ b/tests/presses/test_cam_press_fixes.py @@ -0,0 +1,11 @@ +# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest + +from kvpress import AdaKVPress, CAMPress, SnapKVPress + + +def test_cam_rejects_adakv_base_press(): + with pytest.raises(ValueError, match="requires a ScorerPress"): + CAMPress(base_press=AdaKVPress(SnapKVPress())) From 564c08d6168ca8483788c200e7c7e769fa7e70b3 Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 03:01:18 +0200 Subject: [PATCH 41/57] =?UTF-8?q?Derive=20the=20CAMPress=20kept=20and=20ev?= =?UTF-8?q?icted=20tokens=20from=20one=20sort=20=F0=9F=A4=96=F0=9F=A4=96?= =?UTF-8?q?=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The two independent topk calls could select a token as both kept and merged, or neither, when scores are tied. Without ties the selected tokens are unchanged. Signed-off-by: Maximilian Jeblick --- kvpress/presses/cam_press.py | 9 ++++---- tests/presses/test_cam_press_fixes.py | 33 ++++++++++++++++++++++++++- 2 files changed, 36 insertions(+), 6 deletions(-) diff --git a/kvpress/presses/cam_press.py b/kvpress/presses/cam_press.py index 2e4766a58..fd0c659c9 100644 --- a/kvpress/presses/cam_press.py +++ b/kvpress/presses/cam_press.py @@ -140,8 +140,10 @@ def compress( mean_scores = scores.mean(dim=1) # [bsz, seq_len] — aggregate across KV heads - evict_indices = mean_scores.topk(n_to_evict, dim=-1, largest=False).indices - evict_indices = torch.sort(evict_indices, dim=-1).values + # A single sort makes the kept and evicted tokens complementary, even with tied scores + sorted_indices = mean_scores.argsort(dim=-1, descending=True, stable=True) + kept_indices = torch.sort(sorted_indices[:, : self.target_size], dim=-1).values + evict_indices = torch.sort(sorted_indices[:, self.target_size :], dim=-1).values evict_scores = mean_scores.gather(-1, evict_indices) # Flip so later sequence positions come first; stable sort preserves this order for ties @@ -150,9 +152,6 @@ def compress( merge_indices = evict_indices.gather(-1, n_to_evict - 1 - order) merge_indices = torch.sort(merge_indices, dim=-1).values - kept_indices = mean_scores.topk(self.target_size, dim=-1).indices - kept_indices = torch.sort(kept_indices, dim=-1).values - n_to_merge = merge_indices.shape[1] target_starts = torch.searchsorted(kept_indices, merge_indices, right=True) diff --git a/tests/presses/test_cam_press_fixes.py b/tests/presses/test_cam_press_fixes.py index fbf3380b1..ec578f169 100644 --- a/tests/presses/test_cam_press_fixes.py +++ b/tests/presses/test_cam_press_fixes.py @@ -1,11 +1,42 @@ # SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import types +from dataclasses import dataclass, field + import pytest +import torch + +from kvpress import AdaKVPress, CAMPress, ScorerPress, SnapKVPress + + +@dataclass +class FixedScorePress(ScorerPress): + scores: list = field(default_factory=list) -from kvpress import AdaKVPress, CAMPress, SnapKVPress + def score(self, module, hidden_states, keys, values, attentions, kwargs): + return torch.tensor(self.scores, dtype=keys.dtype).expand(keys.shape[:3]) def test_cam_rejects_adakv_base_press(): with pytest.raises(ValueError, match="requires a ScorerPress"): CAMPress(base_press=AdaKVPress(SnapKVPress())) + + +def test_cam_merged_tokens_are_not_kept_with_tied_scores(): + scores = [3.0, 1.0, 1.0, 1.0, 1.0, 2.0, 2.0, 1.0, 1.0, 1.0, 0.0, 5.0] + seq_len, target_size = len(scores), 8 + press = CAMPress(base_press=FixedScorePress(scores=scores), target_size=target_size, merge_budget=2) + module = types.SimpleNamespace(layer_idx=0) + press.layer_step_counts[0] = seq_len - target_size # all evicted tokens are merge candidates + attentions = torch.ones(1, 1, seq_len) # merge probability of 1 + press._running_attn_sum[0] = attentions.clone() + + keys = torch.arange(seq_len, dtype=torch.float32).view(1, 1, seq_len, 1).expand(-1, -1, -1, seq_len) + values = torch.eye(seq_len).view(1, 1, seq_len, seq_len) # one-hot values identify merged tokens + new_keys, new_values = press.compress(module, torch.randn(1, 4, 8), keys, values, attentions, {}) + + kept = new_keys[0, 0, :, 0].long() + merged = torch.nonzero(new_values[0, 0] * (1 - torch.eye(seq_len)[kept]))[:, 1].unique() + assert len(kept) == target_size and len(merged) > 0 + assert not set(merged.tolist()) & set(kept.tolist()) From f25bb45719f3f94645ce8f7ea767f1cfa84d9f89 Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 03:03:49 +0200 Subject: [PATCH 42/57] =?UTF-8?q?Accumulate=20the=20CAMPress=20attention?= =?UTF-8?q?=20sum=20in=20float32=20=F0=9F=A4=96=F0=9F=A4=96=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit In bfloat16 the running sum of small attention weights stalls (2000 steps of 2e-3 sum to 1.0 instead of 4.0). Merge contributions are cast back to the cache dtype. Signed-off-by: Maximilian Jeblick --- kvpress/presses/cam_press.py | 7 ++--- tests/presses/test_cam_press_fixes.py | 39 ++++++++++++++++++++++++++- 2 files changed, 42 insertions(+), 4 deletions(-) diff --git a/kvpress/presses/cam_press.py b/kvpress/presses/cam_press.py index fd0c659c9..cc78c31ad 100644 --- a/kvpress/presses/cam_press.py +++ b/kvpress/presses/cam_press.py @@ -197,7 +197,7 @@ def compress( -1, num_key_value_heads, -1, head_dim ) - values.scatter_add_(2, scatter_idx, contributions) + values.scatter_add_(2, scatter_idx, contributions.to(values.dtype)) # Physical Pruning kept_indices_expand = kept_indices.view(bsz, 1, self.target_size, 1).expand( @@ -265,7 +265,8 @@ def forward_hook( attn_squeezed = attentions.squeeze(2) if layer_idx not in self._running_attn_sum: - self._running_attn_sum[layer_idx] = attn_squeezed.clone() + # Accumulate in float32, small attention weights vanish when summed in half precision + self._running_attn_sum[layer_idx] = attn_squeezed.to(torch.float32, copy=True) else: # Pad running sum for the new token growth prev_len = self._running_attn_sum[layer_idx].shape[-1] @@ -273,7 +274,7 @@ def forward_hook( if pad_len > 0: pad = torch.zeros( - (bsz, num_key_value_heads, pad_len), device=attn_squeezed.device, dtype=attn_squeezed.dtype + (bsz, num_key_value_heads, pad_len), device=attn_squeezed.device, dtype=torch.float32 ) self._running_attn_sum[layer_idx] = torch.cat([self._running_attn_sum[layer_idx], pad], dim=-1) diff --git a/tests/presses/test_cam_press_fixes.py b/tests/presses/test_cam_press_fixes.py index ec578f169..859a83abc 100644 --- a/tests/presses/test_cam_press_fixes.py +++ b/tests/presses/test_cam_press_fixes.py @@ -6,8 +6,9 @@ import pytest import torch +from transformers import DynamicCache, LlamaConfig, LlamaForCausalLM -from kvpress import AdaKVPress, CAMPress, ScorerPress, SnapKVPress +from kvpress import AdaKVPress, CAMPress, KnormPress, ScorerPress, SnapKVPress @dataclass @@ -40,3 +41,39 @@ def test_cam_merged_tokens_are_not_kept_with_tied_scores(): merged = torch.nonzero(new_values[0, 0] * (1 - torch.eye(seq_len)[kept]))[:, 1].unique() assert len(kept) == target_size and len(merged) > 0 assert not set(merged.tolist()) & set(kept.tolist()) + + +@torch.no_grad() +def test_cam_accumulates_attention_in_float32(): + press = CAMPress(base_press=KnormPress(), compression_interval=10**6, target_size=4) + module = types.SimpleNamespace(layer_idx=0) + cache = DynamicCache() + keys = torch.zeros(1, 1, 8, 4, dtype=torch.bfloat16) + cache.update(keys, keys.clone(), layer_idx=0) + kwargs = { + "hidden_states": torch.zeros(1, 1, 4, dtype=torch.bfloat16), + "past_key_values": cache, + "cache_position": torch.tensor([8]), + "position_embeddings": (torch.ones(1, 1, 4), torch.zeros(1, 1, 4)), + } + attentions = torch.full((1, 1, 1, 8), 2e-3, dtype=torch.bfloat16) + for _ in range(2000): + press.forward_hook(module, [], kwargs, (None, attentions)) + torch.testing.assert_close(press._running_attn_sum[0], torch.full((1, 1, 8), 4.0), rtol=1e-2, atol=0) + + +@torch.no_grad() +def test_cam_compresses_half_precision_caches(): + config = LlamaConfig( + hidden_size=64, + intermediate_size=128, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + vocab_size=128, + ) + model = LlamaForCausalLM(config).to(torch.bfloat16).eval() + cache = DynamicCache() + with CAMPress(base_press=KnormPress(), compression_interval=2, target_size=12)(model): + model.generate(torch.randint(0, 128, (1, 16)), past_key_values=cache, max_new_tokens=5, do_sample=False) + assert cache.get_seq_length() == 12 From 94ba4f0f8228df7105b9beb5479952fccbd5a230 Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 03:08:06 +0200 Subject: [PATCH 43/57] =?UTF-8?q?Count=20the=20tokens=20seen=20in=20Compre?= =?UTF-8?q?ssionRatioDecodingPress=20=F0=9F=A4=96=F0=9F=A4=96=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The target size was derived from position_ids, which models fill from cache_position when they are not passed. cache_position counts the compressed cache, so the target shrank geometrically (67, 34, 19, ... instead of ~44) and the documented NotImplementedError was unreachable. The press now counts the prefill and decoding tokens per layer (reset on prefill), starting from the position of the first decoding step when the prefill ran outside of the press context. The tests of the removed position_ids logic are replaced by tests/presses/test_compression_ratio_decoding_press.py. Signed-off-by: Maximilian Jeblick --- .../compression_ratio_decoding_press.py | 38 ++++++++++++------ kvpress/presses/decoding_press.py | 6 +-- .../test_compression_ratio_decoding_press.py | 32 +++++++++++++++ tests/test_decoding_compression.py | 39 ------------------- 4 files changed, 62 insertions(+), 53 deletions(-) create mode 100644 tests/presses/test_compression_ratio_decoding_press.py diff --git a/kvpress/presses/compression_ratio_decoding_press.py b/kvpress/presses/compression_ratio_decoding_press.py index d27017b60..ad0718110 100644 --- a/kvpress/presses/compression_ratio_decoding_press.py +++ b/kvpress/presses/compression_ratio_decoding_press.py @@ -3,6 +3,10 @@ from dataclasses import dataclass, field +import torch +import torch.nn as nn + +from kvpress.presses.base_press import is_prefilling from kvpress.presses.decoding_press import DecodingPress @@ -12,8 +16,7 @@ class CompressionRatioDecodingPress(DecodingPress): A decoding press that keeps a fixed fraction of all tokens seen so far. Unlike `DecodingPress`, which compresses to a fixed absolute `target_size`, - this subclass derives the target size from the full sequence length observed - so far during decoding. + this subclass derives the target size from the number of tokens seen so far. Parameters ---------- @@ -28,8 +31,8 @@ class CompressionRatioDecodingPress(DecodingPress): Notes ----- - This press requires logical `position_ids` to be passed through the model - forward call. Otherwise an exception will be raised. + The press counts the tokens of the prefill and of every decoding step. If the prefill ran outside of the press + context (as in the pipeline), the count starts from the position of the first decoding step. """ target_compression_ratio: float = 0.5 @@ -38,13 +41,26 @@ class CompressionRatioDecodingPress(DecodingPress): def __post_init__(self): super().__post_init__() assert 0 <= self.target_compression_ratio < 1, "target_compression_ratio must be between 0 and 1" + self.tokens_seen: dict[int, int] = {} # Per-layer number of tokens seen - def _resolve_target_size(self, kwargs: dict) -> int: - total_tokens_seen = self._resolve_total_tokens_seen(kwargs) - return max(1, int(total_tokens_seen * (1 - self.target_compression_ratio))) + def forward_hook(self, module: nn.Module, input: list[torch.Tensor], kwargs: dict, output: list): + layer_idx = int(module.layer_idx) + q_len = kwargs["hidden_states"].shape[1] + if is_prefilling(kwargs["cache_position"], q_len): + self.tokens_seen[layer_idx] = q_len + else: + if layer_idx not in self.tokens_seen: + # The prefill was not seen, count the tokens before the first position of this step + positions = kwargs.get("position_ids") + if positions is None: + positions = kwargs["cache_position"] + self.tokens_seen[layer_idx] = int(positions[..., 0].max()) + self.tokens_seen[layer_idx] += q_len + return super().forward_hook(module, input, kwargs, output) - def _resolve_total_tokens_seen(self, kwargs: dict) -> int: - if "position_ids" in kwargs and kwargs["position_ids"] is not None: - return int(kwargs["position_ids"].max().item()) + 1 + def _resolve_target_size(self, layer_idx: int) -> int: + return max(1, int(self.tokens_seen[layer_idx] * (1 - self.target_compression_ratio))) - raise NotImplementedError("CompressionRatioDecodingPress requires logical position_ids in kwargs") + def reset(self): + super().reset() + self.tokens_seen = {} diff --git a/kvpress/presses/decoding_press.py b/kvpress/presses/decoding_press.py index ac3e79e06..ca0e90ea5 100644 --- a/kvpress/presses/decoding_press.py +++ b/kvpress/presses/decoding_press.py @@ -99,7 +99,7 @@ def compress( storing existing scores in a buffer (e.g. KNormPress) and reusing them in subsequent compressions. """ k_len = keys.shape[2] - target_size = self._resolve_target_size(kwargs) + target_size = self._resolve_target_size(module.layer_idx) target_compression_ratio = self._find_target_compression_ratio(k_len, target_size) logger.debug(f"Compressing {k_len} to {target_size} with ratio {target_compression_ratio}") @@ -138,7 +138,7 @@ def forward_hook(self, module: nn.Module, input: list[torch.Tensor], kwargs: dic self.layer_step_counts[layer_idx] += 1 # Apply compression if we've reached the compression step threshold - target_size = self._resolve_target_size(kwargs) + target_size = self._resolve_target_size(layer_idx) if (self.layer_step_counts[layer_idx] >= self.compression_interval) or (q_len >= target_size): logger.debug( f"Applying decoding compression: layer_step_count ({self.layer_step_counts[layer_idx]}) >= compression_steps ({self.compression_interval})" # noqa: E501 @@ -236,5 +236,5 @@ def _find_target_compression_ratio(self, q_len: int, target_tokens: int) -> floa return ratio - def _resolve_target_size(self, kwargs: dict) -> int: + def _resolve_target_size(self, layer_idx: int) -> int: return self.target_size diff --git a/tests/presses/test_compression_ratio_decoding_press.py b/tests/presses/test_compression_ratio_decoding_press.py new file mode 100644 index 000000000..1ae2778d1 --- /dev/null +++ b/tests/presses/test_compression_ratio_decoding_press.py @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch +from transformers import DynamicCache + +from kvpress import CompressionRatioDecodingPress, KnormPress +from tests.fixtures import unit_test_model # noqa: F401 + + +@pytest.mark.parametrize("prefill_in_context", [True, False]) +@pytest.mark.parametrize("logical_position_ids", [True, False]) +@torch.no_grad() +def test_target_size_follows_all_tokens_seen(unit_test_model, prefill_in_context, logical_position_ids): # noqa: F811 + press = CompressionRatioDecodingPress(base_press=KnormPress(), target_compression_ratio=0.5, compression_interval=8) + cache = DynamicCache() + input_ids = torch.randint(0, 1000, (1, 64), device=unit_test_model.device) + if not prefill_in_context: + logits = unit_test_model(input_ids, past_key_values=cache).logits + + with press(unit_test_model): + if prefill_in_context: + logits = unit_test_model(input_ids, past_key_values=cache).logits + for step in range(64): + position_ids = torch.tensor([[64 + step]], device=input_ids.device) if logical_position_ids else None + next_ids = logits[:, -1:].argmax(dim=-1) + logits = unit_test_model(next_ids, past_key_values=cache, position_ids=position_ids).logits + + # 128 tokens were seen when the last compression happened + assert cache.get_seq_length() == 64 + assert press.tokens_seen == {} diff --git a/tests/test_decoding_compression.py b/tests/test_decoding_compression.py index 8d79c729b..70e3e6312 100644 --- a/tests/test_decoding_compression.py +++ b/tests/test_decoding_compression.py @@ -13,7 +13,6 @@ from kvpress import ( CAMPress, CompactorPress, - CompressionRatioDecodingPress, DecodingPress, KnormPress, KVzapPress, @@ -233,44 +232,6 @@ def test_decoding_press_equivalence(press_factory): ) -def test_sequence_ratio_decoding_press_uses_total_tokens_seen(): - press = CompressionRatioDecodingPress(base_press=KnormPress(), target_compression_ratio=0.5) - - target_size = press._resolve_target_size({"position_ids": torch.tensor([[107]])}) - - assert target_size == 54 - assert press._find_target_compression_ratio(58, target_size) == pytest.approx(1 - (54 / 58)) - - -def test_sequence_ratio_decoding_press_prefers_logical_positions_over_compressed_cache_position(): - press = CompressionRatioDecodingPress(base_press=KnormPress(), target_compression_ratio=0.5) - - target_size = press._resolve_target_size( - { - "position_ids": torch.tensor([[107]]), - "cache_position": torch.tensor([57]), - } - ) - - assert target_size == 54 - - -def test_sequence_ratio_decoding_press_requires_logical_position_ids(): - press = CompressionRatioDecodingPress(base_press=KnormPress(), target_compression_ratio=0.5) - - with pytest.raises(NotImplementedError, match="requires logical position_ids"): - press._resolve_target_size({"cache_position": torch.tensor([57])}) - - -def test_sequence_ratio_decoding_press_noops_if_cache_is_already_below_target(): - press = CompressionRatioDecodingPress(base_press=KnormPress(), target_compression_ratio=0.5) - - target_size = press._resolve_target_size({"position_ids": torch.tensor([[107]])}) - - assert target_size == 54 - assert press._find_target_compression_ratio(50, target_size) == 0.0 - - """ E AttributeError: 'QFilterPress' object has no attribute 'q_filters' E Failed: DecodingPress failed with SnapKVPress: shape '[1, 2, 2, 6]' is invalid for input of size 12 From f137547b66ecc156cd1f4386ac60034c5d737522 Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 12:46:48 +0200 Subject: [PATCH 44/57] =?UTF-8?q?Enter=20the=20inner=20press=20contexts=20?= =?UTF-8?q?in=20PrefillDecodingPress=20=F0=9F=A4=96=F0=9F=A4=96?= =?UTF-8?q?=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Only the inner forward hooks were called, never the inner __call__, so presses relying on it failed (e.g. FinchPress: window_size must be provided). Both inner contexts are now entered with an ExitStack, which is enough since BasePress hooks only compress during prefilling and DecodingPress hooks only during decoding. Signed-off-by: Maximilian Jeblick --- kvpress/presses/prefill_decoding_press.py | 9 +++++-- .../test_prefill_decoding_press_fixes.py | 26 +++++++++++++++++++ 2 files changed, 33 insertions(+), 2 deletions(-) create mode 100644 tests/presses/test_prefill_decoding_press_fixes.py diff --git a/kvpress/presses/prefill_decoding_press.py b/kvpress/presses/prefill_decoding_press.py index 973576576..be1c64b57 100644 --- a/kvpress/presses/prefill_decoding_press.py +++ b/kvpress/presses/prefill_decoding_press.py @@ -2,7 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 import logging -from contextlib import contextmanager +from contextlib import ExitStack, contextmanager from dataclasses import dataclass from typing import Optional @@ -82,8 +82,13 @@ def forward_hook(self, module: nn.Module, input: list[torch.Tensor], kwargs: dic @contextmanager def __call__(self, model: PreTrainedModel): + # Entering both contexts is enough: BasePress hooks only compress during prefilling, + # and DecodingPress hooks only during decoding. try: - with super().__call__(model): + with ExitStack() as stack: + for press in (self.prefilling_press, self.decoding_press): + if press is not None: + stack.enter_context(press(model)) yield finally: # Reset decoding press if it exists diff --git a/tests/presses/test_prefill_decoding_press_fixes.py b/tests/presses/test_prefill_decoding_press_fixes.py new file mode 100644 index 000000000..4c8d2a721 --- /dev/null +++ b/tests/presses/test_prefill_decoding_press_fixes.py @@ -0,0 +1,26 @@ +# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import torch +from transformers import DynamicCache + +from kvpress import DecodingPress, FinchPress, KnormPress, PrefillDecodingPress +from tests.fixtures import unit_test_model # noqa: F401 + + +@torch.no_grad() +def test_prefilling_press_context_is_entered(unit_test_model): # noqa: F811 + finch_press = FinchPress(compression_ratio=0.5) + finch_press.delimiter_token_id = unit_test_model.config.eos_token_id + press = PrefillDecodingPress( + prefilling_press=finch_press, + decoding_press=DecodingPress(base_press=KnormPress(), compression_interval=2, target_size=8), + ) + input_ids = torch.arange(10, 30, device=unit_test_model.device) + input_ids[15] = finch_press.delimiter_token_id + + cache = DynamicCache() + with press(unit_test_model): + unit_test_model(input_ids.unsqueeze(0), past_key_values=cache) + # The delimiter is removed and half of the remaining 19 tokens are kept + assert cache.get_seq_length() == 9 From 8f26624ef05392790dd47bc3fae5bf06607c5677 Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 03:16:27 +0200 Subject: [PATCH 45/57] =?UTF-8?q?Reset=20DMSPress=20state=20per=20layer=20?= =?UTF-8?q?on=20prefill=20=F0=9F=A4=96=F0=9F=A4=96=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The state was only reset when layer 0 prefilled, which never happens on Gemma3 where layer 0 is a skipped sliding window layer. Signed-off-by: Maximilian Jeblick --- kvpress/presses/dms_press.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/kvpress/presses/dms_press.py b/kvpress/presses/dms_press.py index 2a2e1c636..d54e19c36 100644 --- a/kvpress/presses/dms_press.py +++ b/kvpress/presses/dms_press.py @@ -78,10 +78,10 @@ def forward_hook(self, module: nn.Module, input: list[torch.Tensor], kwargs: dic # Extract layer index as int for type safety layer_idx: int = module.layer_idx # type: ignore[assignment] - # Reset the scores buffer and compression ratios if we are in prefilling - if prefilling and (layer_idx == 0): - self.scores_buffer.clear() - self.compression_ratios.clear() + # Reset the state of this layer if we are in prefilling (layer 0 is not hooked for every model, e.g. Gemma3) + if prefilling: + self.scores_buffer.pop(layer_idx, None) + self.compression_ratios.pop(layer_idx, None) # Skip compression during decoding if not enabled if not prefilling and not self.decoding: From e98f0c75450ed5133f8b624a1a2b5f0490f8a126 Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 03:18:33 +0200 Subject: [PATCH 46/57] =?UTF-8?q?Fix=20the=20SimLayerKVPress=20recent=20wi?= =?UTF-8?q?ndow=20when=20n=5Frecent=20<=3D=20n=5Flast=20=F0=9F=A4=96?= =?UTF-8?q?=F0=9F=A4=96=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit keys[:, :, -n_recent + n_last:] becomes keys[:, :, 0:] when n_recent == n_last (the config of tests/default_presses.py), keeping the whole cache plus duplicated initial tokens. The recent window is now sliced from k_len - max(0, n_recent - n_last), and the reported ratio is computed from the kept length instead of assuming n_last == 1. Signed-off-by: Maximilian Jeblick --- kvpress/presses/simlayerkv_press.py | 9 ++++---- tests/presses/test_simlayerkv_fixes.py | 29 ++++++++++++++++++++++++++ 2 files changed, 34 insertions(+), 4 deletions(-) create mode 100644 tests/presses/test_simlayerkv_fixes.py diff --git a/kvpress/presses/simlayerkv_press.py b/kvpress/presses/simlayerkv_press.py index 3c1148148..2add7356a 100644 --- a/kvpress/presses/simlayerkv_press.py +++ b/kvpress/presses/simlayerkv_press.py @@ -106,10 +106,11 @@ def compress( # Compression if self.is_lazy(module, hidden_states, keys, kwargs["position_embeddings"]): - # If layer is lazy, only keep the initial and recent KV pairs - keys = torch.cat([keys[:, :, : self.n_initial], keys[:, :, -self.n_recent + self.n_last :]], dim=2) - values = torch.cat([values[:, :, : self.n_initial], values[:, :, -self.n_recent + self.n_last :]], dim=2) - self.compression_ratios.append((k_len - self.n_initial - self.n_recent + 1) / k_len) + # If layer is lazy, only keep the initial and the n_recent - n_last most recent KV pairs + recent_start = k_len - max(0, self.n_recent - self.n_last) + keys = torch.cat([keys[:, :, : self.n_initial], keys[:, :, recent_start:]], dim=2) + values = torch.cat([values[:, :, : self.n_initial], values[:, :, recent_start:]], dim=2) + self.compression_ratios.append(1 - keys.shape[2] / k_len) else: self.compression_ratios.append(0.0) diff --git a/tests/presses/test_simlayerkv_fixes.py b/tests/presses/test_simlayerkv_fixes.py new file mode 100644 index 000000000..1ffb7c2f7 --- /dev/null +++ b/tests/presses/test_simlayerkv_fixes.py @@ -0,0 +1,29 @@ +# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch +from transformers import DynamicCache + +from kvpress import SimLayerKVPress +from tests.fixtures import unit_test_model # noqa: F401 + + +@pytest.mark.parametrize("n_recent, n_last", [(1, 1), (1, 2), (4, 1), (6, 2)]) +@torch.no_grad() +def test_lazy_layers_keep_initial_and_recent_tokens(unit_test_model, n_recent, n_last): # noqa: F811 + input_ids = torch.randint(0, 1000, (1, 32), device=unit_test_model.device) + reference_cache = DynamicCache() + unit_test_model(input_ids, past_key_values=reference_cache) + + # lazy_threshold=0 makes every layer lazy + press = SimLayerKVPress(lazy_threshold=0.0, n_initial=2, n_recent=n_recent, n_last=n_last) + cache = DynamicCache() + with press(unit_test_model): + unit_test_model(input_ids, past_key_values=cache) + + n_recent_kept = max(0, n_recent - n_last) + kept_positions = list(range(2)) + list(range(32 - n_recent_kept, 32)) + for layer, reference_layer in zip(cache.layers, reference_cache.layers): + torch.testing.assert_close(layer.keys, reference_layer.keys[:, :, kept_positions]) + assert press.compression_ratio == pytest.approx(1 - len(kept_positions) / 32) From 065d854aef810ec0b8b546c376cace5ca5d341be Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 12:47:34 +0200 Subject: [PATCH 47/57] =?UTF-8?q?Compute=20MergingPress=20similarities=20i?= =?UTF-8?q?n=20chunks=20of=20evicted=20tokens=20=F0=9F=A4=96=F0=9F=A4=96?= =?UTF-8?q?=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The float32 (batch, heads, n_evict, n_kept) similarity matrix needed about 8.6 GB per layer at 32k context. It is now computed for chunk_size=1024 evicted tokens at a time, which gives the same nearest survivors (identical for up to chunk_size evicted tokens, up to float32 rounding of the matrix product otherwise). Signed-off-by: Maximilian Jeblick --- kvpress/presses/merging_press.py | 20 +++++++++---- tests/presses/test_merging_press_chunks.py | 34 ++++++++++++++++++++++ 2 files changed, 49 insertions(+), 5 deletions(-) create mode 100644 tests/presses/test_merging_press_chunks.py diff --git a/kvpress/presses/merging_press.py b/kvpress/presses/merging_press.py index 0c0c15136..a614cda8b 100644 --- a/kvpress/presses/merging_press.py +++ b/kvpress/presses/merging_press.py @@ -38,16 +38,20 @@ class MergingPress(BasePress): merge_fraction : float, default=1.0 Fraction of evicted tokens (ranked by similarity) that are merged. Task-dependent: 1.0 wins on retrieval, 0.75 wins on extraction. + chunk_size : int, default=1024 + Number of evicted tokens whose similarities to the survivors are computed at once. + Bounds the memory of the similarity matrix, results only change by floating point rounding. """ press: ScorerPress = None # type: ignore[assignment] similarity_threshold: float = 0.0 merge_fraction: float = 1.0 + chunk_size: int = 1024 def __post_init__(self): - assert isinstance(self.press, ScorerPress), ( - f"MergingPress requires a ScorerPress, got {type(self.press).__name__}" - ) + assert isinstance( + self.press, ScorerPress + ), f"MergingPress requires a ScorerPress, got {type(self.press).__name__}" assert 0.0 <= self.similarity_threshold <= 1.0 assert 0.0 < self.merge_fraction <= 1.0, "merge_fraction must be in (0, 1]" @@ -134,10 +138,16 @@ def merge( kept_values = values.gather(2, keep_idx) evict_values = values.gather(2, evict_idx) - # Cosine similarity → nearest survivor (per evicted token, batched over B, H) + # Cosine similarity → nearest survivor (per evicted token, batched over B, H), + # computed for chunks of evicted tokens to avoid materializing the (n_evict, n_kept) matrix kept_keys = kept_keys / kept_keys.norm(dim=-1, keepdim=True).clamp(min=_EPS) evict_keys = evict_keys / evict_keys.norm(dim=-1, keepdim=True).clamp(min=_EPS) - max_sim, target = (evict_keys @ kept_keys.transpose(-2, -1)).max(dim=-1) + nearest = [ + (evict_keys[:, :, start : start + self.chunk_size] @ kept_keys.transpose(-2, -1)).max(dim=-1) + for start in range(0, n_evict, self.chunk_size) + ] + max_sim = torch.cat([chunk.values for chunk in nearest], dim=-1) + target = torch.cat([chunk.indices for chunk in nearest], dim=-1) # Threshold gate merge_ok = max_sim >= self.similarity_threshold diff --git a/tests/presses/test_merging_press_chunks.py b/tests/presses/test_merging_press_chunks.py new file mode 100644 index 000000000..de1609a8c --- /dev/null +++ b/tests/presses/test_merging_press_chunks.py @@ -0,0 +1,34 @@ +# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch + +from kvpress import KnormPress +from kvpress.presses.merging_press import MergingPress + + +@pytest.mark.parametrize("similarity_threshold, merge_fraction", [(0.0, 1.0), (0.1, 0.75)]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_chunked_merge_matches_unchunked_merge(similarity_threshold, merge_fraction, dtype): + torch.manual_seed(0) + bsz, num_heads, seq_len, head_dim, n_kept = 2, 3, 50, 8, 20 + keys, values = torch.randn(2, bsz, num_heads, seq_len, head_dim, dtype=dtype).unbind(0) + indices = torch.rand(bsz, num_heads, seq_len).topk(n_kept, dim=-1).indices + + def merge(chunk_size): + press = MergingPress( + KnormPress(compression_ratio=0.6), + similarity_threshold=similarity_threshold, + merge_fraction=merge_fraction, + chunk_size=chunk_size, + ) + return press.merge(keys, values, indices) + + # Matrix products of different shapes may round differently, hence assert_close + reference_keys, reference_values = merge(chunk_size=seq_len) + for chunk_size in [1, 3, 7]: + new_keys, new_values = merge(chunk_size) + assert torch.equal(new_keys, reference_keys) + torch.testing.assert_close(new_values, reference_values) + assert torch.equal(merge(chunk_size=seq_len - n_kept)[1], reference_values) From 72ccb20877fa3abe595c7939cfe1fdde840760ba Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 03:25:07 +0200 Subject: [PATCH 48/57] =?UTF-8?q?Always=20keep=20the=20FinchPress=20window?= =?UTF-8?q?=20in=20chunked=20compression=20=F0=9F=A4=96=F0=9F=A4=96?= =?UTF-8?q?=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Chunks covered the whole sequence, so a remainder chunk made only of window (question) tokens evicted some of them. Chunks now only cover the context and the window is always kept. Signed-off-by: Maximilian Jeblick --- kvpress/presses/finch_press.py | 8 ++++--- tests/presses/test_finch_press_fixes.py | 31 +++++++++++++++++++++++++ 2 files changed, 36 insertions(+), 3 deletions(-) create mode 100644 tests/presses/test_finch_press_fixes.py diff --git a/kvpress/presses/finch_press.py b/kvpress/presses/finch_press.py index af2219fde..22d9157c1 100644 --- a/kvpress/presses/finch_press.py +++ b/kvpress/presses/finch_press.py @@ -101,13 +101,15 @@ def compress(self, module, hidden_states, keys, values, attentions, kwargs): n_kept = compute_n_kept(k_len, self.compression_ratio) indices = scores.topk(n_kept, dim=-1).indices else: - assert self.chunk_length > self.window_size / (1 - self.compression_ratio) + # Chunks only cover the context, the window is always kept + context_length = k_len - self.window_size indices = [] - for i in range(0, k_len, self.chunk_length): - chunk_scores = scores[:, :, i : i + self.chunk_length] + for i in range(0, context_length, self.chunk_length): + chunk_scores = scores[:, :, i : min(i + self.chunk_length, context_length)] n_kept = compute_n_kept(chunk_scores.shape[2], self.compression_ratio) chunk_indices = i + chunk_scores.topk(n_kept, dim=-1).indices indices.append(chunk_indices) + indices.append(torch.arange(context_length, k_len, device=keys.device).expand(*scores.shape[:2], -1)) indices = torch.cat(indices, dim=-1) if self.rerotate_keys: indices = torch.sort(indices, dim=2).values diff --git a/tests/presses/test_finch_press_fixes.py b/tests/presses/test_finch_press_fixes.py new file mode 100644 index 000000000..64b7b59c4 --- /dev/null +++ b/tests/presses/test_finch_press_fixes.py @@ -0,0 +1,31 @@ +# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import types + +import pytest +import torch + +from kvpress import FinchPress +from kvpress.utils import compute_n_kept + + +def make_module(num_heads): + return types.SimpleNamespace(head_dim=1, config=types.SimpleNamespace(num_attention_heads=num_heads)) + + +@pytest.mark.parametrize("context_length, window_size, chunk_length", [(10, 3, 5), (12, 4, 5), (7, 5, 3)]) +def test_chunked_compression_keeps_the_window(context_length, window_size, chunk_length): + torch.manual_seed(0) + k_len = context_length + window_size + press = FinchPress(compression_ratio=0.5, chunk_length=chunk_length, rerotate_keys=False) + press.window_size = window_size + keys = torch.arange(k_len, dtype=torch.float32).view(1, 1, k_len, 1) + attentions = torch.rand(1, 1, k_len, k_len).softmax(dim=-1) + + new_keys, _ = press.compress(make_module(1), None, keys, keys.clone(), attentions, {}) + + kept = new_keys[0, 0, :, 0].long().tolist() + assert set(range(context_length, k_len)) <= set(kept) + chunk_lengths = [min(chunk_length, context_length - i) for i in range(0, context_length, chunk_length)] + assert len(kept) == sum(compute_n_kept(length, 0.5) for length in chunk_lengths) + window_size From f4ef15e477c6552e7cfeb87fb305b8660844df7c Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 03:26:04 +0200 Subject: [PATCH 49/57] =?UTF-8?q?Reset=20the=20FinchPress=20window=20size?= =?UTF-8?q?=20at=20each=20prefill=20=F0=9F=A4=96=F0=9F=A4=96=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A sample without delimiter silently reused the window size of the previous sample. It now fails with 'window_size must be provided'. Signed-off-by: Maximilian Jeblick --- kvpress/presses/finch_press.py | 3 +++ tests/presses/test_finch_press_fixes.py | 14 ++++++++++++++ 2 files changed, 17 insertions(+) diff --git a/kvpress/presses/finch_press.py b/kvpress/presses/finch_press.py index 22d9157c1..41a565bf4 100644 --- a/kvpress/presses/finch_press.py +++ b/kvpress/presses/finch_press.py @@ -127,6 +127,9 @@ def embed_token_forward_hook(self, module, input, output): """ Forward hook to detect a delimiter token between the context and the window """ + if input[0].shape[1] > 1: # prefilling + # Reset the window size so that an input without delimiter can't reuse the previous one + self.window_size = None if input[0].shape[1] > 1 and self.delimiter_token_id in input[0][0]: # prefilling assert len(input[0]) == 1, "Only batch size 1 is supported." # Find the delimiter token and compute the window size diff --git a/tests/presses/test_finch_press_fixes.py b/tests/presses/test_finch_press_fixes.py index 64b7b59c4..f6858c9c0 100644 --- a/tests/presses/test_finch_press_fixes.py +++ b/tests/presses/test_finch_press_fixes.py @@ -8,6 +8,7 @@ from kvpress import FinchPress from kvpress.utils import compute_n_kept +from tests.fixtures import unit_test_model # noqa: F401 def make_module(num_heads): @@ -29,3 +30,16 @@ def test_chunked_compression_keeps_the_window(context_length, window_size, chunk assert set(range(context_length, k_len)) <= set(kept) chunk_lengths = [min(chunk_length, context_length - i) for i in range(0, context_length, chunk_length)] assert len(kept) == sum(compute_n_kept(length, 0.5) for length in chunk_lengths) + window_size + + +@torch.no_grad() +def test_window_size_is_not_reused_across_samples(unit_test_model): # noqa: F811 + press = FinchPress(compression_ratio=0.5) + press.delimiter_token_id = unit_test_model.config.eos_token_id + input_ids = torch.arange(10, 30, device=unit_test_model.device) + input_ids_with_delimiter = input_ids.clone() + input_ids_with_delimiter[15] = press.delimiter_token_id + with press(unit_test_model): + unit_test_model(input_ids_with_delimiter.unsqueeze(0)) + with pytest.raises(AssertionError, match="window_size must be provided"): + unit_test_model(input_ids.unsqueeze(0)) From 46c7c2da27aa7115b02265774eee2bd5819db506 Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 03:26:45 +0200 Subject: [PATCH 50/57] =?UTF-8?q?Compute=20FinchPress=20normalized=20score?= =?UTF-8?q?s=20in=20float32=20=F0=9F=A4=96=F0=9F=A4=96=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The attention weights multiplied by the number of non-zero weights overflow float16 for long contexts (0.9 * 80000 = inf). The formula is unchanged. Signed-off-by: Maximilian Jeblick --- kvpress/presses/finch_press.py | 3 ++- tests/presses/test_finch_press_fixes.py | 11 +++++++++++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/kvpress/presses/finch_press.py b/kvpress/presses/finch_press.py index 41a565bf4..4fc67e000 100644 --- a/kvpress/presses/finch_press.py +++ b/kvpress/presses/finch_press.py @@ -72,7 +72,8 @@ def score(self, module, hidden_states, keys, values, attentions, kwargs): if self.normalize_scores: non_zero_counts = torch.arange(k_len - self.window_size, k_len)[None, None, :, None] non_zero_counts = non_zero_counts.to(attn_weights.device) - attn_weights = attn_weights * non_zero_counts + # float32 as the products overflow float16 for long contexts + attn_weights = attn_weights.float() * non_zero_counts # Average per group scores = attn_weights.mean(dim=-2) diff --git a/tests/presses/test_finch_press_fixes.py b/tests/presses/test_finch_press_fixes.py index f6858c9c0..2cab8a75d 100644 --- a/tests/presses/test_finch_press_fixes.py +++ b/tests/presses/test_finch_press_fixes.py @@ -32,6 +32,17 @@ def test_chunked_compression_keeps_the_window(context_length, window_size, chunk assert len(kept) == sum(compute_n_kept(length, 0.5) for length in chunk_lengths) + window_size +def test_normalized_scores_do_not_overflow_in_float16(): + k_len = 80_001 + press = FinchPress(compression_ratio=0.5) + press.window_size = 1 + keys = torch.zeros(1, 1, k_len, 1, dtype=torch.float16) + attentions = torch.full((1, 1, 1, k_len), 0.9, dtype=torch.float16) + scores = press.score(make_module(1), None, keys, keys, attentions, {}) + assert torch.isfinite(scores).all() + torch.testing.assert_close(scores[0, 0, 0], torch.tensor(0.9 * 80_000), rtol=1e-3, atol=0) + + @torch.no_grad() def test_window_size_is_not_reused_across_samples(unit_test_model): # noqa: F811 press = FinchPress(compression_ratio=0.5) From 7c8d0647086993bc65bd7cbf8d29ba8b500ee5bb Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 03:30:13 +0200 Subject: [PATCH 51/57] =?UTF-8?q?Resize=20FinchPress=20embeddings=20only?= =?UTF-8?q?=20when=20needed=20and=20use=20get=5Finput=5Fembeddings=20?= =?UTF-8?q?=F0=9F=A4=96=F0=9F=A4=96=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit resize_token_embeddings(len(tokenizer)) shrank the embeddings of models whose vocabulary is padded (Qwen2.5: 151,936 -> 151,666 rows) on the shared model, so they are only resized when the delimiter id does not fit. model.model.embed_tokens does not exist on Gemma3ForConditionalGeneration, model.get_input_embeddings() is used instead. Signed-off-by: Maximilian Jeblick --- kvpress/presses/finch_press.py | 7 ++- tests/presses/test_finch_press_fixes.py | 84 +++++++++++++++++++++++++ 2 files changed, 88 insertions(+), 3 deletions(-) diff --git a/kvpress/presses/finch_press.py b/kvpress/presses/finch_press.py index 4fc67e000..81d5cbe04 100644 --- a/kvpress/presses/finch_press.py +++ b/kvpress/presses/finch_press.py @@ -152,8 +152,9 @@ def update_model_and_tokenizer(self, model, tokenizer, delimiter_token: str = "< if delimiter_token not in tokenizer.get_vocab(): tokenizer.add_special_tokens({"additional_special_tokens": [delimiter_token]}) self.delimiter_token_id = tokenizer.convert_tokens_to_ids(delimiter_token) # type: ignore - # update model embeddings - model.resize_token_embeddings(len(tokenizer)) + # update model embeddings, len(tokenizer) can be smaller than the embeddings (e.g. Qwen2.5) + if self.delimiter_token_id >= model.get_input_embeddings().num_embeddings: + model.resize_token_embeddings(len(tokenizer)) return tokenizer @contextmanager @@ -167,7 +168,7 @@ def __call__(self, model): with super().__call__(model): try: - hook = model.model.embed_tokens.register_forward_hook(self.embed_token_forward_hook) + hook = model.get_input_embeddings().register_forward_hook(self.embed_token_forward_hook) yield finally: hook.remove() diff --git a/tests/presses/test_finch_press_fixes.py b/tests/presses/test_finch_press_fixes.py index 2cab8a75d..0092741df 100644 --- a/tests/presses/test_finch_press_fixes.py +++ b/tests/presses/test_finch_press_fixes.py @@ -5,6 +5,18 @@ import pytest import torch +from transformers import ( + AutoTokenizer, + DynamicCache, + Gemma3Config, + Gemma3ForConditionalGeneration, + Gemma3TextConfig, + LlamaConfig, + LlamaForCausalLM, + Qwen2Config, + Qwen2ForCausalLM, + SiglipVisionConfig, +) from kvpress import FinchPress from kvpress.utils import compute_n_kept @@ -54,3 +66,75 @@ def test_window_size_is_not_reused_across_samples(unit_test_model): # noqa: F81 unit_test_model(input_ids_with_delimiter.unsqueeze(0)) with pytest.raises(AssertionError, match="window_size must be provided"): unit_test_model(input_ids.unsqueeze(0)) + + +def test_update_model_and_tokenizer_does_not_shrink_embeddings(): + tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct") + config = Qwen2Config( + hidden_size=16, intermediate_size=32, num_hidden_layers=1, num_attention_heads=2, num_key_value_heads=1 + ) + vocab_size = config.vocab_size + assert vocab_size > len(tokenizer) + 1 + model = Qwen2ForCausalLM(config) + + press = FinchPress() + press.update_model_and_tokenizer(model, tokenizer) + assert model.get_input_embeddings().num_embeddings == vocab_size + assert press.delimiter_token_id == len(tokenizer) - 1 + + +def test_update_model_and_tokenizer_grows_embeddings_when_needed(): + tokenizer = AutoTokenizer.from_pretrained("MaxJeblick/llama2-0b-unit-test") + vocab_size = len(tokenizer) + config = LlamaConfig( + hidden_size=16, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=1, + vocab_size=vocab_size, + ) + model = LlamaForCausalLM(config) + + press = FinchPress() + press.update_model_and_tokenizer(model, tokenizer) + assert model.get_input_embeddings().num_embeddings == len(tokenizer) == vocab_size + 1 + + +@torch.no_grad() +def test_finch_press_with_gemma3_for_conditional_generation(): + text_config = Gemma3TextConfig( + hidden_size=64, + intermediate_size=128, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=16, + vocab_size=300, + layer_types=["sliding_attention", "full_attention"], + sliding_window=8, + ) + vision_config = SiglipVisionConfig( + hidden_size=32, intermediate_size=64, num_hidden_layers=1, num_attention_heads=2, image_size=28, patch_size=14 + ) + config = Gemma3Config( + text_config=text_config, + vision_config=vision_config, + mm_tokens_per_image=4, + image_token_index=299, + boi_token_index=297, + eoi_token_index=298, + ) + model = Gemma3ForConditionalGeneration(config).eval() + assert not hasattr(model.model, "embed_tokens") + + press = FinchPress(compression_ratio=0.5, rerotate_keys=False) + press.delimiter_token_id = 296 + input_ids = torch.arange(10, 30) + input_ids[15] = press.delimiter_token_id + cache = DynamicCache() + with press(model): + model(input_ids.unsqueeze(0), past_key_values=cache) + # The delimiter is removed and only the full attention layer is compressed + assert cache.layers[0].keys.shape[2] == 19 + assert cache.layers[1].keys.shape[2] == 9 From d20740c4fe61178c4aa8f19eb9b53ef1125fe304 Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 04:09:49 +0200 Subject: [PATCH 52/57] =?UTF-8?q?Reject=20keys=20masked=20during=20prefill?= =?UTF-8?q?ing=20in=20decoding=20compression=20=F0=9F=A4=96=F0=9F=A4=96?= =?UTF-8?q?=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit DecodingPress and CAMPress remove tokens from the cache, which shifts the positions of keys masked by a head-wise press applied during prefilling (e.g. PrefillDecodingPress(prefilling_press=AdaKVPress(...))). This raised an IndexError in attention_patch or masked the wrong keys; it now raises a clear ValueError. Signed-off-by: Maximilian Jeblick --- kvpress/presses/cam_press.py | 1 + kvpress/presses/decoding_press.py | 12 ++++++++++++ tests/presses/test_decoding_press_fixes.py | 15 ++++++++++++++- 3 files changed, 27 insertions(+), 1 deletion(-) diff --git a/kvpress/presses/cam_press.py b/kvpress/presses/cam_press.py index cc78c31ad..3d09a868d 100644 --- a/kvpress/presses/cam_press.py +++ b/kvpress/presses/cam_press.py @@ -287,6 +287,7 @@ def forward_hook( q_len >= self.target_size ): + self._check_no_masked_keys(module) # Apply compression using cumulative attention scores and buffered hidden states attn_squeezed = self._running_attn_sum[layer_idx] buffered_hidden_states, buffered_kwargs = self._get_buffered_inputs(layer_idx, kwargs) diff --git a/kvpress/presses/decoding_press.py b/kvpress/presses/decoding_press.py index ca0e90ea5..91549d96a 100644 --- a/kvpress/presses/decoding_press.py +++ b/kvpress/presses/decoding_press.py @@ -64,6 +64,17 @@ def __post_init__(self): def post_init_from_model(self, model): self.base_press.post_init_from_model(model) + def _check_no_masked_keys(self, module: nn.Module): + """ + Removing tokens from the cache would shift the positions of keys masked by a head-wise press + (e.g. AdaKVPress, KVzipPress or DMSPress) during prefilling. + """ + if not isinstance(self.base_press, AdaKVPress) and getattr(module, "masked_key_indices", None) is not None: + raise ValueError( + f"{type(self).__name__} removes tokens from the cache and cannot be combined with a head-wise press " + "applied during prefilling, such as AdaKVPress, KVzipPress or DMSPress." + ) + def compress( self, module: nn.Module, @@ -144,6 +155,7 @@ def forward_hook(self, module: nn.Module, input: list[torch.Tensor], kwargs: dic f"Applying decoding compression: layer_step_count ({self.layer_step_counts[layer_idx]}) >= compression_steps ({self.compression_interval})" # noqa: E501 ) + self._check_no_masked_keys(module) keys, values = extract_keys_and_values(cache, module.layer_idx) # Apply compression using buffered hidden states for this layer. Attention weights are not used diff --git a/tests/presses/test_decoding_press_fixes.py b/tests/presses/test_decoding_press_fixes.py index 5d400b0d0..6fbdd4d1c 100644 --- a/tests/presses/test_decoding_press_fixes.py +++ b/tests/presses/test_decoding_press_fixes.py @@ -7,7 +7,7 @@ import torch from transformers import DynamicCache -from kvpress import CAMPress, DecodingPress, KnormPress, ScorerPress +from kvpress import AdaKVPress, CAMPress, DecodingPress, KnormPress, PrefillDecodingPress, ScorerPress from tests.fixtures import unit_test_model, unit_test_model_output_attention # noqa: F401 @@ -55,3 +55,16 @@ def test_state_is_reset_on_prefill(unit_test_model): # noqa: F811 unit_test_model(input_ids, past_key_values=DynamicCache()) assert set(press.layer_step_counts.values()) == {0} assert all(len(buffer) == 0 for buffer in press.hidden_states_buffer.values()) + + +@pytest.mark.parametrize("press_cls", [DecodingPress, CAMPress]) +@torch.no_grad() +def test_decoding_compression_rejects_keys_masked_during_prefilling(unit_test_model, press_cls): # noqa: F811 + press = PrefillDecodingPress( + prefilling_press=AdaKVPress(KnormPress(compression_ratio=0.5)), + decoding_press=press_cls(base_press=KnormPress(), compression_interval=2, target_size=8), + ) + input_ids = torch.randint(0, 1000, (1, 16), device=unit_test_model.device) + with pytest.raises(ValueError, match="head-wise press"): + with press(unit_test_model): + unit_test_model.generate(input_ids, max_new_tokens=4, do_sample=False) From cdcb00e67378f1517b0b36c4d82fb29f70bac3da Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 04:09:49 +0200 Subject: [PATCH 53/57] =?UTF-8?q?Make=20the=20CAMPress=20half-precision=20?= =?UTF-8?q?test=20deterministic=20=F0=9F=A4=96=F0=9F=A4=96=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The randomly initialized model could stop generating early, changing the number of decoding steps. Signed-off-by: Maximilian Jeblick --- tests/presses/test_cam_press_fixes.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/tests/presses/test_cam_press_fixes.py b/tests/presses/test_cam_press_fixes.py index 859a83abc..554051bef 100644 --- a/tests/presses/test_cam_press_fixes.py +++ b/tests/presses/test_cam_press_fixes.py @@ -64,6 +64,7 @@ def test_cam_accumulates_attention_in_float32(): @torch.no_grad() def test_cam_compresses_half_precision_caches(): + torch.manual_seed(0) config = LlamaConfig( hidden_size=64, intermediate_size=128, @@ -75,5 +76,11 @@ def test_cam_compresses_half_precision_caches(): model = LlamaForCausalLM(config).to(torch.bfloat16).eval() cache = DynamicCache() with CAMPress(base_press=KnormPress(), compression_interval=2, target_size=12)(model): - model.generate(torch.randint(0, 128, (1, 16)), past_key_values=cache, max_new_tokens=5, do_sample=False) + model.generate( + torch.randint(0, 128, (1, 16)), + past_key_values=cache, + max_new_tokens=5, + min_new_tokens=5, + do_sample=False, + ) assert cache.get_seq_length() == 12 From 2b3d0794ac6413ec2f1afba49e709ab93c5ed6af Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 12:49:45 +0200 Subject: [PATCH 54/57] =?UTF-8?q?Fix=20the=20kvzap=20AIME=20evaluation=20s?= =?UTF-8?q?cript=20for=20transformers=205=20=F0=9F=A4=96=F0=9F=A4=96?= =?UTF-8?q?=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit transformers 5 returns a BatchEncoding from apply_chat_template, so the script crashed on the first question. Request return_dict=True, call model.generate(**inputs) and slice the output with the input_ids length. Signed-off-by: Maximilian Jeblick --- kvzap/evaluate_aime.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/kvzap/evaluate_aime.py b/kvzap/evaluate_aime.py index 8625b87b8..07709e3fd 100644 --- a/kvzap/evaluate_aime.py +++ b/kvzap/evaluate_aime.py @@ -76,15 +76,17 @@ def evaluate( # Tokenize question messages = [{"role": "user", "content": row["question"]}] - tokens = tokenizer.apply_chat_template(messages, return_tensors="pt", add_generation_prompt=True) - tokens = tokens.to(model.device) + inputs = tokenizer.apply_chat_template( + messages, return_tensors="pt", return_dict=True, add_generation_prompt=True + ) + inputs = inputs.to(model.device) with press(model): # Generation config from model card: https://huggingface.co/Qwen/Qwen3-32B output_tokens = model.generate( - tokens, temperature=0.6, top_p=0.95, top_k=20, min_p=0.0, max_new_tokens=max_new_tokens + **inputs, temperature=0.6, top_p=0.95, top_k=20, min_p=0.0, max_new_tokens=max_new_tokens ) - answer = tokenizer.decode(output_tokens[0, tokens.shape[1] :]) + answer = tokenizer.decode(output_tokens[0, inputs["input_ids"].shape[1] :]) df.loc[idx, "predicted_answer"] = answer if isinstance(press, DMSPress): df.loc[idx, "compression_ratio"] = press.compression_ratio From ea3f2c43f54d5d14f063b7ca2b80e852df5f8663 Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 04:30:45 +0200 Subject: [PATCH 55/57] =?UTF-8?q?Skip=20the=20fake=20key=20search=20when?= =?UTF-8?q?=20no=20key=20is=20masked=20=F0=9F=A4=96=F0=9F=A4=96?= =?UTF-8?q?=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit With an empty mask (e.g. a head-wise press at compression_ratio=0), search_hyperplane ran on every decoding step for nothing and could fail to find a hyperplane on some models. Signed-off-by: Maximilian Jeblick --- kvpress/attention_patch.py | 2 +- tests/test_attention_patch.py | 17 ++++++++++++++++- 2 files changed, 17 insertions(+), 2 deletions(-) diff --git a/kvpress/attention_patch.py b/kvpress/attention_patch.py index 90e472feb..c3fd10816 100644 --- a/kvpress/attention_patch.py +++ b/kvpress/attention_patch.py @@ -82,7 +82,7 @@ def wrapper(module, query, key, value, attention_mask, dropout, **kwargs): if query.shape[2] == key.shape[2]: # Prefilling module.masked_key_indices = None - elif getattr(module, "masked_key_indices", None) is not None: + elif getattr(module, "masked_key_indices", None) is not None and module.masked_key_indices[-1].numel() > 0: # Decoding: build fake keys k s.t. exp() = 0 bsz, num_heads, seq_len, head_dim = query.shape num_key_value_heads = key.shape[1] diff --git a/tests/test_attention_patch.py b/tests/test_attention_patch.py index 1e3dccdbf..ae2c8c6cd 100644 --- a/tests/test_attention_patch.py +++ b/tests/test_attention_patch.py @@ -1,10 +1,13 @@ # SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import types + import pytest import torch -from kvpress.attention_patch import search_hyperplane +import kvpress.attention_patch as attention_patch_module +from kvpress.attention_patch import attention_patch, search_hyperplane def attention_weights(X, K, attention_scaling): @@ -50,3 +53,15 @@ def test_search_hyperplane_raises_for_non_separable_queries(): X = torch.cat([X, -X], dim=1) # no hyperplane has both q and -q on its positive side with pytest.raises(ValueError, match="Could not find fake keys"): search_hyperplane(X, max_iter=10) + + +def test_attention_patch_skips_the_fake_key_search_without_masked_keys(monkeypatch): + def fail(*args, **kwargs): + raise AssertionError("search_hyperplane should not run when no key is masked") + + monkeypatch.setattr(attention_patch_module, "search_hyperplane", fail) + wrapped = attention_patch(lambda module, query, key, value, attention_mask, dropout, **kwargs: key) + empty = torch.empty(0, dtype=torch.long) + module = types.SimpleNamespace(masked_key_indices=(empty, empty, empty)) + query, key = torch.randn(1, 2, 1, 4), torch.randn(1, 2, 5, 4) + torch.testing.assert_close(wrapped(module, query, key.clone(), key.clone(), None, 0.0), key) From 49b5b75a5cf3aa123222fa00921a502684dfc6b9 Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 05:42:54 +0200 Subject: [PATCH 56/57] =?UTF-8?q?Load=20the=20RestoreKV=20adapter=20on=20t?= =?UTF-8?q?he=20model=20device=20=F0=9F=A4=96=F0=9F=A4=96=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit In transformers 5.2, load_adapter keeps the adapter weights on the CPU unless a device map is given, so RestoreKVPress failed with a device mismatch on single-device CUDA models. Signed-off-by: Maximilian Jeblick --- kvpress/presses/restorekv_press.py | 4 +++- tests/presses/test_restorekv_fixes.py | 22 ++++++++++++++++++++++ 2 files changed, 25 insertions(+), 1 deletion(-) create mode 100644 tests/presses/test_restorekv_fixes.py diff --git a/kvpress/presses/restorekv_press.py b/kvpress/presses/restorekv_press.py index 6fe3e4a10..8adcb2f2e 100644 --- a/kvpress/presses/restorekv_press.py +++ b/kvpress/presses/restorekv_press.py @@ -53,7 +53,9 @@ def post_init_from_model(self, model: PreTrainedModel): embeddings_path = hf_hub_download(restore_model_name, "restore_embeddings.safetensors") self.restore_embeddings = load_file(embeddings_path)["restore_embeddings"].to(model.device, dtype=model.dtype) if self.adapter_name not in getattr(model, "peft_config", {}): - model.load_adapter(restore_model_name, adapter_name=self.adapter_name) + # Without a device map, load_adapter keeps the adapter weights on the CPU + load_kwargs = {} if hasattr(model, "hf_device_map") else {"device_map": {"": model.device}} + model.load_adapter(restore_model_name, adapter_name=self.adapter_name, **load_kwargs) model.disable_adapters() self.restore_model_name = restore_model_name logger.info("Loaded %s with %d restore tokens", restore_model_name, self.num_restore_tokens) diff --git a/tests/presses/test_restorekv_fixes.py b/tests/presses/test_restorekv_fixes.py new file mode 100644 index 000000000..19fa992bb --- /dev/null +++ b/tests/presses/test_restorekv_fixes.py @@ -0,0 +1,22 @@ +# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import torch +from safetensors.torch import save_file + +from kvpress import RestoreKVPress +from kvpress.presses import restorekv_press +from tests.fixtures import unit_test_model # noqa: F401 + + +def test_restorekv_loads_the_adapter_on_the_model_device(unit_test_model, monkeypatch, tmp_path): # noqa: F811 + embeddings_path = tmp_path / "restore_embeddings.safetensors" + save_file({"restore_embeddings": torch.zeros(4, unit_test_model.config.hidden_size)}, str(embeddings_path)) + monkeypatch.setattr(restorekv_press, "hf_hub_download", lambda *args, **kwargs: str(embeddings_path)) + load_adapter_kwargs = [] + monkeypatch.setattr(unit_test_model, "load_adapter", lambda *args, **kwargs: load_adapter_kwargs.append(kwargs)) + monkeypatch.setattr(unit_test_model, "disable_adapters", lambda: None) + + RestoreKVPress(compression_ratio=0.5).post_init_from_model(unit_test_model) + + assert load_adapter_kwargs[0]["device_map"] == {"": unit_test_model.device} From ef334401750a3d5b5cb3cf558aa47e30b433ab90 Mon Sep 17 00:00:00 2001 From: Maximilian Jeblick Date: Tue, 6 Oct 2026 13:29:09 +0200 Subject: [PATCH 57/57] =?UTF-8?q?Format=20the=20files=20that=20fail=20the?= =?UTF-8?q?=20black=20check=20on=20main=20=F0=9F=A4=96=F0=9F=A4=96?= =?UTF-8?q?=F0=9F=A4=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit make style fails on main because black 24 formats these three files differently. No code changes. Signed-off-by: Maximilian Jeblick --- evaluation/benchmarks/loogle/create_huggingface_dataset.py | 4 +--- kvpress/presses/entropy_gated_chunkkv_press.py | 6 +++--- tests/presses/test_merging_press.py | 6 +++--- 3 files changed, 7 insertions(+), 9 deletions(-) diff --git a/evaluation/benchmarks/loogle/create_huggingface_dataset.py b/evaluation/benchmarks/loogle/create_huggingface_dataset.py index d53a53148..335e1c2bc 100644 --- a/evaluation/benchmarks/loogle/create_huggingface_dataset.py +++ b/evaluation/benchmarks/loogle/create_huggingface_dataset.py @@ -38,9 +38,7 @@ if task == "shortdep_cloze": df["answer"] = df["answer"].apply( - lambda values: json.dumps( - {f"": value for i, value in enumerate(values)}, ensure_ascii=False - ) + lambda values: json.dumps({f"": value for i, value in enumerate(values)}, ensure_ascii=False) ) df["context"] = df["context"].apply(lambda x: context_prompt[task].format(input=x)) diff --git a/kvpress/presses/entropy_gated_chunkkv_press.py b/kvpress/presses/entropy_gated_chunkkv_press.py index 218a026c7..7c485c77e 100644 --- a/kvpress/presses/entropy_gated_chunkkv_press.py +++ b/kvpress/presses/entropy_gated_chunkkv_press.py @@ -54,9 +54,9 @@ class EntropyGatedChunkKVPress(ChunkKVPress): def __post_init__(self): super().__post_init__() - assert self.chunk_length > self.low_entropy_chunk_length >= 1, ( - "EntropyGatedChunkKVPress requires chunk_length > low_entropy_chunk_length >= 1" - ) + assert ( + self.chunk_length > self.low_entropy_chunk_length >= 1 + ), "EntropyGatedChunkKVPress requires chunk_length > low_entropy_chunk_length >= 1" def compress( self, diff --git a/tests/presses/test_merging_press.py b/tests/presses/test_merging_press.py index e222f3e11..507008af4 100644 --- a/tests/presses/test_merging_press.py +++ b/tests/presses/test_merging_press.py @@ -49,9 +49,9 @@ def test_keys_unchanged(unit_test_model): # noqa: F811 unit_test_model(input_ids.clone(), past_key_values=cache_merge) for i in range(len(cache_hard.layers)): - assert torch.equal(cache_hard.layers[i].keys, cache_merge.layers[i].keys), ( - f"Layer {i}: keys must not be modified" - ) + assert torch.equal( + cache_hard.layers[i].keys, cache_merge.layers[i].keys + ), f"Layer {i}: keys must not be modified" def test_merge_preserves_more_info(unit_test_model): # noqa: F811