# E22i: dictionary-scale 3-bit KV cache, real-model KL, simulated quantize-dequantize (HF transformers, Qwen3.5-4B NF4 weights).
# Quantization is applied to post-RoPE K and to V as they enter the cache (cache.update), only on the 8 full-attention layers.
# Reference = same NF4 model with the unquantized (bf16) cache. Dictionaries fit on held-out calib text (bench/calib_v3.txt), frozen.
import os, sys, json, time, math, torch, numpy as np
import torch.nn.functional as Fn
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
SMOKE = len(sys.argv) > 1 and sys.argv[1] == 'smoke'
MP = os.environ.get('E22I_MODEL', 'Qwen/Qwen3.5-4B')
ROOT = os.path.dirname(os.path.abspath(__file__)); OUT = os.path.join(ROOT, 'smoke' if SMOKE else 'run'); os.makedirs(OUT, exist_ok=True)
BENCH = os.path.join(ROOT, '..', '..', 'bench')
CTX = 4096 if SMOKE else 65536; NSC = 512 if SMOKE else 4096; PCH = 1024; SCH = 256
SPANS = [72000] if SMOKE else [72000,108000,144000,180000,216000,252000]
ARMS = ['ref','q4_g32','q3_g32','q3_g32_rot','d5_rot','d6_rot','d5_norot','d6_norot']
dev = 'cuda:0'; torch.cuda.set_per_process_memory_fraction(0.92)
def log(m):
    s = f"{time.strftime('%Y-%m-%d %H:%M:%S')}  {m}"; print(s, flush=True); open(os.path.join(OUT,'run.log'),'a',encoding='utf8').write(s+'\n')
def hadamard(n):
    H = torch.ones(1,1)
    while H.shape[0] < n: H = torch.cat([torch.cat([H,H],1), torch.cat([H,-H],1)],0)
    return H/math.sqrt(n)
HM = hadamard(256).to(dev)
def f16(x): return x.half().float()
def absmax_q(s, lo, hi):  # s (N,32): ggml-style signed absmax scale d = m/lo, fp16 scale
    i = s.abs().argmax(1, keepdim=True); m = s.gather(1, i); d = f16(m/lo); d = torch.where(d == 0, torch.ones_like(d), d)
    return torch.clamp(torch.round(s/d), lo, hi)*d
def dict_q(s, D):
    best = torch.full((s.shape[0],), float('inf'), device=s.device); out = torch.empty_like(s)
    for dv in D.tolist():
        c = torch.clamp(torch.round(s/dv), -4, 3)*dv; e = ((s-c)**2).sum(1); m = e < best; best = torch.where(m, e, best); out[m] = c[m]
    return out
DICTS = {}   # (bits, rot, layer, kind) -> tensor
STATE = {'arm': 'ref', 'cap': None}
FULL = None
def quant(x, layer, kind):
    arm = STATE['arm']
    if arm == 'ref' or layer not in FULL: return x
    if arm == 'capture': STATE['cap'].setdefault((layer,kind), []).append(x.detach().float().cpu()); return x
    dt = x.dtype; y = x.float(); rot = arm.endswith('_rot')
    if rot: y = y @ HM
    s = y.reshape(-1, 32)
    if arm == 'q4_g32': q = absmax_q(s, -8, 7)
    elif arm.startswith('q3_g32'): q = absmax_q(s, -4, 3)
    else:
        b = int(arm[1]); q = dict_q(s, DICTS[(b, rot, layer, kind)])
    y = q.reshape(y.shape)
    if rot: y = y @ HM.T
    return y.to(dt)
def kmeans1d(v, k, it=100):
    c = torch.quantile(v, torch.linspace(0.5/k, 1-0.5/k, k, dtype=v.dtype))
    for _ in range(it):
        a = (v[:,None]-c[None]).abs().argmin(1)
        for j in range(k):
            m = a == j
            if m.any(): c[j] = v[m].mean()
    return c
def fit_dict(x, n):  # mirrors the round 6 CPU test fit_dict (per layer, per K/V, all heads)
    s = x.reshape(-1,32).cpu(); i = s.abs().argmax(1, keepdim=True); m = s.gather(1,i).squeeze(1); d = f16(m/-4.); d = d[d != 0]
    lv = torch.log(d.abs()); g = torch.Generator().manual_seed(0)
    if len(lv) > 100000: lv = lv[torch.randperm(len(lv), generator=g)[:100000]]
    mag = f16(torch.exp(torch.sort(kmeans1d(lv.double(), n//2)).values.float()))
    return torch.cat([-mag.flip(0), mag]).to(dev)
def main():
    global FULL
    log(f"=== E22i start smoke={SMOKE} model={MP} CTX={CTX} NSC={NSC} spans={SPANS} arms={ARMS}")
    tok = AutoTokenizer.from_pretrained(MP)
    bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type='nf4', bnb_4bit_compute_dtype=torch.bfloat16)
    model = AutoModelForCausalLM.from_pretrained(MP, quantization_config=bnb, dtype=torch.bfloat16, device_map=dev, attn_implementation='sdpa').eval()
    tc = getattr(model.config, 'text_config', model.config)
    FULL = set(i for i, t in enumerate(tc.layer_types) if t == 'full_attention')
    log(f"loaded; full-attn layers {sorted(FULL)}; vram {torch.cuda.memory_allocated()/2**30:.2f} GiB")
    with torch.no_grad():
        o = model(input_ids=torch.tensor([[1,2,3,4]], device=dev), use_cache=True, logits_to_keep=1)
    C = type(o.past_key_values); log(f"cache class {C.__module__}.{C.__name__}")
    orig = C.update
    def upd(self, k, v, layer_idx, *a, **kw):
        return orig(self, quant(k, layer_idx, 'K'), quant(v, layer_idx, 'V'), layer_idx, *a, **kw)
    C.update = upd
    def prefill(ids, keep_last=False):
        cache = None; out = None
        for s in range(0, ids.shape[1], PCH):
            out = model(input_ids=ids[:, s:s+PCH], past_key_values=cache, use_cache=True, logits_to_keep=1); cache = out.past_key_values
        return cache, out
    # calibration: held-out text, never the scored spans
    ctoks = tok(open(os.path.join(BENCH,'calib_v3.txt'), encoding='utf8', errors='ignore').read(), return_tensors='pt').input_ids[:, :8192].to(dev)
    STATE['arm'] = 'capture'; STATE['cap'] = {}
    with torch.no_grad(): prefill(ctoks)
    for (l,k), xs in STATE['cap'].items():
        x = torch.cat(xs, dim=2).to(dev)  # (1,H,T,256) post-RoPE K / V
        for b, n in ((5,32),(6,64)):
            DICTS[(b, False, l, k)] = fit_dict(x, n); DICTS[(b, True, l, k)] = fit_dict(x @ HM, n)
        log(f"calib layer {l} {k}: T={x.shape[2]} d6_rot range {DICTS[(6,True,l,k)].abs().min().item():.4g}..{DICTS[(6,True,l,k)].abs().max().item():.4g}")
    torch.save({str(k): v.cpu() for k, v in DICTS.items()}, os.path.join(OUT, 'dicts.pt')); STATE['cap'] = None; del ctoks
    wtoks = tok(open(os.path.join(BENCH,'wiki_test.txt'), encoding='utf8', errors='ignore').read(), return_tensors='pt').input_ids
    log(f"wiki_test HF tokens {wtoks.shape[1]}")
    n = CTX + NSC + 1
    res = open(os.path.join(OUT, 'results.jsonl'), 'a')
    for E in SPANS:
        assert E <= wtoks.shape[1]
        w = wtoks[:, E-n:E].to(dev); P = CTX
        ref = []
        for arm in ARMS:
            STATE['arm'] = arm; t0 = time.time(); torch.cuda.reset_peak_memory_stats()
            with torch.no_grad():
                cache = prefill(w[:, :P])[0]; kl = []; agree = 0; nll = 0.
                for j, s in enumerate(range(P, P+NSC, SCH)):
                    out = model(input_ids=w[:, s:s+SCH], past_key_values=cache, use_cache=True); cache = out.past_key_values
                    lg = out.logits[0].float(); tgt = w[0, s+1:s+1+SCH]
                    lp = Fn.log_softmax(lg, -1); nll += -lp.gather(1, tgt[:,None]).sum().item()
                    if arm == 'ref': ref.append(lg.cpu()); kl.append(torch.zeros(lg.shape[0])); agree += lg.shape[0]
                    else:
                        lr = Fn.log_softmax(ref[j].to(dev), -1); kl.append((lr.exp()*(lr-lp)).sum(-1).cpu()); agree += (lr.argmax(-1) == lp.argmax(-1)).sum().item()
                    del out, lg, lp
            kl = torch.cat(kl); r = dict(span=E, ctx=CTX, scored=NSC, arm=arm, mean_kl=kl.mean().item(), p99_kl=kl.quantile(0.99).item(), top1=agree/NSC,
                     ppl=math.exp(nll/NSC), wall_s=round(time.time()-t0,1), peak_gib=round(torch.cuda.max_memory_allocated()/2**30,2))
            res.write(json.dumps(r)+'\n'); res.flush(); log(f"result {json.dumps(r)}")
            del cache; import gc; gc.collect(); torch.cuda.empty_cache()
        if SMOKE:  # chunked-prefill consistency: one-shot forward of the whole window vs chunked ref logits
            STATE['arm'] = 'ref'
            with torch.no_grad(): lg = model(input_ids=w[:, :P+SCH], use_cache=False, logits_to_keep=SCH).logits[0].float()
            d = (Fn.log_softmax(lg,-1) - Fn.log_softmax(ref[0].to(dev),-1)).abs().max().item()
            lr = Fn.log_softmax(ref[0].to(dev),-1); kk = (lr.exp()*(lr-Fn.log_softmax(lg,-1))).sum(-1).mean().item()
            log(f"smoke consistency: one-shot vs chunked first scored chunk max|dlogp|={d:.4g} meanKL={kk:.4g}")
        del ref
    log("=== E22iDONE ===")
if __name__ == '__main__': main()