# this is benchmarked on FA3 at this commit: https://github.com/Dao-AILab/flash-attention/tree/4f0640d534888c579a448fd89c2d4e064905d798

import torch

from triton.testing import do_bench, do_bench_cudagraph

from einops import rearrange

from flash_attn_interface import flash_attn_with_kvcache

try:
    from flash_attn.utils.benchmark import pytorch_profiler
except ImportError:
    pytorch_profiler = None

device = "cuda"
dtype = torch.bfloat16
seqlen = 64 * 1024
nheads    = 16
nheads_kv = 16
headdim   = 128
headdim_v = 128
has_qv = False

def run_benchmark(seqlens, seqlen_q, num_splits=0):
    print(f' ----------- starting seq_lengths: {seqlens} seqlen_q: {seqlen_q} num_splits: {num_splits} -----------')
    # page_size = None
    page_size = 1

    batch_size = len(seqlens)

    torch.manual_seed(0)

    cache_seqlens = torch.tensor([seqlens[i] for i in range(batch_size)], device=device, dtype=torch.int)

    q       = torch.randn(batch_size, seqlen_q, nheads, headdim, dtype=dtype, device=device)
    v_cache = torch.randn(batch_size, seqlen, nheads_kv, headdim_v, dtype=dtype, device=device)
    k_cache = torch.randn(batch_size, seqlen, nheads_kv, headdim, dtype=dtype, device=device)
    if page_size is not None:
        assert seqlen % page_size == 0
        k_cache, v_cache = [rearrange(x, "b (n p) h d -> (b n) p h d", p=page_size) for x in [k_cache, v_cache]]
        page_table = rearrange(torch.arange(batch_size * seqlen // page_size, device=device, dtype=torch.int32),
                            "(b s) -> b s", s=seqlen // page_size)
    else:
        page_table = None
    qv = torch.randn(batch_size, seqlen_q, nheads, headdim_v, dtype=dtype, device=device) if has_qv else None

    # Time in ms
    fn = lambda: flash_attn_with_kvcache(q, k_cache, v_cache, cache_seqlens=cache_seqlens, num_splits=num_splits, qv=qv, page_table=page_table, causal=True)
    t0 = do_bench(fn, warmup=1, rep=10)

    mem_io = cache_seqlens.sum().item() * nheads_kv * (headdim + headdim_v) * 2
    flops  = seqlen_q * cache_seqlens.float().sum().item() * nheads * (headdim + headdim_v * 2) * 2
    print(f"Time: {t0 * 1e3:.1f} us, {mem_io * 1e-9 / (t0 * 1e-3):.0f} GB/s, {flops * 1e-12 / (t0 * 1e-3):.0f} TFLOPS/s")


if __name__ == "__main__":
    run_benchmark([4641, 45118, 65536, 256], 4)
    run_benchmark([4641, 45118, 65536, 256], 2)
    run_benchmark([4641, 45118, 65536, 256], 1)
    run_benchmark([256, 512, 1024, 512, 256], 1)
    run_benchmark([4000, 4000, 4000, 4000, 200], 1)
    run_benchmark([100, 100, 100, 100, 16000], 1)