import torch
import bitsandbytes as bnb
from transformers.cache_utils import DynamicCache

# Ensure that DynamicCache and set are safe during serialization
torch.serialization.add_safe_globals([DynamicCache, set])

def quantize_and_save(past_key_values, save_path="cache_nf4.pth"):
    quantized_past_key_values = []
    for k, v in past_key_values:
        # Convert tensors to float16 for quantization
        k, v = k.to(torch.float16), v.to(torch.float16)
        
        # Perform 4-bit quantization using NF4
        k_quant, k_absmax = bnb.functional.quantize_4bit(k, quant_type="nf4")
        v_quant, v_absmax = bnb.functional.quantize_4bit(v, quant_type="nf4")
        
        # Extract float values from QuantState
        # Assuming `QuantState` has an attribute or method to retrieve the absmax value
        # This may vary based on bitsandbytes' implementation
        # Replace `.absmax()` with the correct method if different
        k_absmax_value = k_absmax.absmax().item() if hasattr(k_absmax, 'absmax') else float(k_absmax)
        v_absmax_value = v_absmax.absmax().item() if hasattr(v_absmax, 'absmax') else float(v_absmax)
        
        # Append quantized tensors and their corresponding absmax values
        quantized_past_key_values.append(((k_quant, k_absmax_value), (v_quant, v_absmax_value)))

    # Save the quantized data to disk
    torch.save(quantized_past_key_values, save_path)
    print("Cache saved successfully with NF4 quantization and absmax.")

def load_and_dequantize(save_path="cache_nf4.pth"):
    try:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        loaded_past_key_values = torch.load(save_path, map_location=device)
        print("NF4-quantized cache loaded successfully.")
    except FileNotFoundError:
        print("No NF4 quantized cache file found.")
        loaded_past_key_values = []

    past_key_values = DynamicCache()

    if loaded_past_key_values:
        for layer_idx, ((k_quant, k_absmax), (v_quant, v_absmax)) in enumerate(loaded_past_key_values):
            # Convert absmax values back to tensors
            k_absmax_tensor = torch.tensor(k_absmax, dtype=torch.float32, device=device)
            v_absmax_tensor = torch.tensor(v_absmax, dtype=torch.float32, device=device)

            # Define output shapes based on quantized tensor dimensions
            key_out_shape = (k_quant.shape[0] * 8, k_quant.shape[1])
            value_out_shape = (v_quant.shape[0] * 8, v_quant.shape[1])

            # Dequantize the key and value tensors
            key_states = bnb.functional.dequantize_4bit(
                k_quant,
                absmax=k_absmax_tensor,
                out=torch.empty(key_out_shape, dtype=torch.float16, device=device),
                quant_type="nf4"
            )
            value_states = bnb.functional.dequantize_4bit(
                v_quant,
                absmax=v_absmax_tensor,
                out=torch.empty(value_out_shape, dtype=torch.float16, device=device),
                quant_type="nf4"
            )

            # Update the DynamicCache with the dequantized tensors
            past_key_values.update(key_states, value_states, layer_idx)

    print("DynamicCache populated with dequantized float16 past_key_values.")
    return past_key_values


quantize_and_save(past_key_values)

# Load the quantized cache
past_key_values = load_and_dequantize()







---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[11], line 74
     70     print("DynamicCache populated with dequantized float16 past_key_values.")
     71     return past_key_values
---> 74 quantize_and_save(past_key_values)
     76 # Load the quantized cache
     77 past_key_values = load_and_dequantize()

Cell In[11], line 22, in quantize_and_save(past_key_values, save_path)
     16 v_quant, v_absmax = bnb.functional.quantize_4bit(v, quant_type="nf4")
     18 # Extract float values from QuantState
     19 # Assuming `QuantState` has an attribute or method to retrieve the absmax value
     20 # This may vary based on bitsandbytes' implementation
     21 # Replace `.absmax()` with the correct method if different
---> 22 k_absmax_value = k_absmax.absmax().item() if hasattr(k_absmax, 'absmax') else float(k_absmax)
     23 v_absmax_value = v_absmax.absmax().item() if hasattr(v_absmax, 'absmax') else float(v_absmax)
     25 # Append quantized tensors and their corresponding absmax values

TypeError: 'Tensor' object is not callable