Anneal

__main__.py

@@ -47,6 +47,8 @@ def _config_from_args(args) -> HPCConfig:
         cfg.optimizer = "mse"
     if args.no_spectral:
         cfg.use_spectral = False
+    if args.thermal is not None:
+        cfg.therm_passes = int(args.thermal)
     if args.no_q4_fallback:
         cfg.enable_q4_fallback = False
     if args.keep_embd:
@@ -277,6 +279,13 @@ def main() -> None:
     parser.add_argument("--device", default=None, help="compute device: auto (default), cuda, or cpu")
     parser.add_argument("--mse", action="store_true", help="disable HPC sequence search")
     parser.add_argument("--no-spectral", action="store_true", help="disable spectral/DC objective")
+    parser.add_argument(
+        "--thermal",
+        type=int,
+        default=None,
+        metavar="PASSES",
+        help="thermal refinery heat-bath passes after Phase 4 (default 0 = off; seeded, bit-exact per seed)",
+    )
     parser.add_argument(
         "--format",
         choices=("mixed", "mixed_q1", "custom", "q1_0", "q3_k", "q4_0", "q2_k"),

anneal.py

@@ -0,0 +1,556 @@
+"""Thermal refinery: seeded pair heat-bath annealing on the Q2_K lattice.
+
+With ``fold_depth == 1`` the extended objective for one block is exactly::
+
+    E = sum_i w_i e_i^2
+        + (vw_lambda / n) * sum_p (e_p + e_{p+n/2})^2
+        + (dc_lambda / n) * (sum_i e_i + carry)^2
+
+so the 256 weights decompose into n/2 vesica pairs that interact with each
+other only through the scalar DC residual: with every error affine in its
+code, each pair's conditional energy over its joint states is exactly
+quadratic and closed-form.  ``gibbs_pair_pass`` resamples the pairs from
+those conditionals at a temperature (heat bath), fully batched across blocks
+over any uniform lattice size (``levels``: 4 = Q2_K, 16 = Q4_0, 2 = Q1_0,
+8 = Q3_K K-quants).  ``greedy=True`` resamples at the conditional argmin
+instead, which is a monotone descent.  ``thermal_refine_codes`` /
+``thermal_refine_uniform_codes`` schedule temperatures, then keep the
+elementwise best against the pre-anneal state, so the pass can never regress
+a block.  All randomness comes from a seeded generator, and a fixed seed
+reproduces bit-identical quantizations.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+from .config import HPCConfig
+from .fp16 import nearest_int
+from .q2k import QK_K
+
+EPS_D = 1.0e-15
+
+
+def pair_energy(
+    e: np.ndarray,
+    w: np.ndarray | None,
+    lam_vw: float,
+    lam_dc: float,
+    carry: np.ndarray,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    """Block energies plus the (sse, v, dc) live state for exact deltas."""
+    half = e.shape[1] // 2
+    ww = np.ones_like(e) if w is None else np.asarray(w, dtype=np.float64)
+    we_sq = ww * e * e
+    sse = np.sum(we_sq, axis=1)
+    v = e[:, :half] + e[:, half:]
+    dc = np.sum(e, axis=1) + np.asarray(carry, dtype=np.float64)
+    energy = sse + lam_vw * np.sum(v * v, axis=1) + lam_dc * dc * dc
+    return energy, we_sq, v, dc
+
+
+def gibbs_pair_pass(
+    X_f64: np.ndarray,
+    w_full: np.ndarray | None,
+    d_sub: np.ndarray,
+    m_sub: np.ndarray,
+    codes: np.ndarray,
+    e: np.ndarray,
+    v: np.ndarray,
+    dc: np.ndarray,
+    sse: np.ndarray,
+    temperature: np.ndarray,
+    lam_vw: float,
+    lam_dc: float,
+    rng: np.random.Generator,
+    greedy: bool = False,
+    levels: int = 4,
+) -> None:
+    """One exact sequential Gibbs sweep over all vesica pairs, in place.
+
+    Updates (codes, e, v, dc, sse) for the batch of blocks.  With
+    ``greedy=False`` each pair is resampled from its exact conditional
+    energy via a heat bath at ``temperature`` (per-block vector); with
+    ``greedy=True`` the conditional argmin is taken and only strict
+    energy decreases are accepted.  ``levels`` selects the code lattice
+    size (4 = Q2_K 2-bit, 16 = Q4_0 nibble, 2 = Q1_0 bit, 8 = Q3_K tri);
+    ``d_sub`` holds the per-column error step so that any code move ``q``
+    produces the exact error delta ``d_sub * (q_cur - q)`` (formula is
+    shared across formats; Q1_0's sign trick is ``d_sub = 2*d``), and
+    ``m_sub`` is accepted for signature symmetry but unused here.
+    """
+    n_blocks, n = X_f64.shape
+    half = n // 2
+    q_grid = np.arange(levels, dtype=np.int64)
+    w_full = np.ones_like(e) if w_full is None else np.asarray(w_full, dtype=np.float64)
+    w_lo, w_hi = w_full[:, :half], w_full[:, half:]
+    ds_lo, ds_hi = d_sub[:, :half], d_sub[:, half:]
+    active_lo, active_hi = np.abs(ds_lo) >= EPS_D, np.abs(ds_hi) >= EPS_D
+    idx = np.arange(n_blocks)
+
+    for p in range(half):
+        e_lo = e[:, p].copy()
+        e_hi = e[:, half + p].copy()
+        q_lo = codes[:, p].astype(np.int64)
+        q_hi = codes[:, half + p].astype(np.int64)
+        v_p = v[:, p]
+        dc_p = dc
+
+        # candidate error deltas (affine in q): de(q) = d_sub * (q_cur - q)
+        de_lo = ds_lo[:, p : p + 1] * (q_lo[:, None] - q_grid[None, :])
+        de_hi = ds_hi[:, p : p + 1] * (q_hi[:, None] - q_grid[None, :])
+        de_lo = np.where(active_lo[:, p : p + 1], de_lo, 0.0)
+        de_hi = np.where(active_hi[:, p : p + 1], de_hi, 0.0)
+        de_sum = de_lo[:, :, None] + de_hi[:, None, :]  # (n_blocks, L, L)
+
+        d_e = (
+            w_lo[:, p, None, None]
+            * (2.0 * e_lo[:, None, None] * de_lo[:, :, None] + de_lo[:, :, None] ** 2)
+            + w_hi[:, p, None, None]
+            * (2.0 * e_hi[:, None, None] * de_hi[:, None, :] + de_hi[:, None, :] ** 2)
+            + lam_vw * (2.0 * v_p[:, None, None] * de_sum + de_sum * de_sum)
+            + lam_dc * (2.0 * dc_p[:, None, None] * de_sum + de_sum * de_sum)
+        )
+
+        if greedy:
+            flat = d_e.reshape(n_blocks, -1)
+            cur = q_lo * levels + q_hi
+            pick = np.argmin(flat, axis=1)
+            improve = flat[idx, pick] < -1.0e-18
+            pick = np.where(improve, pick, cur)
+        else:
+            cur = q_lo * levels + q_hi
+            scale = np.maximum(temperature, 1.0e-300)
+            logits = -d_e / scale[:, None, None]
+            logits -= logits.max(axis=(1, 2), keepdims=True)
+            probs = np.exp(logits).reshape(n_blocks, -1)
+            # dead pairs (both sub-scales ~0): codes carry no energy, keep current
+            dead = ~(active_lo[:, p] | active_hi[:, p])
+            if np.any(dead):
+                probs[dead] = 0.0
+                probs[dead, cur[dead]] = 1.0
+            cum = np.cumsum(probs, axis=1)
+            cum /= cum[:, -1:]
+            pick = np.sum(cum < rng.random((n_blocks, 1)), axis=1)
+            pick = np.minimum(pick, levels * levels - 1)
+
+        ia, ib = pick // levels, pick % levels
+        upd_lo = np.take_along_axis(de_lo, ia[:, None], axis=1)[:, 0]
+        upd_hi = np.take_along_axis(de_hi, ib[:, None], axis=1)[:, 0]
+        upd_sum = upd_lo + upd_hi
+        sse += (
+            w_lo[:, p] * (2.0 * e_lo * upd_lo + upd_lo * upd_lo)
+            + w_hi[:, p] * (2.0 * e_hi * upd_hi + upd_hi * upd_hi)
+        )
+        e[:, p] = e_lo + upd_lo
+        e[:, half + p] = e_hi + upd_hi
+        v[:, p] = v_p + upd_sum
+        dc = dc + upd_sum
+        codes[:, p] = ia
+        codes[:, half + p] = ib
+
+
+def codes_recompute_for(
+    X_f64: np.ndarray, d: np.ndarray, m: np.ndarray, ls: np.ndarray, lm: np.ndarray
+) -> np.ndarray:
+    """Nearest rounding of codes, matching q2k.codes_from_params semantics."""
+    x32 = X_f64.astype(np.float32)
+    d32 = np.asarray(d, dtype=np.float32)
+    m32 = np.asarray(m, dtype=np.float32)
+    d_sub = np.repeat(d32[:, None] * ls.astype(np.float32), 16, axis=1)
+    m_sub = np.repeat(m32[:, None] * lm.astype(np.float32), 16, axis=1)
+    safe = d_sub >= np.float32(EPS_D)
+    safe_d = np.where(safe, d_sub, np.float32(1.0))
+    q = np.clip(nearest_int((x32 + m_sub) / safe_d), 0, 3)
+    return np.where(safe, q, 0).astype(np.int64)
+
+
+def thermal_refine_codes(
+    X: np.ndarray,
+    w: np.ndarray | None,
+    d: np.ndarray,
+    m: np.ndarray,
+    ls: np.ndarray,
+    lm: np.ndarray,
+    codes: np.ndarray,
+    carry: np.ndarray,
+    cfg: HPCConfig,
+    dm_solve=None,
+    passes: int | None = None,
+    seed: int | None = None,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, int]:
+    """Seeded heat-bath annealing of the code lattice for a batch of blocks.
+
+    Returns ``(d, m, ls, lm, codes, n_improved)``.  Per block the returned
+    state never has a higher extended energy than the input state: every
+    end-of-pass candidate is compared against the running elementwise best
+    and only strict decreases are kept.
+    """
+    n_blocks = X.shape[0]
+    n_passes = cfg.therm_passes if passes is None else int(passes)
+    lam_vw = (float(cfg.vw_lambda) / QK_K) if cfg.use_spectral else 0.0
+    lam_dc = (float(cfg.dc_lambda) / QK_K) if cfg.use_spectral else 0.0
+    if n_blocks == 0 or n_passes == 0 or (lam_vw == 0.0 and lam_dc == 0.0):
+        return d, m, ls, lm, codes, 0
+    if cfg.fold_depth > 1:
+        raise ValueError("thermal refinement supports fold_depth <= 1")
+
+    X_f64 = np.asarray(X, dtype=np.float64)
+    w_f64 = None if w is None else np.asarray(w, dtype=np.float64)
+    carry_f64 = np.asarray(carry, dtype=np.float64)
+    ls = np.asarray(ls, dtype=np.int64)
+    lm = np.asarray(lm, dtype=np.int64)
+    d = np.asarray(d, dtype=np.float64)
+    m = np.asarray(m, dtype=np.float64)
+    codes = np.asarray(codes, dtype=np.int64).copy()
+    rng = np.random.default_rng(cfg.therm_seed if seed is None else seed)
+
+    def energies(dd, mm, cc):
+        d_sub = np.repeat(dd[:, None] * ls, 16, axis=1)
+        m_sub = np.repeat(mm[:, None] * lm, 16, axis=1)
+        e = X_f64 - (d_sub * cc - m_sub)
+        return pair_energy(e, w_f64, lam_vw, lam_dc, carry_f64)[0]
+
+    energy_best = energies(d, m, codes)
+    energy0 = energy_best.copy()
+    best_d, best_m, best_codes = d.copy(), m.copy(), codes.copy()
+    d_cur, m_cur, codes_cur = d.copy(), m.copy(), codes.copy()
+
+    for t in range(n_passes):
+        temperature = cfg.therm_t0 * np.maximum(energy_best, 1.0e-30) / QK_K * (cfg.therm_t_decay**t)
+        for sweep in range(cfg.therm_sweeps + 2):
+            greedy = sweep >= cfg.therm_sweeps
+            d_sub = np.repeat(d_cur[:, None] * ls, 16, axis=1)
+            m_sub = np.repeat(m_cur[:, None] * lm, 16, axis=1)
+            e = X_f64 - (d_sub * codes_cur - m_sub)
+            energy_cur, wesq, v, dc = pair_energy(e, w_f64, lam_vw, lam_dc, carry_f64)
+            sse = energy_cur - lam_vw * np.sum(v * v, axis=1) - lam_dc * dc * dc
+            gibbs_pair_pass(
+                X_f64,
+                w_f64,
+                d_sub,
+                m_sub,
+                codes_cur,
+                e,
+                v,
+                dc,
+                sse,
+                temperature,
+                lam_vw,
+                lam_dc,
+                np.random.default_rng(rng.integers(0, 2**62)),
+                greedy=greedy,
+            )
+            if greedy:
+                energy_now, _, _, _ = pair_energy(
+                    X_f64 - (d_sub * codes_cur - m_sub), w_f64, lam_vw, lam_dc, carry_f64
+                )
+                if not np.any(energy_now < energy_cur - 1.0e-18):
+                    break
+        codes_cur = np.clip(codes_cur, 0, 3)
+        energy_a = energies(d_cur, m_cur, codes_cur)
+        if dm_solve is not None:
+            dm_w = w_f64 if w_f64 is not None else np.ones_like(X_f64)
+            d2, m2 = dm_solve(X_f64, dm_w, ls, lm, codes_cur, d_cur, m_cur, carry_f64, cfg)
+            codes2 = codes_recompute_for(X_f64, d2, m2, ls, lm)
+            energy_b = energies(d2, m2, codes2)
+            take_b = energy_b < energy_a - 1.0e-18
+            d_cur = np.where(take_b, d2, d_cur)
+            m_cur = np.where(take_b, m2, m_cur)
+            codes_cur = np.where(take_b[:, None], codes2, codes_cur)
+        energy_now = energies(d_cur, m_cur, codes_cur)
+        take = energy_now < energy_best - 1.0e-18
+        best_d = np.where(take, d_cur, best_d)
+        best_m = np.where(take, m_cur, best_m)
+        best_codes = np.where(take[:, None], codes_cur, best_codes)
+        energy_best = np.where(take, energy_now, energy_best)
+
+    n_improved = int(np.sum(energy_best < energy0 - 1.0e-18))
+    return best_d, best_m, ls.copy(), lm.copy(), best_codes.astype(np.uint8), n_improved
+
+
+def thermal_refine_tensor(
+    X: np.ndarray,
+    w: np.ndarray | None,
+    d: np.ndarray,
+    m: np.ndarray,
+    ls: np.ndarray,
+    lm: np.ndarray,
+    codes: np.ndarray,
+    carry: np.ndarray,
+    cfg: HPCConfig,
+    dm_solve=None,
+    chunk: int = 2048,
+    progress=None,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, int]:
+    """Chunked wrapper so peak memory stays bounded on large tensors."""
+    n_blocks = X.shape[0]
+    step = max(1, int(chunk))
+    total_improved = 0
+    for start in range(0, n_blocks, step):
+        stop = min(start + step, n_blocks)
+        sel = slice(start, stop)
+        d_c, m_c, ls_c, lm_c, c_c, improved = thermal_refine_codes(
+            X[sel],
+            None if w is None else w[sel],
+            d[sel],
+            m[sel],
+            ls[sel],
+            lm[sel],
+            codes[sel],
+            carry[sel],
+            cfg,
+            dm_solve=dm_solve,
+            seed=cfg.therm_seed ^ (0x9E3779B97F4A7C15 * (start // step) & 0x3FFFFFFFFFFFFFFF),
+        )
+        d[sel], m[sel], ls[sel], lm[sel], codes[sel] = d_c, m_c, ls_c, lm_c, c_c
+        total_improved += improved
+        if progress is not None:
+            progress(stop, n_blocks, improved)
+    return d, m, ls, lm, codes, total_improved
+
+
+def thermal_refine_uniform_codes(
+    X: np.ndarray,
+    w: np.ndarray | None,
+    codes: np.ndarray,
+    step: np.ndarray,
+    m_const: np.ndarray,
+    carry: np.ndarray,
+    cfg: HPCConfig,
+    lam_vw: float,
+    lam_dc: float,
+    levels: int,
+    passes: int | None = None,
+    seed: int | None = None,
+) -> tuple[np.ndarray, int]:
+    """Heat-bath annealing of a level-index code lattice with FIXED step scales.
+
+    Uniform variant of ``thermal_refine_codes`` for the format quantizers
+    (Q4_0 / Q1_0 / Q3_K): the per-code error step ``step`` (shape (B, n)) and
+    the affine constant ``m_const`` (shape (B, n)) are fixed inputs — the
+    dequantized value of level index ``q`` is ``step*q - m_const`` so the
+    error is ``e = X - (step*q - m_const)`` and every move ``q -> q'`` has
+    the exact closed-form delta again affine in the level grid.  ``codes``
+    are integer level indices in ``[0, levels)``.
+
+    Returns ``(codes, n_improved)``.  Per block the returned state never has
+    a higher extended energy than the input state and its SSE never exceeds
+    the input SSE by more than ``cfg.sse_budget`` (extended-nonincreasing
+    with an SSE guard, matching the pipeline's shaper semantics).
+    """
+    n_blocks, n = X.shape
+    if n_blocks == 0:
+        return codes.astype(np.int64).copy(), 0
+    if n % 2 != 0:
+        raise ValueError("thermal refinement needs an even row width")
+    n_passes = cfg.therm_passes if passes is None else int(passes)
+    if n_passes == 0 or (lam_vw == 0.0 and lam_dc == 0.0):
+        return codes.astype(np.int64).copy(), 0
+    if cfg.fold_depth > 1:
+        raise ValueError("thermal refinement supports fold_depth <= 1")
+
+    X_f64 = np.asarray(X, dtype=np.float64)
+    w_f64 = None if w is None else np.asarray(w, dtype=np.float64)
+    carry_f64 = np.asarray(carry, dtype=np.float64)
+    step = np.asarray(step, dtype=np.float64)
+    m_b = np.asarray(m_const, dtype=np.float64)
+    codes = np.asarray(codes, dtype=np.int64).copy()
+    rng = np.random.default_rng(cfg.therm_seed if seed is None else seed)
+
+    def energies(cc):
+        e = X_f64 - (step * cc - m_b)
+        return pair_energy(e, w_f64, lam_vw, lam_dc, carry_f64)
+
+    energy_best, we_sq, v0, dc0 = energies(codes)
+    energy0 = energy_best.copy()
+    sse0 = np.sum(we_sq, axis=1)
+    sse_cap = sse0 * (1.0 + float(cfg.sse_budget))
+    best_codes = codes.copy()
+    codes_cur = codes.copy()
+
+    for t in range(n_passes):
+        temperature = cfg.therm_t0 * np.maximum(energy_best, 1.0e-30) / n * (cfg.therm_t_decay**t)
+        for sweep in range(cfg.therm_sweeps + 2):
+            greedy = sweep >= cfg.therm_sweeps
+            e = X_f64 - (step * codes_cur - m_b)
+            energy_cur, we_sq, v, dc = pair_energy(e, w_f64, lam_vw, lam_dc, carry_f64)
+            sse = np.sum(we_sq, axis=1)
+            gibbs_pair_pass(
+                X_f64,
+                w_f64,
+                step,
+                m_b,
+                codes_cur,
+                e,
+                v,
+                dc,
+                sse,
+                temperature,
+                lam_vw,
+                lam_dc,
+                np.random.default_rng(rng.integers(0, 2**62)),
+                greedy=greedy,
+                levels=levels,
+            )
+            if greedy:
+                e = X_f64 - (step * codes_cur - m_b)
+                energy_now, _, _, _ = pair_energy(e, w_f64, lam_vw, lam_dc, carry_f64)
+                if not np.any(energy_now < energy_cur - 1.0e-18):
+                    break
+        energy_now, we_sq, v_now, dc_now = energies(codes_cur)
+        sse_now = np.sum(we_sq, axis=1)
+        take = (energy_now < energy_best - 1.0e-18) & (sse_now <= sse_cap)
+        best_codes = np.where(take[:, None], codes_cur, best_codes)
+        energy_best = np.where(take, energy_now, energy_best)
+        codes_cur = best_codes.copy()
+
+    n_improved = int(np.sum(energy_best < energy0 - 1.0e-18))
+    return best_codes, n_improved
+
+
+def _refine_uniform_driver(
+    X: np.ndarray,
+    w: np.ndarray | None,
+    codes: np.ndarray,
+    build_state,
+    carry: np.ndarray | None,
+    lam_vw: float,
+    lam_dc: float,
+    levels: int,
+    cfg: HPCConfig,
+    chunk: int = 2048,
+    progress=None,
+) -> tuple[np.ndarray, int]:
+    """Chunked uniform annealing so peak memory stays bounded on big tensors."""
+    n_blocks = X.shape[0]
+    codes_out = np.asarray(codes, dtype=np.int64).copy()
+    step = max(1, int(chunk))
+    total_improved = 0
+    for start in range(0, n_blocks, step):
+        stop = min(start + step, n_blocks)
+        sel = slice(start, stop)
+        step_c, m_c = build_state(start, stop)
+        carry_c = np.zeros(stop - start) if carry is None else carry[sel]
+        c_c, improved = thermal_refine_uniform_codes(
+            X[sel],
+            None if w is None else w[sel],
+            codes_out[sel],
+            step_c,
+            m_c,
+            carry_c,
+            cfg,
+            lam_vw,
+            lam_dc,
+            levels,
+            seed=cfg.therm_seed ^ (0x9E3779B97F4A7C15 * (start // step) & 0x3FFFFFFFFFFFFFFF),
+        )
+        codes_out[sel] = c_c
+        total_improved += improved
+        if progress is not None:
+            progress(stop, n_blocks, improved)
+    return codes_out, total_improved
+
+
+def thermal_refine_q4_0(
+    X: np.ndarray,
+    W: np.ndarray | None,
+    d: np.ndarray,
+    q: np.ndarray,
+    cfg: HPCConfig,
+    chunk: int = 2048,
+    progress=None,
+) -> tuple[np.ndarray, int]:
+    """Anneal Q4_0 nibble codes against fixed fp16 super-block scales.
+
+    ``X``: raw block targets (B, 32); ``W``: importance weights (B, 32) or
+    None; ``d``: fp16-precision block scales (B,), possibly negative
+    (sign-flip polish outputs); ``q``: nibble codes (B, 32) with
+    dequantization ``d*(q-8)``.  The Gibbs step grid is the signed slope
+    ``d`` itself, so negatives fall out naturally with no code remapping.
+    Returns refined uint8 codes and the number of strictly improved blocks;
+    a no-op when the refinery or spectrals are disabled.
+    """
+    n_blocks, n = X.shape
+    lam_vw = (float(cfg.vw_lambda) / n) if cfg.use_spectral else 0.0
+    lam_dc = (float(cfg.dc_lambda) / n) if cfg.use_spectral else 0.0
+    if n_blocks == 0 or cfg.therm_passes == 0 or (lam_vw == 0.0 and lam_dc == 0.0):
+        return q.copy(), 0
+    d64 = np.asarray(d, dtype=np.float64)
+
+    def build_state(start, stop):
+        ds = np.repeat(d64[start:stop][:, None], n, axis=1)
+        return ds, 8.0 * ds
+
+    codes_out, improved = _refine_uniform_driver(
+        X, W, np.asarray(q, dtype=np.int64), build_state, None,
+        lam_vw, lam_dc, 16, cfg, chunk=chunk, progress=progress,
+    )
+    return codes_out.astype(np.uint8), improved
+
+
+def thermal_refine_q1_0(
+    X: np.ndarray,
+    W: np.ndarray | None,
+    d: np.ndarray,
+    bits: np.ndarray,
+    cfg: HPCConfig,
+    chunk: int = 2048,
+    progress=None,
+) -> tuple[np.ndarray, int]:
+    """Anneal Q1_0 bits against fixed fp16 super-block magnitudes.
+
+    ``bits``: binary codes (B, 32) with dequantization ``d*(2b-1)``; the
+    Gibbs step grid is ``2*d`` with affine constant ``d``.  Returns refined
+    uint8 bits and the number of strictly improved blocks.
+    """
+    n_blocks, n = X.shape
+    lam_vw = (float(cfg.vw_lambda) / n) if cfg.use_spectral else 0.0
+    lam_dc = (float(cfg.dc_lambda) / n) if cfg.use_spectral else 0.0
+    if n_blocks == 0 or cfg.therm_passes == 0 or (lam_vw == 0.0 and lam_dc == 0.0):
+        return bits.copy(), 0
+    d64 = np.asarray(d, dtype=np.float64)
+
+    def build_state(start, stop):
+        ds = np.repeat(2.0 * np.abs(d64[start:stop])[:, None], n, axis=1)
+        return ds, np.repeat(np.abs(d64[start:stop])[:, None], n, axis=1)
+
+    codes_out, improved = _refine_uniform_driver(
+        X, W, np.asarray(bits, dtype=np.int64), build_state, None,
+        lam_vw, lam_dc, 2, cfg, chunk=chunk, progress=progress,
+    )
+    return codes_out.astype(np.uint8), improved
+
+
+def thermal_refine_q3_k(
+    X: np.ndarray,
+    W: np.ndarray | None,
+    scale: np.ndarray,
+    codes: np.ndarray,
+    cfg: HPCConfig,
+    chunk: int = 2048,
+    progress=None,
+) -> tuple[np.ndarray, int]:
+    """Anneal Q3_K signed 3-bit codes against fixed per-superblock scales.
+
+    ``scale``: per-superblock scale values (B, 16) broadcast over the 16
+    codes each; ``codes``: source-signed values (B, 256) in [-4, 3].  The
+    Gibbs grid covers all eight signed levels; negative scales (sign-flip
+    polish outputs) are handled exactly via signed steps.
+    """
+    n_blocks, n = X.shape
+    lam_vw = (float(cfg.vw_lambda) / n) if cfg.use_spectral else 0.0
+    lam_dc = (float(cfg.dc_lambda) / n) if cfg.use_spectral else 0.0
+    if n_blocks == 0 or cfg.therm_passes == 0 or (lam_vw == 0.0 and lam_dc == 0.0):
+        return np.asarray(codes).copy(), 0
+    scale64 = np.asarray(scale, dtype=np.float64)
+
+    def build_state(start, stop):
+        ds = np.repeat(scale64[start:stop], 16, axis=1)
+        return ds, 4.0 * ds
+
+    codes_out, improved = _refine_uniform_driver(
+        X, W, np.asarray(codes, dtype=np.int64) + 4, build_state, None,
+        lam_vw, lam_dc, 8, cfg, chunk=chunk, progress=progress,
+    )
+    return codes_out - 4, improved

config.py

@@ -74,6 +74,16 @@ class HPCConfig:
     quantize_embd_q4: bool = True
     promote_attn_q4: bool = False
 
+    # ---- thermal refinement (pair heat-bath annealer, opt-in) ------------
+    # V2-only knob (0 = off, bit-identical to the deterministic pipeline).
+    # Runs a seeded checkerboard Gibbs sampler over the vesica-pair lattice
+    # after Phase 4; late passes are monotone so output never regresses.
+    therm_passes: int = 0
+    therm_sweeps: int = 2
+    therm_t0: float = 0.25
+    therm_t_decay: float = 0.35
+    therm_seed: int = 271828
+
     # ---- seed clamp -----------------------------------------------------
     clamp_kurtosis_off: float = 6.0
     clamp_kurtosis_soft: float = 20.0
@@ -119,6 +129,14 @@ class HPCConfig:
             raise ValueError("Viterbi weights must be non-negative")
         if self.workers < 0:
             raise ValueError("workers must be >= 0 (0 = auto CPU count)")
+        if self.therm_passes < 0:
+            raise ValueError("therm_passes must be >= 0 (0 disables thermal refinement)")
+        if self.therm_sweeps < 1:
+            raise ValueError("therm_sweeps must be >= 1")
+        if self.therm_t0 <= 0.0:
+            raise ValueError("therm_t0 must be positive")
+        if not 0.0 < self.therm_t_decay < 1.0:
+            raise ValueError("therm_t_decay must lie in (0, 1)")
         if self.device not in ("auto", "cuda", "cpu") and not self.device.startswith("cuda:"):
             raise ValueError("device must be 'auto', 'cuda', 'cuda:N', or 'cpu'")
         if self.quant_format not in ("mixed", "mixed_q1", "custom", "q1_0", "q3_k", "q4_0", "q2_k"):

gui.py

@@ -181,6 +181,7 @@ class HPCQuantizeGUI:
         self.vw_lambda_var = tk.DoubleVar(value=1.0)
         self.dc_decay_var = tk.DoubleVar(value=0.85)
         self.sse_budget_var = tk.DoubleVar(value=0.0005)
+        self.thermal_var = tk.IntVar(value=0)
 
         self.cli_preview_var = tk.StringVar()
         self.status_var = tk.StringVar(value="Ready.")
@@ -213,6 +214,7 @@ class HPCQuantizeGUI:
             self.embd_q4_var,
             self.promote_attn_q4_var,
             self.quant_format_var,
+            self.thermal_var,
         ):
             v.trace_add("write", lambda *_: self._update_cli_preview())
 
@@ -468,6 +470,10 @@ class HPCQuantizeGUI:
         ttk.Entry(spec_f, textvariable=self.dc_decay_var, width=8).grid(row=1, column=1, sticky=tk.EW, padx=4, pady=2)
         ttk.Label(spec_f, text="SSE Budget:").grid(row=1, column=2, sticky=tk.W, padx=4, pady=2)
         ttk.Entry(spec_f, textvariable=self.sse_budget_var, width=8).grid(row=1, column=3, sticky=tk.EW, padx=4, pady=2)
+        ttk.Label(spec_f, text="Thermal Passes (0=off):").grid(row=2, column=0, sticky=tk.W, padx=4, pady=2)
+        ttk.Spinbox(
+            spec_f, from_=0, to=16, increment=1, textvariable=self.thermal_var, width=8
+        ).grid(row=2, column=1, sticky=tk.EW, padx=4, pady=2)
 
         # Right: GGUF Tensor Browser
         browser_frame = ttk.LabelFrame(right_pane, text="3. GGUF Tensor Inspector & Live Status", padding=8)
@@ -794,6 +800,7 @@ class HPCQuantizeGUI:
             vw_lambda=float(self.vw_lambda_var.get()),
             dc_decay=float(self.dc_decay_var.get()),
             sse_budget=float(self.sse_budget_var.get()),
+            therm_passes=max(0, int(self.thermal_var.get())),
             quant_format=self.quant_format_var.get(),
             custom_quant_map=self._custom_map_for_plan() or {},
             enable_q4_fallback=bool(self.q4_fallback_var.get()) and self.quant_format_var.get() not in ("q2_k", "q3_k"),
@@ -841,6 +848,8 @@ class HPCQuantizeGUI:
             cmd.append("--keep-embd")
         if self.promote_attn_q4_var.get():
             cmd.append("--promote-attn-q4")
+        if self.thermal_var.get() > 0:
+            cmd += ["--thermal", str(int(self.thermal_var.get()))]
         cmd.append("--verbose")
         return " ".join(shlex.quote(part) for part in cmd)
 

pipeline.py

@@ -23,6 +23,7 @@ import time
 
 import numpy as np
 
+from .anneal import thermal_refine_tensor
 from .config import HPCConfig
 from .factored import (
     CAND_TO_QUHIT,
@@ -69,6 +70,7 @@ class TensorStats:
     temperature: float
     factor_floor_used: int
     candidate_floor_used: int
+    therm_improved: int = 0
 
 
 @dataclass(slots=True)
@@ -1070,6 +1072,72 @@ class Q2KQuantizer:
                 eta=eta,
             )
 
+        therm_improved = 0
+        if cfg.therm_passes > 0 and n_blocks > 0:
+            _check_stop()
+            t_therm = time.time()
+            _emit(
+                "Phase 4.5: Thermal Refinery (Pair Heat-Bath)",
+                0,
+                n_blocks,
+                f"annealing {n_blocks} blocks: passes={cfg.therm_passes} "
+                f"sweeps={cfg.therm_sweeps} T0={cfg.therm_t0} decay={cfg.therm_t_decay} "
+                f"seed={cfg.therm_seed}",
+                sigma=sigma,
+                temperature=temperature,
+            )
+            therm_d = np.empty(n_blocks, dtype=np.float64)
+            therm_m = np.empty(n_blocks, dtype=np.float64)
+            therm_ls = np.empty((n_blocks, NS), dtype=np.uint8)
+            therm_lm = np.empty((n_blocks, NS), dtype=np.uint8)
+            therm_codes = np.empty((n_blocks, QK_K), dtype=np.uint8)
+            for i, b in enumerate(blocks):
+                therm_d[i] = b.d_float
+                therm_m[i] = b.dmin_float
+                therm_ls[i], therm_lm[i] = b.unpack_scales()
+                therm_codes[i] = unpack_codes(b.qs)
+            last_therm = [0.0, 0]
+
+            def _therm_progress(done, total, improved):
+                now = time.time()
+                if (now - last_therm[0]) >= 0.25 or done == total:
+                    last_therm[0] = now
+                    _emit(
+                        "Phase 4.5: Thermal Refinery (Pair Heat-Bath)",
+                        done,
+                        total,
+                        f"blk {done}/{total} ({max(done, 1)/max(now - t_therm, 1.0e-6):.1f} blk/s) "
+                        f"improved={improved}",
+                        sigma=sigma,
+                        temperature=temperature,
+                    )
+
+            therm_d, therm_m, therm_ls, therm_lm, therm_codes, therm_improved = thermal_refine_tensor(
+                X,
+                None if W is None else W,
+                therm_d,
+                therm_m,
+                therm_ls,
+                therm_lm,
+                therm_codes,
+                carry,
+                cfg,
+                dm_solve=_batch_extended_dm_solve_and_ulp,
+                progress=_therm_progress,
+            )
+            blocks = [
+                make_block(therm_d[i], therm_m[i], therm_ls[i], therm_lm[i], therm_codes[i])
+                for i in range(n_blocks)
+            ]
+            _emit(
+                "Phase 4.5: Thermal Refinery (Pair Heat-Bath)",
+                n_blocks,
+                n_blocks,
+                f"done in {time.time() - t_therm:.1f}s, blocks improved={therm_improved}/{n_blocks}",
+                sigma=sigma,
+                temperature=temperature,
+            )
+
         if cfg.apply_dc_carry and (cfg.dc_lambda > 0.0 or cfg.dc_decay > 0.0):
             _check_stop()
             _emit(
@@ -1103,6 +1171,7 @@ class Q2KQuantizer:
             temperature=temperature,
             factor_floor_used=factor_floor_used,
             candidate_floor_used=candidate_floor_used,
+            therm_improved=therm_improved,
         )
         _emit(
             "Phase 5: Tensor Complete",
@@ -1116,6 +1185,7 @@ class Q2KQuantizer:
             temperature=temperature,
             cand_floors=candidate_floor_used,
             fac_floors=factor_floor_used,
+            therm_improved=therm_improved,
         )
         scales, qs, dm = _stack_blocks(blocks)
         factor_objectives = np.array([r.objective for r in standalone], dtype=np.float64)
@@ -1148,8 +1218,9 @@ def _stack_blocks(blocks):
 
 
 def _report(stats: TensorStats) -> None:
+    extra = f" therm={stats.therm_improved}" if stats.therm_improved else ""
     print(
         f"Q2_K V2: blocks={stats.n_blocks} rmse={stats.rmse:.6e} "
         f"rmse/sigma={stats.rmse_over_sigma:.5f} temp={stats.temperature:.3e} "
-        f"floors(cand={stats.candidate_floor_used}, factor={stats.factor_floor_used})"
+        f"floors(cand={stats.candidate_floor_used}, factor={stats.factor_floor_used}){extra}"
     )

q1_0.py

@@ -18,6 +18,10 @@ per 18-byte block = 1.125 bpw, matching ``block_q1_0`` in ``llama.cpp``):
    Polishes ``d`` over ``+/-8`` fp16 ULPs of both ``+d`` and ``-d``, and
    applies closed-form 1-D quadratic Vesica + DC bit-flip projection bounded by
    ``sse_budget`` plus a strict monotone candidate floor.
+5. **Optional Thermal Refinery (Pair Heat-Bath, opt-in)**:
+   With ``cfg.therm_passes > 0`` the finished bits are re-swept by the seeded
+   pair heat-bath annealer (Anneal) on the 2-level Q1_0 lattice in the same
+   companded target space, non-increasing under the ``sse_budget`` cap.
 """
 
 from __future__ import annotations
@@ -30,6 +34,7 @@ try:
 except ImportError:  # pragma: no cover
     torch = None  # type: ignore[assignment]
 
+from .anneal import thermal_refine_q1_0
 from .config import HPCConfig
 from .cuda_backend import cuda_is_available
 from .factored import MULTS_D
@@ -251,6 +256,13 @@ def _quantize_q1_0_hpc_cpu(
         accept = (sse_new <= sse_cap) & ((sse_new + spec_new) < (sse_live + spec_old) - 1.0e-15)
         best_q = np.where(accept[:, None], q_new, best_q)
 
+    # 5. Optional Thermal Refinery (Pair Heat-Bath, non-regressing):
+    if cfg.therm_passes > 0:
+        thermal_q, _ = thermal_refine_q1_0(
+            X_target.astype(np.float32), W_f32, best_d, best_q, cfg
+        )
+        best_q = thermal_q.astype(np.int64)
+
     best_q_u8 = best_q.astype(np.uint8)
     out = np.empty((B, Q1_0_BLOCK_BYTES), dtype=np.uint8)
     out[:, :2] = best_d16.astype("<u2").view(np.uint8).reshape(B, 2)
@@ -375,6 +387,17 @@ def _quantize_q1_0_hpc_cuda(
             accept = (sse_new <= sse_cap) & ((sse_new + spec_new) < (sse_live + spec_old) - 1.0e-12)
             best_q = torch.where(accept[:, None], q_new, best_q)
 
+        # 5. Optional Thermal Refinery (Pair Heat-Bath, non-regressing):
+        if cfg.therm_passes > 0:
+            thermal_q, _ = thermal_refine_q1_0(
+                X_target.cpu().numpy(),
+                None if W_np is None else W.cpu().numpy(),
+                best_d.cpu().numpy(),
+                best_q.cpu().numpy(),
+                cfg,
+            )
+            best_q = torch.from_numpy(thermal_q.astype(np.int64)).to(dev)
+
         best_q_u8 = best_q.to(torch.uint8)
         deq_fin = best_d[:, None] * (2.0 * best_q_u8.to(torch.float32) - 1.0)
         raw_sse = float(torch.sum((X - deq_fin) ** 2, dtype=torch.float64).item())

q3_k.py

@@ -48,6 +48,7 @@ try:
 except ImportError:  # pragma: no cover
     torch = None  # type: ignore[assignment]
 
+from .anneal import thermal_refine_q3_k
 from .config import HPCConfig
 from .cuda_backend import cuda_is_available
 from .factored import MULTS_D
@@ -646,6 +647,12 @@ def _fit_q3_k_blocks(
         accept = (sse_new <= sse_cap) & ((sse_new + spec_new) < (sse_live + spec_old) - 1.0e-15)
         code = np.where(accept[:, None], code_new, code_f).reshape(B, Q3_K_N_SUBBLOCKS, Q3_K_SUBBLOCK)
 
+    # 5. Optional Thermal Refinery (Pair Heat-Bath, non-regressing):
+    if cfg.therm_passes > 0:
+        scale_rp = best_d[:, None] * sc.astype(np.float32)
+        thermal_c, _ = thermal_refine_q3_k(Xf, Wf, scale_rp, code.reshape(B, QK3_K), cfg)
+        code = thermal_c.astype(np.int16).reshape(B, Q3_K_N_SUBBLOCKS, Q3_K_SUBBLOCK)
+
     return best_d16.astype(np.uint16), sc.astype(np.int16), code.astype(np.int16)
 
 
@@ -806,6 +813,20 @@ def _fit_q3_k_blocks_cuda(
             accept = (sse_new <= sse_cap) & ((sse_new + spec_new) < (sse_live + spec_old) - 1.0e-12)
             code = torch.where(accept[:, None], code_new, code_f).reshape(B, Q3_K_N_SUBBLOCKS, Q3_K_SUBBLOCK)
 
+        # 5. Optional Thermal Refinery (Pair Heat-Bath, non-regressing):
+        if cfg.therm_passes > 0:
+            scale_rp = (best_d[:, None] * sc.to(torch.float32)).cpu().numpy()
+            thermal_c, _ = thermal_refine_q3_k(
+                Xf.cpu().numpy(),
+                Wf.cpu().numpy(),
+                scale_rp,
+                code.reshape(B, QK3_K).cpu().numpy(),
+                cfg,
+            )
+            code = torch.from_numpy(
+                thermal_c.astype(np.int16).reshape(B, Q3_K_N_SUBBLOCKS, Q3_K_SUBBLOCK)
+            ).to(dev)
+
         return (
             best_d16.cpu().numpy().astype(np.uint16),
             sc.cpu().numpy().astype(np.int16),

q4_0.py

@@ -16,9 +16,13 @@ per 18-byte block = 4.5 bpw):
    Viterbi dynamic program with free-energy bin reduction and Manhattan
    transition penalty across the block chain.
 4. **Asymmetric Sign-Flip fp16-ULP Polish & Closed-Form Vesica/DC Shaping**:
-   Polishes ``d`` over ``+/-8`` fp16 ULPs of both ``+d`` and ``-d``, and
-   applies closed-form 1-D quadratic Vesica + DC code projection bounded by
-   ``sse_budget`` plus a strict monotone candidate floor.
+    Polishes ``d`` over ``+/-8`` fp16 ULPs of both ``+d`` and ``-d``, and
+    applies closed-form 1-D quadratic Vesica + DC code projection bounded by
+    ``sse_budget`` plus a strict monotone candidate floor.
+5. **Optional Thermal Refinery (Pair Heat-Bath, opt-in)**:
+    With ``cfg.therm_passes > 0`` the finished nibble codes are re-swept by
+    the seeded pair heat-bath annealer (Anneal) on the 16-level Q4_0 lattice
+    with extended-nonincreasing acceptance under the ``sse_budget`` cap.
 """
 
 from __future__ import annotations
@@ -31,6 +35,7 @@ try:
 except ImportError:  # pragma: no cover
     torch = None  # type: ignore[assignment]
 
+from .anneal import thermal_refine_q4_0
 from .config import HPCConfig
 from .cuda_backend import cuda_is_available
 from .factored import CAND_TO_QUHIT, MULTS_D
@@ -248,6 +253,11 @@ def _quantize_q4_0_hpc_cpu(
         accept = (sse_new <= sse_cap) & ((sse_new + spec_new) < (sse_live + spec_old) - 1.0e-15)
         best_q = np.where(accept[:, None], q_new, best_q)
 
+    # 5. Optional Thermal Refinery (Pair Heat-Bath, non-regressing):
+    if cfg.therm_passes > 0:
+        thermal_q, _ = thermal_refine_q4_0(X_f32, W_f32, best_d, best_q, cfg)
+        best_q = thermal_q.astype(np.int64)
+
     best_q_u8 = best_q.astype(np.uint8)
     out = np.empty((B, Q4_0_BLOCK_BYTES), dtype=np.uint8)
     out[:, :2] = best_d16.astype("<u2").view(np.uint8).reshape(B, 2)
@@ -380,6 +390,17 @@ def _quantize_q4_0_hpc_cuda(
             accept = (sse_new <= sse_cap) & ((sse_new + spec_new) < (sse_live + spec_old) - 1.0e-12)
             best_q = torch.where(accept[:, None], q_new, best_q)
 
+        # 5. Optional Thermal Refinery (Pair Heat-Bath, non-regressing):
+        if cfg.therm_passes > 0:
+            thermal_q, _ = thermal_refine_q4_0(
+                X.cpu().numpy(),
+                None if W_np is None else W.cpu().numpy(),
+                best_d.cpu().numpy(),
+                best_q.cpu().numpy(),
+                cfg,
+            )
+            best_q = torch.from_numpy(thermal_q.astype(np.int64)).to(dev)
+
         best_q_u8 = best_q.to(torch.uint8)
         deq_fin = best_d[:, None] * (best_q_u8.to(torch.float32) - 8.0)
         raw_sse = float(torch.sum((X - deq_fin) ** 2, dtype=torch.float64).item())