Adjust Positional Embedding For Compression

[giulio98]

Feature

Adjust the positional encoding when compressing the cache to improve output quality for long sequences. Based on our experiments, this adjustment significantly reduces gibberish output and enhances performance for extended contexts. The proposed feature integrates seamlessly with existing caching mechanisms and dynamically adjusts cosine and sine positional embeddings during cache compression.

Motivation

Long sequences often result in incoherent outputs when the cache is compressed without considering the positional encodings. This issue arises due to misalignment of the positional embeddings, which leads to inaccurate attention calculations. By adjusting the positional embeddings during compression, we observed a dramatic improvement in model performance and output quality.

This proposal aligns with our experimental findings and can be implemented by extending the DynamicCache class from transformers. Below is an example implementation called FinchCache, which adjusts the positional encodings during cache compression:

from transformers import DynamicCache
import torch

class FinchCache(DynamicCache):
    def __init__(self) -> None:
        super().__init__()
        self._cos_cache = None
        self._sin_cache = None
        self.key_cache = []
        self.value_cache = []

    @staticmethod
    def _rotate_half(x):
        x1 = x[..., : x.shape[-1] // 2]
        x2 = x[..., x.shape[-1] // 2:]
        return torch.cat((-x2, x1), dim=-1)

    def _apply_key_rotary_pos_emb(self, key_states, cos, sin):
        rotated_key_states = (key_states * cos) + (self._rotate_half(key_states) * sin)
        return rotated_key_states

    @staticmethod
    def _rerotate_cos_sin(x, inv_freq, important_pos_batch):
        batch_size, seq_len = important_pos_batch.shape
        idx = torch.arange(0, seq_len, dtype=torch.long, device=important_pos_batch.device).unsqueeze(0).expand(batch_size, -1)
        delta_pos = idx - important_pos_batch
        inv_freq_expanded = inv_freq[None, :, None].float().expand(batch_size, -1, 1)
        delta_pos_expanded = delta_pos[:, None, :].float()
        freqs = (inv_freq_expanded.float() @ delta_pos_expanded.float()).transpose(1, 2)
        emb = torch.cat((freqs, freqs), dim=-1)
        cos = emb.cos()
        sin = emb.sin()
        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)

    @staticmethod
    def gather_important_tokens(states, indices):
        return torch.gather(states, 2, indices.unsqueeze(1).unsqueeze(-1).expand(-1, states.size(1), -1, states.size(3)))

    def compress_cache(self, important_indices, inv_freq):
        new_length = important_indices[0].size(1)
        for layer_idx in range(len(self.key_cache)):
            important_pos = important_indices[layer_idx]
            new_cos, new_sin = self._rerotate_cos_sin(self.key_cache[layer_idx], inv_freq, important_pos)
            gathered_keys = self.gather_important_tokens(self.key_cache[layer_idx], important_pos)
            gathered_values = self.gather_important_tokens(self.value_cache[layer_idx], important_pos)
            self.key_cache[layer_idx] = self._apply_key_rotary_pos_emb(gathered_keys, new_cos, new_sin)
            self.value_cache[layer_idx] = gathered_values
        self._cos_cache = self._cos_cache[:new_length, :]
        self._sin_cache = self._sin_cache[:new_length, :]
        self._seen_tokens = new_length

    def update(
        self,
        key_states: torch.Tensor,
        value_states: torch.Tensor,
        layer_idx: int,
        cache_kwargs: Optional[Dict[str, Any]] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Updates the cache with the new `key_states` and `value_states` for the layer `layer_idx`.

        Parameters:
            key_states (`torch.Tensor`):
                The new key states to cache.
            value_states (`torch.Tensor`):
                The new value states to cache.
            layer_idx (`int`):
                The index of the layer to cache the states for.
            cache_kwargs (`Dict[str, Any]`, `optional`):
                Additional arguments for the cache subclass. No additional arguments are used in `Cache`.

        Return:
            A tuple containing the updated key and value states.
        """
        sin = cache_kwargs.get("sin")
        cos = cache_kwargs.get("cos")
        using_rope = cos is not None and sin is not None
        if layer_idx == 0:
            self._seen_tokens += key_states.shape[-2]
        
        if using_rope and layer_idx == 0:
            # BC: some models still pass `sin`/`cos` with 2 dims. In those models, they are the full sin/cos. Remove
            # after all RoPE models have a llama-like cache utilization.
            if cos.dim() == 2:
                self._cos_cache = cos
                self._sin_cache = sin
            else:
                if self._cos_cache is None:
                    self._cos_cache = cos[0, ...]
                    self._sin_cache = sin[0, ...]
                else:
                    self._cos_cache = torch.cat([self._cos_cache, cos[0, ...]], dim=0)
                    self._sin_cache = torch.cat([self._sin_cache, sin[0, ...]], dim=0)
        
        
        if len(self.key_cache) <= layer_idx:
            # There may be skipped layers, fill them with empty lists
            for _ in range(len(self.key_cache), layer_idx):
                self.key_cache.append([])
                self.value_cache.append([])
            self.key_cache.append(key_states)
            self.value_cache.append(value_states)
        elif len(self.key_cache[layer_idx]) == 0:  # fills previously skipped layers; checking for tensor causes errors
            self.key_cache[layer_idx] = key_states
            self.value_cache[layer_idx] = value_states
        else:
            self.key_cache[layer_idx] = torch.cat([self.key_cache[layer_idx], key_states], dim=-2)
            self.value_cache[layer_idx] = torch.cat([self.value_cache[layer_idx], value_states], dim=-2)

        return self.key_cache[layer_idx], self.value_cache[layer_idx]

The compress_cache function has to be called after the topk indices selection and can be integrated with any presses: performs cache adjustment by taking the indices (from a selection process) and the inverse frequencies of the original position embeddings (model.model.layers[0].self_attn.rotary_emb.inv_freq).

This proposal addresses a significant limitation in handling long sequences and can benefit all users working with compressed cache scenarios.

[maxjeblick]

Thanks a lot for proposing this interesting feature!

I have a few follow-up questions regarding the implementation/usage of FinchCache:

1. Do you have a code example/stub that uses FinchCache? This would help to ensure that an implementation in this repo matches the original version.
2. Could you elaborate a bit more on the effect of _rerotate_cos_sin, followed by _apply_key_rotary_pos_emb on gathered keys? Would this be equivalent to
   a. Compute keys before RoPE (i.e. keys_no_rope = attn.q_proj(hidden_states), followed by reshaping to (keys_no_rope.shape[0], keys_no_rope.shape[1], attn.num_heads, attn.head_dim).transpose(1, 2))
   b. Apply gathered_keys, i.e. gathered_keys = self.gather_important_tokens(keys_no_rope, important_pos)
   c. Apply cos, sin on the gathered gathered_keys, with cos, sin being adjusted to gathered_keys sequence length?
3. When using FinchCache, do you adjust model.generate? We are using position_ids w.r.t. the original context length, see here. This ensures that queries' position ids are compatible with the key's position ids.

[giulio98]

Thank you for your interest in FinchCache and for your thoughtful questions! Let me address them point by point:

1. Code Example for FinchCache Implementation
   This is an example how the finchcache is used by overwriting the generate function of mistral (N.B. is not necessary to overwrite the generate, this is just an example)

import math
import time
from typing import Optional, Tuple, Dict, Any
import torch
from transformers.cache_utils import DynamicCache


class FinchCache(DynamicCache):
    def __init__(self) -> None:
        super().__init__()
        self._cos_cache = None
        self._sin_cache = None
        self.key_cache = []
        self.value_cache = []

    @staticmethod
    def _rotate_half(x):
        x1 = x[..., : x.shape[-1] // 2]
        x2 = x[..., x.shape[-1] // 2:]
        return torch.cat((-x2, x1), dim=-1)

    def _apply_key_rotary_pos_emb(
        self, key_states: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
    ) -> torch.Tensor:
        rotated_key_states = (key_states * cos) + (self._rotate_half(key_states) * sin)
        return rotated_key_states

    @staticmethod
    def _rerotate_cos_sin(x, inv_freq, important_pos_batch):
        batch_size, seq_len = important_pos_batch.shape
        idx = torch.arange(0, seq_len, dtype=torch.long, device=important_pos_batch.device).unsqueeze(0).expand(batch_size, -1)
        delta_pos = idx - important_pos_batch  # [batch_size, seq_len]
        inv_freq_expanded = inv_freq[None, :, None].float().expand(batch_size, -1, 1)
        delta_pos_expanded = delta_pos[:, None, :].float()
        device_type = x.device.type
        freqs = (inv_freq_expanded.float() @ delta_pos_expanded.float()).transpose(1, 2)
        emb = torch.cat((freqs, freqs), dim=-1)
        cos = emb.cos()
        sin = emb.sin()
        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)

    @staticmethod
    def gather_important_tokens(states, indices):
        return torch.gather(states, 2,
                            indices.unsqueeze(1).unsqueeze(-1).expand(-1, states.size(1), -1, states.size(3)))

    def update_rope(self, layer_index, key_states, important_pos, inv_freq):
        new_length = important_pos.size(1)
        new_cos, new_sin = self._rerotate_cos_sin(key_states, inv_freq, important_pos)
        self.key_cache[layer_index] = self._apply_key_rotary_pos_emb(
            self.gather_important_tokens(self.key_cache[layer_index], important_pos),
            new_cos,
            new_sin
        )
        self.value_cache[layer_index] = self.gather_important_tokens(self.value_cache[layer_index], important_pos)
        self._cos_cache = self._cos_cache[:new_length, :]
        self._sin_cache = self._cos_cache[:new_length, :]
        self._seen_tokens = new_length
        return self.key_cache[layer_index], self.value_cache[layer_index]

    def update(
        self,
        key_states: torch.Tensor,
        value_states: torch.Tensor,
        layer_idx: int,
        cache_kwargs: Optional[Dict[str, Any]] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Updates the cache with the new `key_states` and `value_states` for the layer `layer_idx`.

        Parameters:
            key_states (`torch.Tensor`):
                The new key states to cache.
            value_states (`torch.Tensor`):
                The new value states to cache.
            layer_idx (`int`):
                The index of the layer to cache the states for.
            cache_kwargs (`Dict[str, Any]`, `optional`):
                Additional arguments for the cache subclass. No additional arguments are used in `Cache`.

        Return:
            A tuple containing the updated key and value states.
        """
        sin = cache_kwargs.get("sin")
        cos = cache_kwargs.get("cos")
        using_rope = cos is not None and sin is not None
        if layer_idx == 0:
            self._seen_tokens += key_states.shape[-2]
        
        if using_rope and layer_idx == 0:
            # BC: some models still pass `sin`/`cos` with 2 dims. In those models, they are the full sin/cos. Remove
            # after all RoPE models have a llama-like cache utilization.
            if cos.dim() == 2:
                self._cos_cache = cos
                self._sin_cache = sin
            else:
                if self._cos_cache is None:
                    self._cos_cache = cos[0, ...]
                    self._sin_cache = sin[0, ...]
                else:
                    self._cos_cache = torch.cat([self._cos_cache, cos[0, ...]], dim=0)
                    self._sin_cache = torch.cat([self._sin_cache, sin[0, ...]], dim=0)
        
        
        if len(self.key_cache) <= layer_idx:
            # There may be skipped layers, fill them with empty lists
            for _ in range(len(self.key_cache), layer_idx):
                self.key_cache.append([])
                self.value_cache.append([])
            self.key_cache.append(key_states)
            self.value_cache.append(value_states)
        elif len(self.key_cache[layer_idx]) == 0:  # fills previously skipped layers; checking for tensor causes errors
            self.key_cache[layer_idx] = key_states
            self.value_cache[layer_idx] = value_states
        else:
            self.key_cache[layer_idx] = torch.cat([self.key_cache[layer_idx], key_states], dim=-2)
            self.value_cache[layer_idx] = torch.cat([self.value_cache[layer_idx], value_states], dim=-2)

        return self.key_cache[layer_idx], self.value_cache[layer_idx]


import transformers.cache_utils

transformers.cache_utils.DynamicCache = FinchCache

from transformers import MistralForCausalLM
from transformers.utils import logging
from transformers.models.mistral.modeling_mistral import MistralConfig
from transformers import AutoTokenizer
import pandas as pd
from io import StringIO
logger = logging.get_logger(__name__)

class MistralForCompressedCausalLM(MistralForCausalLM):
    def __init__(self, config: MistralConfig, mode, compression_factor, split_size, target_token, condition="question", normalize=True, is_full=False, distance_metric=None, is_table=False, num_tuples=0):  # changed by GC
        config._attn_implementation = "sdpa"
        super().__init__(config)
        self.compression_factor = compression_factor  # added by GC
        self.split_size = split_size  # added by GC
        self.segment_length = self.split_size
        self.mode = mode
        self.target_token = target_token
        self.condition = condition
        self.normalize = normalize




    def generate(
        self,
        accelerator,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.LongTensor] = None,
        question_ids: Optional[torch.LongTensor] = None,
        question_attention_mask: Optional[torch.FloatTensor] = None,
        context_ids: Optional[torch.LongTensor] = None,
        context_attention_mask: Optional[torch.LongTensor] = None,
        **generate_kwargs
    ):

        past_key_values = FinchCache()
        total_context_length = context_ids.size(1)
        if total_context_length + input_ids.size(-1) > self.target_token and not self.target_token == 32768:
            start_processing_time = time.time()

            segment_length = self.split_size
            self.compression_factor = int(math.ceil(total_context_length / (self.target_token - input_ids.size(-1))))

            context_ids_list = torch.split(context_ids, segment_length, dim=1)
            context_attention_mask_list = torch.split(context_attention_mask, segment_length, dim=1)
 
            past_attention_mask = torch.zeros(attention_mask.size(0), 0, dtype=attention_mask.dtype, device=attention_mask.device)
            for step, (segment_context_ids, segment_attention_mask) in enumerate(zip(context_ids_list, context_attention_mask_list)):
                segment_attention_mask = torch.cat([past_attention_mask, segment_attention_mask], dim=1)
                past_cache_len = past_key_values._seen_tokens
                current_ids = torch.cat([segment_context_ids, question_ids], dim=1)
                current_attention_mask = torch.cat([segment_attention_mask, question_attention_mask], dim=1)
                with torch.no_grad():
                    output_question_aware = self.model(
                        input_ids=current_ids,
                        attention_mask=current_attention_mask,
                        output_attentions=True,
                        use_cache=True,
                        past_key_values=past_key_values
                    )
                current_seq_length = segment_context_ids.size(1)
                k = int(current_seq_length // self.compression_factor) + past_cache_len
                for layer_idx, layer_attention in enumerate(output_question_aware.attentions):
                    summed_attention = layer_attention.sum(dim=1)
                    tot_seq_len = summed_attention.size(2)
                    context_attention = summed_attention[:, current_seq_length:, :current_seq_length + past_cache_len]
                    non_zero_counts = torch.arange(1, tot_seq_len + 1, device=context_attention.device)
                    non_zero_counts = non_zero_counts[current_seq_length + past_cache_len:]
                    normalization_factors = non_zero_counts.float() / tot_seq_len
                    context_attention = context_attention * normalization_factors[None, :, None]
                    aggregated_attention = context_attention.sum(dim=1)
                    _, important_tokens = torch.topk(aggregated_attention , k=k, dim=-1, largest=True)
                    important_tokens, _ = torch.sort(important_tokens, dim=-1, descending=False)
                    inv_freq = self.model.layers[0].self_attn.rotary_emb.inv_freq
                    past_key_values.update_rope(layer_idx, past_key_values[layer_idx][0][:, :, :current_seq_length + past_cache_len],
                                                            important_tokens, inv_freq)

                past_attention_mask = torch.ones(segment_attention_mask.size(0), k, device=segment_attention_mask.device, dtype=segment_attention_mask.dtype)
            compressed_length = past_key_values._seen_tokens
            new_input_ids = torch.cat([torch.ones(size=(1, compressed_length), device=input_ids.device, dtype=torch.long), input_ids], dim=-1).to(input_ids.device)
            new_attention_mask = torch.cat([torch.ones(size=(1, compressed_length), device=input_ids.device), attention_mask], dim=-1).to(input_ids.device)
            model_output = super().generate(input_ids=new_input_ids, attention_mask=new_attention_mask, use_cache=True, past_key_values=past_key_values, **generate_kwargs)
            return model_output[:, compressed_length:, ...]

2. Effect of _rerotate_cos_sin and _apply_key_rotary_pos_emb
   The key cache (key_cache) in FinchCache contains keys that are already rotated by positional embeddings (RoPE). The _rerotate_cos_sin function calculates delta values (delta_pos = idx - important_pos_batch) to determine the offset needed for the cos and sin values. These are used to re-rotate the keys into their correct positions with _apply_key_rotary_pos_emb.
   This mechanism generalizes and extends the approach used in SinkCache, as seen in HuggingFace's implementation here.

3.Adjustments to generate and Position IDs
When using FinchCache, you don't need to adjust the model's generate function or pass custom position IDs. HuggingFace's Transformers library handles this internally. You can follow the clean approach described here.
In my example, I use dummy input IDs of length compressed_length to ensure attention masks align properly, though the actual generation relies on the cache. For reference:

new_input_ids = torch.cat([torch.ones(size=(1, compressed_length), device=input_ids.device, dtype=torch.long), input_ids], dim=-1).to(input_ids.device)
new_attention_mask = torch.cat([torch.ones(size=(1, compressed_length), device=input_ids.device), attention_mask], dim=-1).to(input_ids.device)
model_output = super().generate(input_ids=new_input_ids, attention_mask=new_attention_mask, use_cache=True, past_key_values=past_key_values, **generate_kwargs)
return model_output[:, compressed_length:, ...]

Let me know if further clarifications are needed or if you have additional questions!

[maxjeblick]

Thanks for the detailed answer. To clarify if I understand your idea correctly, These are used to re-rotate the keys into their correct positions with _apply_key_rotary_pos_emb, would imply the following:

1. Assume important_tokens indices are the same for all layers and all heads.
2. This implies, the pruning would be equivalent to deleting entire tokens from the input.
3. Pruning the tokens directly from the input sequence (before feeding it to the model) and then doing a normal forward pass (without compression) would then be equivalent to your pruning method using the entire sequence as input. I.e. the RoPE values associated with the pruned keys do not have any "holes".

[giulio98]

No, the two methods are not equivalent. Here’s why:
1. Pruning Tokens Directly from the Input Sequence:
If tokens are pruned directly from the input sequence before being fed into the model, they are completely excluded from the computation. This means that the keys and values generated at every subsequent layer will lack any information about the pruned tokens. In this case, the pruned tokens contribute no signal to the latent representations.
2. Pruning Keys and Values in the Latent Space:
In the pruning method applied in the latent space, the keys and values from earlier layers (including the first layer) already encode information about the full input sequence, including the pruned tokens. When pruning occurs in this latent space, the residual information from the pruned tokens persists in the compressed key and value states. This ensures that the attention mechanism in later layers can still access some of the contextual information from the pruned tokens, even though they are no longer directly represented.
3. Attention Sinks vs. Input Truncation:
In streaming llm (or similar approaches), the compressed keys and values still carry latent information aggregated from the entire input sequence. This is fundamentally different from simply truncating the input (e.g., keeping the first 4 tokens and the last few tokens) and feeding it to the model. Input truncation entirely removes the contribution of the pruned tokens, while key-value compression retains their processed representation, albeit in a reduced form.

The rerotation process serves the sole purpose of adjusting the positional encodings after pruning keys. Without rerotation, the model’s positional encoding system may interpret the pruned positions incorrectly, leading to degraded performance or even nonsensical outputs.

Let’s consider a concrete example:
1. Original Key-Position Pairing:
Before pruning, the keys and their associated positions look like this:
• Keys: [k0, k1, k2, k3, k4, k5, k6, k7, k8]
• Positions: [0, 1, 2, 3, 4, 5, 6, 7, 8]
2. After Pruning k2, k3, and k4:
The remaining keys and positions look like this:
• Keys: [k0, k1, k5, k6, k7, k8]
• Positions (without rerotation): [0, 1, 5, 6, 7, 8]
Here lies the problem: the positional encoding values no longer match the true sequential order of the tokens. For example:
• The model might interpret k5 as being at position 5, which implies it is further along in the sequence than it actually is.
• If the model’s maximum window size is 8, it may mistakenly assume the end of the sequence has been reached, causing it to generate incoherent outputs.
3. Rerotation to Adjust Positions:
To correct this, rerotation adjusts the positions to ensure the remaining keys have continuous, correctly aligned positions:
• Keys: [k0, k1, k5, k6, k7, k8]
• Positions (after rerotation): [0, 1, 2, 3, 4, 5]
Now, the positional encodings are consistent with the logical order of the keys. This ensures the model processes the sequence correctly and continues to generate meaningful outputs.

[maxjeblick]

If tokens are pruned directly from the input sequence before being fed into the model, they are completely excluded from the computation.

I simplified the example too much and -very obviously- these two methods are not equivalent.
My intention was to highlight the continuous nature of the pruned keys, i.e. your point:
Positions (after rerotation): [0, 1, 2, 3, 4, 5]

I would suggest the following workflow:

Implement Finch as a wrapper that accepts any press and modifies the forward hook as follows:

• Recompute keys for the total sequence without applying RoPE (e.g. keys = attn.k_proj(hidden_states))
• Prune keys
• Apply RoPE on pruned sequence -> pruned keys will have continuous representation.

This implementation creates some additional computational overhead (as keys will be recomputed), but it will be easy to implement.
I will start working on this.

[maxjeblick]

I've created a branch with the implementation mentioned above here.

[giulio98]

Thanks, I don't understand if the rerotation is applied by default for any presses. Also in streaming llm the rerotation happens both in the original implentation and in huggingface implementation (see SinkCache in cache_utils.py). The adjustment of the positional encoding according to our experiments gives dramatic improvements and we recommend setting it as default.

[maxjeblick]

I ran some (non exhaustive) benchmarks on max/rerotate_keys branch.
On this sample, it seems that rerotation does not improve metrics. I will check my branch for potential bugs, as
microsoft/Phi-3.5-mini-instruct scores deviate a lot (could be due to incorrect logic in qkv = module.qkv_proj(hidden_states)) .

As noticed, this benchmark is not representative, and there may be models/tasks where rerotating improves the score.

---

The adjustment of the positional encoding according to our experiments gives dramatic improvements

I don't see this as a contradiction. Using the example of sink cache

past_key_values = SinkCache(window_length=256, num_sink_tokens=4)
out = model.generate(**inputs, do_sample=False, max_new_tokens=30, past_key_values=past_key_values)
tokenizer.batch_decode(out, skip_special_tokens=True)[0]

it seems crucial that keys are rerotated, as model.generate uses (in the forward method)

past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
cache_position = torch.arange(past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device)

and past_key_values.get_seq_length() will be the compressed sequence length.

In kvpress, we use context_length = context_ids.shape[1] (see here), which is the uncompressed context length.

So, there are two potential solutions:

• Rerotate keys. Original position ids are not maintained. Start generating w.r.t. compressed sequence length.
• Keep original keys. Cache position ids are not modified, but contain "holes" where keys have been pruned. Start generating w.r.t. uncompressed sequence length.

---

In terms of moving forward, we will first merge this refactor PR and then implement the cache rerotation feature as a new pruner.
It will then be possible to use

from kvpress import RerotateKeysPruner, ExpectedAttentionScorer
press = RerotateKeysPruner(compression_ratio=0.1, scorer=ExpectedAttentionScorer())

# Modifications in pipeline
  answer = self.generate_answer(
      question_ids=question_ids.to(self.model.device),
      cache=cache,
      context_length=context_length if not isinstance(press, RerotateKeysPruner) else cache.get_seq_length(),
      max_new_tokens=max_new_tokens,
  )

[giulio98]

Thanks for the benchmark, now I see the problem: when adjusting the position embeddings, you have to pass to the generate function the compressed sequence length.
Here is why:
Considering the example above you will have this situation
• Keys: [k0, k1, k5, k6, k7, k8]
• Positions (after rerotation): [0, 1, 2, 3, 4, 5]

the last key is at position 5, but if you ask the model to generate from position 8, it will likely break the output generation.

[maxjeblick]

you have to pass to the generate function the compressed sequence length

Yes, exactly, this is what I'm doing here:

kvpress/kvpress/pipeline.py

Line 184 in 3b4338c

context_length if "apply_key_rerotation" not in press.wrappers_applied else cache.get_seq_length()

The press.wrappers_applied code is quite hacky, hence the refactoring PR that will allow for a straigthforward implementation.

[giulio98]

I see. The benchmark you posted earlier were done without passing the compressed length?

[maxjeblick]

The columns are as follows:

• average score no rerotation: Results from current main branch
• average score with rerotation: Results from current rerotate keys branch (with adjusted context length, see line 184 above).

[giulio98]

Did you use a Quantized Cache?
because in that case get_seq_length will return the length based on the _seen_tokens, which i think was not set when after applying the rerotation, see https://github.com/huggingface/transformers/blob/98e8062df3cf82b367e8a243963e4afb0d9d3407/src/transformers/cache_utils.py#L743