import torch
from kernels import get_kernel


kernel = get_kernel("Efficient-Large-Model/Sol-Attn", version=1)

q = torch.randn(1, 4096, 16, 128, device="cuda", dtype=torch.bfloat16)
k = torch.randn_like(q)
v = torch.randn_like(q)

out = kernel.sol_attn(
    q,
    k,
    v,
    tau=1.0,
    thresh_type="exact",
)
print(out.shape)