Delete lattice_repair.py

lattice_repair.py

@@ -1,711 +0,0 @@
-#!/usr/bin/env python3
-"""lattice_repair - whole-lattice cartography + drag of a quantized GGUF onto its F16/BF16 source.
-
-Sibling of lattice_drag.py: same chunked, weight-aware, byte-surgical
-machinery (frames frozen, only outlier codes move, everything else stays
-byte-for-byte identical), extended to the WHOLE lattice:
-
-1.  Whole lattice: every tensor of the quant GGUF is enumerated and matched
-    to the source GGUF by name - the Q2_K/Q3_K/Q4_0 tensors are worked on;
-    all other tensors are measured against the source and dumped too (F32/F16/BF16
-    compared directly; other quant types such as Q8_0/Q6_K are dequantized, so
-    their real displacement is reported, not assumed zero).
-2.  Position difference per weight: for every weight the displacement
-    d[i] = source[i] - dequantized[i] between its position on the source
-    lattice and its position on the quant lattice is computed and dumped
-    in full, BEFORE the drag (.delta.before.f16) and AFTER (.delta.after.f16).
-3.  The drag (the pull): outlier weights (|d| > --thresh code steps of
-    their sub-block) are moved to the nearest code their EXISTING block
-    frame allows, only when that reduces |d| - never regress; scales,
-    minima, fp16 d values and non-outlier codes stay untouched.
-    Every moved weight goes to the --catalog CSV with its source position,
-    old/new lattice positions and old/new codes.
-
-Dumps per tensor (float16 raw files):  <name>.src.f16 (source lattice),
-<name>.qnt.f16 (quant lattice after the drag), <name>.delta.before.f16 /
-<name>.delta.after.f16 (per-weight displacement before / after).
-Kernel/matrix ops add-ons (rope_freq etc.) and all metadata ride along
-unchanged: output is a copy of the quant GGUF with data regions patched
-in place - loads directly in llama.cpp.
-
-    python3 lattice_repair.py --source model-f16.gguf --quant model-Q3.gguf \\
-        --output model-Q3-repaired.gguf --thresh 1.0 --catalog moved.csv
-
-Use --dry-run first: it reports everything and writes nothing.
---dump-dir writes the full source and quant lattices (large).
-
-Requires: numpy, gguf (pip install gguf).
-"""
-
-from __future__ import annotations
-
-import argparse
-import shutil
-import sys
-import time
-from pathlib import Path
-
-import numpy as np
-
-# the pip gguf package; local gguf.py must not shadow it
-sys.path = [p for p in sys.path if p not in ("", str(Path.cwd()))]
-from gguf import GGMLQuantizationType as QT
-from gguf import GGML_QUANT_SIZES, GGUFReader
-from gguf import quants as gq
-
-# bytes per block, elements per block, code range
-FORMATS = {
-    QT.Q2_K: dict(name="Q2_K", be=256, bb=84, cmin=0, cmax=3),
-    QT.Q3_K: dict(name="Q3_K", be=256, bb=110, cmin=-4, cmax=3),
-    QT.Q4_0: dict(name="Q4_0", be=32, bb=18, cmin=0, cmax=15),
-}
-
-
-# ---------------------------------------------------------------------------
-# block layouts.  Every format is expressed as   deq = scale * code + offset
-# per element, so one drag routine serves all three.
-# ---------------------------------------------------------------------------
-def _f16(b):
-    return np.ascontiguousarray(b).view("<f2").reshape(-1).astype(np.float32)
-
-
-def unpack_q2k(raw):
-    B = raw.shape[0]
-    sc, qs = raw[:, :16], raw[:, 16:80]
-    d, dmin = _f16(raw[:, 80:82]), _f16(raw[:, 82:84])
-    codes = np.empty((B, 256), np.int16)
-    for half in range(2):
-        q = qs[:, 32 * half : 32 * half + 32]
-        for j in range(4):
-            codes[:, 128 * half + 32 * j : 128 * half + 32 * j + 32] = (q >> (2 * j)) & 3
-    scale = d[:, None] * (sc & 15).astype(np.float32)
-    off = -dmin[:, None] * (sc >> 4).astype(np.float32)
-    return codes, np.repeat(scale, 16, 1), np.repeat(off, 16, 1)
-
-
-def pack_q2k(raw, codes):
-    B = raw.shape[0]
-    qs = np.zeros((B, 64), np.uint8)
-    for half in range(2):
-        for j in range(4):
-            c = codes[:, 128 * half + 32 * j : 128 * half + 32 * j + 32].astype(np.uint8)
-            qs[:, 32 * half : 32 * half + 32] |= c << (2 * j)
-    out = raw.copy()
-    out[:, 16:80] = qs
-    return out
-
-
-def _q3k_scales(s12):
-    s = s12.astype(np.int32)
-    out = np.empty((s.shape[0], 16), np.int32)
-    for k in range(4):
-        out[:, k] = (s[:, k] & 15) | (((s[:, 8 + k] >> 0) & 3) << 4)
-        out[:, 4 + k] = (s[:, 4 + k] & 15) | (((s[:, 8 + k] >> 2) & 3) << 4)
-        out[:, 8 + k] = (s[:, k] >> 4) | (((s[:, 8 + k] >> 4) & 3) << 4)
-        out[:, 12 + k] = (s[:, 4 + k] >> 4) | (((s[:, 8 + k] >> 6) & 3) << 4)
-    return out - 32
-
-
-def unpack_q3k(raw):
-    B = raw.shape[0]
-    hm, qs = raw[:, :32], raw[:, 32:96]
-    sc = _q3k_scales(raw[:, 96:108])
-    d = _f16(raw[:, 108:110])
-    codes = np.empty((B, 256), np.int16)
-    for g in range(8):
-        low = (qs[:, 32 * (g // 4) : 32 * (g // 4) + 32] >> (2 * (g % 4))) & 3
-        hi = (hm >> g) & 1
-        codes[:, 32 * g : 32 * g + 32] = low.astype(np.int16) + 4 * hi.astype(np.int16) - 4
-    scale = d[:, None] * sc.astype(np.float32)
-    return codes, np.repeat(scale, 16, 1), np.zeros((B, 256), np.float32)
-
-
-def pack_q3k(raw, codes):
-    B = raw.shape[0]
-    u = (codes.astype(np.int16) + 4).astype(np.uint8)
-    hm = np.zeros((B, 32), np.uint8)
-    qs = np.zeros((B, 64), np.uint8)
-    for g in range(8):
-        ug = u[:, 32 * g : 32 * g + 32]
-        qs[:, 32 * (g // 4) : 32 * (g // 4) + 32] |= (ug & 3) << (2 * (g % 4))
-        hm |= (ug >> 2) << g
-    out = raw.copy()
-    out[:, :32] = hm
-    out[:, 32:96] = qs
-    return out
-
-
-def unpack_q4_0(raw):
-    B = raw.shape[0]
-    d = _f16(raw[:, 0:2])
-    qs = raw[:, 2:18]
-    codes = np.empty((B, 32), np.int16)
-    codes[:, :16] = qs & 15
-    codes[:, 16:] = qs >> 4
-    scale = np.repeat(d[:, None], 32, 1)
-    return codes, scale, -8.0 * scale
-
-
-def pack_q4_0(raw, codes):
-    u = codes.astype(np.uint8)
-    out = raw.copy()
-    out[:, 2:18] = u[:, :16] | (u[:, 16:] << 4)
-    return out
-
-
-UNPACK = {QT.Q2_K: unpack_q2k, QT.Q3_K: unpack_q3k, QT.Q4_0: unpack_q4_0}
-PACK = {QT.Q2_K: pack_q2k, QT.Q3_K: pack_q3k, QT.Q4_0: pack_q4_0}
-
-# ---------------------------------------------------------------------------
-# the drag
-# ---------------------------------------------------------------------------
-def drag_chunk(ftype, raw, x, thresh):
-    """Return (new_raw, stats, moved_mask, deq_before, deq_after, codes_old, codes_new)."""
-    f = FORMATS[ftype]
-    codes, s, o = UNPACK[ftype](raw)
-    deq = s * codes + o
-    e = x - deq
-    step = np.abs(s)
-    live = step > 0
-    outlier = live & (np.abs(e) > thresh * step)
-
-    safe = np.where(live, s, 1.0)
-    near = np.clip(np.rint((x - o) / safe), f["cmin"], f["cmax"]).astype(np.int16)
-    new_deq = s * near + o
-    new_e = x - new_deq
-    move = outlier & (np.abs(new_e) < np.abs(e))
-    codes2 = np.where(move, near, codes)
-    deq2 = np.where(move, new_deq, deq)
-    e2 = x - deq2
-
-    flim = (outlier & (np.abs(e2) > thresh * step)).any(axis=1)   # (B,)
-    st = dict(
-        weights=x.size,
-        outliers=int(outlier.sum()),
-        moved=int(move.sum()),
-        frame_limited=int((outlier & (np.abs(e2) > thresh * step)).sum()),
-        blocks_frame_limited=int(flim.sum()),
-        dead_frame=int((~live).sum()),
-        sse_before=float(np.sum(e.astype(np.float64) ** 2)),
-        sse_after=float(np.sum(e2.astype(np.float64) ** 2)),
-        max_before=float(np.max(np.abs(e))) if e.size else 0.0,
-        max_after=float(np.max(np.abs(e2))) if e.size else 0.0,
-        dc_before=float(np.sum(np.abs(e.sum(axis=1)))),
-        dc_after=float(np.sum(np.abs(e2.sum(axis=1)))),
-        blocks_changed=int(move.any(axis=1).sum()),
-        blocks=raw.shape[0],
-    )
-    new_raw = PACK[ftype](raw, codes2) if st["moved"] else raw
-    return new_raw, st, move, flim, deq, deq2, codes, codes2
-
-
-def stats_add(into, st):
-    if into is None:
-        return dict(st)
-    for k, v in st.items():
-        into[k] = max(into.get(k, 0.0), v) if k.startswith("max_") else into.get(k, 0) + v
-    return into
-
-
-# ---------------------------------------------------------------------------
-# frame re-fit for frame-limited outliers (--fix-frames).
-# Only blocks that still hold frame-limited outlier weights are touched.
-# Each block's frame (d, sc) is re-solved on its own grid from the SOURCE:
-# closed-form least squares over the existing codes plus an fp16-ULP window
-# probe; codes are optionally re-projected to the new frame.  A candidate is
-# accepted per block ONLY when it lowers the block cost SSE + dc_lambda*DC^2;
-# every non-accepting block and every untouched block stays byte-identical.
-# ---------------------------------------------------------------------------
-def _bits16(v):
-    """float values -> fp16 bit patterns (uint16)."""
-    return np.asarray(v, np.float32).astype(np.float16).view("<u2").astype(np.uint16)
-
-
-def _val16(b):
-    """fp16 bit patterns (uint16) -> float32 values (same shape)."""
-    return np.ascontiguousarray(b).view("<f2").astype(np.float32)
-
-
-def _ulp_window(b, rad):
-    """fp16 bit windows +-rad ULPs around the magnitudes of patterns b."""
-    mag = (b & 0x7FFF).astype(np.int32)
-    win = np.clip(mag[:, None] + np.arange(-rad, rad + 1), 0, 0x7BFF).astype(np.uint16)
-    return win | (b & 0x8000)[:, None]
-
-
-def _better(cost_new, cost_old):
-    """strictly better beyond float noise (relative 1e-9)."""
-    return cost_new < cost_old * (1.0 - 1e-9) - 1e-30
-
-
-def _block_cost(e, dc_lam):
-    """per-block cost: SSE + dc_lam * DC^2 (reduced over the last axis)."""
-    e = np.asarray(e, np.float64)
-    dc = e.sum(axis=-1)
-    cost = np.sum(e * e, axis=-1) + dc_lam * dc * dc
-    return np.where(np.isfinite(cost), cost, np.inf)
-
-
-def refit_q4_0(raw, x, dc_lam, rad):
-    """d re-solve on Q4_0 blocks; codes optionally re-projected.  Returns
-    (new_raw, improved_mask, deq).  Never mutates the input."""
-    codes, s, o = unpack_q4_0(raw)
-    d16 = np.ascontiguousarray(raw[:, 0] | (raw[:, 1].astype(np.uint16) << 8))
-    cp = codes.astype(np.float64) - 8.0              # deq = d * cp
-    B = raw.shape[0]
-    deq0 = s * codes + o
-    cost0 = _block_cost(x - deq0, dc_lam)
-
-    win_bits = _ulp_window(d16, rad)                 # (B, W)
-    Sxc = np.sum(x.astype(np.float64) * cp, axis=1)
-    Scc = np.sum(cp * cp, axis=1)
-    d_ls = np.where(Scc > 1e-30, Sxc / np.maximum(Scc, 1e-30), 0.0)
-    b_ls = np.where((d_ls > 0) & (d_ls < 6e4), _bits16(d_ls), 0)
-    all_bits = np.concatenate([win_bits, b_ls[:, None]], axis=1)   # (B, W+1)
-    dv = _val16(all_bits)
-    dv = np.where(dv > 0, dv, np.inf)                # invalid candidates lose
-    W = all_bits.shape[1]
-
-    deq_keep = dv[:, :, None] * cp[:, None, :]
-    cost_keep = _block_cost(x[:, None, :] - deq_keep, dc_lam)
-    u_ref = np.clip(np.rint(x[:, None, :] / dv[:, :, None] + 8.0), 0, 15)
-    deq_ref = dv[:, :, None] * (u_ref - 8.0)
-    cost_ref = _block_cost(x[:, None, :] - deq_ref, dc_lam)
-    # choice table: candidates x [keep codes, refit codes]
-    table_bits = np.concatenate([all_bits, all_bits], axis=1)      # (B, 2W)
-    table_cost = np.concatenate([cost_keep, cost_ref], axis=1)
-    pick = np.argmin(table_cost, axis=1)
-    cost_new = table_cost[np.arange(B), pick]
-    mode_ref = pick >= W
-    within = pick % W
-    improved = _better(cost_new, cost0)
-    rows = np.where(improved)[0]
-    if rows.size == 0:
-        return raw, improved, deq0
-    deq = deq0.copy()
-    deq[rows] = np.where(mode_ref[rows, None],
-                         deq_ref[rows, within[rows]],
-                         deq_keep[rows, within[rows]])
-    out = raw.copy()
-    out[rows, :2] = table_bits[np.arange(B), pick][rows].astype("<u2").view(np.uint8).reshape(-1, 2)
-    ref_rows = rows[mode_ref[rows]]
-    if ref_rows.size:
-        ui = u_ref[ref_rows, within[ref_rows]]        # (R, 32)
-        out[ref_rows] = pack_q4_0(out[ref_rows], ui.astype(np.int16))
-    return out, improved, deq
-
-
-def pack_q3_scales(sc):
-    """int subscales (B,16) in [-32,31] -> the 12 scale bytes per block."""
-    u = (np.asarray(sc, np.int32) + 32).astype(np.uint8)           # 0..63
-    B = u.shape[0]
-    s12 = np.zeros((B, 12), np.uint8)
-    for k in range(4):
-        s12[:, k] = (u[:, k] & 15) | ((u[:, 8 + k] & 15) << 4)
-        s12[:, 4 + k] = (u[:, 4 + k] & 15) | ((u[:, 12 + k] & 15) << 4)
-        s12[:, 8 + k] = (((u[:, k] >> 4) & 3)
-                         | (((u[:, 4 + k] >> 4) & 3) << 2)
-                         | (((u[:, 8 + k] >> 4) & 3) << 4)
-                         | (((u[:, 12 + k] >> 4) & 3) << 6))
-    return s12
-
-
-def assemble_q3_k(raw, codes, sc, d_bits):
-    """rebuild Q3_K block bytes from given codes, sub-scale ints and fp16 d bits."""
-    out = pack_q3k(raw, codes)
-    out[:, 96:108] = pack_q3_scales(sc)
-    out[:, 108:110] = np.asarray(d_bits, "<u2").view(np.uint8).reshape(-1, 2)
-    return out
-
-
-def refit_q3_k(raw, x, dc_lam, rad):
-    """(d, sc) re-solve on Q3_K blocks; codes optionally re-projected.
-    Returns (new_raw, improved_mask, deq).  Never mutates the input."""
-    codes, s, o = unpack_q3k(raw)
-    B = raw.shape[0]
-    d16 = np.ascontiguousarray(raw[:, 108] | (raw[:, 109].astype(np.uint16) << 8))
-    deq0 = s * codes + o
-    cost0 = _block_cost(x - deq0, dc_lam)
-
-    x16 = x.astype(np.float64).reshape(B, 16, 16)
-    c16 = codes.astype(np.float64).reshape(B, 16, 16)
-    Sxc = np.sum(x16 * c16, axis=2)                  # (B,16)
-    Scc = np.sum(c16 * c16, axis=2)
-    csum = np.sum(c16, axis=2)
-    axg = np.sum(x16 * x16, axis=2)
-    xtot = x16.reshape(B, 256).sum(axis=1)
-
-    win_bits = _ulp_window(d16, rad)                 # (B, W)
-    # multiplicative ladder: joint d*sc rescales beyond the ULP window
-    d0 = _val16(d16)
-    lad_bits = _bits16(d0[:, None] * (2.0 ** np.arange(-8, 9)))
-    lad_bits = np.where(np.isfinite(_val16(lad_bits)) & (_val16(lad_bits) > 0),
-                        lad_bits, d16[:, None])
-    d_ls = np.where(Scc.sum(1) > 1e-30,
-                    Sxc.sum(1) / np.maximum(Scc.sum(1), 1e-30), 0.0)
-    b_ls = np.where((d_ls > 0) & (d_ls < 6e4), _bits16(d_ls), 0)
-    all_bits = np.concatenate([win_bits, lad_bits, b_ls[:, None]], axis=1)   # (B, W+18)
-    dv = _val16(all_bits)                            # (B, W+18)
-    valid = (dv > 0) & np.isfinite(dv)
-    dv = np.where(valid, dv, np.inf)
-    W = all_bits.shape[1]
-
-    with np.errstate(invalid="ignore"):
-        sc_c = np.clip(np.rint(Sxc[:, None, :] /
-                               np.maximum(dv[:, :, None] * Scc[:, None, :], 1e-30)),
-                       -32, 31).astype(np.int64)    # closed-form sc* per (block, cand)
-    ds = dv[:, :, None] * sc_c
-    sse_c = np.sum(axg[:, None, :] - 2.0 * ds * Sxc[:, None, :]
-                   + ds * ds * Scc[:, None, :], axis=2)
-    dc_c = xtot[:, None] - np.einsum("bdg,bg->bd", sc_c, csum) * dv
-    cost_keep = sse_c + dc_lam * dc_c * dc_c         # codes frozen variant
-
-    # codes re-projected to the candidate frame (vec per group is (B,W+1,16))
-    safe_frame = np.maximum(np.abs(ds), 1e-30)
-    c_ref = np.clip(np.rint(x16[:, None, :, :] / safe_frame[:, :, :, None]), -4, 3)
-    deq_ref = dv[:, :, None, None] * sc_c[:, :, :, None] * c_ref     # (B,W+1,16,16)
-    deq_ref = deq_ref.reshape(B, W, 256)
-    cost_ref = _block_cost(x[:, None, :] - deq_ref, dc_lam)
-    cost_keep = np.where(~valid, np.inf, cost_keep)
-    cost_ref = np.where(~valid, np.inf, cost_ref)
-
-    table_sc = np.concatenate([sc_c, sc_c], axis=1)                  # (B, 2W, 16)
-    table_bits = np.concatenate([all_bits, all_bits], axis=1)
-    table_cost = np.concatenate([cost_keep, cost_ref], axis=1)
-    pick = np.argmin(table_cost, axis=1)
-    cost_new = table_cost[np.arange(B), pick]
-    mode_ref = pick >= W
-    within = pick % W
-    improved = _better(cost_new, cost0)
-    rows = np.where(improved)[0]
-    if rows.size == 0:
-        return raw, improved, deq0
-    deq = deq0.copy()
-    out = raw.copy()
-    for i in rows:
-        p = pick[i]
-        b = table_bits[i, p]
-        sc_int = table_sc[i, p].astype(np.int32)
-        if mode_ref[i]:
-            cs = c_ref[i, within[i]].reshape(256).astype(np.int16)
-            deq[i] = deq_ref[i, within[i]]
-        else:
-            cs = codes[i].astype(np.int16)
-            deq[i] = (dv[i, within[i]] * sc_int.astype(np.float64)
-                      .repeat(16, axis=0).reshape(256) * codes[i].astype(np.float64))
-        out[i] = assemble_q3_k(out[i][None, :], cs[None, :], sc_int[None, :],
-                               np.array([b]))[0]
-    return out, improved, deq
-
-
-def refit_q2_k(raw, x, dc_lam, rad):
-    """(d, dmin) closed-form 2x2 LS re-solve on Q2_K blocks (codes and sub-scale
-    nibbles frozen); codes never change, so only the frame moves.  Returns
-    (new_raw, improved_mask, deq).  Never mutates the input."""
-    codes, s, o = unpack_q2k(raw)
-    B = raw.shape[0]
-    d16 = np.ascontiguousarray(raw[:, 80] | (raw[:, 81].astype(np.uint16) << 8))
-    m16 = np.ascontiguousarray(raw[:, 82] | (raw[:, 83].astype(np.uint16) << 8))
-    sc = raw[:, :16].astype(np.int32)
-    u_e = ((sc & 15).repeat(16, axis=1) * codes).astype(np.float64)  # deq = d*u - m*v
-    v_e = np.repeat((sc >> 4), 16, axis=1).astype(np.float64)        # v = lm't per group
-    deq0 = s * codes + o
-    cost0 = _block_cost(x - deq0, dc_lam)
-
-    Saa = np.sum(u_e * u_e, axis=1)
-    Sab = np.sum(u_e * v_e, axis=1)
-    Sbb = np.sum(v_e * v_e, axis=1)
-    Sxa = np.sum(x.astype(np.float64) * u_e, axis=1)
-    Sxb = np.sum(x.astype(np.float64) * v_e, axis=1)
-    det = Saa * Sbb - Sab * Sab
-    ok = (det > 1e-30) & np.isfinite(det)
-    safe = np.where(ok, det, 1.0)
-    d_star = (Sbb * Sxa - Sab * Sxb) / safe
-    m_star = (Sab * Sxa - Saa * Sxb) / safe
-    ok &= (d_star > 0) & (m_star > 0) & np.isfinite(d_star) & np.isfinite(m_star)
-
-    db_closed = np.where(ok, _bits16(np.where(ok, d_star, 0.0)), 0)
-    mb_closed = np.where(ok, _bits16(np.where(ok, m_star, 0.0)), 0)
-    # candidate families: (d bits column, dmin bits column, width)
-    # [0] closed-form pair, [1] d window (m frozen), [2] m window (d frozen)
-    fams = [
-        (db_closed[:, None], mb_closed[:, None], 1),
-        (_ulp_window(d16, rad), np.broadcast_to(m16[:, None], (B, 2 * rad + 1)), 2 * rad + 1),
-        (np.broadcast_to(d16[:, None], (B, 2 * rad + 1)), _ulp_window(m16, rad), 2 * rad + 1),
-    ]
-    table_db = np.concatenate([f[0] for f in fams], axis=1)
-    table_mb = np.concatenate([f[1] for f in fams], axis=1)
-    offs = np.cumsum([0] + [f[2] for f in fams])
-    dv_t = _val16(np.ascontiguousarray(table_db))
-    mv_t = _val16(np.ascontiguousarray(table_mb))
-    bad = (dv_t <= 0) | (mv_t < 0) | ~np.isfinite(dv_t) | ~np.isfinite(mv_t)
-    dv_t = np.where(bad, np.inf, dv_t)
-    deq_t = dv_t[:, :, None] * u_e[:, None, :] - mv_t[:, :, None] * v_e[:, None, :]
-    cost_t = np.where(bad, np.inf, _block_cost(x[:, None, :] - deq_t, dc_lam))
-    pick = np.argmin(cost_t, axis=1)
-    improved = _better(cost_t[np.arange(B), pick], cost0)
-    rows = np.where(improved)[0]
-    if rows.size == 0:
-        return raw, improved, deq0
-    deq = deq0.copy()
-    out = raw.copy()
-    for i in rows:
-        p = pick[i]
-        dv_i = _val16(np.array([table_db[i, p]]))[0]
-        mv_i = _val16(np.array([table_mb[i, p]]))[0]
-        deq[i] = dv_i * u_e[i] - mv_i * v_e[i]
-        out[i, 80] = table_db[i, p] & 0xFF
-        out[i, 81] = (table_db[i, p] >> 8) & 0xFF
-        out[i, 82] = table_mb[i, p] & 0xFF
-        out[i, 83] = (table_mb[i, p] >> 8) & 0xFF
-    return out, improved, deq
-
-
-REFIT = {QT.Q2_K: refit_q2_k, QT.Q3_K: refit_q3_k, QT.Q4_0: refit_q4_0}
-
-
-def refit_batched(fn, raw, x, dc_lam, rad, batch):
-    """run a re-fit in slices of `batch` blocks (peak memory ~ batch * candidates * 256 * 8 B)."""
-    outs, imps, deqs = [], [], []
-    for a in range(0, raw.shape[0], batch):
-        o, i, d = fn(raw[a : a + batch], x[a : a + batch], dc_lam, rad)
-        outs.append(o)
-        imps.append(i)
-        deqs.append(d)
-    return np.concatenate(outs), np.concatenate(imps), np.concatenate(deqs)
-
-
-# ---------------------------------------------------------------------------
-# source access
-# ---------------------------------------------------------------------------
-def source_elements(t, a, b):
-    """float32 view of elements [a, b) of a source tensor (F32 / F16 / BF16)."""
-    data = np.asarray(t.data)
-    tt = t.tensor_type
-    if tt == QT.BF16:
-        flat = data.reshape(-1).view(np.uint16)[a:b].astype(np.uint32) << 16
-        return flat.view(np.float32)
-    flat = data.reshape(-1)[a:b]
-    if tt in (QT.F16, QT.F32):
-        return flat.astype(np.float32)
-    raise ValueError(f"source tensor type {tt.name} is not F32/F16/BF16")
-
-
-def passthrough_chunks(t, st_t, chunk_elems):
-    """Yield (source, quant) float32 flat chunks for a tensor we do not drag.
-
-    F32/F16/BF16 quant tensors are read directly; any other type gguf-py can
-    decode (Q8_0, Q6_K, ...) is dequantized, so its REAL displacement from the
-    source is measured instead of assumed to be zero.  Raises ValueError /
-    NotImplementedError when a type cannot be decoded."""
-    ne = int(t.n_elements)
-    tt = t.tensor_type
-    if tt in (QT.F32, QT.F16, QT.BF16):
-        for a in range(0, ne, chunk_elems):
-            b = min(a + chunk_elems, ne)
-            yield source_elements(st_t, a, b), source_elements(t, a, b)
-        return
-    bs, ts = GGML_QUANT_SIZES[tt]
-    qdata = np.asarray(t.data).reshape(-1)
-    per = max(1, chunk_elems // bs)
-    for b0 in range(0, ne // bs, per):
-        b1 = min(b0 + per, ne // bs)
-        raw = np.array(qdata[b0 * ts : b1 * ts]).reshape(b1 - b0, ts)
-        yield (source_elements(st_t, b0 * bs, b1 * bs),
-               gq.dequantize(raw, tt).reshape(-1).astype(np.float32))
-
-
-# ---------------------------------------------------------------------------
-def parse_args(argv=None):
-    p = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
-    p.add_argument("--source", required=True, help="F32/F16/BF16 GGUF")
-    p.add_argument("--quant", required=True, help="quantized GGUF (Q2_K/Q3_K/Q4_0 tensors)")
-    p.add_argument("--output", default="", help="default: <quant>-repaired.gguf")
-    p.add_argument("--thresh", type=float, default=1.0,
-                   help="outlier threshold in code steps (default 1.0)")
-    p.add_argument("--dry-run", action="store_true", help="report only, write nothing")
-    p.add_argument("--force", action="store_true", help="overwrite an existing --output")
-    p.add_argument("--fix-frames", action="store_true",
-                   help="frame re-fit on blocks with frame-limited outliers "
-                        "(guarded per block; only blocks that improve change)")
-    p.add_argument("--ulp-radius", type=int, default=8,
-                   help="fp16 ULP window radius for the frame re-fit")
-    p.add_argument("--refit-batch", type=int, default=1024,
-                   help="blocks per re-fit slice (memory control)")
-    p.add_argument("--dc-lambda", type=float, default=0.0,
-                   help="DC^2 coupling in the frame re-fit cost")
-    p.add_argument("--tensor-substr", default="", help="comma-separated name filters")
-    p.add_argument("--catalog", default="", help="CSV of every moved weight")
-    p.add_argument("--dump-dir", default="", help="write full src/quant lattices + displacement maps (float16)")
-    p.add_argument("--chunk-elems", type=int, default=1 << 23)
-    return p.parse_args(argv)
-
-
-def main(argv=None):
-    args = parse_args(argv)
-    src_path, q_path = Path(args.source), Path(args.quant)
-    out_path = Path(args.output) if args.output else q_path.with_name(q_path.stem + "-repaired.gguf")
-
-    src = GGUFReader(str(src_path))
-    qr = GGUFReader(str(q_path))
-    src_by_name = {t.name: t for t in src.tensors}
-    subs = [s for s in args.tensor_substr.split(",") if s]
-
-    # whole lattice: every tensor of the quant file, in file order
-    todo = [t for t in qr.tensors if not subs or any(s in t.name for s in subs)]
-    n_drag = sum(1 for t in todo if t.tensor_type in FORMATS)
-    n_pass = len(todo) - n_drag
-    print(f"whole lattice: {len(todo)} tensors ({n_drag} quantized to drag, "
-          f"{n_pass} passthrough)")
-
-    if not args.dry_run:
-        if out_path.exists() and not args.force:
-            sys.exit(f"{out_path} exists (use --force to overwrite)")
-        print(f"copying {q_path.name} -> {out_path.name} ...")
-        shutil.copyfile(q_path, out_path)
-    out_f = None if args.dry_run else open(out_path, "r+b")
-
-    cat = open(args.catalog, "w") if args.catalog else None
-    if cat:
-        cat.write("tensor,block,idx,fmt,source,old,new,old_code,new_code\n")
-    dump = Path(args.dump_dir) if args.dump_dir else None
-    if dump:
-        dump.mkdir(parents=True, exist_ok=True)
-
-    grand = None
-    pt_w, pt_sse = 0, 0.0
-    t0 = time.time()
-    for t in todo:
-        safe = t.name.replace("/", "__")
-        st_t = src_by_name.get(t.name)
-        if st_t is None:
-            print(f"  skip {t.name}: not in source")
-            continue
-        if int(st_t.n_elements) != int(t.n_elements):
-            print(f"  skip {t.name}: element count differs")
-            continue
-
-        if t.tensor_type not in FORMATS:
-            # not dragged, but still measured: the lattices only coincide for
-            # F32/F16/BF16 - other quant types (Q8_0, Q6_K...) are decoded.
-            fhs = []
-            if dump:
-                fhs = [open(dump / f"{safe}.{k}", "wb") for k in
-                       ("src.f16", "qnt.f16", "delta.before.f16", "delta.after.f16")]
-            n = sse = mx = 0
-            try:
-                for x, q in passthrough_chunks(t, st_t, args.chunk_elems):
-                    d = x - q
-                    n += d.size
-                    sse += float(np.sum(d.astype(np.float64) ** 2))
-                    mx = max(mx, float(np.max(np.abs(d))) if d.size else 0.0)
-                    if fhs:
-                        x.astype(np.float16).tofile(fhs[0])
-                        q.astype(np.float16).tofile(fhs[1])
-                        d.astype(np.float16).tofile(fhs[2])
-                        d.astype(np.float16).tofile(fhs[3])
-            except (ValueError, NotImplementedError, KeyError) as e:
-                print(f"  untouched {t.name} [{t.tensor_type.name}]: cannot measure ({e})")
-            else:
-                pt_w += n
-                pt_sse += sse
-                if mx > 0:
-                    print(f"  untouched {t.name} [{t.tensor_type.name}]: "
-                          f"rmse {(sse / max(n, 1)) ** 0.5:.4e} max|d| {mx:.3e} (not modified)")
-            finally:
-                for fh in fhs:
-                    fh.close()
-            continue
-
-        f = FORMATS[t.tensor_type]
-        be, bb = f["be"], f["bb"]
-        nblocks = int(t.n_elements) // be
-        qdata = np.asarray(t.data).reshape(-1)
-        blocks_per_chunk = max(1, args.chunk_elems // be)
-        tot = None
-        fs = fq = fd_b = fd_a = None
-        if dump:
-            fs = open(dump / f"{safe}.src.f16", "wb")
-            fq = open(dump / f"{safe}.qnt.f16", "wb")
-            fd_b = open(dump / f"{safe}.delta.before.f16", "wb")
-            fd_a = open(dump / f"{safe}.delta.after.f16", "wb")
-        for b0 in range(0, nblocks, blocks_per_chunk):
-            b1 = min(b0 + blocks_per_chunk, nblocks)
-            raw = np.array(qdata[b0 * bb : b1 * bb]).reshape(b1 - b0, bb)
-            x = source_elements(st_t, b0 * be, b1 * be).reshape(b1 - b0, be)
-            new_raw, st, move, flim, deq, deq2, c_old, c_new = drag_chunk(t.tensor_type, raw, x, args.thresh)
-            if args.fix_frames and flim.any():
-                sub_raw = new_raw[flim]
-                ftype = t.tensor_type
-                rr, improved, deq_r = refit_batched(REFIT[ftype], sub_raw, x[flim],
-                                                    args.dc_lambda, args.ulp_radius,
-                                                    args.refit_batch)
-                if improved.any():
-                    new_raw = new_raw.copy()
-                    new_raw[flim] = rr
-                    e_fix = x[flim] - deq_r          # only improved rows differ
-                    deq2 = deq2.copy()
-                    deq2[flim] = deq_r
-                    # merge the re-fit stats into the chunk stats
-                    e2f = x - deq2
-                    st["sse_after"] = float(np.sum(e2f.astype(np.float64) ** 2))
-                    st["max_after"] = float(np.max(np.abs(e2f)))
-                    st["dc_after"] = float(np.sum(np.abs(e2f.sum(axis=1))))
-                    st["blocks_frame_fixed"] = st.get("blocks_frame_fixed", 0) + int(improved.sum())
-                    # catalogue every weight whose position changed (drag OR frame re-fit)
-                    move = deq2 != deq
-                    c_new = UNPACK[ftype](new_raw)[0]
-            if (st["moved"] or st.get("blocks_frame_fixed", 0)) and out_f is not None:
-                out_f.seek(int(t.data_offset) + b0 * bb)
-                out_f.write(new_raw.tobytes())
-            if cat and (st["moved"] or st.get("blocks_frame_fixed", 0)):
-                bi, ei = np.nonzero(move)
-                for i, j in zip(bi, ei):
-                    cat.write(f"{t.name},{b0 + int(i)},{int(j)},{f['name']},"
-                              f"{x[i, j]:.8g},{deq[i, j]:.8g},{deq2[i, j]:.8g},"
-                              f"{int(c_old[i, j])},{int(c_new[i, j])}\n")
-            if dump:
-                x.astype(np.float16).tofile(fs)
-                deq2.astype(np.float16).tofile(fq)
-                (x - deq).astype(np.float16).tofile(fd_b)
-                (x - deq2).astype(np.float16).tofile(fd_a)
-            tot = stats_add(tot, st)
-        for fh in (fs, fq, fd_b, fd_a):
-            if fh:
-                fh.close()
-        if tot is None:
-            continue
-        rms0 = (tot["sse_before"] / tot["weights"]) ** 0.5
-        rms1 = (tot["sse_after"] / tot["weights"]) ** 0.5
-        print(f"  {t.name} [{f['name']}] outliers {tot['outliers']}/{tot['weights']} "
-              f"moved {tot['moved']} (blocks {tot['blocks_changed']}/{tot['blocks']}) "
-              f"frame-limited {tot['frame_limited']} in {tot.get('blocks_frame_limited', 0)} blocks"
-              + (f", frames fixed {tot.get('blocks_frame_fixed', 0)}"
-                 if tot.get('blocks_frame_fixed', 0) else "")
-              + f" | rmse {rms0:.4e}->{rms1:.4e} "
-              f"max|e| {tot['max_before']:.3e}->{tot['max_after']:.3e} "
-              f"mean|DC| {tot['dc_before']/tot['blocks']:.3e}->{tot['dc_after']/tot['blocks']:.3e}")
-        grand = stats_add(grand, tot)
-
-    if cat:
-        cat.close()
-    if out_f:
-        out_f.close()
-    if grand:
-        print(f"\nDRAGGED: {grand['weights']} weights | "
-              f"{grand['outliers']} outliers, {grand['moved']} moved, "
-              f"{grand['frame_limited']} frame-limited"
-              + (f", {grand.get('blocks_frame_fixed', 0)} blocks frame-fixed"
-                 if grand.get('blocks_frame_fixed', 0) else "")
-              + f" | rmse {(grand['sse_before']/grand['weights'])**0.5:.4e} -> "
-              f"{(grand['sse_after']/grand['weights'])**0.5:.4e} "
-              f"max|d| {grand['max_before']:.3e} -> {grand['max_after']:.3e} "
-              f"({time.time()-t0:.0f}s)")
-    if pt_w:
-        print(f"UNTOUCHED: {pt_w} weights in non-dragged tensors, "
-              f"rmse vs source {(pt_sse/pt_w)**0.5:.4e}")
-    print("dry run: nothing written." if args.dry_run else f"done -> {out_path}")
-
-
-if __name__ == "__main__":
-    main()