Fix quantized cache

@@ -720,6 +720,9 @@ class QuantizedLayer(DynamicLayer):     is set as a maximum capacity for the original precision cache. When the length goes beyond maximum capacity, the original     precision cache is discarded and moved into the quantized cache. The quantization is done per-channel with a set `q_group_size`     for both Keys and Values, in contrast to what was described in the paper.++    Note that a beam reordering is only applied to the quantized states on the next `update`, so they may be left in a+    stale batch order in between, see `reorder_cache`.     """      def __init__(@@ -737,6 +740,8 @@ def __init__(         self.q_group_size = q_group_size         self.residual_length = residual_length         self.cumulative_length = 0+        # Beam reordering of the quantized states is deferred until the next `update`, see `reorder_cache`+        self._pending_beam_idx = None      def update(         self, key_states: torch.Tensor, value_states: torch.Tensor, *args, **kwargs@@ -762,13 +767,18 @@ def update(          dequant_keys = self._dequantize(self._quantized_keys)         dequant_values = self._dequantize(self._quantized_values)+        # Beam idx reordering is deferred to here, to avoid a lossy round trip around quantization+        dequant_keys, dequant_values = self._apply_pending_reorder(dequant_keys, dequant_values)+         keys_to_return = torch.cat([dequant_keys, self.keys, key_states], dim=-2)         values_to_return = torch.cat([dequant_values, self.values, value_states], dim=-2)         if self.keys.dim() == 4 and self.keys.shape[-2] + 1 >= self.residual_length:             self._quantized_keys = self._quantize(keys_to_return.contiguous(), axis=self.axis_key)             self._quantized_values = self._quantize(values_to_return.contiguous(), axis=self.axis_value)             self.keys = torch.tensor([], dtype=key_states.dtype, device=key_states.device)             self.values = torch.tensor([], dtype=key_states.dtype, device=key_states.device)+            # The reordering is now baked into the quantized states+            self._pending_beam_idx = None         else:             self.keys = torch.cat([self.keys, key_states], dim=-2)             self.values = torch.cat([self.values, value_states], dim=-2)@@ -781,6 +791,37 @@ def _quantize(self, tensor, axis): ...     @abstractmethod     def _dequantize(self, q_tensor): ... +    def reorder_cache(self, beam_idx: torch.LongTensor) -> None:+        """Reorders this layer's cache for beam search."""+        if not self.is_initialized:+            return++        # Deferred to the next `update`, which dequantizes the states anyway. It is cloned as `beam_idx` belongs to+        # the caller, and composed with any pending one to cover several reorders in a row.+        beam_idx = beam_idx.to(self.device)+        self._pending_beam_idx = (+            beam_idx.clone() if self._pending_beam_idx is None else self._pending_beam_idx.index_select(0, beam_idx)+        )++        # Optional, as the residual cache is emptied whenever it is flushed into the quantized states+        if self.keys.numel() > 0:+            self.keys = self.keys.index_select(0, beam_idx.to(self.keys.device))+            self.values = self.values.index_select(0, beam_idx.to(self.values.device))++    def _apply_pending_reorder(self, keys: torch.Tensor, values: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:+        """Applies the reordering deferred by `reorder_cache` to freshly dequantized states."""+        if self._pending_beam_idx is None:+            return keys, values+        beam_idx = self._pending_beam_idx.to(keys.device)+        return keys.index_select(0, beam_idx), values.index_select(0, beam_idx)++    def reset(self) -> None:+        """Resets the cache values while preserving the objects."""+        super().reset()+        # The quantized states are dropped instead of zeroed, so that the next `update` quantizes from scratch+        self._quantized_keys = self._quantized_values = None+        self._pending_beam_idx = None+     def get_seq_length(self) -> int:         """Returns the sequence length of the cached states."""         return self.cumulative_length

@@ -2293,7 +2293,10 @@ def _prepare_cache_for_generation(                     "cache, please open an issue and tag @zucchini-nlp."                 ) -            cache_config = generation_config.cache_config if generation_config.cache_config is not None else {}+            # Copied, as `generate` may be called several times with the same `cache_config`, whose keys are consumed+            cache_config = (+                copy.deepcopy(generation_config.cache_config) if generation_config.cache_config is not None else {}+            )             cache_config.setdefault("config", self.config.get_text_config(decoder=True))             backend = cache_config.pop("backend", "quanto")             model_kwargs[cache_name] = QuantizedCache(backend=backend, **cache_config)

@@ -30,6 +30,7 @@     cleanup,     get_gpu_count,     is_torch_available,+    require_optimum_quanto,     require_torch,     require_torch_accelerator,     require_torch_gpu,@@ -69,6 +70,7 @@         LinearAttentionAndFullAttentionLayer,         LinearAttentionAndSlidingWindowAttentionLayer,         LinearAttentionLayer,+        QuantoQuantizedLayer,         StaticLayer,         StaticSlidingWindowLayer,     )@@ -181,6 +183,50 @@ def test_max_cache_len_ignores_linear_attention_layers(self):         cache = StaticCache(config=config, max_cache_len=8)         self.assertEqual(cache.get_max_length(), 8) +    @require_optimum_quanto+    def test_quantized_layer_beam_reorder(self):+        """+        `reorder_cache` must reorder the quantized states as well, not only the residual cache. It used to index the+        (emptied) residual cache while guarding on the total sequence length, so it either raised or silently left+        the quantized states in the previous beam order.+        """+        layer = QuantoQuantizedLayer(nbits=4, q_group_size=16, residual_length=4)+        # Row `i` of the states holds the constant `i + 1`, so the beam order can be read back off the cache+        row_values = torch.arange(1, 5, dtype=torch.float32).view(4, 1, 1, 1)++        def states(order, seq_len):+            # Row `i` is filled with the constant of the original row `order[i]`, as the model would produce it+            return row_values.index_select(0, order).expand(4, 2, seq_len, 8).clone()++        order = torch.arange(4)+        layer.update(states(order, 5), states(order, 5))++        # Reorder on every step, with duplicated indices (`beam_idx` is a gather, several beams may pick the same+        # source beam) and crossing the residual cache flushes, which bake the pending reordering in.+        for beam_idx in [[1, 0, 3, 2], [2, 2, 0, 1], [3, 1, 2, 0], [0, 1, 2, 3], [1, 3, 0, 2], [2, 0, 1, 3]]:+            beam_idx = torch.tensor(beam_idx)+            layer.reorder_cache(beam_idx)+            order = order.index_select(0, beam_idx)+            keys, values = layer.update(states(order, 1), states(order, 1))+            # Every cached position of a row must hold that row's constant, whatever the beam shuffling+            expected = row_values.index_select(0, order).expand_as(keys)+            torch.testing.assert_close(keys, expected, rtol=0, atol=1e-2)+            torch.testing.assert_close(values, expected, rtol=0, atol=1e-2)++    @require_optimum_quanto+    def test_quantized_layer_reset(self):+        """`reset` must also drop the quantized states, which hold most of the cache."""+        layer = QuantoQuantizedLayer(nbits=4, q_group_size=16, residual_length=4)++        # 5 tokens for a residual length of 4: the first flush has happened, so most of the cache is quantized+        layer.update(torch.rand(4, 2, 5, 8), torch.rand(4, 2, 5, 8))+        layer.reset()++        self.assertEqual(layer.get_seq_length(), 0)+        keys, _ = layer.update(torch.rand(4, 2, 3, 8), torch.rand(4, 2, 3, 8))+        self.assertEqual(keys.shape[-2], 3)+        self.assertEqual(layer.get_seq_length(), 3)+     def test_dynamic_cache_uses_per_layer_sliding_windows(self):         config = LlamaConfig(             hidden_size=64,@@ -361,6 +407,14 @@ def _skip_on_failed_cache_prerequisites(test, cache_implementation):                 test.skipTest("Offloaded static caches require exactly 1 accelerator")  +def _set_sliding_window(config, cache_implementation):+    """+    Sets the sliding window on `config`, from which the cache layer types are inferred: the sliding cache+    implementations need one to build sliding layers, and the other ones must not have one.+    """+    config.sliding_window = 256 if cache_implementation in ["sliding_window", "hybrid", "hybrid_chunked"] else None++ class CacheIntegrationTest(unittest.TestCase):     """Fast cache integration tests that share the same small model""" @@ -371,12 +425,17 @@ def setUpClass(cls):         cls.model = AutoModelForCausalLM.from_pretrained(             "HuggingFaceTB/SmolLM2-135M-Instruct", device_map="auto", dtype=torch.float16         )-        cls.model.config.sliding_window = 256  # hack to enable the use of caches with sliding windows++    def setUp(self):+        # The model is shared across tests, so the sliding window one of them opts into through+        # `_set_sliding_window` must not leak into the next one. `SmolLM2` has no sliding window of its own.+        self.model.config.sliding_window = None      @parameterized.expand(TEST_CACHE_IMPLEMENTATIONS)     def test_cache_batched(self, cache_implementation):         """Sanity check: caches' `.update` function expects batched inputs"""         _skip_on_failed_cache_prerequisites(self, cache_implementation)+        _set_sliding_window(self.model.config, cache_implementation)          EXPECTED_GENERATION = ["A sequence: 1, 2, 3, 4, 5, 6, 7, 8,", "A sequence: A, B, C, D, E, F, G, H"] @@ -423,6 +482,7 @@ def test_cache_beam_search(self, cache_implementation):         (an output sequence contains multiple beam indices).         """         _skip_on_failed_cache_prerequisites(self, cache_implementation)+        _set_sliding_window(self.model.config, cache_implementation)         if cache_implementation == "offloaded_hybrid_chunked":             # TODO (joao, cyril): something is off with `offloaded_hybrid_chunked`: the             # output sequence (and the corresponding beam scores, if we add `output_scores=True`) are significantly@@ -453,7 +513,7 @@ def test_cache_beam_search(self, cache_implementation):         decoded = self.tokenizer.decode(gen_out.sequences, skip_special_tokens=True)         self.assertListEqual(decoded, EXPECTED_GENERATION) -    @parameterized.expand([("quanto"), ("HQQ")])+    @parameterized.expand([("quanto"), ("hqq")])     def test_quantized_cache_generation(self, backend):         """Tests that QuantizedCache works as expected for both `quanto` and `hqq` backends."""         if backend == "quanto":@@ -462,7 +522,7 @@ def test_quantized_cache_generation(self, backend):             axis_key, axis_value = 0, 0             # This output is taken from a run with the same parameters, and is known to be correct             expected_generation = ["The cat's whiskers are also a sign of anxiety."]-        elif backend == "HQQ":+        elif backend == "hqq":             if not is_hqq_available():                 self.skipTest("HQQ is not available")             axis_key, axis_value = 1, 1@@ -497,10 +557,35 @@ def test_quantized_cache_generation(self, backend):          # Check that something is actually quantized +    @require_optimum_quanto+    def test_quantized_cache_config_is_not_mutated(self):+        """+        Tests that `generate` does not consume the entries of the `cache_config` it is given, which would silently+        change the cache of any subsequent call sharing that dict.+        """+        inputs = self.tokenizer(["The cat"], return_tensors="pt").to(self.model.device)+        cache_config = {"backend": "quanto", "nbits": 4, "q_group_size": 16, "residual_length": 4}+        expected_cache_config = cache_config.copy()++        for _ in range(2):+            gen_out = self.model.generate(+                **inputs,+                do_sample=False,+                max_new_tokens=3,+                return_dict_in_generate=True,+                cache_implementation="quantized",+                cache_config=cache_config,+                disable_compile=True,+            )+            self.assertIsInstance(gen_out.past_key_values, QuantizedCache)+            self.assertEqual(len(gen_out.past_key_values.layers), self.model.config.num_hidden_layers)+            self.assertEqual(cache_config, expected_cache_config)+     @parameterized.expand(TEST_CACHE_IMPLEMENTATIONS)     def test_cache_extra_left_padding(self, cache_implementation):         """Tests that adding extra left-padding does not affect the generation with the cache"""         _skip_on_failed_cache_prerequisites(self, cache_implementation)+        _set_sliding_window(self.model.config, cache_implementation)          EXPECTED_GENERATION = ["The cat's whiskers are also a sign of anxiety."] @@ -878,9 +963,7 @@ def test_cache_gptj_model(self, cache_implementation):         # pytorch_model.bin in the original repo (wrong wte.weight bytes → wrong golden outputs)         model_id = "hf-internal-testing/tiny-random-GPTJForCausalLM-for-CacheHardIntegrationTest"         pipe = pipeline("text-generation", model=model_id, dtype=torch.bfloat16)-        pipe.model.config.sliding_window = (-            256 if cache_implementation in ["sliding_window", "hybrid", "hybrid_chunked"] else None-        )+        _set_sliding_window(pipe.model.config, cache_implementation)         out = pipe(             "hello world",             cache_implementation=cache_implementation,