compressed_tensors.quantization.quant_metadata

@staticmethod
def all_qparam_names():
    """
    All quantization parameter names that might be registered
    onto a module during lifecycle (excluding serialized parameters)
    """
    return [KVCacheScaleType.KEY.value, KVCacheScaleType.VALUE.value] + [
        f"{base_name}_{suffix}"
        for base_name in ("input", "weight", "output")
        for suffix in (
            "global_scale",
            "scale",
            "shape",
            "zero_point",
        )
    ]

@classmethod
def clear_all_qparams(cls, module: Module):
    """
    Remove all parameters related to quantization that might have
    been registered onto a module previously in lifecycle (excluding
    serialized parameters)

    :param module: Module to clear
    """
    for key in cls.all_qparam_names():
        with disable_onloading():
            has_key = hasattr(module, key)

        if has_key:
            delattr(module, key)

@classmethod
def clear_quantization(cls, module: Module):
    """
    Remove all artifacts of quantization from module, non-recursively.
    Artifacts include any qparams, quantization_scheme, or wrapped
    forward method that might have been altered previously in lifecycle.

    `quantization_status` and `quantization_enabled` are left unchanged.

    :param module: Module to clear
    """
    with unwrap_offload_forward(module):
        # Unwrap forward call
        if hasattr(module.forward, "__wrapped__"):
            module.forward = module.forward.__wrapped__.__get__(module)

        # Clear any qparams
        cls.clear_all_qparams(module)

        # Clear quantization_scheme
        if hasattr(module, "quantization_scheme"):
            delattr(module, "quantization_scheme")