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/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"<mask-{i}>": value for i, value in enumerate(values)}, ensure_ascii=False
-            )
+            lambda values: json.dumps({f"<mask-{i}>": 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/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 02d50c7cf..0e473d596 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}"
@@ -173,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:
@@ -327,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.")
@@ -366,7 +378,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()
@@ -375,6 +387,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/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/attention_patch.py b/kvpress/attention_patch.py
index 270ad9e6f..c3fd10816 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(<q, k>) = 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
@@ -70,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(<q, k>) = 0
             bsz, num_heads, seq_len, head_dim = query.shape
             num_key_value_heads = key.shape[1]
diff --git a/kvpress/pipeline.py b/kvpress/pipeline.py
index eb366a554..f1a70bdad 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:
@@ -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)
@@ -305,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/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 5c73ee0fa..aa7d6f8df 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
 
@@ -200,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/cam_press.py b/kvpress/presses/cam_press.py
index 8b521235a..3d09a868d 100644
--- a/kvpress/presses/cam_press.py
+++ b/kvpress/presses/cam_press.py
@@ -4,19 +4,16 @@
 from __future__ import annotations
 
 import logging
-import math
 from dataclasses import dataclass
 
 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__)
 
@@ -58,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 "
 
@@ -141,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
@@ -151,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)
@@ -199,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(
@@ -249,9 +247,8 @@ 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)
 
-        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
 
@@ -268,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]
@@ -276,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)
 
@@ -289,21 +287,14 @@ 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 = 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
-            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
@@ -331,7 +322,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
 
@@ -341,7 +332,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/cap_press.py b/kvpress/presses/cap_press.py
index 446b7a91b..301862490 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)
@@ -165,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.
@@ -172,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)
@@ -238,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}.")
@@ -260,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/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/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/criticalkv_press.py b/kvpress/presses/criticalkv_press.py
index 2db080439..214f7e6c1 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
@@ -60,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.
@@ -131,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
@@ -144,7 +148,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/decoding_press.py b/kvpress/presses/decoding_press.py
index b492ac5ba..91549d96a 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__)
 
@@ -48,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"
@@ -65,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,
@@ -83,7 +93,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
@@ -100,7 +110,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}")
 
@@ -115,8 +125,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 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
         """
@@ -127,43 +137,35 @@ 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 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
 
         # 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
             )
 
-            cache_layer = cache.layers[module.layer_idx]
+            self._check_no_masked_keys(module)
             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
-            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
@@ -178,6 +180,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)
@@ -235,5 +248,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/kvpress/presses/dms_press.py b/kvpress/presses/dms_press.py
index 29634a30e..d54e19c36 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]
@@ -76,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:
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/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/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/kvpress/presses/expected_attention_press.py b/kvpress/presses/expected_attention_press.py
index 88ffe021a..77e8f6191 100644
--- a/kvpress/presses/expected_attention_press.py
+++ b/kvpress/presses/expected_attention_press.py
@@ -2,13 +2,12 @@
 # SPDX-License-Identifier: Apache-2.0
 
 
-import math
 from dataclasses import dataclass
 
 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
@@ -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
@@ -123,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,
@@ -141,14 +162,21 @@ 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, 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/expected_attention_with_stats.py b/kvpress/presses/expected_attention_with_stats.py
index bd776ca37..915fdbd21 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):
     """
@@ -61,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
@@ -87,7 +93,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 +276,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/kvpress/presses/fastkvzip_press.py b/kvpress/presses/fastkvzip_press.py
index c6897be49..168bcd5d1 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)
 
@@ -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):
         """
@@ -208,12 +209,15 @@ 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
+            # 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)
@@ -267,7 +272,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].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
@@ -282,16 +294,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/kvpress/presses/finch_press.py b/kvpress/presses/finch_press.py
index af2219fde..81d5cbe04 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)
@@ -101,13 +102,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
@@ -125,6 +128,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
@@ -146,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
@@ -161,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/kvpress/presses/kvcompose_press.py b/kvpress/presses/kvcompose_press.py
index ed34a8b11..23ad7665d 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.cache_utils import DynamicCache, QuantizedCache
 from transformers.modeling_utils import PreTrainedModel
-from transformers.cache_utils import DynamicCache
+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
@@ -170,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] = []
@@ -209,7 +210,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.
@@ -240,7 +241,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)
@@ -253,17 +254,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 +304,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 +313,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 +328,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 +354,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,61 +380,88 @@ 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")
+        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()
+
+        def new_forward(
+            self,
+            input_ids,
+            past_key_values,
+            *args,
+            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,
+                )
+            if isinstance(past_key_values, QuantizedCache):
+                raise ValueError("KVComposePress does not support QuantizedCache")
 
-        def new_forward(self,
-                        input_ids,
-                        past_key_values,
-                        *args,
-                        press: KVComposePress,
-                        **kwargs,
-                        ):
             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)
+                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)
-            model.model.forward = types.MethodType(new_forward_with_press, model)
-
-            yield
+            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)
+
+                yield
+            finally:
+                model.model.forward = getattr(model, "original_forward_KVComposePress")
+                delattr(model, "original_forward_KVComposePress")
+
+            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")
-            for forward_hook in hooks:
-                forward_hook.remove()
-            self.prepare_important_masks()
-            self.compress_cache(model)
             self._reset_state()
diff --git a/kvpress/presses/kvzip_press.py b/kvpress/presses/kvzip_press.py
index 3e6bd4d39..122d8c896 100644
--- a/kvpress/presses/kvzip_press.py
+++ b/kvpress/presses/kvzip_press.py
@@ -123,9 +123,22 @@ 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)
+            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)
+            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")
+            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)
 
@@ -140,6 +153,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()
@@ -161,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/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/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/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/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/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/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/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/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/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/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
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)
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_cam_press_fixes.py b/tests/presses/test_cam_press_fixes.py
new file mode 100644
index 000000000..554051bef
--- /dev/null
+++ b/tests/presses/test_cam_press_fixes.py
@@ -0,0 +1,86 @@
+# 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 transformers import DynamicCache, LlamaConfig, LlamaForCausalLM
+
+from kvpress import AdaKVPress, CAMPress, KnormPress, ScorerPress, SnapKVPress
+
+
+@dataclass
+class FixedScorePress(ScorerPress):
+    scores: list = field(default_factory=list)
+
+    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())
+
+
+@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():
+    torch.manual_seed(0)
+    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,
+            min_new_tokens=5,
+            do_sample=False,
+        )
+    assert cache.get_seq_length() == 12
diff --git a/tests/presses/test_cap_press_fixes.py b/tests/presses/test_cap_press_fixes.py
new file mode 100644
index 000000000..66b697b19
--- /dev/null
+++ b/tests/presses/test_cap_press_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 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)
+
+
+@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]
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()
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/presses/test_criticalkv_fixes.py b/tests/presses/test_criticalkv_fixes.py
new file mode 100644
index 000000000..c3fea8662
--- /dev/null
+++ b/tests/presses/test_criticalkv_fixes.py
@@ -0,0 +1,50 @@
+# 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
+
+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(
+        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
+
+
+@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]
diff --git a/tests/presses/test_decoding_press_fixes.py b/tests/presses/test_decoding_press_fixes.py
new file mode 100644
index 000000000..6fbdd4d1c
--- /dev/null
+++ b/tests/presses/test_decoding_press_fixes.py
@@ -0,0 +1,70 @@
+# 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
+
+from kvpress import AdaKVPress, CAMPress, DecodingPress, KnormPress, PrefillDecodingPress, ScorerPress
+from tests.fixtures import unit_test_model, 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)
+
+
+@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())
+
+
+@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)
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)
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"]
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..6b6a1945e
--- /dev/null
+++ b/tests/presses/test_ea_with_stats_fixes.py
@@ -0,0 +1,63 @@
+# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from types import SimpleNamespace
+
+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))
+    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
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))
diff --git a/tests/presses/test_fastkvzip_fixes.py b/tests/presses/test_fastkvzip_fixes.py
new file mode 100644
index 000000000..0772ff04a
--- /dev/null
+++ b/tests/presses/test_fastkvzip_fixes.py
@@ -0,0 +1,124 @@
+# 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, 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, unit_test_model  # noqa: F401
+
+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)
+        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(n_gates)
+        ]
+        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)
+
+
+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)
+
+
+@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):
+        pass
diff --git a/tests/presses/test_finch_press_fixes.py b/tests/presses/test_finch_press_fixes.py
new file mode 100644
index 000000000..0092741df
--- /dev/null
+++ b/tests/presses/test_finch_press_fixes.py
@@ -0,0 +1,140 @@
+# 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 transformers import (
+    AutoTokenizer,
+    DynamicCache,
+    Gemma3Config,
+    Gemma3ForConditionalGeneration,
+    Gemma3TextConfig,
+    LlamaConfig,
+    LlamaForCausalLM,
+    Qwen2Config,
+    Qwen2ForCausalLM,
+    SiglipVisionConfig,
+)
+
+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):
+    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
+
+
+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)
+    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))
+
+
+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
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)
+        )
diff --git a/tests/presses/test_kvcompose_fixes.py b/tests/presses/test_kvcompose_fixes.py
new file mode 100644
index 000000000..7ba4ce600
--- /dev/null
+++ b/tests/presses/test_kvcompose_fixes.py
@@ -0,0 +1,112 @@
+# 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, unit_test_model_output_attention  # 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
+
+
+@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_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):
+            raise RuntimeError("boom")
+
+
+@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
+
+
+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
+
+
+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
+    """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
diff --git a/tests/presses/test_kvzip_fixes.py b/tests/presses/test_kvzip_fixes.py
new file mode 100644
index 000000000..8b46a1569
--- /dev/null
+++ b/tests/presses/test_kvzip_fixes.py
@@ -0,0 +1,107 @@
+# 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 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):
+        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,
+    )
+
+
+@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
+
+
+@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"):
+        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)
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
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)
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))
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]
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
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}
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)
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)
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
diff --git a/tests/test_evaluate_config.py b/tests/test_evaluate_config.py
new file mode 100644
index 000000000..9fca3bdc2
--- /dev/null
+++ b/tests/test_evaluate_config.py
@@ -0,0 +1,134 @@
+# 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
+import yaml
+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"
+
+# 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.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)
+
+    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)
+
+
+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)
+    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]
+
+
+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 934cbdb4a..ea52f7b64 100644
--- a/tests/test_evaluation_utils.py
+++ b/tests/test_evaluation_utils.py
@@ -1,12 +1,59 @@
 # 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.needle_in_haystack.utils import insert_needle_in_haystack
 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"
+)
+
+
+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(
+        {
+            "context": [context],
+            "question": [question],
+            "answer": [np.array(answer, dtype=object)],
+            "task": [task],
+            "predicted_answer": [prediction],
+        }
+    )
+
 
 @pytest.mark.parametrize(
     ("text", "expected"),
@@ -43,3 +90,54 @@ 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"),
+}
+
+
+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):
+    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
diff --git a/tests/test_pipeline_fixes.py b/tests/test_pipeline_fixes.py
new file mode 100644
index 000000000..ac7135346
--- /dev/null
+++ b/tests/test_pipeline_fixes.py
@@ -0,0 +1,75 @@
+# 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 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
+    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)
+
+
+@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
+
+
+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]
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]