Release frozen ThunderWorld FP4 1.5x and 2x weights, calibration provenance and loader (part 2)
18e8cca verified | """Native compact-index repair loops; offline packed weight CSR tiles. | |
| All repairs use the same outer factors as the dense base, permitting one | |
| FP32 accumulator. Residual W storage is only the selected N128/K_B tiles. | |
| """ | |
| import torch,triton | |
| import triton.language as tl | |
| from .linear import _pack | |
| def gemm(AC,AS,WC,WS,WG,AG,AI,AM,WR,WRS,WI,OFF,Y, | |
| M:tl.constexpr,N:tl.constexpr,K:tl.constexpr,MP:tl.constexpr, | |
| NA:tl.constexpr,HASW:tl.constexpr,CROSS:tl.constexpr,B:tl.constexpr,BM:tl.constexpr,INTERLEAVE:tl.constexpr,SPLIT_K:tl.constexpr=1,DENSE_W:tl.constexpr=False,WLOC=None,SHARED_BANK:tl.constexpr=False): | |
| mi=tl.program_id(0)*BM+tl.arange(0,BM);nb=tl.program_id(1);ni=nb*128+tl.arange(0,128) | |
| dk=tl.arange(0,B//2);ds=tl.arange(0,B//16);acc=tl.full((BM,128),0,tl.float32) | |
| part=tl.program_id(2) | |
| if SPLIT_K==1:begin_k=0;end_k=K//B | |
| else:begin_k=part*tl.cdiv(K//B,SPLIT_K);end_k=tl.minimum((part+1)*tl.cdiv(K//B,SPLIT_K),K//B) | |
| for kb in range(begin_k,end_k): | |
| ai=mi[:,None]*(K//2)+kb*(B//2)+dk[None,:];asi=mi[:,None]*(K//16)+kb*(B//16)+ds[None,:] | |
| wi=ni[:,None]*(K//2)+kb*(B//2)+dk[None,:];wsi=ni[:,None]*(K//16)+kb*(B//16)+ds[None,:] | |
| a=tl.load(AC+ai,mi[:,None]<MP,0);s=tl.load(AS+asi,mi[:,None]<MP,0.) | |
| w=tl.load(WC+wi);ws=tl.load(WS+wsi) | |
| acc=tl.dot_scaled(a,s,'e2m1',w.T,ws,'e2m1',acc) | |
| if NA==K//B or (INTERLEAVE and NA>0): | |
| if NA==K//B:selected=True | |
| else:selected=tl.load(AM+(tl.program_id(0)*BM//128)*(K//B)+kb) | |
| if selected: | |
| a=tl.load(AC+MP*(K//2)+ai,mi[:,None]<MP,0);s=tl.load(AS+MP*(K//16)+asi,mi[:,None]<MP,0.) | |
| acc=tl.dot_scaled(a,s,'e2m1',w.T,ws,'e2m1',acc) | |
| if NA>0 and NA<K//B and not INTERLEAVE: | |
| for j in range(NA): | |
| kb=tl.load(AI+(tl.program_id(0)*BM//128)*NA+j) | |
| if SPLIT_K==1 or (kb>=begin_k and kb<end_k): | |
| ai=MP*(K//2)+mi[:,None]*(K//2)+kb*(B//2)+dk[None,:] | |
| asi=MP*(K//16)+mi[:,None]*(K//16)+kb*(B//16)+ds[None,:] | |
| wi=ni[:,None]*(K//2)+kb*(B//2)+dk[None,:];wsi=ni[:,None]*(K//16)+kb*(B//16)+ds[None,:] | |
| a=tl.load(AC+ai,mi[:,None]<MP,0);s=tl.load(AS+asi,mi[:,None]<MP,0.) | |
| w=tl.load(WC+wi);ws=tl.load(WS+wsi) | |
| acc=tl.dot_scaled(a,s,'e2m1',w.T,ws,'e2m1',acc) | |
| if HASW: | |
| if DENSE_W:begin=nb*(K//B);end=(nb+1)*(K//B) | |
| else:begin=tl.load(OFF+nb);end=tl.load(OFF+nb+1) | |
| for slot in range(begin,end): | |
| if DENSE_W:kb=slot-begin | |
| else:kb=tl.load(WI+slot) | |
| if SPLIT_K==1 or (kb>=begin_k and kb<end_k): | |
| storage_slot=tl.load(WLOC+slot) if SHARED_BANK else slot | |
| ai=mi[:,None]*(K//2)+kb*(B//2)+dk[None,:];asi=mi[:,None]*(K//16)+kb*(B//16)+ds[None,:] | |
| a=tl.load(AC+ai,mi[:,None]<MP,0);s=tl.load(AS+asi,mi[:,None]<MP,0.) | |
| w=tl.load(WR+storage_slot*128*(B//2)+tl.arange(0,128)[:,None]*(B//2)+dk[None,:]) | |
| ws=tl.load(WRS+storage_slot*128*(B//16)+tl.arange(0,128)[:,None]*(B//16)+ds[None,:]) | |
| acc=tl.dot_scaled(a,s,'e2m1',w.T,ws,'e2m1',acc) | |
| if CROSS and NA>0: | |
| selected=tl.load(AM+(tl.program_id(0)*BM//128)*(K//B)+kb) | |
| if selected: | |
| a=tl.load(AC+MP*(K//2)+ai,mi[:,None]<MP,0);s=tl.load(AS+MP*(K//16)+asi,mi[:,None]<MP,0.) | |
| acc=tl.dot_scaled(a,s,'e2m1',w.T,ws,'e2m1',acc) | |
| if SPLIT_K==1: | |
| tl.store(Y+mi[:,None]*N+ni[None,:],acc*(tl.load(WG)*tl.load(AG)),mi[:,None]<M) | |
| else: | |
| tl.store(Y+part*M*N+mi[:,None]*N+ni[None,:],acc,mi[:,None]<M) | |
| def reduce_split_k(PARTS,Y,WG,AG,SIZE:tl.constexpr,SPLITS:tl.constexpr,BLOCK:tl.constexpr): | |
| idx=tl.program_id(0)*BLOCK+tl.arange(0,BLOCK) | |
| acc=tl.full((BLOCK,),0,tl.float32) | |
| for part in tl.static_range(SPLITS): | |
| acc+=tl.load(PARTS+part*SIZE+idx,idx<SIZE,0.) | |
| tl.store(Y+idx,acc*(tl.load(WG)*tl.load(AG)),idx<SIZE) | |
| class Native: | |
| def __init__(self,wc,ws,wg,d,am,wm,b=128,cross=False,bm=32,interleave=False,num_warps=4,num_stages=1,split_k=1): | |
| if b not in (64,128,256) or bm not in (32,64,128): | |
| raise ValueError("validated experimental tiles: K64/128/256, M32/64/128") | |
| if not wc.is_cuda or torch.cuda.get_device_capability(wc.device)!=(12,0): | |
| raise ValueError("native experiment requires SM120") | |
| if am.dtype!=torch.bool or wm.dtype!=torch.bool: | |
| raise ValueError("repair supports must be boolean") | |
| if num_warps not in (4,8) or num_stages not in (1,2,3,4): | |
| raise ValueError("supported launch settings: warps 4/8, stages 1/2/3/4") | |
| self.num_warps=num_warps;self.num_stages=num_stages | |
| if type(split_k) is not int or split_k not in (1,2,4,8): | |
| raise ValueError("split_k must be 1,2,4,8") | |
| self.split_k=split_k | |
| self.b=b;self.bm=bm;self.interleave=interleave;self.cross=cross;self.d=d;self.wc=wc[0].contiguous().clone();self.ws=ws[0].contiguous().clone();self.wg=wg | |
| self.n,self.k=self.wc.shape[0],self.wc.shape[1]*2 | |
| # All host inspection and sparse weight gathering occurs offline. | |
| self.am=am.reshape(-1,self.k//b,b).any(-1).contiguous() | |
| counts=self.am.sum(-1) | |
| assert bool((counts==counts[0]).all()) | |
| self.na=int(counts[0]);self.ai=self.am.nonzero()[:,1].reshape(self.am.shape[0],self.na).int().contiguous() | |
| wm=wm.reshape(self.n//128,128,self.k//b,b).any((1,3));pairs=wm.nonzero();self.hasw=bool(pairs.numel()) | |
| self.dense_w=bool(wm.all()) | |
| self.off=torch.cat([torch.zeros(1,device=wc.device,dtype=torch.int32),wm.sum(-1).int().cumsum(0).int()]) | |
| self.wi=pairs[:,1].int().contiguous() | |
| self.wloc=self.wi;self.shared_bank=False | |
| tiles=wc[1].reshape(self.n//128,128,self.k//b,b//2).permute(0,2,1,3) | |
| scales=ws[1].view(torch.uint8).reshape(self.n//128,128,self.k//b,b//16).permute(0,2,1,3) | |
| self.wr=tiles[pairs[:,0],pairs[:,1]].contiguous() | |
| self.wrs=scales[pairs[:,0],pairs[:,1]].contiguous().view(torch.float8_e4m3fn) | |
| self.union=self.am.repeat_interleave(b,1).reshape(self.am.shape[0],self.k//128,128).any(-1).contiguous() | |
| def pack(self,x): | |
| return _pack(x,self.d,self.na>0,self.na==self.k//self.b,self.union) | |
| def run_packed(self,x,packed): | |
| ac,sc,ag,mp=packed;m=x.shape[0] | |
| out=torch.empty((m,self.n),device=x.device,dtype=x.dtype) | |
| parts=out if self.split_k==1 else torch.empty((self.split_k,m,self.n),device=x.device,dtype=torch.float32) | |
| gemm[(triton.cdiv(m,self.bm),self.n//128,self.split_k)](ac,sc,self.wc,self.ws,self.wg,ag,self.ai,self.am,self.wr,self.wrs,self.wi,self.off,parts,m,self.n,self.k,mp,self.na,self.hasw,self.cross,self.b,self.bm,self.interleave,self.split_k,self.dense_w,WLOC=self.wloc,SHARED_BANK=self.shared_bank,num_warps=self.num_warps,num_stages=self.num_stages) | |
| if self.split_k>1: | |
| reduce_split_k[(triton.cdiv(m*self.n,512),)](parts,out,self.wg,ag,m*self.n,self.split_k,512,num_warps=4) | |
| return out | |
| def __call__(self,x): | |
| if x.ndim!=2 or x.shape[1]!=self.k or x.shape[0]<=0 or triton.cdiv(x.shape[0],128)!=self.am.shape[0]: | |
| raise ValueError("activation shape must match the prepared M128 support and K") | |
| if x.dtype!=torch.bfloat16 or x.device!=self.wc.device or not x.is_contiguous(): | |
| raise ValueError("expected contiguous BF16 activation on the packed-weight device") | |
| return self.run_packed(x,self.pack(x)) | |
| def residual_weight_bytes(self):return sum(t.numel()*t.element_size() for t in [self.wr,self.wrs,self.wi,self.off]) if self.hasw else 0 | |