# Copyright (c) 2025 by FlashInfer team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Recurrent KDA (Kimi Delta Attention) Benchmark

Benchmarks the recurrent KDA register-tile kernel with per-K-dimension gating.
T=1 is standard decode; T>1 is fused speculative decode with
``num_spec_tokens=T-1``.
KDA differs from GDN by having gate g[B, T, HV, K] instead of a scalar gate.

Usage:
    python benchmarks/bench_recurrent_kda.py --batch-size 1 4 16 64 128 256
    python benchmarks/bench_recurrent_kda.py --head-size 64 --batch-size 1 32 128
    python benchmarks/bench_recurrent_kda.py --seq-len 1 2 3 4 --batch-size 1 32
"""

import argparse
import numpy as np
import torch
import torch.nn.functional as F

from flashinfer.testing import bench_gpu_time

# Import the recurrent KDA kernel
try:
    from flashinfer.kda_decode import _RECURRENT_KDA_AVAILABLE, recurrent_kda

    RECURRENT_KDA_AVAILABLE = _RECURRENT_KDA_AVAILABLE
except ImportError:
    RECURRENT_KDA_AVAILABLE = False


# ============================================================================
# FLOPs and Bytes Calculation
# ============================================================================


def recurrent_kda_flops(
    batch_size: int,
    num_q_heads: int,
    _num_k_heads: int,
    num_v_heads: int,
    head_size: int,
    seq_len: int = 1,
) -> int:
    """
    Calculate FLOPs for KDA (Kimi Delta Attention) decode.

    8 * K * V FLOPs per token per head:
    1. k @ state (prediction):    2 * K * V
    2. outer(k, v_new) update:    2 * K * V
    3. q @ state (output):        2 * K * V
    4. Per-K gate application:    2 * K * V  (K*V element-wise multiply + K exp() calls)

    Note: K = V = head_size for KDA. State ops are per-HV (value) head.
    """
    total_flops = 8 * seq_len * batch_size * num_v_heads * head_size * head_size
    return total_flops


def recurrent_kda_bytes(
    batch_size: int,
    num_q_heads: int,
    num_k_heads: int,
    num_v_heads: int,
    head_size: int,
    dtype: torch.dtype,
    seq_len: int = 1,
) -> int:
    """
    Calculate memory bytes for recurrent KDA.

    Includes:
    - Q, K, V tensors: [B, T, H, K] - dtype
    - G tensor (per-K gate): [B, T, HV, K] - dtype (extra vs GDN)
    - Beta: [B, T, HV] - dtype
    - State read: [B, HV, V, K] - bf16 storage for logical [K, V]
    - State write: [B*T, HV, V, K] for spec decode, [B, HV, V, K] for T=1
    - Output: [B, T, HV, V] - dtype
    """
    elem_size = torch.tensor([], dtype=dtype).element_size()
    state_dtype_bytes = 2  # BF16 state

    # Input tensors: q/k use H (query heads), v uses HV (value heads)
    q_bytes = batch_size * seq_len * num_q_heads * head_size * elem_size
    k_bytes = batch_size * seq_len * num_k_heads * head_size * elem_size
    v_bytes = batch_size * seq_len * num_v_heads * head_size * elem_size

    # Per-K gate: [B, T, HV, K]
    g_bytes = batch_size * seq_len * num_v_heads * head_size * elem_size

    # Beta: [B, T, HV]
    beta_bytes = batch_size * seq_len * num_v_heads * elem_size

    # Output: [B, T, HV, V]
    o_bytes = batch_size * seq_len * num_v_heads * head_size * elem_size

    # State: read one initial slot per sequence and write one checkpoint per token.
    state_write_slots = batch_size if seq_len == 1 else batch_size * seq_len
    state_bytes = (
        (batch_size + state_write_slots)
        * num_v_heads
        * head_size
        * head_size
        * state_dtype_bytes
    )

    total_bytes = (
        q_bytes + k_bytes + v_bytes + g_bytes + beta_bytes + o_bytes + state_bytes
    )
    return total_bytes


# ============================================================================
# Benchmark Function
# ============================================================================


def bench_recurrent_kda(
    batch_size: int,
    seq_len: int,
    num_q_heads: int,
    num_k_heads: int,
    num_v_heads: int,
    head_size: int,
    dtype: torch.dtype,
    warmup_iters: int = 10,
    bench_iters: int = 100,
):
    """Benchmark recurrent KDA decode. T>1 maps to fused speculative decode."""
    if not RECURRENT_KDA_AVAILABLE:
        raise RuntimeError("recurrent KDA kernel is not available")

    if seq_len < 1:
        raise ValueError(f"seq_len must be >= 1, got {seq_len}")
    if num_k_heads != num_q_heads:
        raise ValueError(
            "recurrent KDA uses one K head per Q head; "
            f"got num_q_heads={num_q_heads}, num_k_heads={num_k_heads}"
        )

    # Create inputs
    T = seq_len
    q = torch.randn(batch_size, T, num_q_heads, head_size, dtype=dtype, device="cuda")
    k = torch.randn(batch_size, T, num_q_heads, head_size, dtype=dtype, device="cuda")
    v = torch.randn(batch_size, T, num_v_heads, head_size, dtype=dtype, device="cuda")

    # KDA-specific: per-K log-space gate [B, T, HV, K]
    g = F.logsigmoid(
        torch.randn(batch_size, T, num_v_heads, head_size, device="cuda")
    ).to(dtype)

    # Beta: [B, T, HV] (pre-sigmoided)
    beta = torch.randn(batch_size, T, num_v_heads, dtype=dtype, device="cuda").sigmoid()

    state_slots = batch_size if T == 1 else batch_size * T
    # State pool storage for logical [K, V] states. Spec decode writes one
    # checkpoint slot per token.
    state = torch.randn(
        state_slots,
        num_v_heads,
        head_size,
        head_size,
        dtype=torch.bfloat16,
        device="cuda",
    )

    # Scale factor
    scale = 1.0 / (head_size**0.5)

    num_spec_tokens = None if T == 1 else T - 1

    # Benchmark with bench_gpu_time (CUPTI for accurate kernel timing)
    kernel_times_ms = bench_gpu_time(
        lambda: recurrent_kda(
            q=q,
            k=k,
            v=v,
            g=g,
            beta=beta,
            initial_state=state,
            scale=scale,
            use_qk_l2norm_in_kernel=True,
            num_spec_tokens=num_spec_tokens,
        ),
        enable_cupti=True,
        dry_run_iters=warmup_iters,
        repeat_iters=bench_iters,
    )

    # Calculate metrics
    kernel_median_ms = np.median(kernel_times_ms)
    flops = recurrent_kda_flops(
        batch_size, num_q_heads, num_k_heads, num_v_heads, head_size, seq_len
    )
    bytes_accessed = recurrent_kda_bytes(
        batch_size, num_q_heads, num_k_heads, num_v_heads, head_size, dtype, seq_len
    )

    kernel_tflops = flops / kernel_median_ms / 1e9 if kernel_median_ms > 0 else 0
    kernel_tb_per_sec = (
        bytes_accessed / kernel_median_ms / 1e9 if kernel_median_ms > 0 else 0
    )

    return {
        "batch_size": batch_size,
        "seq_len": seq_len,
        "mode": "decode" if T == 1 else "spec",
        "kernel_median_us": kernel_median_ms * 1000,
        "kernel_tflops": kernel_tflops,
        "kernel_tb_per_sec": kernel_tb_per_sec,
    }


# ============================================================================
# Runner
# ============================================================================


def run_recurrent_kda_benchmark(args, dtype):
    """Run recurrent KDA decode benchmarks."""
    if not RECURRENT_KDA_AVAILABLE:
        print("Error: recurrent KDA kernel is not available.")
        print("Make sure flashinfer.kda_decode.recurrent_kda is importable.")
        return

    invalid_seq_lens = [t for t in args.seq_len if t < 1]
    if invalid_seq_lens:
        print("Error: --seq-len values must be >= 1")
        return
    valid_seq_lens = args.seq_len

    print("\n" + "=" * 100)
    print(f"Recurrent KDA Benchmark (T={valid_seq_lens})")
    print(
        f"Config: q_heads={args.num_q_heads}, k_heads={args.num_k_heads}, "
        f"v_heads={args.num_v_heads}, head_size={args.head_size}, "
        f"dtype={args.dtype}"
    )
    print("=" * 100)
    print()
    print(
        f"{'batch':>6} {'T':>4} {'mode':>8} {'time(us)':>10} "
        f"{'TFLOPS':>10} {'TB/s':>10}"
    )
    print("-" * 100)

    all_results = []
    for batch_size in args.batch_size:
        for seq_len in valid_seq_lens:
            try:
                result = bench_recurrent_kda(
                    batch_size=batch_size,
                    seq_len=seq_len,
                    num_q_heads=args.num_q_heads,
                    num_k_heads=args.num_k_heads,
                    num_v_heads=args.num_v_heads,
                    head_size=args.head_size,
                    dtype=dtype,
                    warmup_iters=args.warmup,
                    bench_iters=args.iters,
                )
                all_results.append(result)

                print(
                    f"{result['batch_size']:>6} {result['seq_len']:>4} "
                    f"{result['mode']:>8} "
                    f"{result['kernel_median_us']:>10.2f} "
                    f"{result['kernel_tflops']:>10.2f} "
                    f"{result['kernel_tb_per_sec']:>10.2f}"
                )
            except Exception as e:
                mode = "decode" if seq_len == 1 else "spec"
                print(
                    f"{batch_size:>6} {seq_len:>4} {mode:>8} "
                    f"{'ERROR':>10} - {type(e).__name__}: {e}"
                )

    print("-" * 100)
    print()

    # Summary by T value
    for t in valid_seq_lens:
        t_results = [r for r in all_results if r["seq_len"] == t]
        if t_results:
            avg_time = np.mean([r["kernel_median_us"] for r in t_results])
            avg_tflops = np.mean([r["kernel_tflops"] for r in t_results])
            print(
                f"T={t}: Average time={avg_time:.2f}us, Average TFLOPS={avg_tflops:.2f}"
            )


# ============================================================================
# Main
# ============================================================================


def main():
    parser = argparse.ArgumentParser(
        description="Recurrent KDA Benchmark",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python benchmarks/bench_recurrent_kda.py --batch-size 1 4 16 64 128 256
  python benchmarks/bench_recurrent_kda.py --head-size 64 --batch-size 1 32 128
  python benchmarks/bench_recurrent_kda.py --seq-len 1 2 3 4 --batch-size 1 32
""",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        nargs="+",
        default=[1, 4, 16, 64, 128, 256],
        help="Batch sizes to benchmark",
    )
    parser.add_argument("--num-q-heads", type=int, default=16)
    parser.add_argument("--num-k-heads", type=int, default=16)
    parser.add_argument("--num-v-heads", type=int, default=32)
    parser.add_argument("--head-size", type=int, default=128, choices=[64, 128])
    parser.add_argument(
        "--dtype", type=str, choices=["float16", "bfloat16"], default="bfloat16"
    )
    parser.add_argument(
        "--seq-len",
        type=int,
        nargs="+",
        default=[1],
        help="Sequence lengths to benchmark; T=1 is standard decode, T>1 is spec decode",
    )
    parser.add_argument(
        "--warmup",
        type=int,
        default=10,
        help="Number of warmup iterations",
    )
    parser.add_argument(
        "--iters",
        type=int,
        default=100,
        help="Number of benchmark iterations",
    )
    args = parser.parse_args()

    # Resolve dtype
    dtype_map = {"float16": torch.float16, "bfloat16": torch.bfloat16}
    dtype = dtype_map[args.dtype]

    run_recurrent_kda_benchmark(args, dtype)


if __name__ == "__main__":
    main()
