import torch,pytest
from pathlib import Path
from kernels import get_local_kernel
kernel=get_local_kernel(Path(__file__).resolve().parents[1]/'build','yue2_qkv_fused')
@pytest.mark.parametrize('b,capacity',[(1,32),(2,257)])
def test_decode_qkv_cache(b,capacity):
 torch.manual_seed(42);hq,hk,d=16,8,128;x=torch.randn(b,1,(hq+2*hk)*d,device='cuda',dtype=torch.bfloat16);qw=torch.randn(d,device='cuda',dtype=torch.bfloat16);kw=torch.randn(d,device='cuda',dtype=torch.bfloat16)
 a=torch.randn(b,1,d//2,device='cuda');c=a.cos();s=a.sin();positions=torch.tensor([capacity-1-i for i in range(b)],device='cuda');keys=torch.randn(b,capacity,hk,d,device='cuda',dtype=torch.bfloat16);values=keys.clone();expected_k=keys.clone();expected_v=values.clone()
 q,k,v=x.split((hq*d,hk*d,hk*d),-1);q=q.reshape(b,1,hq,d);k=k.reshape(b,1,hk,d);v=v.reshape(b,1,hk,d)
 def norm(z,w):return z*torch.rsqrt(z.float().pow(2).mean(-1,keepdim=True)+1e-6).bfloat16()*w
 def rope(z):
  x1,x2=z.chunk(2,-1);cb=c.unsqueeze(2).bfloat16();sb=s.unsqueeze(2).bfloat16();return torch.cat([x1*cb-x2*sb,x2*cb+x1*sb],-1)
 q=rope(norm(q,qw));k=rope(norm(k,kw));slots=positions[:,None,None,None].expand(b,1,hk,d);expected_k.scatter_(1,slots,k);expected_v.scatter_(1,slots,v)
 actual=kernel.qkv_decode(x,qw,kw,c,s,positions,keys,values,hq,1e-6)
 torch.testing.assert_close(actual,q,rtol=0,atol=0);torch.testing.assert_close(keys,expected_k,rtol=0,atol=0);torch.testing.assert_close(values,expected_v,rtol=0,atol=0)
@pytest.mark.parametrize('rows,dim',[(1,2048),(16,128),(2,2048),(32,128),(67,2048),(1,6144)])
def test_exact_rmsnorm(rows,dim):
 torch.manual_seed(456);x=torch.randn(rows,dim,device='cuda',dtype=torch.bfloat16);w=torch.randn(dim,device='cuda',dtype=torch.bfloat16)
 ref=x*torch.rsqrt(x.float().pow(2).mean(-1,keepdim=True)+1e-6).bfloat16()*w
 torch.testing.assert_close(kernel.rmsnorm_exact(x,w,1e-6),ref,rtol=0,atol=0)