Release frozen ThunderWorld FP4 1.5x and 2x weights, calibration provenance and loader (part 2)
18e8cca verified | """Online Q/K/V packing plus native FP4 tile-streaming attention. | |
| Packing is performed afresh for every input. No weights, old activations, | |
| scores or probability matrices are cached. P is quantized only on chip. | |
| """ | |
| import triton | |
| import triton.language as tl | |
| from ._quantize import quantize, quantize_p | |
| def pack_qk(X, CODES, SCALES, FACTORS, M: tl.constexpr, H: tl.constexpr, | |
| SX: tl.constexpr, SH: tl.constexpr, MP: tl.constexpr, REPAIR: tl.constexpr): | |
| r=tl.program_id(0)*16+tl.arange(0,16) | |
| h=tl.program_id(1) | |
| d=tl.arange(0,128) | |
| x=tl.load(X+r[:,None]*SX+h*SH+d[None,:],r[:,None]<M,0).to(tl.float32) | |
| p,s,rp,rs,f=quantize(x,16,128,REPAIR) | |
| idx=(h*MP+r[:,None])*64+tl.arange(0,64)[None,:] | |
| si=(h*MP+r[:,None])*8+tl.arange(0,8)[None,:] | |
| tl.store(CODES+idx,p,r[:,None]<MP);tl.store(SCALES+si,s,r[:,None]<MP) | |
| if REPAIR: | |
| tl.store(CODES+H*MP*64+idx,rp,r[:,None]<MP) | |
| tl.store(SCALES+H*MP*8+si,rs,r[:,None]<MP) | |
| tl.store(FACTORS+h*MP+r,f,r<MP) | |
| def pack_v(X, CODES, SCALES, FACTORS, ENERGY, N:tl.constexpr,H:tl.constexpr, | |
| SX:tl.constexpr,SH:tl.constexpr,NT:tl.constexpr,REPAIR:tl.constexpr,NEED_ENERGY:tl.constexpr,TRANSPOSE:tl.constexpr=False): | |
| di=tl.program_id(0)*16+tl.arange(0,16) | |
| nt=tl.program_id(1);h=tl.program_id(2) | |
| ni=nt*128+tl.arange(0,128) | |
| x=tl.load(X+ni[None,:]*SX+h*SH+di[:,None],ni[None,:]<N,0).to(tl.float32) | |
| if NEED_ENERGY:tl.store(ENERGY+(h*NT+nt)*8+tl.program_id(0),tl.sum(tl.sum(x*x,1),0)/16384.) | |
| p,s,rp,rs,f=quantize(x,16,128,REPAIR) | |
| row=(h*NT+nt)*128+di | |
| idx=row[:,None]*64+tl.arange(0,64)[None,:] | |
| si=row[:,None]*8+tl.arange(0,8)[None,:] | |
| if TRANSPOSE: | |
| idx=(h*NT+nt)*128*64+tl.arange(0,64)[None,:]*128+di[:,None] | |
| si=(h*NT+nt)*128*8+tl.arange(0,8)[None,:]*128+di[:,None] | |
| tl.store(CODES+idx,p);tl.store(SCALES+si,s) | |
| if REPAIR: | |
| tl.store(CODES+H*NT*128*64+idx,rp) | |
| tl.store(SCALES+H*NT*128*8+si,rs) | |
| tl.store(FACTORS+row,f) | |
| def streaming_attention(QC,QS,QF,KC,KS,KF,VC,VS,VF,FLAGS,O,STATS, | |
| M:tl.constexpr,N:tl.constexpr,HQ:tl.constexpr,HK:tl.constexpr, | |
| MP:tl.constexpr,NP:tl.constexpr, | |
| BM:tl.constexpr,BN:tl.constexpr,REPAIR:tl.constexpr,FULL:tl.constexpr,TERMS:tl.constexpr,SPLITS:tl.constexpr,LAZY_P:tl.constexpr,FAST_P:tl.constexpr=False,TRANSPOSE_V:tl.constexpr=False,RELOAD_Q:tl.constexpr=False,PVFLAGS=None,SEPARATE:tl.constexpr=False,PV_REPAIR:tl.constexpr=False,PV_FULL:tl.constexpr=False,PV_TERMS:tl.constexpr=3): | |
| mi=tl.program_id(0)*BM+tl.arange(0,BM);h=tl.program_id(1);kh=h//(HQ//HK) | |
| dd=tl.arange(0,128);dk=tl.arange(0,64);ds=tl.arange(0,8) | |
| qi=(h*MP+mi[:,None])*64+dk[None,:] | |
| qsi=(h*MP+mi[:,None])*8+ds[None,:] | |
| q0=tl.load(QC+qi,mi[:,None]<MP,0);qs=tl.load(QS+qsi,mi[:,None]<MP,0.0) | |
| qf=tl.load(QF+h*MP+mi,mi<MP,0) | |
| if REPAIR and not RELOAD_Q: | |
| q1=tl.load(QC+HQ*MP*64+qi,mi[:,None]<MP,0) | |
| qr=tl.load(QS+HQ*MP*8+qsi,mi[:,None]<MP,0.0) | |
| acc=tl.full((BM,128),0,tl.float32) | |
| mx=tl.full((BM,),-float('inf'),tl.float32);denom=tl.full((BM,),0,tl.float32) | |
| split=tl.program_id(2) | |
| for start in range(split*tl.cdiv(tl.cdiv(N,BN),SPLITS),tl.minimum((split+1)*tl.cdiv(tl.cdiv(N,BN),SPLITS),tl.cdiv(N,BN))): | |
| ni=start*BN+tl.arange(0,BN) | |
| ki=(kh*NP+ni[:,None])*64+dk[None,:] | |
| ksi=(kh*NP+ni[:,None])*8+ds[None,:] | |
| k0=tl.load(KC+ki);ks=tl.load(KS+ksi);kf=tl.load(KF+kh*NP+ni) | |
| score=tl.dot_scaled(q0,qs,'e2m1',k0.T,ks,'e2m1') | |
| if REPAIR: | |
| if FULL:selected=True | |
| else:selected=tl.load(FLAGS+(h*tl.cdiv(M,128)+tl.program_id(0)*BM//128)*tl.cdiv(N,128)+start*BN//128)>1 | |
| if selected: | |
| if RELOAD_Q: | |
| q1=tl.load(QC+HQ*MP*64+qi,mi[:,None]<MP,0) | |
| qr=tl.load(QS+HQ*MP*8+qsi,mi[:,None]<MP,0.0) | |
| score=tl.dot_scaled(q1,qr,'e2m1',k0.T,ks,'e2m1',score) | |
| if TERMS>=3: | |
| k1=tl.load(KC+HK*NP*64+ki);kr=tl.load(KS+HK*NP*8+ksi) | |
| score=tl.dot_scaled(q0,qs,'e2m1',k1.T,kr,'e2m1',score) | |
| if TERMS==4:score=tl.dot_scaled(q1,qr,'e2m1',k1.T,kr,'e2m1',score) | |
| score=score*qf[:,None]*kf[None,:]*0.08838834764831845 | |
| score=tl.where(ni[None,:]<N,score,-float('inf')) | |
| nm=tl.maximum(mx,tl.max(score,1));alpha=tl.exp(mx-nm) | |
| p=tl.exp(score-nm[:,None]);denom=denom*alpha+tl.sum(p,1) | |
| if SEPARATE: | |
| if PV_FULL:pv_selected=True | |
| elif PV_REPAIR:pv_selected=tl.load(PVFLAGS+(h*tl.cdiv(M,128)+tl.program_id(0)*BM//128)*tl.cdiv(N,128)+start*BN//128)>1 | |
| else:pv_selected=False | |
| elif REPAIR:pv_selected=selected | |
| else:pv_selected=False | |
| if (PV_REPAIR if SEPARATE else REPAIR) and LAZY_P and not (PV_FULL if SEPARATE else FULL): | |
| if pv_selected: | |
| p0,ps,p1,pr,pf=quantize_p(p,BM,BN,True,FAST_P) | |
| else: | |
| p0,ps,p1,pr,pf=quantize_p(p,BM,BN,False,FAST_P) | |
| else: | |
| p0,ps,p1,pr,pf=quantize_p(p,BM,BN,PV_REPAIR if SEPARATE else REPAIR,FAST_P) | |
| vi=((kh*tl.cdiv(N,128)+start*BN//128)*128+dd[:,None])*64+(start*BN%128)//2+tl.arange(0,BN//2)[None,:] | |
| vsi=((kh*tl.cdiv(N,128)+start*BN//128)*128+dd[:,None])*8+(start*BN%128)//16+tl.arange(0,BN//16)[None,:] | |
| if TRANSPOSE_V: | |
| vi=(kh*tl.cdiv(N,128)+start*BN//128)*128*64+((start*BN%128)//2+tl.arange(0,BN//2)[None,:])*128+dd[:,None] | |
| vsi=(kh*tl.cdiv(N,128)+start*BN//128)*128*8+((start*BN%128)//16+tl.arange(0,BN//16)[None,:])*128+dd[:,None] | |
| v0=tl.load(VC+vi);vs=tl.load(VS+vsi) | |
| vf=tl.load(VF+(kh*tl.cdiv(N,128)+start*BN//128)*128+dd) | |
| tile=tl.dot_scaled(p0,ps,'e2m1',v0.T,vs,'e2m1') | |
| if PV_REPAIR if SEPARATE else REPAIR: | |
| if pv_selected: | |
| tile=tl.dot_scaled(p1,pr,'e2m1',v0.T,vs,'e2m1',tile) | |
| if (PV_TERMS if SEPARATE else TERMS)>=3: | |
| v1=tl.load(VC+HK*NP*64+vi);vr=tl.load(VS+HK*NP*8+vsi) | |
| tile=tl.dot_scaled(p0,ps,'e2m1',v1.T,vr,'e2m1',tile) | |
| if (PV_TERMS if SEPARATE else TERMS)==4:tile=tl.dot_scaled(p1,pr,'e2m1',v1.T,vr,'e2m1',tile) | |
| acc=acc*alpha[:,None]+tile*pf[:,None]*vf[None,:];mx=nm | |
| if SPLITS==1: | |
| tl.store(O+(mi[:,None]*HQ+h)*128+dd[None,:],acc/denom[:,None],mi[:,None]<M) | |
| else: | |
| oi=((split*HQ+h)*M+mi[:,None])*128+dd[None,:] | |
| tl.store(O+oi,acc,mi[:,None]<M) | |
| si=((split*HQ+h)*M+mi)*2 | |
| tl.store(STATS+si,mx,mi<M);tl.store(STATS+si+1,denom,mi<M) | |
| def combine(PARTS,STATS,OUT,M:tl.constexpr,H:tl.constexpr,SPLITS:tl.constexpr,BS:tl.constexpr): | |
| r=tl.program_id(0)*4+tl.arange(0,4);h=tl.program_id(1) | |
| # Keep split reduction local to each output element. A cross-warp tl.sum | |
| # lowering in the first prototype produced shared-memory WAW race reports. | |
| top=tl.full((4,),-float('inf'),tl.float32) | |
| for s in tl.static_range(SPLITS): | |
| mx=tl.load(STATS+((s*H+h)*M+r)*2,r<M,-float('inf')) | |
| top=tl.maximum(top,mx) | |
| top=tl.where(r<M,top,0.) | |
| di=tl.arange(0,128) | |
| accum=tl.full((4,128),0,tl.float32);total=tl.full((4,),0,tl.float32) | |
| for s in tl.static_range(SPLITS): | |
| si=((s*H+h)*M+r)*2 | |
| mx=tl.load(STATS+si,r<M,-float('inf'));den=tl.load(STATS+si+1,r<M,0.) | |
| alpha=tl.exp(mx-top);total=total+den*alpha | |
| idx=((s*H+h)*M+r[:,None])*128+di[None,:] | |
| value=tl.load(PARTS+idx,r[:,None]<M,0.) | |
| accum=accum+value*alpha[:,None] | |
| tl.store(OUT+(r[:,None]*H+h)*128+di[None,:],accum/total[:,None],r[:,None]<M) | |
| def pack_inputs(q,k,v,repair=True,need_energy=False,terms=3,transpose_v=False,defer_residual=False): | |
| import torch | |
| m,n,h,hk=q.shape[1],k.shape[1],q.shape[2],k.shape[2] | |
| mp=triton.cdiv(m,16)*16;np=triton.cdiv(n,128)*128;planes=2 if repair else 1 | |
| outputs=[];energy=torch.empty((hk,np//128,8),device=q.device,dtype=torch.float32) if need_energy else q | |
| for kind,x,hh,nn,padded in [('q',q,h,m,mp),('k',k,hk,n,np),('v',v,hk,n,np)]: | |
| local_repair=repair and (kind=='q' or terms>=3) | |
| planes=2 if local_repair else 1 | |
| codes=torch.empty((planes,hh*padded,64),device=q.device,dtype=torch.uint8) | |
| scales=torch.empty((planes,hh*padded,8),device=q.device,dtype=torch.float8_e4m3fn) | |
| factors=torch.empty((hh*padded,),device=q.device,dtype=torch.float32) | |
| if kind=='v': | |
| pack_v[(8,np//128,hh)](x,codes,scales,factors,energy,nn,hh,x.stride(1),x.stride(2),np//128,local_repair and not defer_residual,need_energy,transpose_v,num_warps=4) | |
| else: | |
| pack_qk[(triton.cdiv(padded,16),hh)](x,codes,scales,factors,nn,hh,x.stride(1),x.stride(2),padded,local_repair and not defer_residual,num_warps=4) | |
| outputs.extend([codes,scales,factors]) | |
| return (outputs,mp,np,energy) if need_energy else (outputs,mp,np) | |
| def select_tiles(Q,K,ENERGY,FLAGS,M:tl.constexpr,N:tl.constexpr,HQ:tl.constexpr,HK:tl.constexpr, | |
| Q_S:tl.constexpr,Q_H:tl.constexpr,K_S:tl.constexpr,K_H:tl.constexpr, | |
| COUNT:tl.constexpr,TERMS:tl.constexpr,BT:tl.constexpr): | |
| h=tl.program_id(0);mt=tl.program_id(1);kh=h//(HQ//HK) | |
| # Sixteen MMA rows duplicate the eight fixed query samples. | |
| qi=tl.minimum(mt*128+(tl.arange(0,16)%8)*16,M-1) | |
| sample=tl.arange(0,BT*8);ki=sample//8*128+sample%8*16 | |
| dd=tl.arange(0,128) | |
| q=tl.load(Q+qi[:,None]*Q_S+h*Q_H+dd[None,:]) | |
| k=tl.load(K+tl.minimum(ki[:,None],N-1)*K_S+kh*K_H+dd[None,:]) | |
| sim=tl.dot(q,k.T)*0.08838834764831845 | |
| sim=tl.where(ki[None,:]<N,sim,-float('inf')) | |
| weights=tl.exp(sim-tl.max(sim,1)[:,None]);weights=weights/tl.sum(weights,1)[:,None] | |
| mass=tl.sum(tl.sum(tl.reshape(weights,(16,BT,8)),2),0)/16. | |
| nt=tl.cdiv(N,128);ti=tl.arange(0,BT) | |
| energy=tl.sqrt(tl.sum(tl.load(ENERGY+(kh*nt+ti[:,None])*8+tl.arange(0,8)[None,:],ti[:,None]<nt,0.),1)) | |
| score=tl.where(ti<nt,mass*energy,-float('inf')) | |
| # Rank the small number of key tiles; deterministic lower-index tie break. | |
| rank=tl.sum(((score[None,:]>score[:,None])|((score[None,:]==score[:,None])&(ti[None,:]<ti[:,None]))).to(tl.int32),1) | |
| tl.store(FLAGS+(h*tl.cdiv(M,128)+mt)*nt+ti,tl.where(rank<COUNT,TERMS,1),ti<nt) | |
| def select_tiles_head(Q,K,ENERGY,FLAGS,M:tl.constexpr,N:tl.constexpr,HQ:tl.constexpr,HK:tl.constexpr, | |
| Q_S:tl.constexpr,Q_H:tl.constexpr,K_S:tl.constexpr,K_H:tl.constexpr, | |
| COUNTS,TERMS:tl.constexpr,BT:tl.constexpr,PVFLAGS=None,PVCOUNTS=None,SEPARATE:tl.constexpr=False,PVTERMS:tl.constexpr=3): | |
| h=tl.program_id(0);mt=tl.program_id(1);kh=h//(HQ//HK) | |
| COUNT=tl.load(COUNTS+h) | |
| # Sixteen MMA rows duplicate the eight fixed query samples. | |
| qi=tl.minimum(mt*128+(tl.arange(0,16)%8)*16,M-1) | |
| sample=tl.arange(0,BT*8);ki=sample//8*128+sample%8*16 | |
| dd=tl.arange(0,128) | |
| q=tl.load(Q+qi[:,None]*Q_S+h*Q_H+dd[None,:]) | |
| k=tl.load(K+tl.minimum(ki[:,None],N-1)*K_S+kh*K_H+dd[None,:]) | |
| sim=tl.dot(q,k.T)*0.08838834764831845 | |
| sim=tl.where(ki[None,:]<N,sim,-float('inf')) | |
| weights=tl.exp(sim-tl.max(sim,1)[:,None]);weights=weights/tl.sum(weights,1)[:,None] | |
| mass=tl.sum(tl.sum(tl.reshape(weights,(16,BT,8)),2),0)/16. | |
| nt=tl.cdiv(N,128);ti=tl.arange(0,BT) | |
| energy=tl.sqrt(tl.sum(tl.load(ENERGY+(kh*nt+ti[:,None])*8+tl.arange(0,8)[None,:],ti[:,None]<nt,0.),1)) | |
| score=tl.where(ti<nt,mass*energy,-float('inf')) | |
| # Rank the small number of key tiles; deterministic lower-index tie break. | |
| rank=tl.sum(((score[None,:]>score[:,None])|((score[None,:]==score[:,None])&(ti[None,:]<ti[:,None]))).to(tl.int32),1) | |
| tl.store(FLAGS+(h*tl.cdiv(M,128)+mt)*nt+ti,tl.where(rank<COUNT,TERMS,1),ti<nt) | |
| if SEPARATE: | |
| pv_count=tl.load(PVCOUNTS+h) | |
| tl.store(PVFLAGS+(h*tl.cdiv(M,128)+mt)*nt+ti,tl.where(rank<pv_count,PVTERMS,1),ti<nt) | |