break out function

Signed-off-by: Kyle Sayers <[email protected]>

Author: kylesayrs

src/compressed_tensors/quantization/lifecycle/apply.py

@@ -28,6 +28,7 @@
     initialize_module_for_quantization,
     is_attention_module,
 )
+from compressed_tensors.quantization.quant_args import QuantizationArgs
 from compressed_tensors.quantization.quant_config import (
     QuantizationConfig,
     QuantizationStatus,
@@ -128,21 +129,11 @@ def apply_quantization_config(
     # force zero points during initialization
     force_zero_point = config.quantization_status != QuantizationStatus.COMPRESSED
 
-    # apply kv cache quantization before any attention quantization
-    # because attention quantization is a superset of kv cache quantization
+    # apply and initialize kv cache quantization
     if config.kv_cache_scheme is not None:
-        scheme = QuantizationScheme(
-            targets=[".*self_attn$"], input_activations=config.kv_cache_scheme
+        _apply_kv_cache_scheme(
+            model, config.kv_cache_scheme, config.quantization_status, force_zero_point
         )
-        for submodule in model.modules():
-            if is_attention_module(submodule):
-                submodule.quantization_scheme = scheme
-                initialize_hooked_kv_cache(model, submodule)
-                initialize_module_for_quantization(
-                    submodule,
-                    force_zero_point=force_zero_point,
-                )
-                submodule.quantization_status = config.quantization_status
 
     # build mapping of targets to schemes for easier matching
     # use ordered dict to preserve target ordering in config
@@ -191,6 +182,29 @@ def apply_quantization_config(
         submodule.quantization_status = config.quantization_status
 
 
+def _apply_kv_cache_scheme(
+    model: torch.nn.Module,
+    kv_cache_scheme: QuantizationArgs,
+    status: QuantizationStatus,
+    force_zero_point: bool,
+):
+    # applies and initializes kv cache quantization
+    # this step cannot come after attention apply/initialize
+    # otherwise it will override the attention qparams
+    scheme = QuantizationScheme(
+        targets=[".*self_attn$"], input_activations=kv_cache_scheme
+    )
+    for submodule in model.modules():
+        if is_attention_module(submodule):
+            submodule.quantization_scheme = scheme
+            initialize_hooked_kv_cache(model, submodule)
+            initialize_module_for_quantization(
+                submodule,
+                force_zero_point=force_zero_point,
+            )
+            submodule.quantization_status = status
+
+
 def _load_quant_args_from_mapping(
     base_name: str, module_name: str, module: Module, mapping: Dict
 ):