# E21w: 3-bit fixed-width KV quant variants on real Qwen3.5-4B f16 KV (E21t dumps, first 6000 tokens/layer).
# CPU only, deterministic (seed 0). Prior art: KVQuant, KIVI, QServe SmoothAttention, KVTuner, QuaRot.
import numpy as np, json, hashlib, itertools
np.random.seed(0); RNG=np.random.default_rng(0)
DOCS=['wikiA','wikiB','mathA']; NT=6000; L=8; H=4; HD=256; G=32
NQ=128          # proxy queries per kv head
QPOS=np.sort(RNG.choice(np.arange(3000,NT),NQ,replace=False))
def load(doc,kind):
    return np.fromfile(f'data/{doc}_{kind}.bin',dtype=np.float16).reshape(L,NT,H,HD).astype(np.float32)
def had(n):
    M=np.array([[1.]])
    while M.shape[0]<n: M=np.block([[M,M],[M,-M]])
    return (M/np.sqrt(n)).astype(np.float32)
HM=had(HD)
def q_abs(x,b,g=G):
    # q4_0-style symmetric: fp16 scale per group of g along channel; levels -2^(b-1)..2^(b-1)-1
    s=x.reshape(-1,g); i=np.argmax(np.abs(s),1); m=s[np.arange(len(s)),i][:,None]
    d=(m/-(2**(b-1))).astype(np.float16).astype(np.float32); d[d==0]=1e-12
    q=np.clip(np.round(s/d),-(2**(b-1)),2**(b-1)-1)*d
    return q.reshape(x.shape)
def bits(b,g=G): return b+16/g
# x: (NT,H,HD) one layer
def scheme(x,name,b,cs=None,g=G):
    if name=='A': return q_abs(x,b,g)
    if name=='B': return q_abs(x@HM,b,g)@HM.T
    if name=='C': return q_abs(x/cs,b,g)*cs
    if name=='D': return (q_abs((x/cs)@HM,b,g)@HM.T)*cs
def relmse(xq,x): return float(((xq-x)**2).sum()/(x**2).sum())
def calib(x,alpha):  # x: (L,NT,H,HD) of calibration text -> per-channel scales (L,H,HD)
    a=np.abs(x).max(1)+1e-6; s=a**alpha; return s/np.exp(np.log(s).mean(-1,keepdims=True))
def attn(K,V,Qs):
    # K,V (NT,H,HD); Qs proxy queries (H,NQ,HD) at QPOS; causal over strictly earlier keys
    out=[];P=[]
    for h in range(H):
        S=(Qs[h]@K[:,h].T)/np.sqrt(HD)
        mask=np.arange(NT)[None,:]>=QPOS[:,None]; S[mask]=-np.inf
        S-=S.max(1,keepdims=True); p=np.exp(S); p/=p.sum(1,keepdims=True)
        out.append(p@V[:,h]); P.append(p)
    return np.stack(out),np.stack(P)
def aerr(Kq,Vq,K,V,o,p):
    oq,pq=attn(Kq,Vq,K[QPOS].transpose(1,0,2))
    e=np.linalg.norm(oq-o,axis=-1)/np.linalg.norm(o,axis=-1)   # (H,NQ)
    kl=np.where(p>0,p*(np.log(np.maximum(p,1e-30))-np.log(np.maximum(pq,1e-30))),0).sum(-1)
    return e.mean(1),kl.mean(1)
# --- calibration on wikiA (alpha chosen on wikiA only, by 3-bit K err after D) ---
KA=load('wikiA','K'); VA=load('wikiA','V')
alphas=[0.5,0.75,1.0]; best={}
for kind,X in (('K',KA),('V',VA)):
    sc={a:calib(X,a) for a in alphas}
    errs={a:np.mean([relmse(scheme(X[l],'C',3,sc[a][l]),X[l]) for l in range(L)]) for a in alphas}
    a=min(errs,key=errs.get); best[kind]=dict(alpha=a,scales=sc[a],calib_err=errs)
del KA,VA
CS={k:best[k]['scales'] for k in 'KV'}
VARS=['A','B','C','D']; R={'config':dict(NT=NT,G=G,NQ=NQ,qpos_min=int(QPOS[0]),
     alpha_K=best['K']['alpha'],alpha_V=best['V']['alpha'],calib_text='wikiA',
     query='PROXY: post-RoPE keys at 128 random positions in [3000,6000) used as queries, 1 per KV head (no matched Q captured)'),
     'calib_err':{k:{str(a):float(v) for a,v in best[k]['calib_err'].items()} for k in 'KV'},'units':[],'alloc_table':{}}
OPTS=[(2,32),(3,32),(4,32),(2,128),(3,128),(4,128)]
for doc in DOCS:
    K=load(doc,'K'); V=load(doc,'V')
    for l in range(L):
        o,p=attn(K[l],V[l],K[l][QPOS].transpose(1,0,2))
        cfgs=[(v,3) for v in VARS]+[('B',4),('A',4)]
        for v,b in cfgs:
            Kq=scheme(K[l],v,b,CS['K'][l]); Vq=scheme(V[l],v,b,CS['V'][l])
            ea,kl=aerr(Kq,Vq,K[l],V[l],o,p)
            ek=[relmse(Kq[:,h],K[l][:,h]) for h in range(H)]; ev=[relmse(Vq[:,h],V[l][:,h]) for h in range(H)]
            # K-only quant (V exact) attention error to isolate K effect
            eak,klk=aerr(Kq,V[l],K[l],V[l],o,p)
            for h in range(H):
                R['units'].append(dict(doc=doc,layer=l,head=h,var=v,bits=b,bpv=bits(b),K_err=ek[h],V_err=ev[h],
                    attn_err=float(ea[h]),KL=float(kl[h]),attn_err_Konly=float(eak[h]),KL_Konly=float(klk[h])))
        # allocation table: per layer, K opt x V opt, using schemes B (rotation) and D (scale+rotation)
        for v in ('B','D'):
            Kq={o_:scheme(K[l],v,o_[0],CS['K'][l],o_[1]) for o_ in OPTS}
            Vq={o_:scheme(V[l],v,o_[0],CS['V'][l],o_[1]) for o_ in OPTS}
            for ok,ov in itertools.product(OPTS,OPTS):
                ea,kl=aerr(Kq[ok],Vq[ov],K[l],V[l],o,p)
                R['alloc_table'][f'{doc}|{v}|{l}|{ok[0]},{ok[1]}|{ov[0]},{ov[1]}']=[float(ea.mean()),float(kl.mean())]
        print(doc,l,flush=True)
json.dump(R,open('results_raw.json','w'),indent=0,sort_keys=True)