Upload 5 files

anneal.py

@@ -563,3 +563,180 @@ def thermal_refine_q3_k(
         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,
+) -> 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,
+    )
+    return codes_out - 4, improved

q3_k.py

@@ -48,7 +48,7 @@ try:
 except ImportError:  # pragma: no cover
     torch = None  # type: ignore[assignment]
 
-from .anneal import thermal_refine_q3_k
+from .anneal import thermal_refine_q3_k, thermal_refine_q3_k_cuda
 from .config import HPCConfig
 from .cuda_backend import cuda_is_available
 from .factored import MULTS_D
@@ -814,18 +814,17 @@ def _fit_q3_k_blocks_cuda(
             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):
+        #    On CUDA the anneal keeps every tensor on-device — the CPU/numpy
+        #    chain previously forced a full D2H→CPU→D2H round-trip per chunk
+        #    and made the GPU idle for the whole refinery.
         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,
+            scale_rp = best_d[:, None] * sc.to(torch.float32)   # (B, 16) per-superblock
+            thermal_c, _ = thermal_refine_q3_k_cuda(
+                Xf, None if Wf is None else Wf, scale_rp,
+                code.reshape(B, QK3_K), cfg,
             )
-            code = torch.from_numpy(
-                thermal_c.astype(np.int16).reshape(B, Q3_K_N_SUBBLOCKS, Q3_K_SUBBLOCK)
-            ).to(dev)
+            # thermal_refine_q3_k_cuda already returns source-signed codes [-4, 3]
+            code = thermal_c.reshape(B, Q3_K_N_SUBBLOCKS, Q3_K_SUBBLOCK).to(torch.int16)
 
         return (
             best_d16.cpu().numpy().astype(np.uint16),