from thunder.benchmarks import LitGPTSDPABenchmark, Benchmark
import thunder
import torch

bench: Benchmark = LitGPTSDPABenchmark(
        config="Llama-2-7b-hf", device="cuda:0", dtype=thunder.bfloat16, requires_grad=True
    )

args, kwargs = bench.make_batch()
torch.cuda.synchronize()
fn = bench.fn()

jfn = thunder.jit(fn)
jfn(*args, **kwargs)

print(thunder.last_traces(jfn)[-1])

# Constructed by Delete Last Used (took 0 milliseconds)
import torch
from thunder.executors.torchex import no_autocast

@torch.no_grad()
@no_autocast
def augmented_forward_fn(q, k, v):
  # q: "cuda:0 bf16[16, 32, 4096, 128]"
  # k: "cuda:0 bf16[16, 32, 4096, 128]"
  # v: "cuda:0 bf16[16, 32, 4096, 128]"
  (t0, t1, t2, t3, _, _, t4, t5, _) = sdpafx_grad_forward_scaled_dot_product_efficient_attention(q, k, v, 0.1, True, scale=None)
  return {'output': t0, 'flat_args': [q, k, v], 'flat_output': (t0,)}, ((k, q, t0, t1, t2, t3, t4, t5, v), (True, 0.1))