Release frozen ThunderWorld FP4 1.5x and 2x weights, calibration provenance and loader (part 2)
18e8cca verified | """Portable native FP4 FFN/projection kernels; weights are already packed.""" | |
| import triton | |
| import triton.language as tl | |
| from ._quantize import quantize_linear | |
| def scaled_amax(X, D, PARTIAL, M:tl.constexpr, K:tl.constexpr, B:tl.constexpr): | |
| idx=tl.program_id(0)*B+tl.arange(0,B) | |
| x=tl.load(X+idx,idx<M*K,0).to(tl.float32) | |
| d=tl.load(D+idx%K) | |
| x=(x*d).to(tl.bfloat16).to(tl.float32) | |
| tl.store(PARTIAL+tl.program_id(0),tl.max(tl.abs(x),0)) | |
| def finish_amax(PARTIAL, FACTOR, COUNT:tl.constexpr, B:tl.constexpr): | |
| idx=tl.arange(0,B);value=tl.max(tl.load(PARTIAL+idx,idx<COUNT,0),0) | |
| tl.store(FACTOR,tl.where(value>0,value/2688.,1.)) | |
| def activation_rank(X,D,RANK,M:tl.constexpr,K:tl.constexpr,BT:tl.constexpr): | |
| mb=tl.program_id(0);kt=tl.arange(0,BT) | |
| # Eight sampled rows/columns within each 128x128 activation block. | |
| rows=tl.minimum(mb*128+tl.arange(0,8)*16,M-1) | |
| cols=kt[:,None]*128+tl.arange(0,8)[None,:]*16 | |
| x=tl.load(X+rows[None,:,None]*K+cols[:,None,:],cols[:,None,:]<K,0).to(tl.float32) | |
| d=tl.load(D+cols,cols<K,0) | |
| x=(x*d[:,None,:]).to(tl.bfloat16).to(tl.float32) | |
| energy=tl.sum(tl.sum(x*x,2),1) | |
| energy=tl.where(kt<K//128,energy,-float('inf')) | |
| rank=tl.sum(((energy[None,:]>energy[:,None])|((energy[None,:]==energy[:,None])&(kt[None,:]<kt[:,None]))).to(tl.int32),1) | |
| tl.store(RANK+mb*(K//128)+kt,rank,kt<K//128) | |
| def pack_activation(X,D,FACTOR,UNION,CODES,SCALES,M:tl.constexpr,K:tl.constexpr,MP:tl.constexpr, | |
| REPAIR:tl.constexpr,FULL:tl.constexpr): | |
| mi=tl.program_id(0)*16+tl.arange(0,16);ki=tl.program_id(1)*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) | |
| factor=tl.load(FACTOR) | |
| if REPAIR and not FULL: | |
| active=tl.load(UNION+(tl.program_id(0)*16//128)*(K//128)+tl.program_id(1)) | |
| if active:p,s,rp,rs,_=quantize_linear(x,factor,16,128,True) | |
| else: | |
| p,s,rp,rs,_=quantize_linear(x,factor,16,128,False) | |
| rp=tl.full((16,64),0,tl.uint8);rs=tl.full((16,8),0,tl.float8e4nv) | |
| else:p,s,rp,rs,_=quantize_linear(x,factor,16,128,REPAIR) | |
| pi=mi[:,None]*(K//2)+tl.program_id(1)*64+tl.arange(0,64)[None,:] | |
| si=mi[:,None]*(K//16)+tl.program_id(1)*8+tl.arange(0,8)[None,:] | |
| tl.store(CODES+pi,p,mi[:,None]<MP);tl.store(SCALES+si,s,mi[:,None]<MP) | |
| if REPAIR: | |
| tl.store(CODES+MP*(K//2)+pi,rp,mi[:,None]<MP) | |
| tl.store(SCALES+MP*(K//16)+si,rs,mi[:,None]<MP) | |
| def linear_gemm(AC,AS,WC,WS,AF,WG,BIAS,MASK,Y,M:tl.constexpr,N:tl.constexpr,K:tl.constexpr, | |
| MP:tl.constexpr,BIAS_ON:tl.constexpr,REPAIR:tl.constexpr, | |
| FULL:tl.constexpr,BM:tl.constexpr): | |
| mi=tl.program_id(0)*BM+tl.arange(0,BM);ni=tl.program_id(1)*128+tl.arange(0,128) | |
| dk=tl.arange(0,64);ds=tl.arange(0,8);acc=tl.full((BM,128),0,tl.float32) | |
| for kt in range(K//128): | |
| ai=mi[:,None]*(K//2)+kt*64+dk[None,:] | |
| asi=mi[:,None]*(K//16)+kt*8+ds[None,:] | |
| wi=ni[:,None]*(K//2)+kt*64+dk[None,:] | |
| wsi=ni[:,None]*(K//16)+kt*8+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 REPAIR: | |
| if FULL:selected=True | |
| else:selected=tl.load(MASK+((tl.program_id(0)*BM//128)*(N//128)+tl.program_id(1))*(K//128)+kt) | |
| if selected: | |
| r=tl.load(AC+MP*(K//2)+ai,mi[:,None]<MP,0) | |
| rs=tl.load(AS+MP*(K//16)+asi,mi[:,None]<MP,0.) | |
| acc=tl.dot_scaled(r,rs,'e2m1',w.T,ws,'e2m1',acc) | |
| y=acc*(tl.load(AF)*tl.load(WG)) | |
| if BIAS_ON:y=y+tl.load(BIAS+ni)[None,:].to(tl.float32) | |
| tl.store(Y+mi[:,None]*N+ni[None,:],y,mi[:,None]<M) | |