Release frozen ThunderWorld FP4 1.5x and 2x weights, calibration provenance and loader (part 2)
18e8cca verified | """Current-residual weighted A selection. Only selected residual codes are written. | |
| The base pass emits 16-row residual summaries; no full residual-code materialization | |
| or selector GEMM. A second pass re-encodes selected blocks with the same factor. | |
| """ | |
| import torch | |
| import triton | |
| import triton.language as tl | |
| from ._linear_kernels import scaled_amax,finish_amax | |
| from ._calibrated import quantize_calibrated | |
| def base_summary(X,D,F,SENS,C,S,E,M:tl.constexpr,K:tl.constexpr,MP:tl.constexpr,BASE:tl.constexpr, | |
| RES:tl.constexpr=1.,GAIN_SCORE:tl.constexpr=False): | |
| rb=tl.program_id(0);kb=tl.program_id(1) | |
| mi=rb*16+tl.arange(0,16);ki=kb*128+tl.arange(0,128) | |
| x=tl.load(X+mi[:,None]*K+ki[None,:],mi[:,None]<M,0).to(tl.float32) | |
| d=tl.load(D+ki);x=(x*d[None,:]).to(tl.bfloat16).to(tl.float32) | |
| f=tl.load(F) | |
| p,s,r,rs,_=quantize_calibrated(x,f,16,128,GAIN_SCORE,BASE,RES,1.) | |
| code=tl.reshape(tl.join(p&15,p>>4),(16,128)).to(tl.int32) | |
| mag=code&7 | |
| value=tl.where(mag<4,mag*.5,tl.where(mag<6,mag-2.,2.*mag-8.)) | |
| value=tl.where(code>=8,-value,value) | |
| decoded=tl.reshape(tl.reshape(value,(16,8,16))*s.to(tl.float32)[:,:,None],(16,128))*f | |
| residual=x-decoded | |
| benefit=residual*residual | |
| if GAIN_SCORE: | |
| # Evaluate the represented residual, not a hypothetical exact repair. | |
| # Codes stay in registers; only selected residual codes are written. | |
| rc=tl.reshape(tl.join(r&15,r>>4),(16,128)).to(tl.int32) | |
| rm=rc&7 | |
| rv=tl.where(rm<4,rm*.5,tl.where(rm<6,rm-2.,2.*rm-8.)) | |
| rv=tl.where(rc>=8,-rv,rv) | |
| represented=tl.reshape(tl.reshape(rv,(16,8,16))*rs.to(tl.float32)[:,:,None],(16,128))*f | |
| remaining=residual-represented | |
| benefit=tl.maximum(benefit-remaining*remaining,0.) | |
| sensitivity=tl.load(SENS+ki) | |
| score=tl.sum(tl.sum(benefit*sensitivity[None,:],1),0) | |
| tl.store(E+rb*(K//128)+kb,score) | |
| tl.store(C+mi[:,None]*(K//2)+kb*64+tl.arange(0,64)[None,:],p,mi[:,None]<MP) | |
| tl.store(S+mi[:,None]*(K//16)+kb*8+tl.arange(0,8)[None,:],s,mi[:,None]<MP) | |
| def choose(E,MASK,M:tl.constexpr,K:tl.constexpr,BT:tl.constexpr,NA:tl.constexpr,ENERGY:tl.constexpr=1.): | |
| mb=tl.program_id(0);kt=tl.arange(0,BT) | |
| rows=mb*8+tl.arange(0,8) | |
| scores=tl.load(E+rows[:,None]*(K//128)+kt[None,:],(rows[:,None]<tl.cdiv(M,16))&(kt[None,:]<K//128),0.) | |
| energy=tl.where(kt<K//128,tl.sum(scores,0),-float('inf')) | |
| rank=tl.sum(((energy[None,:]>energy[:,None])|((energy[None,:]==energy[:,None])&(kt[None,:]<kt[:,None]))).to(tl.int32),1) | |
| active=rank<NA | |
| if ENERGY < 1.: | |
| # Stop at the smallest stable prefix explaining ENERGY of the current | |
| # residual sensitivity proxy, subject to the existing per-tile cap. | |
| # Zero residual needs no repair. No host decision or additional GEMM. | |
| values=tl.where(kt<K//128,energy,0.) | |
| before=tl.sum(tl.where(rank[None,:]<rank[:,None],values[None,:],0.),1) | |
| active=active & (before < ENERGY*tl.sum(values,0)) | |
| tl.store(MASK+mb*(K//128)+kt,active,kt<K//128) | |
| def residual_only(X,D,F,MASK,C,S,M:tl.constexpr,K:tl.constexpr,MP:tl.constexpr,BASE:tl.constexpr,RES:tl.constexpr): | |
| rb=tl.program_id(0);kb=tl.program_id(1) | |
| active=tl.load(MASK+(rb//8)*(K//128)+kb) | |
| if active: | |
| mi=rb*16+tl.arange(0,16);ki=kb*128+tl.arange(0,128) | |
| x=tl.load(X+mi[:,None]*K+ki[None,:],mi[:,None]<M,0).to(tl.float32) | |
| d=tl.load(D+ki);x=(x*d[None,:]).to(tl.bfloat16).to(tl.float32) | |
| _,_,r,rs,_=quantize_calibrated(x,tl.load(F),16,128,True,BASE,RES,1.) | |
| tl.store(C+MP*(K//2)+mi[:,None]*(K//2)+kb*64+tl.arange(0,64)[None,:],r,mi[:,None]<MP) | |
| tl.store(S+MP*(K//16)+mi[:,None]*(K//16)+kb*8+tl.arange(0,8)[None,:],rs,mi[:,None]<MP) | |
| def pack_sparse(x,d,sensitivity,mask,count,*,base=1.,res=1.,factor=None,energy=1.,score='residual_energy'): | |
| if score not in ('residual_energy','represented_gain'):raise ValueError('Unknown sparse activation score') | |
| m,k=x.shape;mp=triton.cdiv(m,16)*16 | |
| if not 0<count<k//128:raise ValueError('Sparse packing requires a partial integer quota') | |
| if factor is None: | |
| partial=torch.empty((triton.cdiv(m*k,4096),),device=x.device,dtype=torch.float32) | |
| factor=torch.empty((),device=x.device,dtype=torch.float32) | |
| scaled_amax[(partial.numel(),)](x,d,partial,m,k,4096,num_warps=4) | |
| finish_amax[(1,)](partial,factor,partial.numel(),triton.next_power_of_2(partial.numel()),num_warps=4) | |
| c=torch.empty((2,mp,k//2),device=x.device,dtype=torch.uint8) | |
| s=torch.empty((2,mp,k//16),device=x.device,dtype=torch.float8_e4m3fn) | |
| summary=torch.empty((triton.cdiv(m,16),k//128),device=x.device,dtype=torch.float32) | |
| base_summary[(triton.cdiv(m,16),k//128)](x,d,factor,sensitivity,c,s,summary,m,k,mp,base,res, | |
| score=='represented_gain',num_warps=4,enable_fp_fusion=False) | |
| choose[(triton.cdiv(m,128),)](summary,mask,m,k,triton.next_power_of_2(k//128),count,energy,num_warps=4) | |
| residual_only[(triton.cdiv(m,16),k//128)](x,d,factor,mask,c,s,m,k,mp,base,res,num_warps=4) | |
| return c,s,factor,mp | |