Add DuoAttention on the fly

Signed-off-by: SimJeg <sjegou@nvidia.com>

Author: SimJeg

evaluation/evaluate.py

@@ -65,6 +65,7 @@
     "think": ThinKPress(),
     "tova": TOVAPress(),
     "duo_attention": DuoAttentionPress(),
+    "duo_attention_on_the_fly": DuoAttentionPress(on_the_fly_scoring=True),
     "chunkkv": ChunkKVPress(press=SnapKVPress(), chunk_length=20),
 }
 

kvpress/presses/duo_attention_press.py

@@ -1,13 +1,18 @@
 # SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from time import time
 from contextlib import contextmanager
 from dataclasses import dataclass, field
 from io import StringIO
 
+
 import numpy as np
 import requests  # type: ignore[import-untyped]
 import torch
+from datasets import load_dataset
+from transformers import AutoTokenizer
+from transformers.models.llama.modeling_llama import apply_rotary_pos_emb
 
 from kvpress.presses.base_press import BasePress
 
@@ -29,12 +34,15 @@ class DuoAttentionPress(BasePress):
     Splits attention heads into two types:
     - Retrieval heads: use the full KV cache
     - Streaming heads: use only sink and recent tokens.
-
-    Head classification is based on scores loaded from https://github.com/mit-han-lab/duo-attention/
     The higher the head_compression_ratio, the more streaming heads are used.
+
+    Head classification is based on scores.
+    - If on_the_fly_scoring=False, scores are loaded from https://github.com/mit-han-lab/duo-attention/
+    - (experimental) If on_the_fly_scoring=True, scores are computed using duo_attention_on_the_fly 
     """
 
     head_compression_ratio: float = 0.0
+    on_the_fly_scoring: bool = False
     compression_ratio_: float = field(init=False, default=None)
     recent_size: int = field(init=False, default=None)
     sink_size: int = field(init=False, default=None)
@@ -44,15 +52,19 @@ def __post_init_from_model__(self, model):
         """
         Initialize sink_size, recent_size, and streaming_mask from a model
         """
-        # Load attention pattern from the DuoAttention repo
-        self.sink_size, self.recent_size, head_scores = self.load_attention_pattern(model)
-
-        # Define retrieval and streaming heads through a binary mask
-        n_pruned = round(head_scores.size * self.head_compression_ratio)
-        self.streaming_mask = torch.zeros(head_scores.shape, dtype=bool, device=model.device)
-        if n_pruned > 0:
-            indices = np.argsort(head_scores, axis=None)[:n_pruned]
-            self.streaming_mask[np.unravel_index(indices, head_scores.shape)] = True
+        if getattr(self, "_post_init_model_name", None) != model.config.name_or_path:
+            # Load attention pattern from the DuoAttention repo
+            self.sink_size, self.recent_size, head_scores = self.load_attention_pattern(model)
+            if self.on_the_fly_scoring:
+                head_scores = duo_attention_on_the_fly(model)
+
+            # Define retrieval and streaming heads through a binary mask
+            n_pruned = round(head_scores.size * self.head_compression_ratio)
+            self.streaming_mask = torch.zeros(head_scores.shape, dtype=bool, device=model.device)
+            if n_pruned > 0:
+                indices = np.argsort(head_scores, axis=None)[:n_pruned]
+                self.streaming_mask[np.unravel_index(indices, head_scores.shape)] = True
+            self._post_init_model_name = model.config.name_or_path
 
     @property
     def compression_ratio(self) -> float:
@@ -108,3 +120,70 @@ def __call__(self, model):
         self.__post_init_from_model__(model)
         with super().__call__(model):
             yield
+
+
+def duo_attention_on_the_fly(model, num_samples=50, q_len=500):
+    """
+    New experimental method to quickly compute DuoAttention scores:
+    - Compute the mean query and key on num_samples random samples from BookSum
+    - Repeat the mean query and key q_len times and apply RoPE to get (Q, K)
+    - Compute the attention weights for (Q[-1], K) and compute the "area under the cumulated attention curve"
+    These scores could also be saved to avoid recomputing them but this method is still experimental
+    """
+
+    start = time()
+    print(f"Starting computation of DuoAttention scores based on {num_samples} samples")
+    tokenizer = AutoTokenizer.from_pretrained(model.config.name_or_path)
+    num_heads = model.config.num_attention_heads
+    num_key_value_heads = model.config.num_key_value_heads
+    num_key_value_groups = num_heads // num_key_value_heads
+
+    # Load data
+    dataset = load_dataset("kmfoda/booksum", split="train").to_pandas()
+    texts = dataset.sample(num_samples, random_state=42)["chapter"].tolist()
+
+    # Initialize variables
+    position_ids = torch.arange(q_len).unsqueeze(0)
+    scores = torch.zeros((model.config.num_hidden_layers, num_key_value_heads))
+
+    # Compute scores
+    for text in texts:
+        with torch.no_grad():
+            # Compute hidden states
+            inputs = tokenizer(text, return_tensors="pt").to(model.device)
+            hidden_states = list(model(**inputs, output_hidden_states=True).hidden_states[:-1])
+
+            for layer_idx, h in enumerate(hidden_states):
+                module = model.model.layers[layer_idx]
+                d = module.self_attn.head_dim
+                h = module.input_layernorm(h)
+
+                # Mean query
+                q = module.self_attn.q_proj(h)
+                q = q.view(1, q.shape[1], -1, d)
+                q = q.mean(dim=1, keepdim=True)
+                q = q.repeat(1, q_len, 1, 1).transpose(1, 2)
+
+                # Mean key
+                k = module.self_attn.k_proj(h)
+                k = k.view(1, k.shape[1], -1, d)
+                k = k.mean(dim=1, keepdim=True)
+                k = k.repeat(1, q_len, 1, 1).transpose(1, 2)
+
+                # Apply RoPE
+                cos, sin = model.model.rotary_emb(h, position_ids.to(h.device))
+                q, k = apply_rotary_pos_emb(q, k, cos, sin)
+                k = k.repeat_interleave(num_key_value_groups, dim=1)
+
+                # Compute attention weights for the last token
+                attn_weights = torch.matmul(q[:, :, -1:, :], k.transpose(2, 3)) / (d**0.5)
+                attn_weights = attn_weights.softmax(dim=-1, dtype=torch.float32).squeeze()
+
+                # Compute score: area under the cumulated attention curve
+                s = torch.cumsum(attn_weights, dim=1).mean(1).cpu()
+                s = s.view(-1, num_key_value_groups).mean(1)
+
+                # Store the scores
+                scores[layer_idx] += s / num_samples
+    print(f"Finished computation of DuoAttention scores in {time() - start:.2f}s")
+    return scores.numpy()

kvpress/presses/qfilter_press.py

@@ -27,9 +27,11 @@ class QFilterPress(ScorerPress):
     """
 
     def __post_init_from_model__(self, model):
-        model_name = model.config.name_or_path.split("/")[-1]
-        self.q_filters = self.load_q_filters(model_name)
-        self.q_filters = self.q_filters.to(model.dtype)
+        if getattr(self, "_post_init_model_name", None) != model.config.name_or_path:
+            model_name = model.config.name_or_path.split("/")[-1]
+            self.q_filters = self.load_q_filters(model_name)
+            self.q_filters = self.q_filters.to(model.dtype)
+            self._post_init_model_name = model.config.name_or_path
 
     @staticmethod
     def load_q_filters(model_name):