Repository: NVIDIA/kvpress

Issue:
Request for New Press Support: Integrating CriticalKV with Existing Methods
### Feature

Support for a new work: CriticalKV [Openreview](https://openreview.net/forum?id=lRTDMGYCpy) [ArXiv](https://arxiv.org/abs/2502.03805). This method introduces a novel metric for KV cache entry selection from an output perturbation perspective, which can be seamlessly integrated to enhance existing techniques.

### Motivation

CriticalKV is fully compatible with most existing methods and requires only a few additional lines of code to compute the new selection metric. I implemented a demo version of CriticalKV in [repo](https://github.com/FFY0/AdaKV-in-NVIDIA-kvpress/tree/criticalkv) and applied it to enhance four existing methods—SnapKV, Ada-SnapKV, Expected-Attention, and Ada-Expected-Attention. Across different compression rates on the Ruler benchmark with Llama-3.1-8B-Instruct,  CriticalKV consistently improved their performance significantly.

If you have any suggestions, I’d be happy to further refine and improve this implementation!

![Image](https://github.com/user-attachments/assets/06f39494-52a0-480d-99bd-abac93c97f18)

Hints from issue comments:
@SimJeg: This looks great, congrats for the new work and results. In
ExpectedAttentionPress, we already multiply by the value norm, did you
disable (use_vnorm=False) before applying your new method?

Le ven. 7 févr. 2025, 06:35, Yuan Feng ***@***.***> a écrit :

> Feature
>
> Support for a new work: CriticalKV Openreview
> <https://openreview.net/forum?id=lRTDMGYCpy> ArXiv
> <https://arxiv.org/abs/2502.03805>. This method introduces a novel metric
> for KV cache entry selection from an output perturbation perspective, which
> can be seamlessly integrated to enhance existing techniques.
> Motivation
>
> CriticalKV is fully compatible with most existing methods and requires
> only a few additional lines of code to compute the new selection metric. I
> implemented a demo version of CriticalKV in repo
> <https://github.com/FFY0/AdaKV-in-NVIDIA-kvpress/tree/criticalkv> and
> applied it to enhance four existing methods—SnapKV, Ada-SnapKV,
> Expected-Attention, and Ada-Expected-Attention. Across different
> compression rates on the Ruler benchmark, CriticalKV consistently improved
> their performance significantly.
>
> If you have any suggestions, I’d be happy to further refine and improve
> this implementation!
>
> all_cases_average_score.png (view on web)
> <https://github.com/user-attachments/assets/06f39494-52a0-480d-99bd-abac93c97f18>
>
> —
> Reply to this email directly, view it on GitHub
> <https://github.com/NVIDIA/kvpress/issues/45>, or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/ADE64VIKDNBPVVOCA6ECPP32ORAYNAVCNFSM6AAAAABWVFHEWOVHI2DSMVQWIX3LMV43ASLTON2WKOZSHAZTOMRXG44TQMQ>
> .
> You are receiving this because you are subscribed to this thread.Message
> ID: ***@***.***>
>
@FFY0: Yes, I noticed that. When integrated into ExpectedAttentionPress, the `use_vnorm` operation is disabled in the `__post_init__` function. Therefore, the improvement achieved by CriticalKV demonstrates a more effective way than directly multiplying the value norm.@SimJeg: Hi @FFY0,

Please find below a proposed updated version for your code. It includes:
- fix linting (flake8 compatible)
- proper use of dataclass
- minor refacto of the code and docstrings for clarity
- improved `vwl1norm` (3x faster for same output, using torch.matmul) + use of static method
- no forcing of use_vnorm=False

Could you please: 
- refactor your code to have the code below in a single script called `criticalkv_press.py` 
- update init.py to import the 2 classes from this unified script
- update evaluate.py to explicitly use v_norm=False for ExpectedAttentionPress
- update the README to add the following line at the end of the Available presses sections:

`CriticalKVPress` and `CriticalAdaKVPress` ([paper](https://arxiv.org/abs/2502.03805)): rescales the scores by the L1 norm of Wo @ values

- add the 2 presses in the tests (l45, in `test_press.py`, add it as a wrapper press)
- open a PR
- sign off your commits using "git commit -s"
- Share a plot here with everything superposed as well as raw results


```python
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

import logging
from dataclasses import dataclass

import torch
from transformers.models.llama.modeling_llama import repeat_kv

from kvpress.presses.base_press import BasePress
from kvpress.presses.scorer_press import ScorerPress
from kvpress.presses.expected_attention_press import ExpectedAttentionPress

logger = logging.getLogger(__name__)


class CriticalKVPress(ScorerPress):
    """
    CriticalKV (https://arxiv.org/abs/2502.03805) rescales the scores of a ScorerPress by
    the L1 norm of Wo @ values
    """

    def __init__(self, press: ScorerPress, epsilon: float = 1e-4, first_stage_ratio: float = 0.5):
        self.press = press
        self.epsilon = epsilon
        self.first_stage_ratio = first_stage_ratio

        assert isinstance(self.press, ScorerPress), "CriticalAdaKVPress requires a ScorerPress as input"
        if isinstance(self.press, ExpectedAttentionPress) and self.press.use_vnorm:
            logger.warning("use_vnorm should be disabled for CriticalAdaKVPress")

    @property
    def compression_ratio(self):
        return self.press.compression_ratio

    @compression_ratio.setter
    def compression_ratio(self, value):
        self.press.compression_ratio = value

    @staticmethod
    def vwl1norm(values, module):
        bsz, num_key_value_heads, q_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)
        V = repeat_kv(values, num_key_value_groups)
        WoV_norm = torch.matmul(V, Wo).norm(p=1, dim=-1)
        WoV_norm = WoV_norm.view(bsz, num_key_value_heads, num_key_value_groups, q_len).mean(dim=2)
        return WoV_norm

    def score(self, module, hidden_states, keys, values, attentions, kwargs):

        # Stage 1
        scores = self.press.score(module, hidden_states, keys, values, attentions, kwargs)
        q_len = keys.shape[2]
        selection_budget = int((1 - self.compression_ratio) * q_len * self.first_stage_ratio)
        top_k_index = torch.topk(scores, selection_budget, sorted=True, dim=-1).indices

        # Stage 2
        projected_norm = self.vwl1norm(values, module)
        scores = (scores + self.epsilon) * projected_norm

        # Merge the two stages
        scores.scatter_(-1, top_k_index, torch.finfo(scores.dtype).max)

        return scores


@dataclass
class CriticalAdaKVPress(BasePress):
    """
    CriticalAdaKV (https://arxiv.org/abs/2502.03805) rescales the scores of a ScorerPress by
    the L1 norm of Wo @ values and combines it with AdaKV (https://arxiv.org/abs/2407.11550).
    """

    press: ScorerPress
    alpha_safeguard: float = 0.20
    epsilon: float = 1e-4
    first_stage_ratio: float = 0.5

    def __post_init__(self):
        assert 0 <= self.alpha_safeguard <= 1, "alpha_safeguard should be in 0, 1]"
        assert isinstance(self.press, ScorerPress), "CriticalAdaKVPress requires a ScorerPress as input"
        if isinstance(self.press, ExpectedAttentionPress) and self.press.use_vnorm:
            logger.warning("use_vnorm should be disabled for CriticalAdaKVPress")

    @property
    def compression_ratio(self):
        return self.press.compression_ratio

    @compression_ratio.setter
    def compression_ratio(self, value):
        self.press.compression_ratio = value

    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"

        # Compute scores
        scores = self.press.score(module, hidden_states, keys, values, attentions, kwargs)
        bsz, num_key_value_heads, q_len = scores.shape

        # Make sure to keep at least alpha * (1 - compression_ratio) KV pairs per head
        n_kept = int(q_len * (1 - self.compression_ratio))  # ScorerPress definition
        n_safe = int(n_kept * self.alpha_safeguard)
        top_indices = torch.topk(scores, n_safe, dim=-1).indices
        scores.scatter_(-1, top_indices, torch.finfo(scores.dtype).max)

        ############################
        # Start of CriticalKV code #
        ############################

        # Budget allocation
        budget_scores = scores.scatter(-1, top_indices, torch.finfo(scores.dtype).max)
        budget_scores = budget_scores.reshape(bsz, -1)
        top_indices = torch.topk(budget_scores, n_kept * num_key_value_heads, dim=-1).indices
        top_indices_head_idx = top_indices // q_len
        head_budgets = torch.zeros(num_key_value_heads, device=keys.device, dtype=torch.int64)
        head_budgets.scatter_add_(0, top_indices_head_idx.flatten(), torch.ones_like(top_indices_head_idx.flatten()))

        # Stage 1
        head_selection_budget_1st = (head_budgets * self.first_stage_ratio).to(torch.int64).tolist()
        top_k_index = torch.topk(scores, max(head_selection_budget_1st), sorted=True, dim=-1).indices
        for head_idx in range(num_key_value_heads):
            phase1_budget = head_selection_budget_1st[head_idx]
            scores[:, head_idx, :].scatter_(-1, top_k_index[:, head_idx, :phase1_budget], torch.finfo(scores.dtype).max)

        # Stage 2
        projected_norm = CriticalKVPress.vwl1norm(values, module)
        scores = (scores + self.epsilon) * projected_norm
        top_k_index = torch.topk(scores, max(head_budgets), sorted=True, dim=-1).indices
        for head_idx in range(num_key_value_heads):
            budget = head_budgets[head_idx]
            scores[:, head_idx, :].scatter_(-1, top_k_index[:, head_idx, :budget], torch.finfo(scores.dtype).max)

        ##########################
        # End of CriticalKV code #
        ##########################

        # Compute bottom-k across heads
        n_pruned = num_key_value_heads * (q_len - n_kept)
        indices = torch.topk(-scores.reshape(bsz, -1), n_pruned, dim=1).indices.flatten()

        # Save indices to mask during the attention mechanism. Please refer to attention_patch.py for more details
        batch_indices = torch.arange(bsz).repeat_interleave(n_pruned)
        head_indices = indices // q_len
        seq_indices = indices % q_len
        module.masked_key_indices = (batch_indices, head_indices, seq_indices)
        return keys, values

```@SimJeg: I made some experiments using `ExpectedAttentionPress` and 50% compression: 
- use $\Vert V \Vert$ or $\Vert Wo \cdot V \Vert$: as expected $\Vert Wo \cdot V \Vert$ is slightly better (0.5 to 1 %)
- use L1 vs L2: it's basically the same (as in Figure 13 of the paper)
- use $\epsilon = 0$ vs $\epsilon = 10^{-4}$: interestingly this is where I see the biggest improvement (1.5 to 2%). My interpretation: below a certain threshold, the value norm is more reliable than the expected attention weight.

| value   | norm | epsilon | score     |
|---------|------|---------|-----------|
| WoV | L2   | 1e-4    | 79.252308 |
| WoV| L1   | 1e-4    | 79.178462 |
| V   | L1   | 1e-4    | 78.436923 |
| V   | L2   | 1e-4    | 78.349231 |
| WoV | L2   | 0       | 77.252308 |
| WoV | L1   | 0       | 77.200769 |
| V   | L2   | 0       | 76.603077 |
| V   | L1   | 0       | 76.443077 |

@FFY0, could you share the raw results for `expected_attention` and `critical_expected_attention` at 50% compression to see if they align with above results ? Would be interesting to understand if first stage is useful for this press.@FFY0: Hello @SimJeg. In my case, the `criti_expected_attention` is `79.59076923`, while the `expected_attention` is `76.89307692`. The purpose of this two-stage selection ablation is to ensure that the assumption in the paper hold, especially in scenarios involving small cache budgets. This is particularly crucial in compression scenarios with questions, where the cache budget often drops below 0.1.  

Below are the previous ablation results of the **first-stage selection** on the 16 datasets from LongBench with llama-3.1-8B and Mistral-v0.2-7b. The "**w/ ours ablation case**" refers to the setting where the **first stage is removed**. The results generally align with our theoretical expectations—that is, the **smaller the budget, the more critical the first stage becomes**.

Llama
![Image](https://github.com/user-attachments/assets/5c3562b6-d2b3-4910-a83d-ed5a94af4435)
Mistral
![Image](https://github.com/user-attachments/assets/e5a90ff7-fe40-4145-83cd-f43fa148bb05)@FFY0: In addition, the `vwl1norm` method, which utilizes `torch.matmul` along with a static method, indeed accelerates computation for short sequences. However, when applied to long sequences, it incurs a high memory overhead due to intermediate variables. There are two solutions to this issue:  

1. **Per-head computation**: Our previous implementation processes each head separately. While this approach is slower, it is easier to understand. Morever, since its computation time is still negligible compared to LLM inference costs, we initially adopted it.

2. **Kernel fusion**: By fusing the projection and normalization operations into a single kernel, we can eliminate intermediate variables, achieving optimal speed and memory efficiency. However, this approach may increase the complexity of maintaining this codebase.  

In practice, given that compression cost is negligible compared to the inference costs of LLMs, solution 1 seems quite acceptable. However, if you're interested in solution 2, we would be happy to contribute our kernel.@SimJeg: Thanks for your detailed answers ! My point was not to challenge the paper but rather to better understand where the gain comes from. 

For the vwl1norm operation, I missed that point indeed that the dimension of WoV is 32 times bigger than the one of V ! You can switch back to your head-wise for loop in the code I shared but please add a comment to explain why you are not using matmul.

Note: there was a small bug in the init code I shared, I just fixed it@FFY0: Thanks for your response. Based on previous evaluations, I think the gain of different operations seems to vary across different scenarios, such as changes in cache size or different compression settings (with or without compression). Fine-tuning these parameters for specific cases could further enhance performance.

Additionally, segmenting `vwl1norm` along the sequence length might be a better approach. This would maintain constant memory usage while improving computational efficiency by choosing a suitable block size. I'm currently working on implementing and testing this modification.😊@SimJeg: > Additionally, segmenting vwl1norm along the sequence length might be a better approach. 

To me the head approach is fine and it avoids adding another parameter the user has to understand. Also it has the same effect on memory as when you compute Q, K or V@FFY0: > To me the head approach is fine and it avoids adding another parameter the user has to understand. 

That makes sense. The simple head-wise approach should be sufficient for most cases. I'll revert to using head-wise computation.@SimJeg: @FFY0 feel free to open the PR with your current branch whenever you want

Solution patch:
diff --git a/README.md b/README.md
index 3885909..f45c60d 100644
--- a/README.md
+++ b/README.md
@@ -76,6 +76,7 @@ Finally we provide wrapper presses that can be combined with other presses:
 - `ComposedPress` ([source](kvpress/presses/composed_press.py)): compose multiple presses together by chaining their forward hooks
 - `KeyRerotationPress` ([source](kvpress/presses/key_rerotation_press.py)): rerotate pruned keys to have continuous RoPE embeddings
 - `ChunkPress` ([source](kvpress/presses/chunk_press.py), [paper](https://direct.mit.edu/tacl/article/doi/10.1162/tacl_a_00716/125280)): compress the KV cache on each sequence chunk separately. This can yield to more uniform compression across long sequences
+- `CriticalKVPress` and `CriticalAdaKVPress` ([source](kvpress/presses/criticalkv_press.py), [paper](https://arxiv.org/abs/2502.03805)): refine the scores using the L1 norm of Wo @ values, coupled with a two-stage selection.
 
 For a detailed list of existing KV cache compression methods, check [Awesome-KV-Cache-Compression](https://github.com/October2001/Awesome-KV-Cache-Compression) or [Awesome-LLM-Compression](https://github.com/HuangOwen/Awesome-LLM-Compression?tab=readme-ov-file#kv-cache-compression)
 
diff --git a/evaluation/evaluate.py b/evaluation/evaluate.py
index d165fda..00fdbad 100644
--- a/evaluation/evaluate.py
+++ b/evaluation/evaluate.py
@@ -17,6 +17,8 @@ from transformers import pipeline
 from zero_scrolls.calculate_metrics import calculate_metrics as zero_scrolls_scorer
 
 from kvpress import (
+    CriticalKVPress,
+    CriticalAdaKVPress,
     AdaKVPress,
     ExpectedAttentionPress,
     KnormPress,
@@ -45,6 +47,10 @@ SCORER_DICT = {
 }
 
 PRESS_DICT = {
+    "criti_adasnapkv": CriticalAdaKVPress(SnapKVPress()),
+    "criti_ada_expected_attention": CriticalAdaKVPress(ExpectedAttentionPress(use_vnorm=False)),
+    "criti_snapkv": CriticalKVPress(SnapKVPress()),
+    "criti_expected_attention": CriticalKVPress(ExpectedAttentionPress(use_vnorm=False)),
     "adasnapkv": AdaKVPress(SnapKVPress()),
     "ada_expected_attention": AdaKVPress(ExpectedAttentionPress()),
     "expected_attention": ExpectedAttentionPress(),
diff --git a/kvpress/__init__.py b/kvpress/__init__.py
index c5c4d99..54e419c 100644
--- a/kvpress/__init__.py
+++ b/kvpress/__init__.py
@@ -20,11 +20,13 @@ from kvpress.presses.snapkv_press import SnapKVPress
 from kvpress.presses.streaming_llm_press import StreamingLLMPress
 from kvpress.presses.think_press import ThinKPress
 from kvpress.presses.tova_press import TOVAPress
-
+from kvpress.presses.criticalkv_press import CriticalKVPress, CriticalAdaKVPress
 # Patch the attention functions to support head-wise compression
 patch_attention_functions()
 
 __all__ = [
+    "CriticalAdaKVPress",
+    "CriticalKVPress",
     "AdaKVPress",
     "BasePress",
     "ComposedPress",
diff --git a/kvpress/presses/criticalkv_press.py b/kvpress/presses/criticalkv_press.py
new file mode 100644
index 0000000..9574a98
--- /dev/null
+++ b/kvpress/presses/criticalkv_press.py
@@ -0,0 +1,161 @@
+# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+from dataclasses import dataclass
+
+import torch
+from transformers.models.llama.modeling_llama import repeat_kv
+
+from kvpress.presses.base_press import BasePress
+from kvpress.presses.scorer_press import ScorerPress
+from kvpress.presses.expected_attention_press import ExpectedAttentionPress
+
+logger = logging.getLogger(__name__)
+
+
+class CriticalKVPress(ScorerPress):
+    """
+    CriticalKV (https://arxiv.org/abs/2502.03805) rescales the scores of a ScorerPress by
+    the L1 norm of Wo @ values
+    """
+
+    def __init__(self, press: ScorerPress, epsilon: float = 1e-4, first_stage_ratio: float = 0.5):
+        self.press = press
+        self.epsilon = epsilon
+        self.first_stage_ratio = first_stage_ratio
+
+        assert isinstance(self.press, ScorerPress), "CriticalAdaKVPress requires a ScorerPress as input"
+        if isinstance(self.press, ExpectedAttentionPress) and self.press.use_vnorm:
+            logger.warning("use_vnorm should be disabled for CriticalAdaKVPress")
+
+    @property
+    def compression_ratio(self):
+        return self.press.compression_ratio
+
+    @compression_ratio.setter
+    def compression_ratio(self, value):
+        self.press.compression_ratio = value
+
+    @staticmethod
+    def vwl1norm(values, module):
+        bsz, num_key_value_heads, q_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)
+        V = repeat_kv(values, num_key_value_groups)
+
+        # We use head-wise computation instead of direct matmul to reduce the memory usage of WoV.
+        # Future kernel fusion optimization could eliminate this intermediate variables to enhance performance.
+        head_WoV_norm_list = []
+        for head in range(V.size(1)):
+            head_WoV = V[: , head, : , ...].matmul(Wo[head, ...].unsqueeze(0))
+            head_WoV_norm = torch.norm(head_WoV, p=1, dim=-1)
+            head_WoV_norm_list.append(head_WoV_norm)
+
+        # b_size, num_heads, q_len , k_len
+        WoV_norm = torch.stack(head_WoV_norm_list, dim=1)
+        WoV_norm = WoV_norm.view(bsz, num_key_value_heads, module.num_key_value_groups, q_len).mean(dim=2)
+        return WoV_norm
+
+    def score(self, module, hidden_states, keys, values, attentions, kwargs):
+        # Stage 1
+        scores = self.press.score(module, hidden_states, keys, values, attentions, kwargs)
+        q_len = keys.shape[2]
+        selection_budget = int((1 - self.compression_ratio) * q_len * self.first_stage_ratio)
+        top_k_index = torch.topk(scores, selection_budget, sorted=True, dim=-1).indices
+
+        # Stage 2
+        projected_norm = self.vwl1norm(values, module)
+        scores = (scores + self.epsilon) * projected_norm
+
+        # Merge the two stages
+        scores.scatter_(-1, top_k_index, torch.finfo(scores.dtype).max)
+
+        return scores
+
+
+@dataclass
+class CriticalAdaKVPress(BasePress):
+    """
+    CriticalAdaKV (https://arxiv.org/abs/2502.03805) rescales the scores of a ScorerPress by
+    the L1 norm of Wo @ values and combines it with AdaKV (https://arxiv.org/abs/2407.11550).
+    """
+
+    press: ScorerPress
+    alpha_safeguard: float = 0.20
+    epsilon: float = 1e-4
+    first_stage_ratio: float = 0.5
+
+    def __post_init__(self):
+        assert 0 <= self.alpha_safeguard <= 1, "alpha_safeguard should be in 0, 1]"
+        assert isinstance(self.press, ScorerPress), "CriticalAdaKVPress requires a ScorerPress as input"
+        if isinstance(self.press, ExpectedAttentionPress) and self.press.use_vnorm:
+            logger.warning("use_vnorm should be disabled for CriticalAdaKVPress")
+
+    @property
+    def compression_ratio(self):
+        return self.press.compression_ratio
+
+    @compression_ratio.setter
+    def compression_ratio(self, value):
+        self.press.compression_ratio = value
+
+    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"
+
+        # Compute scores
+        scores = self.press.score(module, hidden_states, keys, values, attentions, kwargs)
+        bsz, num_key_value_heads, q_len = scores.shape
+
+        # Make sure to keep at least alpha * (1 - compression_ratio) KV pairs per head
+        n_kept = int(q_len * (1 - self.compression_ratio))  # ScorerPress definition
+        n_safe = int(n_kept * self.alpha_safeguard)
+        top_indices = torch.topk(scores, n_safe, dim=-1).indices
+        scores.scatter_(-1, top_indices, torch.finfo(scores.dtype).max)
+
+        ############################
+        # Start of CriticalKV code #
+        ############################
+
+        # Budget allocation
+        budget_scores = scores.scatter(-1, top_indices, torch.finfo(scores.dtype).max)
+        budget_scores = budget_scores.reshape(bsz, -1)
+        top_indices = torch.topk(budget_scores, n_kept * num_key_value_heads, dim=-1).indices
+        top_indices_head_idx = top_indices // q_len
+        head_budgets = torch.zeros(num_key_value_heads, device=keys.device, dtype=torch.int64)
+        head_budgets.scatter_add_(0, top_indices_head_idx.flatten(), torch.ones_like(top_indices_head_idx.flatten()))
+
+        # Stage 1
+        head_selection_budget_1st = (head_budgets * self.first_stage_ratio).to(torch.int64).tolist()
+        top_k_index = torch.topk(scores, max(head_selection_budget_1st), sorted=True, dim=-1).indices
+        for head_idx in range(num_key_value_heads):
+            phase1_budget = head_selection_budget_1st[head_idx]
+            scores[:, head_idx, :].scatter_(-1, top_k_index[:, head_idx, :phase1_budget], torch.finfo(scores.dtype).max)
+
+        # Stage 2
+        projected_norm = CriticalKVPress.vwl1norm(values, module)
+        scores = (scores + self.epsilon) * projected_norm
+        top_k_index = torch.topk(scores, max(head_budgets), sorted=True, dim=-1).indices
+        for head_idx in range(num_key_value_heads):
+            budget = head_budgets[head_idx]
+            scores[:, head_idx, :].scatter_(-1, top_k_index[:, head_idx, :budget], torch.finfo(scores.dtype).max)
+
+        ##########################
+        # End of CriticalKV code #
+        ##########################
+
+        # Compute bottom-k across heads
+        n_pruned = num_key_value_heads * (q_len - n_kept)
+        indices = torch.topk(-scores.reshape(bsz, -1), n_pruned, dim=1).indices.flatten()
+
+        # Save indices to mask during the attention mechanism. Please refer to attention_patch.py for more details
+        batch_indices = torch.arange(bsz).repeat_interleave(n_pruned)
+        head_indices = indices // q_len
+        seq_indices = indices % q_len
+        module.masked_key_indices = (batch_indices, head_indices, seq_indices)
+        return keys, values


Test patch:
diff --git a/tests/presses/test_presses.py b/tests/presses/test_presses.py
index ee2e0fa..27adab6 100644
--- a/tests/presses/test_presses.py
+++ b/tests/presses/test_presses.py
@@ -8,6 +8,8 @@ from torch import nn
 from transformers import DynamicCache
 
 from kvpress import (
+    CriticalKVPress,
+    CriticalAdaKVPress,
     AdaKVPress,
     ChunkPress,
     ComposedPress,
@@ -42,7 +44,8 @@ def test_chunk_press(unit_test_model):  # noqa: F811
 
 
 @pytest.mark.parametrize("press_dict", default_presses)
-@pytest.mark.parametrize("wrapper_press", [None, ComposedPress, KeyRerotationPress, AdaKVPress, ChunkPress])
+@pytest.mark.parametrize("wrapper_press", [None, ComposedPress, KeyRerotationPress, AdaKVPress, ChunkPress,
+                                           CriticalKVPress, CriticalAdaKVPress])
 def test_presses_run(unit_test_model, press_dict, wrapper_press):  # noqa: F811
     cls = press_dict["cls"]
     for kwargs in press_dict["kwargs"]:
@@ -51,14 +54,13 @@ def test_presses_run(unit_test_model, press_dict, wrapper_press):  # noqa: F811
             press = ComposedPress(presses=[press])
         if isinstance(wrapper_press, KeyRerotationPress):
             press = KeyRerotationPress(press=press)
-        if isinstance(wrapper_press, AdaKVPress):
-            if not isinstance(press, ScorerPress):
-                return
+        if isinstance(wrapper_press, (AdaKVPress, CriticalKVPress, CriticalAdaKVPress)):
+            if isinstance(press, ScorerPress):
+                press = wrapper_press(press=press)
             else:
-                press = AdaKVPress(press=press)
+                return
         if isinstance(wrapper_press, ChunkPress):
             press = ChunkPress(press=press, chunk_length=2)
-
         with press(unit_test_model):
             input_ids = unit_test_model.dummy_inputs["input_ids"]
             unit_test_model(input_ids, past_key_values=DynamicCache()).past_key_values