Fix _remove_answer_from_cache crash with QuantizedCache(backend='hqq')

README.md

@@ -172,7 +172,7 @@ We support KV cache quantization through the transformers `QuantizedCache` class
 ```python
 from transformers import QuantizedCache
 
-cache = QuantizedCache(backend="quanto", nbits=4)
+cache = QuantizedCache(backend="quanto", config=pipe.model.config, nbits=4)
 
 pipe(..., cache=cache)
 ```

kvpress/pipeline.py

@@ -238,32 +238,36 @@ def _forward(
                     context_length = cache.get_seq_length()
 
                 cache_seq_lengths = [cache.get_seq_length(layer_idx) for layer_idx in range(len(cache))]
+                # A quantized cache can't be sliced back to the context, so save its layer states instead
+                layer_states = (
+                    [vars(layer).copy() for layer in cache.layers] if isinstance(cache, QuantizedCache) else None
+                )
                 answer = self.generate_answer(
                     question_ids=question_ids.to(self.model.device),
                     cache=cache,
                     context_length=context_length,
                     max_new_tokens=max_new_tokens,
                 )
-                self._remove_answer_from_cache(cache, cache_seq_lengths)
+                self._remove_answer_from_cache(cache, cache_seq_lengths, layer_states)
 
                 answers.append(answer)
         return answers
 
-    def _remove_answer_from_cache(self, cache: Cache, cache_seq_lengths: list[int]):
+    def _remove_answer_from_cache(
+        self, cache: Cache, cache_seq_lengths: list[int], layer_states: Optional[list[dict]] = None
+    ):
+        if isinstance(cache, QuantizedCache):
+            # A QuantizedLayer flushes its full-precision residual into the quantized storage once it reaches
+            # residual_length, so the question and answer can't be sliced off. Restore the saved states instead.
+            # A shallow copy is enough: QuantizedLayer.update assigns new tensors instead of writing in place.
+            for layer, state in zip(cache.layers, layer_states):
+                vars(layer).update(state)
+            return
 
         for layer_idx, sequence_length in enumerate(cache_seq_lengths):
             cache.layers[layer_idx].keys = cache.layers[layer_idx].keys[:, :, :sequence_length]
             cache.layers[layer_idx].values = cache.layers[layer_idx].values[:, :, :sequence_length]
 
-        if isinstance(cache, QuantizedCache):
-            for layer_idx, sequence_length in enumerate(cache_seq_lengths):
-                cache.layers[layer_idx]._quantized_keys = cache.layers[layer_idx]._quantized_keys[
-                    :, :, :sequence_length
-                ]
-                cache.layers[layer_idx]._quantized_values = cache.layers[layer_idx]._quantized_values[
-                    :, :, :sequence_length
-                ]
-
     def generate_answer(
         self, question_ids: torch.Tensor, cache: Cache, context_length: int, max_new_tokens: int
     ) -> str:

tests/test_pipeline.py

@@ -7,7 +7,7 @@
 import pytest
 import torch
 from transformers import AutoTokenizer, DynamicCache, QuantizedCache
-from transformers.utils import is_flash_attn_2_available, is_optimum_quanto_available
+from transformers.utils import is_flash_attn_2_available, is_hqq_available, is_optimum_quanto_available
 
 from kvpress import ExpectedAttentionPress, KVzipPress
 from kvpress.pipeline import KVPressTextGenerationPipeline
@@ -106,13 +106,22 @@ def test_pipeline_answer_is_correct(danube_500m_model, caplog):  # noqa: F811
     assert "Compressed Context Length: 16" in messages
 
 
-@pytest.mark.skipif(not is_optimum_quanto_available(), reason="Optimum Quanto is not available")
-def test_pipeline_with_quantized_cache(kv_press_danube_pipeline, caplog):  # noqa: F811
+@pytest.mark.parametrize(
+    "backend",
+    [
+        pytest.param(
+            "quanto",
+            marks=pytest.mark.skipif(not is_optimum_quanto_available(), reason="Optimum Quanto is not available"),
+        ),
+        pytest.param("hqq", marks=pytest.mark.skipif(not is_hqq_available(), reason="HQQ is not available")),
+    ],
+)
+def test_pipeline_with_quantized_cache(backend, kv_press_danube_pipeline, caplog):  # noqa: F811
     with caplog.at_level(logging.DEBUG):
         context = "This is a test article. It was written on 2022-01-01."
         questions = ["When was this article written?"]
         press = ExpectedAttentionPress(compression_ratio=0.4)
-        cache = QuantizedCache(backend="quanto", config=kv_press_danube_pipeline.model.config, nbits=4)
+        cache = QuantizedCache(backend=backend, config=kv_press_danube_pipeline.model.config, nbits=4)
         answers = kv_press_danube_pipeline(context, questions=questions, press=press, cache=cache)["answers"]
 
     assert len(answers) == 1
@@ -126,6 +135,36 @@ def test_pipeline_with_quantized_cache(kv_press_danube_pipeline, caplog):  # noq
     assert "Compressed Context Length: 16" in messages
 
 
+@pytest.mark.parametrize(
+    "backend",
+    [
+        pytest.param(
+            "quanto",
+            marks=pytest.mark.skipif(not is_optimum_quanto_available(), reason="Optimum Quanto is not available"),
+        ),
+        pytest.param("hqq", marks=pytest.mark.skipif(not is_hqq_available(), reason="HQQ is not available")),
+    ],
+)
+def test_pipeline_with_quantized_cache_multiple_questions(backend, kv_press_danube_pipeline):  # noqa: F811
+    # residual_length=4 makes the full-precision residual flush into the quantized storage while answering,
+    # so the cache must be restored to the context after each answer, not just sliced.
+    context = "This is a test article. It was written on 2022-01-01."
+    questions = ["When was this article written?", "What is this article?"]
+    press = ExpectedAttentionPress(compression_ratio=0.4)
+
+    def make_cache():
+        return QuantizedCache(backend=backend, config=kv_press_danube_pipeline.model.config, nbits=4, residual_length=4)
+
+    cache = make_cache()
+    answers = kv_press_danube_pipeline(context, questions=questions, press=press, cache=cache)["answers"]
+
+    assert cache.get_seq_length() == 16  # the compressed context, without any question or answer
+
+    for question, answer in zip(questions, answers):
+        single = kv_press_danube_pipeline(context, question=question, press=press, cache=make_cache())["answer"]
+        assert answer == single
+
+
 def test_pipeline_compresses_context(unit_test_model, caplog):  # noqa: F811
     with caplog.at_level(logging.DEBUG):
         answers = generate_answer(unit_test_model)