"""Offline activation-weighted block-scale search for NVFP4 weights."""
import torch
import triton
import triton.language as tl
@triton.jit
def _round(v):
 a=tl.abs(v)
 code=(a>.25).to(tl.int32)+(a>=.75)+(a>1.25)+(a>=1.75)+(a>2.5)+(a>=3.5)+(a>5.)
 val=tl.where(code<4,code*.5,tl.where(code==4,2.,tl.where(code==5,3.,tl.where(code==6,4.,6.))))
 return code|tl.where(v<0,8,0),val*tl.where(v<0,-1.,1.)
@triton.jit
def _quant(W,H,G,Q,S,COUNT:tl.constexpr,K:tl.constexpr,SEARCH:tl.constexpr,B:tl.constexpr):
 b=tl.program_id(0)*B+tl.arange(0,B)
 i=tl.arange(0,16)
 w=tl.load(W+b[:,None]*16+i[None,:],b[:,None]<COUNT,0).to(tl.float32)
 h=tl.load(H+(b[:,None]*16+i[None,:])%K).to(tl.float32)
 g=tl.load(G)
 amax=tl.max(tl.abs(w),1)*g/6
 best=tl.full((B,),float('inf'),tl.float32)
 selected=tl.full((B,),1.,tl.float32)
 for j in range(SEARCH):
  factor=1. if SEARCH==1 else .7+.05*j
  s=tl.minimum(448.,tl.maximum(amax*factor,.001953125)).to(tl.float8e4nv).to(tl.float32)
  code,dq=_round(w*g/s[:,None])
  delta=w-dq*s[:,None]/g
  err=tl.sum(delta*delta*h,1)
  selected=tl.where(err<best,s,selected);best=tl.minimum(err,best)
 code,_=_round(w*g/selected[:,None])
 halves=tl.reshape(code,(B,8,2))
 packed=tl.sum(halves << (tl.arange(0,2)[None,None,:]*4),2)
 tl.store(Q+b[:,None]*8+tl.arange(0,8)[None,:],packed,b[:,None]<COUNT)
 tl.store(S+b,selected,b<COUNT)
def quantize_weight(w,global_scale,hdiag,search=True):
 n,k=w.shape
 assert n%128==0 and k%64==0
 q=torch.empty((n,k//2),device=w.device,dtype=torch.uint8)
 sf=torch.empty((n,k//16),device=w.device,dtype=torch.float8_e4m3fn)
 _quant[(triton.cdiv(n*k//16,64),)](w.contiguous(),hdiag.contiguous(),global_scale,q,sf,n*k//16,k,15 if search else 1,64,num_warps=4)
 sw=sf.view(torch.uint8).reshape(n//128,4,32,k//64,4).permute(0,3,2,1,4).contiguous().flatten()
 return q,sw