Upload lattice_repair.py

lattice_repair.py

@@ -0,0 +1,711 @@
+#!/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()