| 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) |
|
|
