add ALL_QPARAM_KEYS var

Signed-off-by: Brian Dellabetta <[email protected]>

Author: brian-dellabetta

src/compressed_tensors/quantization/lifecycle/initialize.py

@@ -45,6 +45,7 @@
     "initialize_module_for_quantization",
     "is_attention_module",
     "KVCacheScaleType",
+    "ALL_QPARAM_KEYS",
 ]
 
 
@@ -56,6 +57,18 @@ class KVCacheScaleType(Enum):
     VALUE = "v_scale"
 
 
+ALL_QPARAM_KEYS = [KVCacheScaleType.KEY.value, KVCacheScaleType.VALUE.value] + [
+    f"{base_name}_{suffix}"
+    for base_name in ("input", "weight", "output")
+    for suffix in (
+        "global_scale",
+        "scale",
+        "zero_point",
+        "g_idx",
+    )
+]
+
+
 def initialize_module_for_quantization(
     module: Module,
     scheme: Optional[QuantizationScheme] = None,
@@ -146,17 +159,7 @@ def _clear_all_qparams(
 
     :param module: module to clear qparams from
     """
-    keys = [KVCacheScaleType.KEY.value, KVCacheScaleType.VALUE.value] + [
-        f"{base_name}_{suffix}"
-        for base_name in ("input", "weight", "output")
-        for suffix in (
-            "global_scale",
-            "scale",
-            "zero_point",
-            "g_idx",
-        )
-    ]
-    for key in keys:
+    for key in ALL_QPARAM_KEYS:
         if hasattr(module, key):
             delete_offload_parameter(module, key)