🚨Add an FP8 backend to `QuantizedCache`

@@ -240,6 +240,10 @@ A [`StoppingCriteria`] can be used to change when to stop generation (other than     - update     - lazy_initialization +[[autodoc]] Fp8QuantizedLayer+    - update+    - lazy_initialization+ [[autodoc]] Cache     - update     - early_initialization

@@ -176,10 +176,11 @@ responses = tokenizer.batch_decode(out[:,-28:], skip_special_tokens=True)  ## Quantized cache -The [`QuantizedCache`] reduces memory requirements by quantizing the KV values to a lower precision. [`QuantizedCache`] currently supports two quantization backends:+The [`QuantizedCache`] reduces memory requirements by quantizing the KV values to a lower precision. [`QuantizedCache`] currently supports three quantization backends:  - `hqq` supports int1, int2, int3, int4, and int8 datatypes. - `quanto` supports int2 and int4 datatypes. This is the default quantization backend.+- `fp8` stores the KV values in 8-bit floating point (`float8_e4m3fn`). It requires no extra dependency, and quantizes per-tensor rather than per-group, so it rejects `nbits`, `axis_key`, `axis_value`, `q_group_size` and `residual_length`.  > [!WARNING] > Quantizing the cache can harm latency if the context length is short and there is enough GPU memory available for generation without enabling cache quantization. Try to find a balance between memory efficiency and latency.@@ -238,6 +239,21 @@ print(tokenizer.batch_decode(out, skip_special_tokens=True)[0]) I like rock music because it's loud and energetic. It's a great way to express myself and rel ``` +The `fp8` backend does not take any additional parameter, and raises if it is given one of the group-wise ones.++```py+import torch+from transformers import AutoTokenizer, AutoModelForCausalLM++tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-4B-Instruct-2507")+model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-4B-Instruct-2507", dtype=torch.float16, device_map="auto")+inputs = tokenizer("I like rock music because", return_tensors="pt").to(model.device)++out = model.generate(**inputs, do_sample=False, max_new_tokens=20, cache_implementation="quantized", cache_config={"backend": "fp8"})+print(tokenizer.batch_decode(out, skip_special_tokens=True)[0])+I like rock music because it is powerful and emotional. I also like classical music because it is elegant and sophisticated. I enjoy+```+ ## Encoder-decoder cache  [`EncoderDecoderCache`] is designed for encoder-decoder models. It manages both the self-attention and cross-attention caches to ensure storage and retrieval of previous kv pairs. It is possible to individually set a different cache type for the encoder and decoder.

@@ -349,6 +349,7 @@         "DynamicIndexedLayer",         "DynamicLayer",         "EncoderDecoderCache",+        "Fp8QuantizedLayer",         "HQQQuantizedLayer",         "QuantizedCache",         "QuantoQuantizedLayer",@@ -468,6 +469,7 @@     from .cache_utils import DynamicIndexedLayer as DynamicIndexedLayer     from .cache_utils import DynamicLayer as DynamicLayer     from .cache_utils import EncoderDecoderCache as EncoderDecoderCache+    from .cache_utils import Fp8QuantizedLayer as Fp8QuantizedLayer     from .cache_utils import HQQQuantizedLayer as HQQQuantizedLayer     from .cache_utils import QuantizedCache as QuantizedCache     from .cache_utils import QuantoQuantizedLayer as QuantoQuantizedLayer

@@ -689,7 +689,36 @@ def reorder_cache(self, beam_idx: torch.LongTensor) -> None:             self.indexer_keys = self.indexer_keys.index_select(0, beam_idx.to(self.indexer_keys.device))  -class QuantizedLayer(DynamicLayer):+class BaseQuantizedLayer(DynamicLayer):+    """+    Base class for the cache layers that store their key and value states in a quantized form, in+    `_quantized_keys` and `_quantized_values`. Their layout is left to the backend, so the sequence length is+    tracked separately, in `cumulative_length`.+    """++    def __init__(self):+        super().__init__()+        self.cumulative_length = 0+        self._quantized_keys = self._quantized_values = None++    @abstractmethod+    def _quantize(self, tensor): ...++    @abstractmethod+    def _dequantize(self, q_tensor): ...++    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++    def get_seq_length(self) -> int:+        """Returns the sequence length of the cached states."""+        return self.cumulative_length+++class QuantizedLayer(BaseQuantizedLayer):     """     A quantized layer similar to what is described in the [KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache paper](https://huggingface.co/papers/2402.02750).     It allows the model to generate longer sequence length without allocating too much memory for the key and value caches by@@ -702,6 +731,9 @@ class QuantizedLayer(DynamicLayer):      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`.++    What is shared here is the residual window, not the grouping itself: a backend whose quantized states can be+    appended to, as FP8 does, has no use for it and subclasses `BaseQuantizedLayer` directly instead.     """      def __init__(@@ -718,10 +750,12 @@ def __init__(         self.axis_value = axis_value         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 +    @abstractmethod+    def _quantize(self, tensor, axis): ...+     def update(         self, key_states: torch.Tensor, value_states: torch.Tensor, *args, **kwargs     ) -> tuple[torch.Tensor, torch.Tensor]:@@ -765,12 +799,6 @@ def update(          return keys_to_return, values_to_return -    @abstractmethod-    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:@@ -798,14 +826,8 @@ def _apply_pending_reorder(self, keys: torch.Tensor, values: torch.Tensor) -> tu     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-  class QuantoQuantizedLayer(QuantizedLayer):     def __init__(@@ -919,6 +941,76 @@ def _dequantize(self, qtensor):         return tensor  +class Fp8QuantizedLayer(BaseQuantizedLayer):+    """+    A quantized layer storing the key and value states in 8-bit floating point. The per-tensor scale is frozen+    after the first `update`, so new states are appended as-is: no residual cache, no re-quantization.++    A finer-grained variant belongs here rather than under `QuantizedLayer`, as grouping along the head+    dimension keeps states appendable and only turns the stored scale into a tensor.+    """++    fp8_dtype = torch.float8_e4m3fn+    fp8_max = torch.finfo(fp8_dtype).max++    def lazy_initialization(self, key_states: torch.Tensor, value_states: torch.Tensor) -> None:+        super().lazy_initialization(key_states, value_states)+        self._quantized_keys = self._quantize(key_states)+        self._quantized_values = self._quantize(value_states)++    def _quantize(self, tensor, scale=None):+        """Quantize `tensor` to FP8, calibrating the scale on its amplitude if it is not provided."""+        if scale is None:+            # `clamp` guards against an all-zeros tensor, which would give a null scale+            scale = (tensor.abs().amax().float() / self.fp8_max).clamp(min=torch.finfo(torch.float32).tiny)+        return (tensor.float() / scale).to(self.fp8_dtype), scale++    def _dequantize(self, qtensor):+        quant_tensor, scale = qtensor+        return (quant_tensor.float() * scale).to(self.dtype)++    def update(+        self, key_states: torch.Tensor, value_states: torch.Tensor, *args, **kwargs+    ) -> tuple[torch.Tensor, torch.Tensor]:+        """+        Update the key and value caches in-place, and return the necessary keys and value states.++        Args:+            key_states (`torch.Tensor`): The new key states to cache.+            value_states (`torch.Tensor`): The new value states to cache.++        Returns:+            tuple[`torch.Tensor`, `torch.Tensor`]: The key and value states.+        """+        self.cumulative_length += key_states.shape[-2]++        if not self.is_initialized:+            self.lazy_initialization(key_states, value_states)+            return key_states, value_states++        quantized_keys, key_scale = self._quantized_keys+        quantized_values, value_scale = self._quantized_values+        new_keys, _ = self._quantize(key_states, scale=key_scale)+        new_values, _ = self._quantize(value_states, scale=value_scale)+        self._quantized_keys = (torch.cat([quantized_keys, new_keys], dim=-2), key_scale)+        self._quantized_values = (torch.cat([quantized_values, new_values], dim=-2), value_scale)++        keys_to_return = self._dequantize(self._quantized_keys)+        values_to_return = self._dequantize(self._quantized_values)+        return keys_to_return, values_to_return++    def reorder_cache(self, beam_idx: torch.LongTensor) -> None:+        """Reorders the cache for beam search. The per-tensor scale lets the FP8 states be reordered as-is."""+        if not self.is_initialized:+            return++        quantized_keys, key_scale = self._quantized_keys+        quantized_values, value_scale = self._quantized_values+        beam_idx = beam_idx.to(quantized_keys.device)+        self._quantized_keys = (quantized_keys.index_select(0, beam_idx), key_scale)+        self._quantized_values = (quantized_values.index_select(0, beam_idx), value_scale)++ class LinearAttentionCacheLayerMixin(ABC):     """Base, abstract class for a linear attention single layer's cache.""" @@ -1935,7 +2027,7 @@ class QuantizedCache(Cache):      Args:         backend (`str`):-            The quantization backend to use. One of `("quanto", "hqq").+            The quantization backend to use. One of `("quanto", "hqq", "fp8")`.         config (`PreTrainedConfig`):             The config of the model for which this Cache will be used.         nbits (`int`, *optional*, defaults to 4):@@ -1947,26 +2039,52 @@ class QuantizedCache(Cache):         q_group_size (`int`, *optional*, defaults to 64):             Quantization is done per-channel according to a set `q_group_size` for both keys and values.         residual_length (`int`, *optional*, defaults to 128):-            Maximum capacity for the original precision cache+            Maximum capacity for the original precision cache.++    The last five arguments configure the group-wise integer backends. The `fp8` backend quantizes per-tensor and+    rejects them.     """      def __init__(         self,         backend: str,         config: PreTrainedConfig,-        nbits: int = 4,-        axis_key: int = 0,-        axis_value: int = 0,-        q_group_size: int = 64,-        residual_length: int = 128,+        nbits: int | None = None,+        axis_key: int | None = None,+        axis_value: int | None = None,+        q_group_size: int | None = None,+        residual_length: int | None = None,     ):+        group_wise_kwargs = {+            "nbits": nbits,+            "axis_key": axis_key,+            "axis_value": axis_value,+            "q_group_size": q_group_size,+            "residual_length": residual_length,+        }         if backend == "quanto":             layer_class = QuantoQuantizedLayer         elif backend == "hqq":             layer_class = HQQQuantizedLayer+        elif backend == "fp8":+            layer_class = Fp8QuantizedLayer         else:             raise ValueError(f"Unknown quantization backend `{backend}`") +        # A backend either takes all of the group-wise options or none of them. A future group-wise fp8 variant+        # would need this to be per-backend instead.+        if backend == "fp8":+            unsupported = [name for name, value in group_wise_kwargs.items() if value is not None]+            if len(unsupported) > 0:+                raise ValueError(+                    f"The `fp8` backend quantizes per-tensor and does not support {unsupported}. Use the `quanto` or "+                    "`hqq` backend to quantize group-wise."+                )+            layer_kwargs = {}+        else:+            # The unset ones are left out, so that the layer's own defaults apply+            layer_kwargs = {name: value for name, value in group_wise_kwargs.items() if value is not None}+         config = config.get_text_config(decoder=True)         layer_types, _ = get_layer_types_and_kwargs(config)         invalid_layer_types = set(layer_types) - {"full_attention"}@@ -1975,10 +2093,7 @@ def __init__(                 "`QuantizedCache` is only supported for models with only full attention layers. We found the following invalid layer "                 f"types: {invalid_layer_types}"             )-        layers = [-            layer_class(nbits, axis_key, axis_value, q_group_size, residual_length)-            for _ in range(config.num_hidden_layers)-        ]+        layers = [layer_class(**layer_kwargs) for _ in range(config.num_hidden_layers)]         super().__init__(layers=layers)  

@@ -18,7 +18,7 @@  import torch -from ...cache_utils import DynamicCache, QuantizedCache, QuantizedLayer, StaticCache+from ...cache_utils import BaseQuantizedLayer, DynamicCache, QuantizedCache, StaticCache from ...generation import (     GenerateDecoderOnlyOutput,     GenerationConfig,@@ -399,7 +399,7 @@ def _reset_negative_cache_for_audio_start(             mask_4d = diffusion_start_mask.view(-1, 1, 1, 1)             # Reset the KV cache             for layer in negative_model_kwargs["past_key_values"].layers:-                if isinstance(layer, QuantizedLayer):+                if isinstance(layer, BaseQuantizedLayer):                     continue                 if layer.keys is not None and layer.values is not None:                     layer_mask_4d = mask_4d.to(layer.keys.device)

@@ -68,6 +68,7 @@         DynamicIndexedLayer,         DynamicLayer,         DynamicSlidingWindowLayer,+        Fp8QuantizedLayer,         LinearAttentionAndFullAttentionLayer,         LinearAttentionAndSlidingWindowAttentionLayer,         LinearAttentionLayer,@@ -184,16 +185,23 @@ 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):+    @parameterized.expand(["quanto", "fp8"])+    def test_quantized_layer_beam_reorder(self, backend):         """         `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.+        the quantized states in the previous beam order. FP8 has no residual cache, and reorders right away.         """-        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)+        if backend == "quanto":+            if not is_optimum_quanto_available():+                self.skipTest("Quanto is not available")+            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, atol = torch.arange(1, 5, dtype=torch.float32).view(4, 1, 1, 1), 1e-2+        else:+            layer = Fp8QuantizedLayer()+            # Same, with constants exactly representable in `float8_e4m3fn` once divided by the calibrated scale+            row_values, atol = torch.tensor([1.0, 2.0, 4.0, 8.0]).view(4, 1, 1, 1), 0          def states(order, seq_len):             # Row `i` is filled with the constant of the original row `order[i]`, as the model would produce it@@ -211,8 +219,8 @@ def states(order, seq_len):             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)+            torch.testing.assert_close(keys, expected, rtol=0, atol=atol)+            torch.testing.assert_close(values, expected, rtol=0, atol=atol)      @parameterized.expand(         [@@ -242,19 +250,61 @@ def test_quantized_layer_chunked_updates(self, _, chunk_lengths, expected_residu                 layer._dequantize(layer._quantized_keys).shape[-2], sequence_length - expected_residual_length             ) -    @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)+    @parameterized.expand(["quanto", "fp8"])+    def test_quantized_layer_reset(self, backend):+        """`reset` must also drop the quantized states, which hold most of the cache, and any calibrated scale."""+        if backend == "quanto":+            if not is_optimum_quanto_available():+                self.skipTest("Quanto is not available")+            layer = QuantoQuantizedLayer(nbits=4, q_group_size=16, residual_length=4)+        else:+            layer = Fp8QuantizedLayer() -        # 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))+        # 5 tokens for a residual length of 4, so the first flush has happened and most of the cache is quantized.+        # The large amplitude makes a stale FP8 scale obvious on the next calibration.+        layer.update(torch.rand(4, 2, 5, 8) * 100, torch.rand(4, 2, 5, 8) * 100)+        calibrated_scale = layer._quantized_keys[1].clone() if backend == "fp8" else None         layer.reset()          self.assertEqual(layer.get_seq_length(), 0)+        self.assertIsNone(layer._quantized_keys)+        self.assertIsNone(layer._quantized_values)+         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)+        if backend == "fp8":+            self.assertLess(layer._quantized_keys[1].item(), calibrated_scale.item())++    def test_fp8_quantized_layer(self):+        """Tests that `Fp8QuantizedLayer` appends new tokens using the scale calibrated on the first `update`."""+        layer = Fp8QuantizedLayer()++        # As in the attention layer, the states come out of a transpose, so they are not contiguous+        keys, values = torch.randn(1, 5, 2, 8).transpose(1, 2), torch.randn(1, 5, 2, 8).transpose(1, 2)+        layer.update(keys, values)+        quantized_keys, calibrated_scale = layer._quantized_keys+        self.assertEqual(quantized_keys.dtype, torch.float8_e4m3fn)+        self.assertEqual(quantized_keys.shape, keys.shape)+        calibrated_scale = calibrated_scale.clone()++        # Even states with a much larger amplitude do not change the scale, which is frozen after calibration+        cached_keys, _ = layer.update(torch.randn(1, 2, 1, 8) * 100, torch.randn(1, 2, 1, 8))+        self.assertEqual(cached_keys.shape, (1, 2, 6, 8))+        self.assertEqual(cached_keys.dtype, keys.dtype)+        quantized_keys, scale = layer._quantized_keys+        self.assertTrue(torch.equal(scale, calibrated_scale))+        self.assertEqual(quantized_keys.shape, (1, 2, 6, 8))++    def test_fp8_quantized_cache_rejects_group_wise_options(self):+        """Tests that the `fp8` backend, which quantizes per-tensor, does not silently drop the group-wise options."""+        config = LlamaConfig(num_hidden_layers=2)+        with self.assertRaisesRegex(ValueError, "does not support"):+            QuantizedCache(backend="fp8", config=config, q_group_size=128, nbits=2)++        layer = QuantizedCache(backend="fp8", config=config).layers[0]+        self.assertIsInstance(layer, Fp8QuantizedLayer)+        self.assertFalse(hasattr(layer, "q_group_size"))      def test_dynamic_cache_uses_per_layer_sliding_windows(self):         config = LlamaConfig(@@ -600,21 +650,26 @@ 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"), ("fp8")])     def test_quantized_cache_generation(self, backend):-        """Tests that QuantizedCache works as expected for both `quanto` and `hqq` backends."""+        """Tests that QuantizedCache works as expected for the `quanto`, `hqq` and `fp8` backends."""+        # The group-wise integer backends share the same options, while `fp8` quantizes per-tensor+        cache_config = {"backend": backend, "nbits": 4, "q_group_size": 16, "residual_length": 4}         if backend == "quanto":             if not is_optimum_quanto_available():                 self.skipTest("Quanto is not available")-            axis_key, axis_value = 0, 0+            cache_config.update({"axis_key": 0, "axis_value": 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":             if not is_hqq_available():                 self.skipTest("HQQ is not available")-            axis_key, axis_value = 1, 1+            cache_config.update({"axis_key": 1, "axis_value": 1})             # HQQ has slightly different numerics             expected_generation = ["The cat's whiskers are also a sign of anxiety."]+        elif backend == "fp8":+            cache_config = {"backend": backend}+            expected_generation = ["The cat's whiskers are also a sign of anxiety."]         else:             return @@ -626,14 +681,7 @@ def test_quantized_cache_generation(self, backend):             max_new_tokens=10,             return_dict_in_generate=True,             cache_implementation="quantized",-            cache_config={-                "backend": backend,-                "nbits": 4,-                "q_group_size": 16,-                "residual_length": 4,-                "axis_key": axis_key,-                "axis_value": axis_value,-            },+            cache_config=cache_config,             disable_compile=True,         )