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"):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    quantized_past_key_values = []
    for k, v in past_key_values:
        # Move tensors to the specified device (GPU if available)
        k, v = k.to(device).to(torch.float16), v.to(device).to(torch.float16)
        
        # Perform 4-bit quantization using NF4
        k_quant, k_quant_state = bnb.functional.quantize_4bit(k, quant_type="nf4")
        v_quant, v_quant_state = bnb.functional.quantize_4bit(v, quant_type="nf4")
        
        # Extract the absolute maximum value from the QuantState object
        k_absmax_value = k_quant_state.absmax.item()
        v_absmax_value = v_quant_state.absmax.item()
        
        # 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

# Example past_key_values for demonstration purposes
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

quantize_and_save(past_key_values)

# Load the quantized cache
past_key_values = load_and_dequantize()












---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[19], line 75
     72 # Example past_key_values for demonstration purposes
     73 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
---> 75 quantize_and_save(past_key_values)
     77 # Load the quantized cache
     78 past_key_values = load_and_dequantize()

Cell In[19], line 21, in quantize_and_save(past_key_values, save_path)
     18 v_quant, v_quant_state = bnb.functional.quantize_4bit(v, quant_type="nf4")
     20 # Extract the absolute maximum value from the QuantState object
---> 21 k_absmax_value = k_quant_state.absmax.item()
     22 v_absmax_value = v_quant_state.absmax.item()
     24 # Append quantized tensors and their corresponding absmax values

RuntimeError: a Tensor with 14048 elements cannot be converted to Scalar