import torch,pytest,json
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('rows,dim',[(1,2048),(2,2048),(16,128),(32,128),(1,6144),(67,2048)])
def test_rmsnorm(rows,dim):
 torch.manual_seed(12)
 x=(torch.randn(rows,dim,device='cuda')*torch.logspace(-3,3,rows,device='cuda')[:,None]).bfloat16();w=torch.randn(dim,device='cuda',dtype=torch.bfloat16)
 ref=x*torch.rsqrt(x.float().pow(2).mean(-1,keepdim=True)+1e-6).to(x.dtype)*w
 out=kernel.rmsnorm(x,w,1e-6)
 torch.testing.assert_close(out,ref,rtol=.008,atol=1e-5)
 print('rmsnorm exact fraction',rows,dim,float((out==ref).float().mean()))
@pytest.mark.parametrize('b,t,h',[(1,1,16),(1,1,8),(2,1,16),(1,17,8)])
def test_rope(b,t,h):
 torch.manual_seed(13);x=torch.randn(b,t,h,128,device='cuda',dtype=torch.bfloat16);a=torch.randn(b,t,1,64,device='cuda');c=a.cos();s=a.sin();cb=c.bfloat16();sb=s.bfloat16();x1,x2=x.chunk(2,-1)
 ref=torch.cat([x1*cb-x2*sb,x2*cb+x1*sb],-1)
 torch.testing.assert_close(kernel.rope(x,c,s),ref,rtol=0,atol=0)
@pytest.mark.parametrize('rows',[1,2,16])
def test_swiglu(rows):
 torch.manual_seed(14);x=torch.randn(rows,12288,device='cuda',dtype=torch.bfloat16);g,u=x.chunk(2,-1)
 torch.testing.assert_close(kernel.swiglu(x),torch.nn.functional.silu(g)*u,rtol=0,atol=0)
def test_invalid_norm_weight():
 with pytest.raises(RuntimeError,match='weight'):kernel.rmsnorm(torch.ones(1,128,device='cuda',dtype=torch.bfloat16),torch.ones(64,device='cuda',dtype=torch.bfloat16),1e-6)