File size: 2,323 Bytes
b35ab1f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | from pathlib import Path
import pytest
import torch
from kernels import get_local_kernel
kernel=get_local_kernel(Path(__file__).resolve().parents[1]/'build','yue2_qkv_fused')
def inputs(rows,scale=1):
torch.manual_seed(51)
x=(torch.randn(rows,2048,device='cuda')*scale).bfloat16()
w=torch.randn(2048,device='cuda').bfloat16()
p=torch.randperm(2048,device='cuda').int()
s=(torch.randint(0,2,(2048,),device='cuda')*2-1).to(torch.int8)
bounds=torch.linspace(-.065,.065,15,device='cuda')
return x,w,p,s,bounds
@pytest.mark.parametrize('rows',[0,1,2,8,32])
@pytest.mark.parametrize('scale',[0,1e-9,1,100])
def test_continuous_formula(rows,scale):
x,w,p,s,b=inputs(rows,scale)
packed,norms=kernel.rmsquant(x,w,p,s,b)
y=x.float()*torch.rsqrt(x.float().square().mean(-1,keepdim=True)+1e-6)*w.float()
expected_norm=y.norm(dim=-1)
y=y/(expected_norm[:,None]+1e-8)
y=y[:,p.long()]*s.float()
width=1
while width<2048:
z=y.reshape(rows,2048//(2*width),2,width)
a,c=z[:,:,0],z[:,:,1]
y=torch.stack([a+c,a-c],dim=2).reshape(rows,2048)
width*=2
codes=torch.bucketize(y/(2048**.5),b)
expected=(codes[:,0::2]|(codes[:,1::2]<<4)).byte()
torch.testing.assert_close(norms,expected_norm,rtol=2e-6,atol=1e-9)
# Only boundaries within FP32 reduction noise may differ.
if rows:
assert (expected!=packed).float().mean().item()<.001
assert torch.isfinite(norms).all()
def test_graph_and_invalid_contract():
args=inputs(2)
expected=kernel.rmsquant(*args)
stream=torch.cuda.Stream();stream.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(stream):
for _ in range(3):kernel.rmsquant(*args)
torch.cuda.current_stream().wait_stream(stream)
graph=torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):result=kernel.rmsquant(*args)
graph.replay();torch.cuda.synchronize()
for a,b in zip(expected,result):torch.testing.assert_close(a,b,rtol=0,atol=0)
with pytest.raises(RuntimeError,match='BF16'):
kernel.rmsquant(args[0].float(),*args[1:])
with pytest.raises(RuntimeError,match='positive'):
kernel.rmsquant(*args,rms_eps=0)
with pytest.raises(RuntimeError,match='permutation'):
kernel.rmsquant(args[0],args[1],args[2].long(),*args[3:])
|