| """Bounded weight preparation and executable current-recipe linears."""
|
| from functools import lru_cache
|
| from pathlib import Path
|
| import torch
|
| from .kernel_build import stable_source
|
| import torch.nn.functional as F
|
| from .numerics import Packed,quantize,bridge_quantize,gemm_release_reference
|
| from .contracts import blend
|
|
|
| @lru_cache(None)
|
| def extension():
|
| from torch.utils.cpp_extension import load
|
| return load(name='fidel_compatibility_packed_reference_v1',
|
| sources=[str(stable_source(Path(__file__).with_name('packed_gemm_reference.cu')))],
|
| extra_cuda_cflags=['-O3','--fmad=false'],verbose=False)
|
|
|
| def packed_product(a,b):
|
| if a.data.is_cuda:
|
| return extension().gemm(a.data,b.data,a.scales,b.scales,a.global_scale,b.global_scale)
|
| return gemm_release_reference(a,b)
|
|
|
| def bridge_values(x):
|
| if x.is_cuda:return extension().bridge(x.float().contiguous())
|
| return bridge_quantize(x)
|
|
|
|
|
|
|
|
|
| DENSE_F32_PARTITIONS={
|
| (16,1536):24,(32,1536):24,(384,1536):24,(1536,1536):24,
|
| (512,1536):16,(256,1536):16,(12288,1536):1,(1536,6144):24,
|
| (1536,512):1,(1536,256):1,(1536,32):1,(384,32):1,
|
| (128,1536):24,(64,1536):24,(1536,128):1,(1536,64):1,
|
| (131072,1536):1,(131072,256):1,
|
| }
|
|
|
| class LinearEngine:
|
| def __init__(self,weight_rows=512):
|
| self.weight_rows=weight_rows
|
| self.weights={}
|
| self.dense_weights={}
|
| self.calls=0
|
| def clear(self):
|
| self.weights.clear();self.dense_weights.clear()
|
| def prepare(self,name,weight,seed,bridge):
|
| key=(name,seed,bridge,weight.device,weight.dtype,weight.data_ptr(),weight._version,tuple(weight.stride()))
|
| if key in self.weights:return self.weights[key]
|
| global_scale=(weight.float().abs().amax()*(1./(6*448*(16/17)))).clamp_min(torch.finfo(torch.float32).tiny)
|
| data,scales,original,means=[],[],[],[]
|
| for start in range(0,weight.shape[0],self.weight_rows):
|
| block=weight[start:start+self.weight_rows].float().contiguous()
|
| q=quantize(block,eden=True,seed=seed+1,global_scale_override=global_scale,
|
| group_offset=start*(weight.shape[1]//16))
|
| data.append(q.data);scales.append(q.scales)
|
| if bridge:
|
| qb,sb=bridge_values(block)
|
| original.append(qb);means.append(sb.mean(1,keepdim=True))
|
| value=(Packed(torch.cat(data),torch.cat(scales),global_scale),
|
| torch.cat(original) if bridge else None,torch.cat(means) if bridge else None)
|
| self.weights[key]=value
|
| return value
|
| def __call__(self,x,weight,bias=None,*,name,seed=0,bridge=False,activation_row_offset=0):
|
| if x.requires_grad or weight.requires_grad:raise ValueError('Inference only')
|
| if x.shape[-1]!=weight.shape[1]:raise ValueError('Linear width mismatch')
|
| if type(activation_row_offset) is not int or activation_row_offset<0 or activation_row_offset%128:
|
| raise ValueError('Activation offsets must start on a 128-row tile')
|
| self.calls+=1
|
| shape=x.shape;rows=x.reshape(-1,shape[-1]);live=rows.shape[0]
|
| padded=F.pad(rows,(0,0,0,(-live)%128))
|
| qw,bridge_w,bridge_mean=self.prepare(name,weight,seed,bridge)
|
| qx=quantize(padded.float(),eden=True,seed=seed,rowwise=True,
|
| group_offset=activation_row_offset*(padded.shape[1]//16))
|
| recipe=packed_product(qx,qw)
|
| if bias is not None:recipe=recipe+bias.float()
|
| recipe=recipe[:live].to(x.dtype)
|
| output=[]
|
| for start in range(0,live,128):
|
| block=padded[start:start+128];count=min(128,live-start)
|
| if bridge:
|
| q,s=bridge_values(block)
|
| dot=(q@bridge_w.T).to(torch.bfloat16)
|
| mean=s.mean(1,keepdim=True)
|
| baseline=(dot*mean)*bridge_mean.T if x.dtype==torch.bfloat16 else dot.float()*mean.float()*bridge_mean.T.float()
|
| if bias is not None:baseline=baseline+bias.to(baseline.dtype)
|
| elif block.is_cuda and block.dtype==weight.dtype==torch.bfloat16 and torch.cuda.get_device_capability(block.device)!=(10,0):
|
| from .math_primitives import bf16_product
|
| key=(weight.data_ptr(),weight._version,tuple(weight.shape),tuple(weight.stride()),str(weight.device))
|
| if key not in self.dense_weights:self.dense_weights[key]=weight.T.contiguous()
|
| baseline=bf16_product(block[:((count+15)//16)*16],self.dense_weights[key]).to(block.dtype)
|
| if bias is not None:baseline=baseline+bias.to(baseline.dtype)
|
| elif block.is_cuda and block.dtype==torch.float32 and torch.cuda.get_device_capability(block.device)!=(10,0) and tuple(weight.shape) in DENSE_F32_PARTITIONS:
|
| from .math_primitives import extension as math_extension
|
| partitions=DENSE_F32_PARTITIONS[tuple(weight.shape)]
|
| baseline=math_extension().dense_fma(block[:count].float().contiguous(),weight.float().contiguous(),partitions)
|
| if bias is not None:baseline=baseline+bias.float()
|
| elif block.dtype==weight.dtype:
|
| baseline=F.linear(block,weight,bias)
|
| else:
|
| baseline=F.linear(block.float(),weight.float(),None if bias is None else bias.float()).to(x.dtype)
|
| output.append(blend(baseline[:count],recipe[start:start+count],legacy_bridge=bridge))
|
| return torch.cat(output).reshape(*shape[:-1],weight.shape[0])
|
|
|
