kvpress 0.0.1__py3-none-any.whl

kvpress/__init__.py

@@ -0,0 +1,36 @@
+# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from kvpress.per_layer_compression_wrapper import apply_per_layer_compression
+from kvpress.pipeline import KVPressTextGenerationPipeline
+from kvpress.presses.base_press import BasePress
+from kvpress.presses.expected_attention_press import ExpectedAttentionPress
+from kvpress.presses.knorm_press import KnormPress
+from kvpress.presses.observed_attention_press import ObservedAttentionPress
+from kvpress.presses.random_press import RandomPress
+from kvpress.presses.snapkv_press import SnapKVPress
+from kvpress.presses.streaming_llm_press import StreamingLLMPress
+
+__all__ = [
+    "BasePress",
+    "ExpectedAttentionPress",
+    "KnormPress",
+    "ObservedAttentionPress",
+    "RandomPress",
+    "SnapKVPress",
+    "StreamingLLMPress",
+    "KVPressTextGenerationPipeline",
+    "apply_per_layer_compression",
+]

kvpress/per_layer_compression_wrapper.py

@@ -0,0 +1,59 @@
+# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+from typing import List
+
+import torch
+
+from kvpress.presses.base_press import BasePress
+
+logger = logging.getLogger(__name__)
+
+
+def apply_per_layer_compression(press: BasePress, compression_ratios: List[float]) -> BasePress:
+    """
+    Apply per-layer compression to a given press object.
+    This function wraps the forward hook of the press object to apply per-layer compression.
+
+    Parameters
+    ----------
+    press : BasePress
+        The press object to apply per-layer compression to.
+    compression_ratios : Dict[int, float]
+
+    Returns
+    -------
+    BasePress
+        The press object with per-layer compression applied.
+    """
+    press.compression_ratios = compression_ratios  # type: ignore[attr-defined]
+    press.compression_ratio = None
+
+    logger.warning(
+        "Per layer compression wrapper is an experimental feature and only works with flash attention. "
+        "Please make sure that the model uses flash attention."
+    )
+
+    original_forward_hook = press.forward_hook
+
+    def _forward_hook(module: torch.nn.Module, input: list[torch.Tensor], kwargs: dict, output: list):
+        press.compression_ratio = press.compression_ratios[module.layer_idx]  # type: ignore[attr-defined]
+        output = original_forward_hook(module, input, kwargs, output)
+        press.compression_ratio = None
+        return output
+
+    press.forward_hook = _forward_hook  # type: ignore[method-assign]
+    return press

kvpress/pipeline.py

@@ -0,0 +1,266 @@
+# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import contextlib
+import logging
+from typing import Dict, List, Optional
+
+import torch
+from transformers import AutoModelForCausalLM, DynamicCache, Pipeline
+from transformers.pipelines import PIPELINE_REGISTRY
+from transformers.pipelines.base import GenericTensor
+
+from kvpress.presses.base_press import BasePress
+from kvpress.presses.observed_attention_press import ObservedAttentionPress
+
+logger = logging.getLogger(__name__)
+
+
+class KVPressTextGenerationPipeline(Pipeline):
+    """
+    Pipeline for key-value compression in causal language models.
+    This pipeline allows you to compress a long prompt using a key-value press
+    and then generate answers using greedy decoding.
+    """
+
+    def _sanitize_parameters(
+        self,
+        question: Optional[str] = None,
+        questions: Optional[List[str]] = None,
+        answer_prefix: Optional[str] = None,
+        press: Optional[BasePress] = None,
+        max_new_tokens: int = 50,
+        max_context_length: Optional[int] = None,
+        **kwargs,
+    ):
+        """
+        Sanitize the input parameters for the pipeline.
+        The user can either provide a single question or a list of questions to be asked about the context.
+
+        Parameters
+        ----------
+        question : str, optional
+            The question to be asked about the context. Exclusive with `questions`.
+        questions : List[str], optional
+            A list of questions to be asked about the context. Exclusive with `question`.
+        answer_prefix : str, optional
+            The prefix to be added to the generated answer.
+        press : BasePress, optional
+            The key-value press to use for compression.
+        max_new_tokens : int, optional
+            The maximum number of new tokens to generate for each answer.
+        max_context_length : int, optional
+            The maximum number of tokens in the context. By default will use the maximum length supported by the model.
+        **kwargs : dict
+            Additional keyword arguments, currently ignored.
+
+        Returns
+        -------
+        Tuple[Dict, Dict, Dict]
+            A tuple containing three dictionaries:
+                - preprocess_kwargs: The keyword arguments for the preprocess function.
+                - forward_kwargs: The keyword arguments for the forward function.
+                - postprocess_kwargs: The keyword arguments for the postprocess function.
+        """
+
+        answer_prefix = answer_prefix or ""
+        postprocess_kwargs = {"single_question": questions is None}
+        assert question is None or questions is None, "Either question or questions should be provided, not both."
+        questions = questions or ([question] if question else [""])
+        if max_context_length is None:
+            max_context_length = min(self.tokenizer.model_max_length, int(1e10))  # 1e10 to avoid overflow
+        preprocess_kwargs = {
+            "questions": questions,
+            "answer_prefix": answer_prefix,
+            "max_context_length": max_context_length,
+        }
+        forward_kwargs = {"press": press, "max_new_tokens": max_new_tokens}
+        return preprocess_kwargs, forward_kwargs, postprocess_kwargs
+
+    def preprocess(
+        self,
+        context: str,
+        questions: List[str],
+        answer_prefix: str,
+        max_context_length: int,
+    ):
+        """
+        Apply the chat template to the triplet (context, questions, answer_prefix) and tokenize it.
+
+        Returns
+        -------
+        Dict[str, GenericTensor]
+            A dictionary containing the tokenized context (key: "context_ids") and questions (key: "questions_ids").
+
+        """
+
+        # Apply chat template if available
+        if self.tokenizer.chat_template is None:
+            bos_token = getattr(self.tokenizer, "bos_token", "")
+            context = bos_token + context
+            question_suffix = "\n"  # to separate the question from the answer
+        else:
+            separator = "\n" + "#" * len(context)
+            context = self.tokenizer.apply_chat_template(
+                [{"role": "user", "content": context + separator}], add_generation_prompt=True, tokenize=False
+            )
+            context, question_suffix = context.split(separator)
+
+        # Add question_suffix and answer prefix
+        # e.g. for llama3.1, question_suffix="<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n")
+        questions = [question + question_suffix + answer_prefix for question in questions]
+
+        # Tokenize the context and questions
+        context_ids = self.tokenizer.encode(context, return_tensors="pt", add_special_tokens=False)
+        question_ids = [
+            self.tokenizer.encode(question, return_tensors="pt", add_special_tokens=False) for question in questions
+        ]
+
+        # Truncate context
+        if context_ids.shape[1] > max_context_length:
+            logger.warning(
+                f"Context length has been truncated from {context_ids.shape[1]} to {max_context_length} tokens."
+            )
+            context_ids = context_ids[:, :max_context_length]
+
+        return {"context_ids": context_ids, "questions_ids": question_ids}
+
+    def _forward(
+        self, input_tensors: Dict[str, GenericTensor], max_new_tokens: int = 50, press: Optional[BasePress] = None
+    ):
+        """
+        Forward pass of the kv-press pipeline.
+
+        Parameters
+        ----------
+        input_tensors : Dict[str, GenericTensor]
+            A dictionary containing the tokenized context and questions.
+        max_new_tokens : int, optional
+            The maximum number of new tokens to generate for each answer. Defaults to 50.
+        press : BasePress, optional
+            The key-value press to use for compression. Defaults to None.
+
+        Returns
+        -------
+        List[str]
+            A list of generated answers.
+        """
+
+        context_ids = input_tensors["context_ids"].to(self.model.device)
+        context_length = context_ids.shape[1]
+
+        # Prefilling using the press on the context
+        with press(self.model) if press is not None else contextlib.nullcontext():
+            past_key_values = self.model(
+                input_ids=context_ids,
+                past_key_values=DynamicCache(),
+                output_attentions=isinstance(press, ObservedAttentionPress),
+                num_logits_to_keep=1,
+            ).past_key_values
+
+        logger.debug(f"Context Length: {context_length}")
+        logger.debug(f"Compressed Context Length: {past_key_values.get_seq_length()}")
+
+        # Greedy decoding for each question
+        answers = []
+        for question_ids in input_tensors["questions_ids"]:
+            answer = self.generate_answer(
+                question_ids=question_ids.to(self.model.device),
+                past_key_values=past_key_values,
+                context_length=context_length,
+                max_new_tokens=max_new_tokens,
+            )
+            answers.append(answer)
+
+        return answers
+
+    def postprocess(self, model_outputs, single_question):
+        if single_question:
+            return {"answer": model_outputs[0]}
+        return {"answers": model_outputs}
+
+    def generate_answer(
+        self, question_ids: torch.Tensor, past_key_values: DynamicCache, context_length: int, max_new_tokens: int
+    ) -> str:
+        """
+        Generate an answer to a question using greedy decoding.
+
+        Parameters
+        ----------
+        question_ids : torch.Tensor
+            The tokenized question.
+        past_key_values : DynamicCache
+            The compressed key-value cache.
+        context_length : int
+            The length of the context.
+        max_new_tokens : int
+            The maximum number of new tokens to generate.
+
+        Returns
+        -------
+        str
+            The generated answer.
+        """
+
+        cache_seq_lengths = [
+            past_key_values.get_seq_length(layer_idx=layer_idx) for layer_idx in range(len(past_key_values))
+        ]
+
+        position_ids = torch.arange(
+            context_length, context_length + question_ids.shape[1], device=self.model.device
+        ).unsqueeze(0)
+
+        # if the user doesn't provide a question, skip forward pass
+        outputs = self.model(
+            input_ids=question_ids.to(self.model.device),
+            past_key_values=past_key_values,
+            position_ids=position_ids,
+            num_logits_to_keep=1,
+        )
+
+        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]
+
+        for i in range(max_new_tokens - 1):
+            outputs = self.model(
+                input_ids=generated_ids[-1].unsqueeze(0).unsqueeze(0),
+                past_key_values=outputs.past_key_values,
+                position_ids=position_ids + i,
+            )
+            new_id = outputs.logits[0, -1].argmax()
+            generated_ids.append(new_id)
+            if new_id.item() in should_stop_token_ids:
+                break
+        answer = self.tokenizer.decode(torch.stack(generated_ids), skip_special_tokens=True)
+
+        # remove the generated tokens from the cache
+        past_key_values.key_cache = [
+            key[:, :, :cache_seq_len] for key, cache_seq_len in zip(past_key_values.key_cache, cache_seq_lengths)
+        ]
+        past_key_values.value_cache = [
+            value[:, :, :cache_seq_len] for value, cache_seq_len in zip(past_key_values.value_cache, cache_seq_lengths)
+        ]
+        return answer
+
+
+PIPELINE_REGISTRY.register_pipeline(
+    "kv-press-text-generation",
+    pipeline_class=KVPressTextGenerationPipeline,
+    pt_model=AutoModelForCausalLM,
+)

kvpress/presses/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

kvpress/presses/base_press.py

@@ -0,0 +1,146 @@
+# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+from contextlib import contextmanager
+from typing import Dict, Generator
+
+import torch
+from torch import nn
+from transformers import LlamaForCausalLM, MistralForCausalLM, Phi3ForCausalLM, PreTrainedModel, Qwen2ForCausalLM
+
+logger = logging.getLogger(__name__)
+
+
+class BasePress:
+    """Base class for pruning methods.
+    Each pruning method should implement a `score` method that computes the scores for each KV pair in a layer.
+    This score is used to prune the KV pairs with the lowest scores in the `hook` method
+    The `hook` method is called after the forward pass of a layer and updates the cache with the pruned KV pairs.
+    The press can be applied to a model by calling it with the model as an argument.
+    """
+
+    def __init__(self, compression_ratio: float = 0.0):
+        self.compression_ratio = compression_ratio
+        assert 0 <= compression_ratio < 1, "Compression ratio must be between 0 and 1"
+
+    def score(
+        self,
+        module: nn.Module,
+        hidden_states: torch.Tensor,
+        keys: torch.Tensor,
+        values: torch.Tensor,
+        attentions: torch.Tensor,
+        kwargs,
+    ) -> torch.Tensor:
+        """Compute the scores for each KV pair in the layer.
+
+        Parameters
+        ----------
+        module :
+            Transformer layer, see `hook` method for more details.
+        hidden_states :
+            Hidden states of the layer.
+        keys :
+            Keys of the cache. Note keys are after RoPE.
+        values :
+            Values of the cache.
+        attentions :
+            Attention weights of the layer.
+        kwargs :
+            Keyword arguments, as given to the forward pass of the layer.
+
+        Returns
+        -------
+            Scores for each KV pair in the layer, shape keys.shape[:-1].
+
+        """
+        raise NotImplementedError
+
+    def forward_hook(self, module: nn.Module, input: list[torch.Tensor], kwargs: Dict, output: list):
+        """Cache compression hook called after the forward pass of a decoder layer.
+        The hook is applied only during the pre-filling phase if there is some pruning ratio.
+        The current implementation only allows to remove a constant number of KV pairs.
+
+        Parameters
+        ----------
+        module :
+            Transformer attention layer.
+        input :
+            Input to the hook. This is the input to the forward pass of the layer.
+        kwargs :
+            Keyword arguments, as given to the forward pass of the layer.
+        output :
+            Output of the hook. This is the original output of the forward pass of the layer.
+
+        Returns
+        -------
+            Modified output of the forward pass of the layer.
+
+        """
+        # See e.g. LlamaDecoderLayer.forward for the output structure
+        if len(output) == 3:
+            _, attentions, cache = output
+        else:
+            attentions, cache = None, output[-1]
+
+        hidden_states = kwargs["hidden_states"]
+        q_len = hidden_states.shape[1]
+
+        # Don't compress if the compression ratio is 0 or this is not pre-filling
+        if (self.compression_ratio == 0) or (cache.seen_tokens > q_len):
+            return output
+
+        keys = cache.key_cache[module.layer_idx]
+        values = cache.value_cache[module.layer_idx]
+
+        with torch.no_grad():
+            scores = self.score(module, hidden_states, keys, values, attentions, kwargs)
+
+        # Prune KV pairs with the lowest scores
+        n_kept = int(q_len * (1 - self.compression_ratio))
+        indices = scores.topk(n_kept, dim=-1).indices
+        indices = indices.unsqueeze(-1).expand(-1, -1, -1, module.head_dim)
+
+        # Update cache
+        cache.key_cache[module.layer_idx] = keys.gather(2, indices)
+        cache.value_cache[module.layer_idx] = values.gather(2, indices)
+
+        return output
+
+    @contextmanager
+    def __call__(self, model: PreTrainedModel) -> Generator:
+        """
+        Context manager to apply a compression method to a model.
+        Apply this context manager during the pre-filling phase to compress the context.
+
+        Parameters
+        ----------
+        model : PreTrainedModel
+            Model to apply the compression method to
+        """
+
+        if not isinstance(model, (LlamaForCausalLM, MistralForCausalLM, Phi3ForCausalLM, Qwen2ForCausalLM)):
+            logger.warning(f"Model {type(model)} not tested")
+
+        try:
+            hooks = []
+            for layer in model.model.layers:
+                hooks.append(layer.self_attn.register_forward_hook(self.forward_hook, with_kwargs=True))
+
+            yield
+        finally:
+            for forward_hook in hooks:
+                forward_hook.remove()

kvpress/presses/expected_attention_press.py

@@ -0,0 +1,148 @@
+# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+import math
+from dataclasses import dataclass
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from transformers.models.llama.modeling_llama import repeat_kv
+
+from kvpress.presses.base_press import BasePress
+
+
+@dataclass
+class ExpectedAttentionPress(BasePress):
+    """
+    Compute scores based on the expected attention on next positions. To do so
+        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)
+        5. Rescale the scores by the norm of the values
+    The first n_sink tokens are removed from calculations (sink attention phenomenon).
+    """
+
+    compression_ratio: float = 0.0
+    n_future_positions: int = 512
+    n_sink: int = 4
+    use_covariance: bool = True
+    use_vnorm: bool = True
+
+    def get_query_statistics(self, module: nn.Module, hidden_states: torch.Tensor):
+        """
+        Compute the mean and covariance matrix of the queries
+        """
+
+        bsz, q_len, _ = hidden_states.shape
+        n, d = module.num_heads, module.head_dim
+
+        # Remove first hidden_states that likely contain outliers
+        h = hidden_states[:, self.n_sink :]
+
+        if hasattr(module, "q_proj"):
+            Wq = module.q_proj.weight
+        elif hasattr(module, "qkv_proj"):
+            Wq = module.qkv_proj.weight[: n * d]
+        else:
+            raise NotImplementedError(f"ExpectedAttentionPress not yet implemented for {module.__class__}.")
+
+        # Query mean
+        mean_h = torch.mean(h, dim=1, keepdim=True)
+        mu = torch.matmul(mean_h, Wq.T).squeeze(1)
+        mu = mu.view(bsz, n, d)
+
+        # Query covariance
+        cov = None
+        if self.use_covariance:
+            h = h - mean_h
+            cov = torch.matmul(h.transpose(1, 2), h) / h.shape[1]
+            cov = torch.matmul(Wq, torch.matmul(cov, Wq.T))  # TODO: not optimal
+            cov = cov.view(bsz, n, d, n, d).diagonal(dim1=1, dim2=3)
+            cov = cov.permute(0, 3, 1, 2)
+
+        # RoPE rotation matrix on next n_future_positions
+        if "position_ids" in inspect.signature(module.rotary_emb.forward).parameters:
+            position_ids = torch.arange(q_len, q_len + self.n_future_positions).unsqueeze(0).to(mu.device)
+            cos, sin = module.rotary_emb(mu, position_ids)
+            cos, sin = cos[0], sin[0]
+        else:
+            cos, sin = module.rotary_emb(mu, q_len + self.n_future_positions)
+            cos, sin = cos[q_len:], sin[q_len:]
+
+        Id = torch.eye(d, device=cos.device, dtype=cos.dtype)
+        P = torch.zeros((d, d), device=cos.device, dtype=cos.dtype)
+        P[d // 2 :, : d // 2], P[: d // 2, d // 2 :] = torch.eye(d // 2), -torch.eye(d // 2)
+        R = cos.unsqueeze(1) * Id + sin.unsqueeze(1) * P
+
+        # Apply average rotation to the mean and covariance
+        R = R.mean(dim=0)
+        mu = torch.matmul(mu, R.T)
+        if self.use_covariance:
+            cov = torch.matmul(R, torch.matmul(cov, R.T))
+
+        # Instead of using the average rotation matrix, we could use a mixture of gaussian statistics to
+        # estimate mean and covariance. Estimation is better, but end-to-end performance was lower.
+        # mu = torch.einsum("bhj, fij -> bhfi", mu, R)
+        # mean_mu = mu.mean(dim=2, keepdim=True)
+        # if self.use_covariance:
+        #     cov = torch.einsum("fki, bhkl, fjl -> bhfij", R, cov, R)
+        #     cov = cov.mean(dim=2)
+        #     cov += torch.einsum("bhfi, bhfj -> bhji", mu - mean_mu, mu - mean_mu) / self.n_future_positions
+        # mu = mean_mu.squeeze(2)
+
+        return mu, cov
+
+    def score(
+        self,
+        module: nn.Module,
+        hidden_states: torch.Tensor,
+        keys: torch.Tensor,
+        values: torch.Tensor,
+        attentions: torch.Tensor,
+        kwargs,
+    ) -> torch.Tensor:
+
+        # Remove sink tokens
+        assert keys.size(2) > self.n_sink, f"Input should contain more tokens than n_sink={self.n_sink}"
+        keys = keys[:, :, self.n_sink :]
+        values = values[:, :, self.n_sink :]
+
+        # Compute query statistics
+        mean_query, cov_query = self.get_query_statistics(module, hidden_states)
+
+        # Compute scores
+        bsz, num_key_value_heads, q_len, d = keys.shape
+        keys = repeat_kv(keys, module.num_key_value_groups).transpose(2, 3)
+        scores = torch.matmul(mean_query.unsqueeze(2), keys).squeeze(2) / math.sqrt(d)
+        if self.use_covariance:
+            scores += torch.einsum("bhin, bhij, bhjn->bhn", keys, cov_query, keys) / d / 2
+        scores = F.softmax(scores, dim=-1)
+
+        # Average scores across groups
+        scores = scores.view(bsz, num_key_value_heads, module.num_key_value_groups, q_len)
+        scores = scores.mean(dim=2)
+
+        # Rescale scores by the norm of the values
+        if self.use_vnorm:
+            scores = scores * values.norm(dim=-1)
+
+        # Add back the sink tokens. Use max score to make sure they are not pruned.
+        scores = F.pad(scores, (self.n_sink, 0), value=scores.max().item())
+
+        return scores