# E22o: headroom checks for 2-bit-class V (supporting measurement, no novelty claim).
# Test 1 attention-weighted V bits (1a future-attention upper bound, 1b causal window+cumulative attention).
# Test 2 W_o: 2a V error before/after gate and o_proj; 2b W_o right-singular-basis bit allocation vs Hadamard control.
# Test 3 key-cohort error ledger: 3a premise (within-cohort attention correlation vs random), 3b ledger encoder arms.
# Test 4 per-layer gain correction for 2-bit V (variance match, calibration grid).
# Harness lineage: e22i..e22n (HF Qwen3.5-4B NF4, simulated quantize-dequantize at cache insert on the 8 full-attention layers,
# 6 spans of wiki_test, 64K context, 4,096 scored tokens, KL vs the unquantized cache). K = rotated 4-bit G128 in all quantized-V arms.
import os, sys, json, time, math, gc, torch
import torch.nn.functional as Fn
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
import bitsandbytes.functional as BF
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]
LC = 2048 if SMOKE else 8192; NCAL = 256 if SMOKE else 512
GRID = [0.9, 1.0, 1.3] if SMOKE else [0.8,0.85,0.9,0.95,1.0,1.1,1.2,1.3,1.4,1.5,1.6]  # widened below 0.9: smoke showed variance-match gains < 1
NQ = 256 if SMOKE else 1024; WIN = 512; NSEED = 5
FR = {'225': 0.0605, '250': 0.1855}   # 2.125 + 2f + 1/256 flag = 2.25 / 2.5 bits
ARMS = ['ref','k4_v4','k4rg128','A2','A4','t1a_225','t1a_250','t1b_225','t1b_250','t2_wo16','t2_wo80','t2_h16','t2_h80',
        't3_plain','t3_led_sign','t3_led_km'] + [f't3_led_rand{s}' for s in range(NSEED)] + ['t4_gvar','t4_ggrid','t4_led_km_ggrid']
BITS_V = {'k4_v4':4.5,'k4rg128':16,'A2':2.125,'A4':4.125,'t2_wo16':2.25,'t2_wo80':2.5,'t2_h16':2.25,'t2_h80':2.5,'t3_plain':2.0625,
          't4_gvar':2.0625,'t4_ggrid':2.0625}
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):  # ggml-style signed absmax, fp16 scale per row
    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 qrot(x, b, G):
    y = x.float() @ HM; q = absmax_q(y.reshape(-1, G), -(2**(b-1)), 2**(b-1)-1).reshape(y.shape); return q @ HM.T
def qplain(x, b, G=32): return absmax_q(x.float().reshape(-1, G), -(2**(b-1)), 2**(b-1)-1).reshape(x.shape)
def q_alloc(x, R, r):  # R (H,256,256) orthogonal per head; first r coords 3-bit (one fp16 scale), rest 2-bit in 2 groups
    y = torch.einsum('bhtd,hde->bhte', x.float(), R); sh = y.shape[:-1]
    qt = absmax_q(y[..., :r].reshape(-1, r), -4, 3).reshape(*sh, r)
    qr = absmax_q(y[..., r:].reshape(-1, (256-r)//2), -2, 1).reshape(*sh, 256-r)
    return torch.einsum('bhte,hde->bhtd', torch.cat([qt, qr], -1), R)
FULL = None; RW = {}; RH = None; SIGNIDX = {}; KM = {}; GVAR = {}; GGRID = {}
STATE = {'arm': 'ref', 'phase': 'prefill', 'recv': False, 'gain': None}
POS = {}; LEDGER = {}; SHADOW = {}; CAP = {}; HI = {}; RECV = {}; RC = {}; DEC = {}; QS = {}; KREF = {}; E2A = {}; STAT = {}
def lab_sign(kr, l):  # kr (H,T,256) rotated K -> cohort = signs of the 6 highest-variance rotated coords (calibration)
    H, T, _ = kr.shape; b = torch.gather(kr, 2, SIGNIDX[l][:, None, :].expand(H, T, 6)) > 0
    return (b.long() * (2**torch.arange(6, device=dev))).sum(-1)
def lab_km(kr, l): return torch.cdist(kr, KM[l]).argmin(-1)
def ledger_q(xr, lab, E):  # xr (H,T,256) rotated V; lab (H,T); E (H,64,256) per-cohort fp32 error ledger, updated in place
    H, T, D = xr.shape; key = (torch.arange(H, device=dev)[:, None]*64 + lab).reshape(-1); N = key.numel()
    flat = xr.reshape(-1, D); Ef = E.view(-1, D); out = torch.empty_like(flat)
    sk, perm = torch.sort(key, stable=True); cnt = torch.bincount(sk, minlength=H*64); st = torch.cumsum(cnt, 0) - cnt
    rk = torch.empty_like(key); rk[perm] = torch.arange(N, device=dev) - st[sk]
    o2 = torch.argsort(rk*(H*64) + key); rc = torch.bincount(rk).tolist(); s0 = 0
    for c in rc:
        idx = o2[s0:s0+c]; kk = key[idx]; u = flat[idx] + Ef[kk]; d = absmax_q(u, -2, 1); Ef[kk] = u - d; out[idx] = d; s0 += c
    return out.view(H, T, D)
def quant_kv(k, v, l):
    arm = STATE['arm']; T = k.shape[2]; pos0 = POS.get(l, 0); POS[l] = pos0 + T
    if arm == 'ref' or l not in FULL: return k, v
    dt = v.dtype
    if arm == 'capture': CAP.setdefault(l, []).append((k.detach().float().cpu(), v.detach().float().cpu())); return k, v
    if arm == 'k4_v4': return qplain(k, 4).to(dt), qplain(v, 4).to(dt)
    kq = qrot(k, 4, 128).to(dt)
    if arm == 'k4rg128': return kq, v
    if arm == 'A2':
        if l not in SHADOW: SHADOW[l] = torch.empty((v.shape[0], v.shape[1], CTX+NSC+SCH, v.shape[3]), dtype=v.dtype, pin_memory=True)
        SHADOW[l][:, :, pos0:pos0+T] = v.detach().cpu(); vq = qrot(v, 2, 128)  # fp shadow kept in pinned CPU RAM (OOM fix)
    elif arm == 'A4' or arm.startswith('t1b'): vq = qrot(v, 4, 128)
    elif arm.startswith('t1a'):
        sel = HI[l][pos0:pos0+T]; vq = torch.where(sel[None, None, :, None], qrot(v, 4, 128), qrot(v, 2, 128))
    elif arm.startswith('t2_'):
        r = int(arm.split('_')[1].lstrip('woh')); vq = q_alloc(v, RW[l] if arm[3] == 'w' else RH, r)
    else:  # t3 / t4: rotated 2-bit, one fp16 scale per 256 values; optional ledger, optional gain
        g = 1.0 if STATE['gain'] is None else STATE['gain'][l]
        vr = v.float()[0] @ HM
        if '_led_' in arm:
            kr = k.float()[0] @ HM
            lab = lab_sign(kr, l) if 'sign' in arm else lab_km(kr, l)
            if 'rand' in arm:
                gg = torch.Generator().manual_seed(int(arm[-1])*1000003 + l*1009 + pos0); lab = lab[:, torch.randperm(T, generator=gg).to(dev)]
            if l not in LEDGER: LEDGER[l] = torch.zeros(v.shape[1], 64, 256, device=dev)
            d = ledger_q(vr, lab, LEDGER[l])
        else: d = absmax_q(vr.reshape(-1, 256), -2, 1).view(vr.shape)
        vq = ((d*g) @ HM.T)[None]
    return kq, vq.to(dt)
def att_sum(q, k, scaling):  # sum over heads and queries of causal softmax attention each key receives
    H = q.shape[1]; nr = H // k.shape[1]; kv = k.shape[2]; ql = q.shape[2]
    msk = torch.arange(kv, device=dev)[None, :] > (kv - ql + torch.arange(ql, device=dev))[:, None]; tot = torch.zeros(kv, device=dev)
    for h in range(H):
        s = (q[0, h] @ k[0, h//nr].T).float()*scaling; s.masked_fill_(msk, float('-inf')); tot += torch.softmax(s, -1).sum(0); del s
    return tot
ORIG = ALL_ATTENTION_FUNCTIONS['sdpa']; CALLS = {'n': 0}
def attn_wrap(module, q, k, v, mask, **kw):
    out = ORIG(module, q, k, v, mask, **kw)
    l = getattr(module, 'layer_idx', None); CALLS['n'] += 1
    if l not in FULL: return out
    arm = STATE['arm']; kv = k.shape[2]
    with torch.no_grad():
        if arm == 'ref' and STATE['recv']:
            RECV[l][:kv] += att_sum(q, k, module.scaling)
            if STATE['phase'] == 'score': QS.setdefault(l, []).append(q[0].clone()); KREF[l] = k
        elif arm.startswith('t1b'):
            RC[l][:kv] += att_sum(q, k, module.scaling); f = FR[arm[4:]]; cut = kv - WIN; d0 = DEC.get(l, 0)
            if cut > d0:
                n = cut - d0; nk = int(round(f*n)); sc = RC[l][d0:cut]; keep = torch.zeros(n, dtype=torch.bool, device=dev)
                if nk > 0: keep[sc.topk(nk).indices] = True
                dem = d0 + torch.nonzero(~keep).squeeze(1)
                v[0][:, dem, :] = qrot(v[0][:, dem, :], 2, 128).to(v.dtype)
                STAT['kept'] = STAT.get('kept', 0) + nk; STAT['dec'] = STAT.get('dec', 0) + n; DEC[l] = cut
                if SMOKE and 'chk' not in STAT and len(dem): STAT['chk'] = (l, int(dem[0]), v[0][:, dem[0], :].clone())
            if SMOKE and 'chk' in STAT and STAT['chk'][0] == l and 'chk_ok' not in STAT and DEC[l] > STAT['chk'][1] + PCH:
                STAT['chk_ok'] = bool(torch.equal(v[0][:, STAT['chk'][1], :], STAT['chk'][2])); log(f"smoke: demotion persists in cache = {STAT['chk_ok']}")
        elif arm == 'A2' and STATE['phase'] == 'score':
            ql = q.shape[2]; oq = out[0].reshape(1, ql, -1).float(); orf = ORIG(module, q, k, SHADOW[l][:, :, :k.shape[2]].to(k.device), mask, **kw)[0].reshape(1, ql, -1).float()
            gate = module._e22o_qp.view(1, ql, -1, 512)[..., 256:].reshape(1, ql, -1).float(); sg = torch.sigmoid(gate)
            po_q = module.o_proj((oq*sg).to(torch.bfloat16)).float(); po_r = module.o_proj((orf*sg).to(torch.bfloat16)).float()
            a = E2A.setdefault(l, [0.]*6)
            for i, t in enumerate([oq-orf, orf, (oq-orf)*sg, orf*sg, po_q-po_r, po_r]): a[i] += t.pow(2).sum().item()
    return out
ALL_ATTENTION_FUNCTIONS['sdpa'] = attn_wrap
def kmeans(X, k=64, it=25, seed=0):
    g = torch.Generator().manual_seed(seed); C = X[torch.randperm(X.shape[0], generator=g)[:k].to(dev)].clone()
    for _ in range(it):
        a = torch.cdist(X, C).argmin(1); Cn = torch.zeros_like(C).index_add_(0, a, X); c = torch.bincount(a, minlength=k).float(); m = c > 0; C[m] = Cn[m]/c[m, None]
    return C
def pooled(Z, lab):  # pooled mean pairwise Pearson within cohorts; Z (nq, P) columns unit-norm standardized
    S = torch.zeros(64, Z.shape[0], device=dev).index_add_(0, lab, Z.T); n = torch.bincount(lab, minlength=64).float()
    return (((S**2).sum(1) - n).sum() / (n*(n-1)).sum()).item()
def premise(l, P, scaling):
    K = KREF[l][0]; Qs = torch.cat(QS[l], 1); H = K.shape[0]; nr = Qs.shape[0]//H; ns = Qs.shape[1]
    KR = K[:, :P].float() @ HM; LS = lab_sign(KR, l); LK = lab_km(KR, l); del KR
    res = {'sign': [], 'km': [], 'sign_rand': [], 'km_rand': []}; g = torch.Generator().manual_seed(1000 + l)
    for h in range(H):
        cand = Qs[h*nr:(h+1)*nr].reshape(-1, Qs.shape[-1]); cpos = (P + torch.arange(ns, device=dev)).repeat(nr)
        sel = torch.randperm(cand.shape[0], generator=g)[:NQ].to(dev); q = cand[sel]; pq = cpos[sel]
        s = (q @ K[h].T).float()*scaling; s.masked_fill_(torch.arange(K.shape[1], device=dev)[None] > pq[:, None], float('-inf'))
        A = torch.softmax(s, -1)[:, :P].contiguous(); del s
        m = A.mean(0); sd = A.std(0); ok = sd > 0; Z = ((A - m)/(sd*math.sqrt(A.shape[0]-1)).clamp_min(1e-30))[:, ok].contiguous(); del A
        for nm, L in (('sign', LS[h][ok]), ('km', LK[h][ok])):
            res[nm].append(pooled(Z, L))
            rr = []
            for sd_ in range(NSEED):
                gg = torch.Generator().manual_seed(sd_*7919 + l*31 + h); rr.append(pooled(Z, L[torch.randperm(L.numel(), generator=gg).to(dev)]))
            res[nm+'_rand'].append(sum(rr)/len(rr))
        del Z
    return {k: sum(v)/len(v) for k, v in res.items()} | {'per_head': res}
def main():
    global FULL, RH
    log(f"=== E22o 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')
    ATT = {m.layer_idx: m for m in model.modules() if type(m).__name__.endswith('Attention') and getattr(m, 'layer_idx', None) in FULL and hasattr(m, 'o_proj')}
    log(f"loaded; full-attn layers {sorted(FULL)}; attn modules {sorted(ATT)}; vram {torch.cuda.memory_allocated()/2**30:.2f} GiB")
    for l, m in ATT.items():
        def hk(mod, i, o, m=m): m._e22o_qp = o if STATE['arm'] == 'A2' and STATE['phase'] == 'score' else None
        m.q_proj.register_forward_hook(hk)
        W = m.o_proj.weight; W = BF.dequantize_4bit(W.data, W.quant_state).float() if hasattr(W, 'quant_state') else W.float()
        H = tc.num_key_value_heads; nr = tc.num_attention_heads // H; R = []; en = []
        for h in range(H):
            M = torch.cat([W[:, j*256:(j+1)*256] for j in range(h*nr, (h+1)*nr)], 0); U, S, Vh = torch.linalg.svd(M, full_matrices=False)
            R.append(Vh.T); e = S**2/(S**2).sum(); en.append((e[:16].sum().item(), e[:80].sum().item()))
        RW[l] = torch.stack(R); log(f"W_o layer {l}: W {tuple(W.shape)}; energy share top16/top80 per head {[(round(a,3), round(b,3)) for a, b in en]}")
    RH = HM[None].expand(tc.num_key_value_heads, 256, 256).contiguous()
    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__}; attn wrapper calls {CALLS['n']}")
    assert CALLS['n'] > 0, 'attention wrapper not called'
    orig = C.update
    def upd(self, k, v, layer_idx, *a, **kw):
        kq, vq = quant_kv(k, v, layer_idx); return orig(self, kq, vq, layer_idx, *a, **kw)
    C.update = upd
    def reset():
        POS.clear(); LEDGER.clear(); SHADOW.clear(); DEC.clear()
    def run(w, P, nsc, ref=None):
        reset(); STATE['phase'] = 'prefill'; cache = None
        for s in range(0, P, PCH):
            cache = model(input_ids=w[:, s:min(s+PCH, P)], past_key_values=cache, use_cache=True, logits_to_keep=1).past_key_values
        STATE['phase'] = 'score'; kl = []; agree = 0; nll = 0.; outs = []
        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 ref is None: outs.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
        STATE['phase'] = 'prefill'
        return cache, outs, torch.cat(kl), agree, nll
    with torch.no_grad():
        # 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.to(dev)
        lc = min(LC, ctoks.shape[1] - NCAL - 1); log(f"calib tokens {ctoks.shape[1]}, using {lc} context + {NCAL} scored")
        STATE['arm'] = 'capture'; reset(); cache = None
        for s in range(0, lc, PCH): cache = model(input_ids=ctoks[:, s:min(s+PCH, lc)], past_key_values=cache, use_cache=True, logits_to_keep=1).past_key_values
        del cache
        for l, xs in sorted(CAP.items()):
            K = torch.cat([a for a, _ in xs], 2)[0].to(dev); V = torch.cat([b for _, b in xs], 2)[0].to(dev)
            KR = K @ HM; VR = V @ HM
            SIGNIDX[l] = KR.var(1).topk(6, dim=-1).indices
            KM[l] = torch.stack([kmeans(KR[h]) for h in range(KR.shape[0])])
            D = absmax_q(VR.reshape(-1, 256), -2, 1); GVAR[l] = (VR.std()/D.std()).item()
            sz = torch.bincount(lab_km(KR, l).reshape(-1), minlength=64); ss = torch.bincount(lab_sign(KR, l).reshape(-1), minlength=64)
            log(f"calib layer {l}: T={K.shape[1]} gain_var={GVAR[l]:.4f} km cohort size max/min {sz.max().item()}/{sz.min().item()} sign cohorts nonempty {(ss>0).sum().item()} max {ss.max().item()}")
            del K, V, KR, VR, D
        CAP.clear(); STATE['arm'] = 'ref'
        cref = run(ctoks, lc, NCAL)[1]
        def calkl(gains):
            STATE['arm'] = 't4_cal'; STATE['gain'] = gains; r = run(ctoks, lc, NCAL, cref)[2].mean().item(); STATE['gain'] = None; return r
        G = dict(GVAR); log(f"calib KL plain 2-bit G256 (gain 1): {calkl({l: 1.0 for l in FULL}):.5f}; variance-match gains: {calkl(G):.5f}")
        for l in sorted(FULL):
            best = (calkl(G), G[l])  # keep the current (variance-match) gain unless a grid value beats it
            for gv in GRID:
                G2 = dict(G); G2[l] = gv; r = calkl(G2)
                if best is None or r < best[0]: best = (r, gv)
            G[l] = best[1]; log(f"gain grid layer {l}: best g={best[1]} calib KL {best[0]:.5f}")
        GGRID.update(G); log(f"GVAR {GVAR}; GGRID {GGRID}")
        json.dump({'gvar': GVAR, 'ggrid': GGRID, 'signidx': {l: v.tolist() for l, v in SIGNIDX.items()}}, open(os.path.join(OUT, 'calib.json'), 'w'))
        torch.save({l: v.cpu() for l, v in KM.items()}, os.path.join(OUT, 'km.pt')); del ctoks, cref
    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]}")
    res = open(os.path.join(OUT, 'results.jsonl'), 'a')
    for E in SPANS:
        n = CTX + NSC + 1; assert E <= wtoks.shape[1]
        w = wtoks[:, E-n:E].to(dev); P = CTX; ref = None
        for arm in ARMS:
            STATE['arm'] = arm; STAT.clear(); E2A.clear(); t0 = time.time(); torch.cuda.reset_peak_memory_stats()
            STATE['gain'] = GVAR if arm == 't4_gvar' else (GGRID if 'ggrid' in arm else None)
            if arm == 'ref':
                STATE['recv'] = True; QS.clear(); KREF.clear()
                for l in FULL: RECV[l] = torch.zeros(P+NSC, device=dev)
            if arm.startswith('t1a'):
                f = FR[arm[4:]]; nk = int(round(f*(P+NSC)))
                for l in FULL: HI[l] = torch.zeros(P+NSC, dtype=torch.bool, device=dev); HI[l][RECV[l].topk(nk).indices] = True
            if arm.startswith('t1b'):
                for l in FULL: RC[l] = torch.zeros(P+NSC, device=dev)
            with torch.no_grad():
                cache, outs, kl, agree, nll = run(w, P, NSC, ref)
                if arm == 'ref':
                    ref = outs; STATE['recv'] = False
                    for l in sorted(FULL):
                        pr = premise(l, P, ATT[l].scaling)
                        r = dict(span=E, kind='premise3a', layer=l, **pr); res.write(json.dumps(r)+'\n'); res.flush(); log(f"premise {json.dumps(r)}")
                    QS.clear(); KREF.clear()
            bv = BITS_V.get(arm)
            if arm.startswith('t1a'): bv = 2.125 + 2*FR[arm[4:]] + 1/256
            if arm.startswith('t1b'): hf = (STAT.get('kept', 0) + WIN*len(FULL)) / ((P+NSC)*len(FULL)); bv = 2.125 + 2*hf + 1/256
            if '_led_' in arm: bv = 2.0625
            r = dict(span=E, ctx=P, scored=NSC, arm=arm, bits_v=bv, 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))
            if arm == 'A2': r['e2a'] = {l: dict(pre=math.sqrt(a[0]/a[1]), gated=math.sqrt(a[2]/a[3]), post_oproj=math.sqrt(a[4]/a[5])) for l, a in sorted(E2A.items())}
            if arm.startswith('t1b'): r['kept_frac'] = STAT.get('kept', 0)/max(1, STAT.get('dec', 1))
            res.write(json.dumps(r)+'\n'); res.flush(); log(f"result {json.dumps(r)}")
            del cache, kl; SHADOW.clear(); LEDGER.clear(); gc.collect(); torch.cuda.empty_cache()
        del ref, w
    log("=== E22oDONE ===")
if __name__ == '__main__': main()