| import json |
| import os |
| import sys |
| import time |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| import tensorrt as trt |
| from huggingface_hub import snapshot_download |
|
|
| REPO=os.environ.get('CF_RUNTIME_REPO','patdev/Companion-Forge-L4-ONNX') |
| ROOT=Path(snapshot_download(REPO,repo_type='model',token=os.environ.get('HF_TOKEN'),local_dir='/tmp/cf-val',allow_patterns=['engines/l4-sm89/custom-ops/*.plan','plugins/tensorrt/companion_sparse_trt.py','plugins/reference/ops.py'])) |
| sys.path.insert(0,str(ROOT/'plugins/tensorrt'));sys.path.insert(0,str(ROOT/'plugins/reference')) |
| import companion_sparse_trt |
| import ops as ref |
|
|
| LOGGER=trt.Logger(trt.Logger.ERROR) |
| DT={trt.float32:torch.float32,trt.float16:torch.float16,trt.int32:torch.int32,trt.int64:torch.int64,trt.bool:torch.bool} |
| if hasattr(trt,'bfloat16'): DT[trt.bfloat16]=torch.bfloat16 |
|
|
| class OutputAllocator(trt.IOutputAllocator): |
| def __init__(self,dtype): |
| trt.IOutputAllocator.__init__(self);self.dtype=dtype;self.tensor=None;self.shape=None;self.nbytes=0 |
| def reallocate_output(self,tensor_name,memory,size,alignment): |
| self.nbytes=int(size); n=max(1,(int(size)+torch.tensor([],dtype=self.dtype).element_size()-1)//torch.tensor([],dtype=self.dtype).element_size()) |
| self.tensor=torch.empty(n,device='cuda',dtype=self.dtype);return int(self.tensor.data_ptr()) |
| def reallocate_output_async(self,tensor_name,memory,size,alignment,stream): |
| return self.reallocate_output(tensor_name,memory,size,alignment) |
| def notify_shape(self,tensor_name,dims): |
| self.shape=tuple(int(x) for x in dims) |
|
|
| class Runner: |
| def __init__(self,name): |
| blob=(ROOT/'engines/l4-sm89/custom-ops'/f'{name}.plan').read_bytes();self.rt=trt.Runtime(LOGGER);self.eng=self.rt.deserialize_cuda_engine(blob);assert self.eng;self.ctx=self.eng.create_execution_context();self.name=name |
| def run(self,feeds): |
| keep=[] |
| for n,x in feeds.items(): |
| dt=DT[self.eng.get_tensor_dtype(n)];x=x.to(dtype=dt).contiguous();feeds[n]=x;keep.append(x) |
| if self.eng.get_tensor_mode(n)==trt.TensorIOMode.INPUT: |
| self.ctx.set_input_shape(n,tuple(x.shape));self.ctx.set_tensor_address(n,int(x.data_ptr())) |
| outs={};allocs={} |
| for i in range(self.eng.num_io_tensors): |
| n=self.eng.get_tensor_name(i) |
| if self.eng.get_tensor_mode(n)!=trt.TensorIOMode.OUTPUT:continue |
| dt=DT[self.eng.get_tensor_dtype(n)];shape=tuple(int(x) for x in self.ctx.get_tensor_shape(n)) |
| if any(x<0 for x in shape): |
| a=OutputAllocator(dt);allocs[n]=a;self.ctx.set_output_allocator(n,a) |
| else: |
| y=torch.empty(shape if shape else (),device='cuda',dtype=dt);outs[n]=y;keep.append(y);self.ctx.set_tensor_address(n,int(y.data_ptr())) |
| ok=self.ctx.execute_async_v3(stream_handle=torch.cuda.current_stream().cuda_stream);assert ok,f'{self.name} execute failed';torch.cuda.synchronize() |
| for n,a in allocs.items(): |
| if a.tensor is None:raise RuntimeError(f'{self.name}/{n}: allocator not called') |
| if a.shape is None:raise RuntimeError(f'{self.name}/{n}: shape notification missing') |
| num=int(np.prod(a.shape)) if a.shape else 1 |
| outs[n]=a.tensor[:num].view(a.shape) |
| return outs |
|
|
| def diff(a,b): |
| a=a.float();b=b.float();d=(a-b).abs();return {'max_abs':float(d.max()) if d.numel() else 0.0,'mean_abs':float(d.mean()) if d.numel() else 0.0,'cos':float(torch.nn.functional.cosine_similarity(a.reshape(1,-1),b.reshape(1,-1)).item()) if d.numel() else 1.0} |
|
|
| def grid4(): |
| xyz=torch.cartesian_prod(torch.arange(4),torch.arange(4),torch.arange(4)).to(torch.int32);return torch.cat([torch.zeros((64,1),dtype=torch.int32),xyz],1).cuda() |
|
|
| report={} |
| torch.manual_seed(123) |
| |
| f=torch.randn(32,8,device='cuda',dtype=torch.float16);c=grid4()[:32].contiguous();r0,r1=ref.sparse_subdivide(f,c);o=Runner('SparseSubdivide').run({'feats':f,'coords':c});report['SparseSubdivide']={'feats':diff(o['out_feats'],r0),'coords_equal':bool(torch.equal(o['out_coords'],r1))} |
| |
| f=torch.randn(64,8,device='cuda',dtype=torch.float16);c=grid4();r0,r1,ri=ref.sparse_downsample(f,c,(2,2,2));o=Runner('SparseDownsample').run({'feats':f,'coords':c});m=int(o['count'].item());report['SparseDownsample']={'count':m,'ref_count':int(r0.shape[0]),'feats':diff(o['out_feats'][:m],r0),'coords_equal':bool(torch.equal(o['out_coords'][:m],r1)),'inverse_equal':bool(torch.equal(o['inverse'],ri))} |
| |
| f=torch.randn(8,8,device='cuda',dtype=torch.float16);tc=grid4();inv=torch.arange(64,device='cuda',dtype=torch.int32)%8;r0,r1=ref.sparse_upsample(f,tc,inv);o=Runner('SparseUpsample').run({'feats':f,'target_coords':tc,'inverse':inv});report['SparseUpsample']={'feats':diff(o['out_feats'],r0),'coords_equal':bool(torch.equal(o['out_coords'],r1))} |
| |
| qkv=torch.randn(64,3,4,16,device='cuda',dtype=torch.float16);c=grid4();r=ref.sparse_window_attention(qkv,c,8,(0,0,0));o=Runner('SparseWindowAttention').run({'qkv':qkv,'coords':c});report['SparseWindowAttention']=diff(o['out'],r) |
| |
| import spconv.pytorch as spconv |
| f=torch.randn(64,8,device='cuda',dtype=torch.float16);c=grid4();mod=spconv.SubMConv3d(8,16,3,bias=True,algo=spconv.ConvAlgo.Native).cuda().half().eval();st=spconv.SparseConvTensor(f,c,[64,64,64],1);ry=mod(st);w=mod.weight.detach().contiguous();b=mod.bias.detach().contiguous();print('SPCONV_WEIGHT_SHAPE',tuple(w.shape),flush=True) |
| try: |
| o=Runner('SparseConv3D').run({'feats':f,'coords':c,'weight':w,'bias':b});m=int(o['count'].item());report['SparseConv3D']={'count':m,'ref_count':int(ry.features.shape[0]),'feats':diff(o['out_feats'][:m],ry.features),'coords_equal':bool(torch.equal(o['out_coords'][:m],ry.indices.to(torch.int32)))} |
| except Exception as e: |
| report['SparseConv3D']={'error':repr(e),'weight_shape':tuple(w.shape)} |
| |
| for name in ['SparseSubdivide','SparseUpsample','SparseWindowAttention']: |
| pass |
| print('REPORT',json.dumps(report,indent=2),flush=True) |
| Path('/tmp/custom_validation.json').write_text(json.dumps(report,indent=2)) |
|
|
