import torch
import numpy as np
import time

L=2048
S=512
NH=128
HS=1024
dtype=torch.float32

q = torch.randn([1, NH, L, HS], dtype=dtype, device="mps")
k = torch.randn([1, NH, S, HS], dtype=q.dtype, device="mps")
v = torch.randn([1, NH, S, HS], dtype=q.dtype, device="mps")

op = torch.nn.functional.scaled_dot_product_attention

num_warmup = 5
n_stats = 50

for _ in range(num_warmup):
    out = op(q, k, v)
    torch.mps.synchronize()

stats = []
for i in range(n_stats):
    start = time.perf_counter_ns()
    out = op(q, k, v)
    torch.mps.synchronize()
    end = time.perf_counter_ns()
    stats.append((end-start)/1e3)

print(f"Took {np.mean(stats):0.2f}\u00B1{np.std(stats):0.2f} \u03bcs per iter")

L=1, S=72, NH=32, HS=128
# Before change:
Took 310.55±37.74 μs per iter
# After change:
Took 368.98±35.13 μs per iter

L=2048, S=512, NH=128, HS=1024
# Before change:
Took 52839.62±147.31 μs per iter
# After change:
Took 52825.06±132.55 μs per iter

Before change:
----------------------------------------------- benchmark 'hub': 1 tests -----------------------------------------------
Name (time in ms)                Min       Max      Mean  StdDev    Median     IQR  Outliers     OPS  Rounds  Iterations
------------------------------------------------------------------------------------------------------------------------
test_train[hf_GPT2-mps]     398.4452  407.1690  400.6670  3.6971  398.8504  3.3065       1;1  2.4958       5           1
------------------------------------------------------------------------------------------------------------------------

After change:
----------------------------------------------- benchmark 'hub': 1 tests -----------------------------------------------
Name (time in ms)                Min       Max      Mean  StdDev    Median     IQR  Outliers     OPS  Rounds  Iterations
------------------------------------------------------------------------------------------------------------------------
test_train[hf_GPT2-mps]     395.7955  400.3153  398.8589  1.7881  399.2772  1.8033       1;0  2.5072       5           1
------------------------------------------------------------------------------------------------------------------------