# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "helion",
#     "kernels",
#     "torch",
# ]
# ///

from __future__ import annotations

from pathlib import Path

import kernels
import torch
import torch.nn.functional as F


def main() -> None:
    project = Path(__file__).parent
    build = project / "result" if (project / "result").exists() else project / "build"
    linear_attention = kernels.get_local_kernel(build, "cuda")

    b, t, h, d = 2, 512, 8, 64
    q = torch.randn(b, t, h, d, device="cuda", dtype=torch.bfloat16)
    k = F.normalize(
        torch.randn(b, t, h, d, device="cuda", dtype=torch.float32), dim=-1
    ).to(q)
    v = torch.randn_like(q)
    g = -torch.rand(b, t, h, d, device="cuda", dtype=torch.float32) * 0.1
    beta = torch.rand(b, t, h, device="cuda", dtype=torch.bfloat16)

    output, final_state = linear_attention.chunk_kda(
        q,
        k,
        v,
        g,
        beta,
        output_final_state=True,
    )
    print("output:", tuple(output.shape))
    print("final state:", tuple(final_state.shape))


if __name__ == "__main__":
    main()