File size: 5,839 Bytes
7218492 | 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 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | """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)
# Native B200 FP32 dense profiles under the release's fixed 128-row baseline.
# Keys are (output width, reduction width); every admitted profile has retained
# real-operand authority. Other shapes remain behind the model fidelity gate.
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])
|