from transformers import AutoModelForCausalLM, DynamicCache
import torch
from tqdm import tqdm


model_name = "mistralai/Mistral-7B-v0.1"
# model_name = "google/gemma-2-9b-it"
# model_name = "openai/gpt-oss-20b"
device = 0

model = AutoModelForCausalLM.from_pretrained(model_name, device_map=device, torch_dtype=torch.bfloat16)

input_sizes = torch.linspace(50, 8000, 50).tolist()

old_sizes = []
new_sizes = []
for size in tqdm(input_sizes):
    with torch.no_grad():
        input = torch.randint(1000, 3000, (1, int(size)), device=device)
        
        # initializing DynamicCache without config will use only full layers
        old_output = model(input, past_key_values=DynamicCache(), logits_to_keep=1)
        cache = old_output.past_key_values
        tot = sum([layer.keys.numel() * 2 * layer.keys.element_size() for layer in cache.layers])
        old_sizes.append(tot / 1024**3)

        # Initializing it with the config will infer and use the sliding window/hybrid structure
        new_output = model(input, past_key_values=DynamicCache(config=model.config), logits_to_keep=1)
        cache = new_output.past_key_values
        tot = sum([layer.keys.numel() * 2 * layer.keys.element_size() for layer in cache.layers])
        new_sizes.append(tot / 1024**3)

import matplotlib.pyplot as plt

plt.figure()
plt.plot(input_sizes, old_sizes, "r", label="before")
plt.plot(input_sizes, new_sizes, "b", label="now")
plt.xlabel("Cache size [tokens]")
plt.ylabel("Cache memory usage [GiB]")
plt.grid()
plt.legend()
plt.show()