File size: 4,049 Bytes
99323da 9a2611f 99323da | 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 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 | from pathlib import Path
import pytest
import torch
from kernels import get_local_kernel
kernel=get_local_kernel(Path(__file__).resolve().parents[1]/'build','int8_ordered_attention')
def quantize(x):
scale=x.float().abs().amax(-1).clamp_min(1e-12)/127
return (x.float()/scale[...,None]).round().clamp(-127,127).to(torch.int8),scale
def reference(q,k,v,ks,vs,lengths):
out=[]
for b,length in enumerate(lengths):
if not length:out.append(torch.zeros_like(q[b],dtype=torch.float32));continue
key=(k[b,:length].float()*ks[b,:length,None].movedim(1,2)).repeat_interleave(2,1)
value=(v[b,:length].float()*vs[b,:length,None].movedim(1,2)).repeat_interleave(2,1)
scores=torch.einsum('hd,thd->ht',q[b].float(),key)*(128**-.5)
out.append(torch.einsum('ht,thd->hd',scores.softmax(-1),value))
return torch.stack(out)
@pytest.mark.parametrize('parts',[16,32,64,128])
@pytest.mark.parametrize('dtype',[torch.bfloat16,torch.float16])
@pytest.mark.parametrize('capacity,lengths',[(129,[0,129]),(513,[513,129]),(11053,[2053,6053])])
def test_quantized_lengths_and_graph(parts,dtype,capacity,lengths):
torch.manual_seed(995)
q=torch.randn(2,16,128,device='cuda',dtype=dtype)
k,ks=quantize(torch.randn(2,capacity,8,128,device='cuda',dtype=dtype))
v,vs=quantize(torch.randn(2,capacity,8,128,device='cuda',dtype=dtype))
for b,length in enumerate(lengths):ks[b,length:]=float('nan');vs[b,length:]=float('nan')
used=torch.tensor(lengths,device='cuda',dtype=torch.int32)
out=kernel.attention(q,k,v,ks,vs,used,parts)
assert torch.isfinite(out).all()
torch.testing.assert_close(out.float(),reference(q,k,v,ks,vs,lengths),atol=.001,rtol=.01)
stream=torch.cuda.Stream();stream.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(stream):
for _ in range(3):kernel.attention(q,k,v,ks,vs,used,parts)
torch.cuda.current_stream().wait_stream(stream)
graph=torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):captured=kernel.attention(q,k,v,ks,vs,used,parts)
short=[min(x,7) for x in lengths];used.copy_(torch.tensor(short,device='cuda',dtype=torch.int32));graph.replay();torch.cuda.synchronize()
torch.testing.assert_close(captured.float(),reference(q,k,v,ks,vs,short),atol=.001,rtol=.01)
def test_invalid_scale_and_cache_contracts():
q=torch.zeros(1,16,128,device='cuda',dtype=torch.bfloat16)
k=torch.zeros(1,10,8,128,device='cuda',dtype=torch.int8)
s=torch.ones(1,10,8,device='cuda');used=torch.ones(1,device='cuda',dtype=torch.int32)
with pytest.raises(RuntimeError,match='scales'):kernel.attention(q,k,k,s.half(),s,used)
with pytest.raises(RuntimeError,match='scales'):kernel.attention(q,k,k,s[:,:,:7],s,used)
with pytest.raises(RuntimeError,match='KV dtype'):kernel.attention(q,k.bfloat16(),k,s,s,used)
def test_unaligned_cache_view():
torch.manual_seed(109)
q=torch.randn(1,16,128,device='cuda',dtype=torch.bfloat16)
shape=(1,129,8,128);n=129*8*128
backing=torch.randint(-127,128,(n+1,),device='cuda',dtype=torch.int8)
k=backing[1:].view(shape);v=k.clone();scales=torch.full((1,129,8),.02,device='cuda');used=torch.tensor([127],device='cuda',dtype=torch.int32)
out=kernel.attention(q,k,v,scales,scales,used,32)
torch.testing.assert_close(out.float(),reference(q,k,v,scales,scales,[127]),atol=.001,rtol=.01)
@pytest.mark.parametrize('dtype',[torch.bfloat16,torch.float16])
@pytest.mark.parametrize('batch',[1,2])
@pytest.mark.parametrize('parts',[32,64,128])
@pytest.mark.parametrize('length',[129,6053])
def test_original_reduction_order(dtype,batch,parts,length):
control=get_local_kernel(Path(__file__).resolve().parents[2]/'int8-kv/build','int8_kv_runtime')
torch.manual_seed(110)
q=torch.randn(batch,16,128,device='cuda',dtype=dtype)
k,ks=quantize(torch.randn(batch,length+37,8,128,device='cuda',dtype=dtype))
v,vs=quantize(torch.randn(batch,length+37,8,128,device='cuda',dtype=dtype))
used=torch.tensor([length-i*31 for i in range(batch)],device='cuda',dtype=torch.int32)
torch.testing.assert_close(kernel.attention(q,k,v,ks,vs,used,parts),control.attention(q,k,v,ks,vs,used,parts),atol=0,rtol=0)
|