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