| import json |
| from pathlib import Path |
| import pytest |
| import torch |
| from kernels import get_local_kernel |
| from orbitquant.kernels.native_packed_matmul import matmul_packed_w4a4_int8_with_native_kernel as original |
|
|
| kernel=get_local_kernel(Path(__file__).resolve().parents[1] / 'build', 'orbitquant_gemv') |
|
|
| @pytest.mark.parametrize('rows,n,k',[(1,2048,2048),(2,6144,2048),(8,2048,6144),(1,129,128),(2,1024,2048)]) |
| @pytest.mark.parametrize('dtype',[torch.bfloat16,torch.float16]) |
| def test_exact_integer_gemv(rows,n,k,dtype): |
| torch.manual_seed(123) |
| x=torch.randint(0,256,(rows,k//2),device='cuda',dtype=torch.uint8) |
| w=torch.randint(0,256,(n*k//2,),device='cuda',dtype=torch.uint8) |
| xn=torch.rand(rows,device='cuda');wn=torch.rand(n,device='cuda',dtype=torch.bfloat16) |
| ac=torch.arange(-8,8,device='cuda',dtype=torch.int8);wc=ac.flip(0).contiguous() |
| bias=None |
| kw=dict(activation_scale=.03125,weight_scale=.0625,bias=bias,output_dtype=dtype) |
| ref=original(x,w,xn,wn,ac,wc,out_features=n,in_features=k,**kw) |
| out=kernel.gemv(x,w,xn,wn,ac,wc,**kw) |
| torch.testing.assert_close(out,ref,rtol=0,atol=0) |
| stream=torch.cuda.Stream();stream.wait_stream(torch.cuda.current_stream()) |
| with torch.cuda.stream(stream): |
| for _ in range(3):kernel.gemv(x,w,xn,wn,ac,wc,**kw) |
| torch.cuda.current_stream().wait_stream(stream) |
| graph=torch.cuda.CUDAGraph() |
| with torch.cuda.graph(graph): captured=kernel.gemv(x,w,xn,wn,ac,wc,**kw) |
| graph.replay();torch.cuda.synchronize() |
| torch.testing.assert_close(captured,ref,rtol=0,atol=0) |
|
|
| def test_rejects_bad_weight_shape(): |
| x=torch.zeros(1,64,device='cuda',dtype=torch.uint8) |
| with pytest.raises(RuntimeError,match='packed weight size'): |
| kernel.gemv(x,x.flatten(),torch.ones(1,device='cuda'),torch.ones(2,device='cuda',dtype=torch.bfloat16),torch.zeros(16,device='cuda',dtype=torch.int8),torch.zeros(16,device='cuda',dtype=torch.int8),activation_scale=1,weight_scale=1) |
|
|
