File size: 1,960 Bytes
3ea3da7 | 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 | 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)
|