import os
os.environ["TOKENIZERS_PARALLELISM"] = "false"

import time
from transformers import AutoTokenizer, Gemma2ForCausalLM
from transformers.cache_utils import HybridCache
import torch

torch.set_float32_matmul_precision("high")
# catch re-compilations
torch._logging.set_logs(graph_breaks=True, recompiles=True)

tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-9b")
model = Gemma2ForCausalLM.from_pretrained("google/gemma-2-9b", attn_implementation="eager")
model.to("cuda")

model.forward = torch.compile(model.forward, mode="reduce-overhead", fullgraph=True)

input_text = "The theory of special relativity states "
input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")
prompt_length = input_ids.input_ids.shape[1]

model.generation_config.min_new_tokens = model.generation_config.max_new_tokens = 32

past_key_values = HybridCache(
    config=model.config,
    max_batch_size=1,
    max_cache_len=prompt_length + 4 * model.generation_config.max_new_tokens,
    device=model.device,
    dtype=model.dtype
)

# enable passing kv cache
model._supports_cache_class = True
model.generation_config.cache_implementation = None

for i in range(3):
    # two warm-ups
    outputs_1 = model.generate(**input_ids, past_key_values=past_key_values, do_sample=True, temperature=1)
    outputs_2 = model.generate(outputs_1, past_key_values=past_key_values, do_sample=True, temperature=1)

    # one timed run
    torch.cuda.synchronize("cuda")
    start = time.time()
    outputs_3 = model.generate(outputs_2, past_key_values=past_key_values, do_sample=True, temperature=1)
    torch.cuda.synchronize("cuda")
    runtime = time.time() - start
    print(f"Run {i}: {model.generation_config.max_new_tokens / runtime} tok/s")

    past_key_values.reset()