compressed-tensors 0.6.0__py3-none-any.whl → 0.7.0__py3-none-any.whl

compressed_tensors/quantization/cache.py

@@ -0,0 +1,201 @@
+# Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from enum import Enum
+from typing import Any, Dict, List, Optional, Tuple
+
+from compressed_tensors.quantization.observers import Observer
+from compressed_tensors.quantization.quant_args import QuantizationArgs
+from torch import Tensor
+from transformers import DynamicCache as HFDyanmicCache
+
+
+class KVCacheScaleType(Enum):
+    KEY = "k_scale"
+    VALUE = "v_scale"
+
+
+class QuantizedKVParameterCache(HFDyanmicCache):
+
+    """
+    Quantized KV cache used in the forward call based on HF's dynamic cache.
+    Quantization strategy (tensor, group, channel) set from Quantization arg's strategy
+    Singleton, so that the same cache gets reused in all forward call of self_attn.
+    Each time forward is called, .update() is called, and ._quantize(), ._dequantize()
+     gets called appropriately.
+    The size of tensor is
+     `[batch_size, num_heads, seq_len - residual_length, head_dim]`.
+
+
+    Triggered by adding kv_cache_scheme in the recipe.
+
+    Example:
+
+    ```python3
+    recipe = '''
+    quant_stage:
+        quant_modifiers:
+            QuantizationModifier:
+                kv_cache_scheme:
+                    num_bits: 8
+                    type: float
+                    strategy: tensor
+                    dynamic: false
+                    symmetric: true
+    '''
+
+    """
+
+    _instance = None
+    _initialized = False
+
+    def __new__(cls, *args, **kwargs):
+        """Singleton"""
+        if cls._instance is None:
+            cls._instance = super(QuantizedKVParameterCache, cls).__new__(cls)
+        return cls._instance
+
+    def __init__(self, quantization_args: QuantizationArgs):
+        if not self._initialized:
+            super().__init__()
+
+            self.quantization_args = quantization_args
+
+            self.k_observers: List[Observer] = []
+            self.v_observers: List[Observer] = []
+
+            # each index corresponds to layer_idx of the attention layer
+            self.k_scales: List[Tensor] = []
+            self.v_scales: List[Tensor] = []
+
+            self.k_zps: List[Tensor] = []
+            self.v_zps: List[Tensor] = []
+
+            self._initialized = True
+
+    def update(
+        self,
+        key_states: Tensor,
+        value_states: Tensor,
+        layer_idx: int,
+        cache_kwargs: Optional[Dict[str, Any]] = None,
+    ) -> Tuple[Tensor, Tensor]:
+        """
+        Get the k_scale and v_scale and output the
+         fakequant-ed key_states and value_states
+        """
+
+        if len(self.k_observers) <= layer_idx:
+            k_observer = self.quantization_args.get_observer()
+            v_observer = self.quantization_args.get_observer()
+
+            self.k_observers.append(k_observer)
+            self.v_observers.append(v_observer)
+
+        q_key_states = self._quantize(
+            key_states.contiguous(), KVCacheScaleType.KEY, layer_idx
+        )
+        q_value_states = self._quantize(
+            value_states.contiguous(), KVCacheScaleType.VALUE, layer_idx
+        )
+
+        qdq_key_states = self._dequantize(q_key_states, KVCacheScaleType.KEY, layer_idx)
+        qdq_value_states = self._dequantize(
+            q_value_states, KVCacheScaleType.VALUE, layer_idx
+        )
+
+        keys_to_return, values_to_return = qdq_key_states, qdq_value_states
+
+        return keys_to_return, values_to_return
+
+    def get_seq_length(self, layer_idx: Optional[int] = 0) -> int:
+        """
+        Returns the sequence length of the cached states.
+        A layer index can be optionally passed.
+        """
+        if len(self.key_cache) <= layer_idx:
+            return 0
+        # since we cannot get the seq_length of each layer directly and
+        # rely on `_seen_tokens` which is updated every "layer_idx" == 0,
+        # this is a hack to get the actual seq_length for the given layer_idx
+        # this part of code otherwise fails when used to
+        # verify attn_weight shape in some models
+        return self._seen_tokens if layer_idx == 0 else self._seen_tokens - 1
+
+    def reset_states(self):
+        """reset the kv states (used in calibration)"""
+        self.key_cache: List[Tensor] = []
+        self.value_cache: List[Tensor] = []
+        # Used in `generate` to keep tally of how many tokens the cache has seen
+        self._seen_tokens = 0
+        self._quantized_key_cache: List[Tensor] = []
+        self._quantized_value_cache: List[Tensor] = []
+
+    def reset(self):
+        """
+        Reset the instantiation, create new instance on init
+        """
+        QuantizedKVParameterCache._instance = None
+        QuantizedKVParameterCache._initialized = False
+
+    def _quantize(self, tensor, kv_type, layer_idx):
+        """Quantizes a key/value using a defined quantization method."""
+        from compressed_tensors.quantization.lifecycle.forward import quantize
+
+        if kv_type == KVCacheScaleType.KEY:  # key type
+            observer = self.k_observers[layer_idx]
+            scales = self.k_scales
+            zps = self.k_zps
+        else:
+            assert kv_type == KVCacheScaleType.VALUE
+            observer = self.v_observers[layer_idx]
+            scales = self.v_scales
+            zps = self.v_zps
+
+        scale, zp = observer(tensor)
+        if len(scales) <= layer_idx:
+            scales.append(scale)
+            zps.append(zp)
+        else:
+            scales[layer_idx] = scale
+            zps[layer_idx] = scale
+
+        q_tensor = quantize(
+            x=tensor,
+            scale=scale,
+            zero_point=zp,
+            args=self.quantization_args,
+        )
+        return q_tensor
+
+    def _dequantize(self, qtensor, kv_type, layer_idx):
+        """Dequantizes back the tensor that was quantized by `self._quantize()`"""
+        from compressed_tensors.quantization.lifecycle.forward import dequantize
+
+        if kv_type == KVCacheScaleType.KEY:
+            scale = self.k_scales[layer_idx]
+            zp = self.k_zps[layer_idx]
+        else:
+            assert kv_type == KVCacheScaleType.VALUE
+            scale = self.v_scales[layer_idx]
+            zp = self.v_zps[layer_idx]
+
+        qdq_tensor = dequantize(
+            x_q=qtensor,
+            scale=scale,
+            zero_point=zp,
+            args=self.quantization_args,
+        )
+        return qdq_tensor

compressed_tensors/quantization/lifecycle/apply.py

@@ -43,6 +43,7 @@ from compressed_tensors.quantization.utils import (
     infer_quantization_status,
     is_kv_cache_quant_scheme,
     iter_named_leaf_modules,
+    iter_named_quantizable_modules,
 )
 from compressed_tensors.utils.helpers import fix_fsdp_module_name, replace_module
 from compressed_tensors.utils.offload import update_parameter_data
@@ -106,8 +107,8 @@ def load_pretrained_quantization(model: Module, model_name_or_path: str):
 
 
 def apply_quantization_config(
-    model: Module, config: QuantizationConfig, run_compressed: bool = False
-) -> Dict:
+    model: Module, config: Union[QuantizationConfig, None], run_compressed: bool = False
+) -> OrderedDict:
     """
     Initializes the model for quantization in-place based on the given config
 
@@ -116,6 +117,10 @@ def apply_quantization_config(
     :param run_compressed: Whether the model will be run in compressed mode or
         decompressed fully on load
     """
+    # Workaround for when HF Quantizer passes None, see PR #180
+    if config is None:
+        return OrderedDict()
+
     # remove reference to the original `config`
     # argument. This function can mutate it, and we'd
     # like to keep the original `config` as it is.
@@ -135,15 +140,23 @@ def apply_quantization_config(
     # list of submodules to ignore
     ignored_submodules = defaultdict(list)
     # mark appropriate layers for quantization by setting their quantization schemes
-    for name, submodule in iter_named_leaf_modules(model):
+    for name, submodule in iter_named_quantizable_modules(
+        model,
+        include_children=True,
+        include_attn=True,
+    ):  # child modules and attention modules
         # potentially fix module name to remove FSDP wrapper prefix
         name = fix_fsdp_module_name(name)
         if matches := find_name_or_class_matches(name, submodule, config.ignore):
             for match in matches:
                 ignored_submodules[match].append(name)
             continue  # layer matches ignore list, continue
+
         targets = find_name_or_class_matches(name, submodule, target_to_scheme)
+
         if targets:
+            # mark modules to be quantized by adding
+            # quant scheme to the matching layers
             scheme = _scheme_from_targets(target_to_scheme, targets, name)
             if run_compressed:
                 format = config.format
@@ -200,6 +213,9 @@ def process_kv_cache_config(
     :param config: the QuantizationConfig
     :return: the QuantizationConfig with additional "kv_cache" group
     """
+    if targets == KV_CACHE_TARGETS:
+        _LOGGER.info(f"KV cache targets set to default value of: {KV_CACHE_TARGETS}")
+
     kv_cache_dict = config.kv_cache_scheme.model_dump()
     kv_cache_scheme = QuantizationScheme(
         output_activations=QuantizationArgs(**kv_cache_dict),

compressed_tensors/quantization/lifecycle/calibration.py

@@ -56,10 +56,9 @@ def set_module_for_calibration(module: Module, quantize_weights_upfront: bool =
         observer = module.weight_observer
         g_idx = getattr(module, "weight_g_idx", None)
 
-        offloaded = False
-        if is_module_offloaded(module):
+        offloaded = is_module_offloaded(module)
+        if offloaded:
             module._hf_hook.pre_forward(module)
-            offloaded = True
 
         scale, zero_point = observer(module.weight, g_idx=g_idx)
         update_parameter_data(module, scale, "weight_scale")

compressed_tensors/quantization/lifecycle/forward.py

@@ -14,9 +14,10 @@
 
 from functools import wraps
 from math import ceil
-from typing import Optional
+from typing import Callable, Optional
 
 import torch
+from compressed_tensors.quantization.cache import QuantizedKVParameterCache
 from compressed_tensors.quantization.observers.helpers import calculate_range
 from compressed_tensors.quantization.quant_args import (
     QuantizationArgs,
@@ -62,6 +63,7 @@ def quantize(
     :param g_idx: optional mapping from column index to group index
     :return: fake quantized tensor
     """
+
     return _process_quantization(
         x=x,
         scale=scale,
@@ -165,8 +167,8 @@ def _process_quantization(
     x: torch.Tensor,
     scale: torch.Tensor,
     zero_point: torch.Tensor,
-    g_idx: Optional[torch.Tensor],
     args: QuantizationArgs,
+    g_idx: Optional[torch.Tensor] = None,
     dtype: Optional[torch.dtype] = None,
     do_quantize: bool = True,
     do_dequantize: bool = True,
@@ -266,6 +268,7 @@ def wrap_module_forward_quantized(module: Module, scheme: QuantizationScheme):
             return forward_func_orig.__get__(module, module.__class__)(*args, **kwargs)
 
         input_ = args[0]
+
         compressed = module.quantization_status == QuantizationStatus.COMPRESSED
 
         if scheme.input_activations is not None:
@@ -285,9 +288,11 @@ def wrap_module_forward_quantized(module: Module, scheme: QuantizationScheme):
         output = forward_func_orig.__get__(module, module.__class__)(
             input_, *args[1:], **kwargs
         )
-
         if scheme.output_activations is not None:
+
             # calibrate and (fake) quantize output activations when applicable
+            # kv_cache scales updated on model self_attn forward call in
+            # wrap_module_forward_quantized_attn
             output = maybe_calibrate_or_quantize(
                 module, output, "output", scheme.output_activations
             )
@@ -304,6 +309,50 @@ def wrap_module_forward_quantized(module: Module, scheme: QuantizationScheme):
     setattr(module, "forward", bound_wrapped_forward)
 
 
+def wrap_module_forward_quantized_attn(module: Module, scheme: QuantizationScheme):
+    # expects a module already initialized and injected with the parameters in
+    # initialize_module_for_quantization
+    if hasattr(module.forward, "__func__"):
+        forward_func_orig = module.forward.__func__
+    else:
+        forward_func_orig = module.forward.func
+
+    @wraps(forward_func_orig)  # ensures docstring, names, etc are propagated
+    def wrapped_forward(self, *args, **kwargs):
+
+        # kv cache stored under weights
+        if module.quantization_status == QuantizationStatus.CALIBRATION:
+            quantization_args: QuantizationArgs = scheme.output_activations
+            past_key_value: QuantizedKVParameterCache = quantization_args.get_kv_cache()
+            kwargs["past_key_value"] = past_key_value
+
+            # QuantizedKVParameterCache used for obtaining k_scale, v_scale only,
+            # does not store quantized_key_states and quantized_value_state
+            kwargs["use_cache"] = False
+
+            attn_forward: Callable = forward_func_orig.__get__(module, module.__class__)
+
+            past_key_value.reset_states()
+
+            rtn = attn_forward(*args, **kwargs)
+
+            update_parameter_data(
+                module, past_key_value.k_scales[module.layer_idx], "k_scale"
+            )
+            update_parameter_data(
+                module, past_key_value.v_scales[module.layer_idx], "v_scale"
+            )
+
+            return rtn
+
+        return forward_func_orig.__get__(module, module.__class__)(*args, **kwargs)
+
+    # bind wrapped forward to module class so reference to `self` is correct
+    bound_wrapped_forward = wrapped_forward.__get__(module, module.__class__)
+    # set forward to wrapped forward
+    setattr(module, "forward", bound_wrapped_forward)
+
+
 def maybe_calibrate_or_quantize(
     module: Module, value: torch.Tensor, base_name: str, args: "QuantizationArgs"
 ) -> torch.Tensor:

compressed_tensors/quantization/lifecycle/frozen.py

@@ -14,6 +14,7 @@
 
 
 from compressed_tensors.quantization.quant_config import QuantizationStatus
+from compressed_tensors.quantization.utils import is_kv_cache_quant_scheme
 from torch.nn import Module
 
 
@@ -44,7 +45,11 @@ def freeze_module_quantization(module: Module):
         delattr(module, "input_observer")
     if scheme.weights and not scheme.weights.dynamic:
         delattr(module, "weight_observer")
-    if scheme.output_activations and not scheme.output_activations.dynamic:
+    if (
+        scheme.output_activations
+        and not is_kv_cache_quant_scheme(scheme)
+        and not scheme.output_activations.dynamic
+    ):
         delattr(module, "output_observer")
 
     module.quantization_status = QuantizationStatus.FROZEN

compressed_tensors/quantization/lifecycle/helpers.py

@@ -16,62 +16,15 @@
 Miscelaneous helpers for the quantization lifecycle
 """
 
-from typing import Optional
-
-import torch
 from torch.nn import Module
 
 
 __all__ = [
-    "update_layer_weight_quant_params",
     "enable_quantization",
     "disable_quantization",
 ]
 
 
-def update_layer_weight_quant_params(
-    layer: Module,
-    weight: Optional[torch.Tensor] = None,
-    g_idx: Optional[torch.Tensor] = None,
-    reset_obs: bool = False,
-):
-    """
-    Update quantization parameters on layer
-
-    :param layer: input layer
-    :param weight: weight to update quant params with, defaults to layer weight
-    :param g_idx: optional mapping from column index to group index
-    :param reset_obs: reset the observer before calculating quant params,
-        defaults to False
-    """
-    attached_weight = getattr(layer, "weight", None)
-
-    if weight is None:
-        weight = attached_weight
-    scale = getattr(layer, "weight_scale", None)
-    zero_point = getattr(layer, "weight_zero_point", None)
-    if g_idx is None:
-        g_idx = getattr(layer, "weight_g_idx", None)
-    observer = getattr(layer, "weight_observer", None)
-
-    if weight is None or observer is None or scale is None or zero_point is None:
-        # scale, zp, or observer not calibratable or weight not available
-        return
-
-    if reset_obs:
-        observer.reset()
-
-    if attached_weight is not None:
-        weight = weight.to(attached_weight.dtype)
-
-    updated_scale, updated_zero_point = observer(weight)
-
-    # update scale and zero point
-    device = next(layer.parameters()).device
-    scale.data = updated_scale.to(device)
-    zero_point.data = updated_zero_point.to(device)
-
-
 def enable_quantization(module: Module):
     module.quantization_enabled = True
 

compressed_tensors/quantization/lifecycle/initialize.py

@@ -17,8 +17,10 @@ import logging
 from typing import Optional
 
 import torch
+from compressed_tensors.quantization.cache import KVCacheScaleType
 from compressed_tensors.quantization.lifecycle.forward import (
     wrap_module_forward_quantized,
+    wrap_module_forward_quantized_attn,
 )
 from compressed_tensors.quantization.quant_args import (
     ActivationOrdering,
@@ -27,6 +29,7 @@ from compressed_tensors.quantization.quant_args import (
 )
 from compressed_tensors.quantization.quant_config import QuantizationStatus
 from compressed_tensors.quantization.quant_scheme import QuantizationScheme
+from compressed_tensors.quantization.utils import is_kv_cache_quant_scheme
 from compressed_tensors.utils import get_execution_device, is_module_offloaded
 from torch.nn import Module, Parameter
 
@@ -62,72 +65,85 @@ def initialize_module_for_quantization(
         # no scheme passed and layer not targeted for quantization - skip
         return
 
-    if scheme.input_activations is not None:
-        _initialize_scale_zero_point_observer(
-            module, "input", scheme.input_activations, force_zero_point=force_zero_point
-        )
-    if scheme.weights is not None:
-        if hasattr(module, "weight"):
-            weight_shape = module.weight.shape
+    if is_attention_module(module):
+        # wrap forward call of module to perform
+        # quantized actions based on calltime status
+        wrap_module_forward_quantized_attn(module, scheme)
+        _initialize_attn_scales(module)
+
+    else:
+
+        if scheme.input_activations is not None:
             _initialize_scale_zero_point_observer(
                 module,
-                "weight",
-                scheme.weights,
-                weight_shape=weight_shape,
+                "input",
+                scheme.input_activations,
                 force_zero_point=force_zero_point,
             )
-        else:
-            _LOGGER.warning(
-                f"module type {type(module)} targeted for weight quantization but "
-                "has no attribute weight, skipping weight quantization "
-                f"for {type(module)}"
-            )
-    if scheme.output_activations is not None:
-        _initialize_scale_zero_point_observer(
-            module,
-            "output",
-            scheme.output_activations,
-            force_zero_point=force_zero_point,
-        )
-
-    module.quantization_scheme = scheme
-    module.quantization_status = QuantizationStatus.INITIALIZED
-
-    offloaded = False
-    if is_module_offloaded(module):
-        try:
-            from accelerate.hooks import add_hook_to_module, remove_hook_from_module
-            from accelerate.utils import PrefixedDataset
-        except ModuleNotFoundError:
-            raise ModuleNotFoundError(
-                "Offloaded model detected. To use CPU offloading with "
-                "compressed-tensors the `accelerate` package must be installed, "
-                "run `pip install compressed-tensors[accelerate]`"
-            )
-
-        offloaded = True
-        hook = module._hf_hook
-        prefix_dict = module._hf_hook.weights_map
-        new_prefix = {}
-
-        # recreate the prefix dict (since it is immutable)
-        # and add quantization parameters
-        for key, data in module.named_parameters():
-            if key not in prefix_dict:
-                new_prefix[f"{prefix_dict.prefix}{key}"] = data
+        if scheme.weights is not None:
+            if hasattr(module, "weight"):
+                weight_shape = None
+                if isinstance(module, torch.nn.Linear):
+                    weight_shape = module.weight.shape
+                _initialize_scale_zero_point_observer(
+                    module,
+                    "weight",
+                    scheme.weights,
+                    weight_shape=weight_shape,
+                    force_zero_point=force_zero_point,
+                )
             else:
-                new_prefix[f"{prefix_dict.prefix}{key}"] = prefix_dict[key]
-        new_prefix_dict = PrefixedDataset(new_prefix, prefix_dict.prefix)
-        remove_hook_from_module(module)
-
-    # wrap forward call of module to perform quantized actions based on calltime status
-    wrap_module_forward_quantized(module, scheme)
-
-    if offloaded:
-        # we need to re-add the hook for offloading now that we've wrapped forward
-        add_hook_to_module(module, hook)
-        if prefix_dict is not None:
-            module._hf_hook.weights_map = new_prefix_dict
+                _LOGGER.warning(
+                    f"module type {type(module)} targeted for weight quantization but "
+                    "has no attribute weight, skipping weight quantization "
+                    f"for {type(module)}"
+                )
+
+        if scheme.output_activations is not None:
+            if not is_kv_cache_quant_scheme(scheme):
+                _initialize_scale_zero_point_observer(
+                    module, "output", scheme.output_activations
+                )
+
+        module.quantization_scheme = scheme
+        module.quantization_status = QuantizationStatus.INITIALIZED
+
+        offloaded = False
+        if is_module_offloaded(module):
+            try:
+                from accelerate.hooks import add_hook_to_module, remove_hook_from_module
+                from accelerate.utils import PrefixedDataset
+            except ModuleNotFoundError:
+                raise ModuleNotFoundError(
+                    "Offloaded model detected. To use CPU offloading with "
+                    "compressed-tensors the `accelerate` package must be installed, "
+                    "run `pip install compressed-tensors[accelerate]`"
+                )
+
+            offloaded = True
+            hook = module._hf_hook
+            prefix_dict = module._hf_hook.weights_map
+            new_prefix = {}
+
+            # recreate the prefix dict (since it is immutable)
+            # and add quantization parameters
+            for key, data in module.named_parameters():
+                if key not in prefix_dict:
+                    new_prefix[f"{prefix_dict.prefix}{key}"] = data
+                else:
+                    new_prefix[f"{prefix_dict.prefix}{key}"] = prefix_dict[key]
+            new_prefix_dict = PrefixedDataset(new_prefix, prefix_dict.prefix)
+            remove_hook_from_module(module)
+
+        # wrap forward call of module to perform
+        # quantized actions based on calltime status
+        wrap_module_forward_quantized(module, scheme)
+
+        if offloaded:
+            # we need to re-add the hook for offloading now that we've wrapped forward
+            add_hook_to_module(module, hook)
+            if prefix_dict is not None:
+                module._hf_hook.weights_map = new_prefix_dict
 
 
 def _initialize_scale_zero_point_observer(
@@ -156,9 +172,10 @@ def _initialize_scale_zero_point_observer(
             # (output_channels, 1)
             expected_shape = (weight_shape[0], 1)
         elif quantization_args.strategy == QuantizationStrategy.GROUP:
+            num_groups = weight_shape[1] // quantization_args.group_size
             expected_shape = (
                 weight_shape[0],
-                weight_shape[1] // quantization_args.group_size,
+                max(num_groups, 1)
             )
 
     scale_dtype = module.weight.dtype
@@ -189,3 +206,34 @@ def _initialize_scale_zero_point_observer(
             requires_grad=False,
         )
         module.register_parameter(f"{base_name}_g_idx", init_g_idx)
+
+
+def is_attention_module(module: Module):
+    return "attention" in module.__class__.__name__.lower() and (
+        hasattr(module, "k_proj")
+        or hasattr(module, "v_proj")
+        or hasattr(module, "qkv_proj")
+    )
+
+
+def _initialize_attn_scales(module: Module) -> None:
+    """Initlaize k_scale, v_scale for  self_attn"""
+
+    expected_shape = 1  # per tensor
+
+    param = next(module.parameters())
+    scale_dtype = param.dtype
+    device = param.device
+
+    init_scale = Parameter(
+        torch.empty(expected_shape, dtype=scale_dtype, device=device),
+        requires_grad=False,
+    )
+
+    module.register_parameter(KVCacheScaleType.KEY.value, init_scale)
+
+    init_scale = Parameter(
+        torch.empty(expected_shape, dtype=scale_dtype, device=device),
+        requires_grad=False,
+    )
+    module.register_parameter(KVCacheScaleType.VALUE.value, init_scale)

compressed_tensors/quantization/quant_args.py

@@ -122,6 +122,12 @@ class QuantizationArgs(BaseModel, use_enum_values=True):
 
         return Observer.load_from_registry(self.observer, quantization_args=self)
 
+    def get_kv_cache(self):
+        """Get the singleton KV Cache"""
+        from compressed_tensors.quantization.cache import QuantizedKVParameterCache
+
+        return QuantizedKVParameterCache(self)
+
     @field_validator("type", mode="before")
     def validate_type(cls, value) -> QuantizationType:
         if isinstance(value, str):