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