Upload 6 files
49c0661 verified 3 days ago
29.4 kB
"""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 thermal_spectral_lams(cfg: HPCConfig, n: int) -> tuple[float, float]:
    """Per-code-normalized ``(lam_vw, lam_dc)`` with the zero-residual tilt applied.
    ``therm_lambda_tilt`` scales both spectral lambdas during annealing so the
    exact Gibbs chain trades weighted SSE for coherent-residual (DC/vesica)
    drain; ``1.0`` (default) is the shipping objective.
    """
    if not cfg.use_spectral:
        return 0.0, 0.0
    tilt = max(1.0, float(cfg.therm_lambda_tilt))
    return (float(cfg.vw_lambda) / n) * tilt, (float(cfg.dc_lambda) / n) * tilt


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, lam_dc = thermal_spectral_lams(cfg, QK_K)
    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, lam_dc = thermal_spectral_lams(cfg, n)
    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, lam_dc = thermal_spectral_lams(cfg, n)
    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, lam_dc = thermal_spectral_lams(cfg, n)
    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


def thermal_refine_uniform_codes_cuda(
    X: "torch.Tensor",
    w: "torch.Tensor | None",
    codes: "torch.Tensor",
    step: "torch.Tensor",
    m_const: "torch.Tensor",
    carry: "torch.Tensor | None",
    lam_vw: float,
    lam_dc: float,
    levels: int,
    cfg: HPCConfig,
    passes: int | None = None,
    seed: int | None = None,
) -> tuple["torch.Tensor", int]:
    """Torch/CUDA sibling of ``thermal_refine_uniform_codes``.
    Identical algorithm and acceptance structure (extended-nonincreasing
    best-state with an ``sse_budget`` guard, seeded heat-bath sweeps plus a
    greedy endgame), run on-device in float64 so results are deterministic
    per device class (bit-exact across runs at the same seed, but not
    bit-identical to the numpy CPU chain — floating-point ordering differs).
    All tensor inputs live on one CUDA device: ``codes`` int64 (B, n),
    the rest float32/float64.  Returns ``(codes, n_improved)`` with
    ``codes`` still on-device.
    """
    import torch  # local import keeps anneal.py torch-free for CPU installs

    dev = X.device
    f64 = torch.float64
    B, n = X.shape
    half = n // 2
    n_passes = cfg.therm_passes if passes is None else int(passes)
    if B == 0 or n % 2 != 0:
        return codes.to(torch.int64), 0
    if n_passes == 0 or (lam_vw == 0.0 and lam_dc == 0.0):
        return codes.to(torch.int64).clone(), 0
    if cfg.fold_depth > 1:
        raise ValueError("thermal refinement supports fold_depth <= 1")

    Xf = X.to(f64)
    wf = None if w is None else w.to(f64)
    cf = None if carry is None else carry.to(f64)
    if cf is None:
        cf = torch.zeros(B, dtype=f64, device=dev)
    sf = step.to(f64)
    mf_data = m_const.to(f64)
    cc = codes.to(torch.int64).clone()
    grid = torch.arange(levels, device=dev, dtype=torch.int64)
    ones_col = torch.ones(B, dtype=f64, device=dev)

    def energies(c):
        e = Xf - (sf * c.to(f64) - mf_data)
        w2 = torch.sum(e * e, dim=1) if wf is None else torch.sum(wf * e * e, dim=1)
        v = e[:, :half] + e[:, half:]
        dc = torch.sum(e, dim=1) + cf
        return (w2 + lam_vw * torch.sum(v * v, dim=1) + lam_dc * dc * dc), e, w2, v, dc

    energy_best, e, sse0, v, dc = energies(cc)
    energy0 = energy_best.clone()
    sse_cap = sse0 * (1.0 + float(cfg.sse_budget))
    best = cc.clone()
    gen = torch.Generator(device=dev)
    gen.manual_seed(int(cfg.therm_seed if seed is None else seed) & 0x7FFFFFFFFFFFFFFF)

    for t in range(n_passes):
        temperature = (
            cfg.therm_t0
            * torch.maximum(energy_best, torch.full_like(energy_best, 1.0e-30))
            / n
            * (cfg.therm_t_decay ** t)
        )
        broke = False
        for sweep in range(cfg.therm_sweeps + 2):
            greedy = sweep >= cfg.therm_sweeps
            energy_cur, e, _, v, dc = energies(cc)
            for p in range(half):
                e_lo = e[:, p]
                e_hi = e[:, half + p]
                q_lo = cc[:, p]
                q_hi = cc[:, half + p]
                ds_lo = sf[:, p]
                ds_hi = sf[:, half + p]
                act_lo = ds_lo.abs() >= EPS_D
                act_hi = ds_hi.abs() >= EPS_D
                w_lo = ones_col if wf is None else wf[:, p]
                w_hi = ones_col if wf is None else wf[:, half + p]
                de_lo = ds_lo[:, None] * (q_lo[:, None] - grid[None, :])
                de_lo = torch.where(act_lo[:, None], de_lo, torch.zeros_like(de_lo))
                de_hi = ds_hi[:, None] * (q_hi[:, None] - grid[None, :])
                de_hi = torch.where(act_hi[:, None], de_hi, torch.zeros_like(de_hi))
                de_sum = de_lo[:, :, None] + de_hi[:, None, :]
                v_p = v[:, p]
                d_e = (
                    w_lo[:, None, None] * (2.0 * e_lo[:, None, None] * de_lo[:, :, None] + de_lo[:, :, None] ** 2)
                    + w_hi[:, 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[:, None, None] * de_sum + de_sum * de_sum)
                )
                cur = q_lo * levels + q_hi
                if greedy:
                    flat = d_e.reshape(B, -1)
                    pick = flat.argmin(dim=1)
                    improve = torch.gather(flat, 1, pick[:, None]).squeeze(1) < -1.0e-18
                    pick = torch.where(improve, pick, cur)
                else:
                    scc = torch.clamp(temperature, min=1.0e-300)[:, None, None]
                    logits = -d_e / scc
                    logits = logits - logits.amax(dim=(1, 2), keepdim=True)
                    probs = logits.exp().reshape(B, -1)
                    dead = ~(act_lo | act_hi)
                    if bool(dead.any()):
                        probs[dead] = 0.0
                        probs[dead, cur[dead]] = 1.0
                    cum = probs.cumsum(dim=1)
                    cum = cum / cum[:, -1:]
                    u = torch.rand(B, 1, generator=gen, device=dev, dtype=f64)
                    pick = (cum < u).sum(dim=1).clamp(max=levels * levels - 1)
                ia = pick // levels
                ib = pick - ia * levels
                upd_lo = torch.gather(de_lo, 1, ia[:, None]).squeeze(1)
                upd_hi = torch.gather(de_hi, 1, ib[:, None]).squeeze(1)
                usum = upd_lo + upd_hi
                e[:, p] = e_lo + upd_lo
                e[:, half + p] = e_hi + upd_hi
                v[:, p] += usum
                dc = dc + usum
                cc[:, p] = ia
                cc[:, half + p] = ib
            if greedy:
                energy_now, _, _, _, _ = energies(cc)
                if not bool(torch.any(energy_now < energy_cur - 1.0e-18)):
                    broke = True
                    break
        _, _, sse_now, _, _ = energies(cc)
        take = (energy_now < energy_best - 1.0e-18) & (sse_now <= sse_cap)
        best = torch.where(take[:, None], cc, best)
        energy_best = torch.where(take, energy_now, energy_best)
        cc = best.clone()
        if broke and bool(torch.all(energy_now >= energy_best)):
            pass  # keep exact CPU-loop semantics: continue across passes regardless

    n_improved = int((energy_best < energy0 - 1.0e-18).sum().item())
    return best, n_improved


def thermal_refine_q3_k_cuda(
    X: "torch.Tensor",
    W: "torch.Tensor | None",
    scale: "torch.Tensor",
    codes: "torch.Tensor",
    cfg: HPCConfig,
    seed: int | None = None,
) -> tuple["torch.Tensor", int]:
    """On-device anneal for Q3_K signed 3-bit codes (see the CPU twin).
    ``scale``: per-superblock scale values (B, 16) on-device, ``codes``:
    signed values (B, 256) in [-4, 3].  Exactly mirrors
    ``thermal_refine_q3_k`` without any host round-trip.
    """
    import torch
    n_blocks, n = X.shape
    lam_vw, lam_dc = thermal_spectral_lams(cfg, n)
    if n_blocks == 0 or cfg.therm_passes == 0 or (lam_vw == 0.0 and lam_dc == 0.0):
        return codes.to(torch.int64).clone(), 0

    def build_state(start, stop):
        ds_full = scale.repeat_interleave(16, dim=1)   # (B, 256)
        return ds_full, 4.0 * ds_full

    ds_full, m_full = build_state(0, n_blocks)
    codes_out, improved = thermal_refine_uniform_codes_cuda(
        X, W, codes.to(torch.int64) + 4, ds_full, m_full, None,
        lam_vw, lam_dc, 8, cfg, seed=seed,
    )
    return codes_out - 4, improved