Upload 9 files

config.py

@@ -24,10 +24,13 @@ class HPCConfig:
     n_cand_m: int = 6
     factor_pin: int = 3
     factor_iters: int = 1
-    factor_standalone_iters: int = 16
-    factor_starts: int = 3
+    factor_1d_iters: int = 6
+    factor_1d_refine_sweeps: int = 1
+    factor_standalone_iters: int = 8
+    factor_starts: int = 2
     factor_ulp_radius: int = 1
     factor_refine_1d: bool = True
+    factor_portfolio: bool = True
     factor_spectral_refine: bool = True
     factor_exact_emit: bool = True
 
@@ -89,4 +92,18 @@ class HPCConfig:
             raise ValueError("ULP radii must be non-negative")
         if self.refine_passes < 1:
             raise ValueError("refine_passes must be >= 1")
+        if not 1 <= self.subblock_states <= 6:
+            raise ValueError("subblock_states must lie in [1, 6]")
+        if self.factor_1d_iters < 1 or self.factor_standalone_iters < 1:
+            raise ValueError("factor iteration budgets must be >= 1")
+        if self.factor_starts < 1:
+            raise ValueError("factor_starts must be >= 1")
+        if self.sse_budget < 0.0:
+            raise ValueError("sse_budget must be non-negative")
+        if not 0.0 < self.same_penalty <= 1.0:
+            raise ValueError("same_penalty must lie in (0, 1]")
+        if self.log_slack < 0.0:
+            raise ValueError("log_slack must be non-negative")
+        if self.viterbi_alpha < 0.0 or self.viterbi_beta < 0.0:
+            raise ValueError("Viterbi weights must be non-negative")
         return self

factored.py

@@ -201,17 +201,47 @@ def _refine_integer_codes(v, w, d, codes, nmax, sweeps=2):
     return d, codes.astype(np.uint8)
 
 
-def factor_1d(v, w, nmax: int = 15, iters: int = 8, refine: bool = True):
+def _factor_1d_legacy(v, w, nmax: int = 15, iters: int = 6):
+    v = np.maximum(np.asarray(v, dtype=np.float64), 0.0)
+    w = np.maximum(np.asarray(w, dtype=np.float64), 0.0)
+    vmax = float(v.max())
+    if vmax < 1.0e-30:
+        return 0.0, np.zeros(v.size, dtype=np.uint8)
+    d = vmax / nmax
+    for _ in range(iters):
+        codes = np.clip(nearest_int(v / d), 0, nmax)
+        den = float(np.sum(w * codes * codes))
+        if den < 1.0e-30:
+            break
+        dn = float(np.sum(w * v * codes) / den)
+        conv = abs(dn - d) <= 1.0e-7 * abs(d) + 1.0e-30
+        d = dn
+        if conv:
+            break
+    codes = np.clip(nearest_int(v / d), 0, nmax)
+    return d, codes.astype(np.uint8)
+
+
+def factor_1d(v, w, nmax: int = 15, iters: int = 8, refine: bool = True, sweeps: int = 1):
     """Weighted 1-D factorisation ``v ~= d * L`` with ``L in [0, nmax]``."""
     v = np.maximum(np.asarray(v, dtype=np.float64), 0.0)
     w = np.maximum(np.asarray(w, dtype=np.float64), 0.0)
     d, codes = _factor_1d_lloyd(v, w, nmax, iters)
     if refine and d > 0.0:
-        d, codes = _refine_integer_codes(v, w, d, codes, nmax)
+        d, codes = _refine_integer_codes(v, w, d, codes, nmax, sweeps)
     return d, codes
 
 
-def _range_init(x, w, cfg: HPCConfig):
+def _factor_1d_for(v, w, nmax, cfg: HPCConfig, legacy: bool):
+    if legacy:
+        return _factor_1d_legacy(v, w, nmax, 6)
+    return factor_1d(
+        v, w, nmax, cfg.factor_1d_iters, cfg.factor_refine_1d,
+        cfg.factor_1d_refine_sweeps,
+    )
+
+
+def _range_init(x, w, cfg: HPCConfig, legacy_1d: bool = False):
     x = np.asarray(x, dtype=np.float64)
     w = np.asarray(w, dtype=np.float64)
     blocks = x.reshape(NS, 16)
@@ -233,8 +263,8 @@ def _range_init(x, w, cfg: HPCConfig):
                 nearest_int((sub + tj) / sj), 0, 3
             ).astype(np.uint8)
     sw = wblocks.sum(axis=1) if w is not None else np.full(NS, 16.0)
-    d, ls = factor_1d(s, sw, 15, cfg.factor_iters, cfg.factor_refine_1d)
-    m, lm = factor_1d(t, sw, 15, cfg.factor_iters, cfg.factor_refine_1d)
+    d, ls = _factor_1d_for(s, sw, 15, cfg, legacy_1d)
+    m, lm = _factor_1d_for(t, sw, 15, cfg, legacy_1d)
     return d, m, ls, lm, codes
 
 
@@ -246,7 +276,7 @@ def _seeded_init(x, scales, mins, d_seed, m_seed):
     return d, m, ls, lm, codes
 
 
-def _ls_init(x, w, cfg: HPCConfig):
+def _ls_init(x, w, cfg: HPCConfig, legacy_1d: bool = False):
     x = np.asarray(x, dtype=np.float64)
     w = np.asarray(w, dtype=np.float64)
     blocks = x.reshape(NS, 16)
@@ -262,8 +292,8 @@ def _ls_init(x, w, cfg: HPCConfig):
             ).astype(np.uint8)
     s, t = _sub_ls(x, w, codes)
     sw = w.reshape(NS, 16).sum(axis=1)
-    d, ls = factor_1d(s, sw, 15, cfg.factor_iters, cfg.factor_refine_1d)
-    m, lm = factor_1d(t, sw, 15, cfg.factor_iters, cfg.factor_refine_1d)
+    d, ls = _factor_1d_for(s, sw, 15, cfg, legacy_1d)
+    m, lm = _factor_1d_for(t, sw, 15, cfg, legacy_1d)
     codes = codes_from_params(x, d, m, ls, lm)
     return d, m, ls, lm, codes
 
@@ -274,34 +304,23 @@ def _sub_ls(x, w, codes):
     q = np.asarray(codes, dtype=np.float64).reshape(NS, 16)
     xb = x.reshape(NS, 16)
     wb = w.reshape(NS, 16)
-    s = np.zeros(NS)
-    t = np.zeros(NS)
-    for j in range(NS):
-        ww = wb[j]
-        qq = q[j]
-        xx = xb[j]
-        Sqq = float(np.sum(ww * qq * qq))
-        Swq = float(np.sum(ww * qq))
-        Sw = float(np.sum(ww))
-        Sqx = float(np.sum(ww * qq * xx))
-        Sx = float(np.sum(ww * xx))
-        det = Sqq * Sw - Swq * Swq
-        if det > 1.0e-30:
-            sj = (Sqx * Sw - Swq * Sx) / det
-            tj = (sj * Swq - Sx) / Sw
-        else:
-            sj = 0.0
-            tj = 0.0
-        if sj < 0.0:
-            sj = 0.0
-        if tj < 0.0:
-            tj = 0.0
-            sj = Sqx / Sqq if Sqq > 1.0e-30 else 0.0
-            if sj < 0.0:
-                sj = 0.0
-        s[j] = sj
-        t[j] = tj
-    return s, t
+    Sqq = np.sum(wb * q * q, axis=1)
+    Swq = np.sum(wb * q, axis=1)
+    Sw = np.sum(wb, axis=1)
+    Sqx = np.sum(wb * q * xb, axis=1)
+    Sx = np.sum(wb * xb, axis=1)
+    det = Sqq * Sw - Swq * Swq
+    valid = det > 1.0e-30
+    safe_det = np.where(valid, det, 1.0)
+    safe_sw = np.where(Sw > 1.0e-30, Sw, 1.0)
+    sj = np.where(valid, (Sqx * Sw - Swq * Sx) / safe_det, 0.0)
+    tj = np.where(valid, (sj * Swq - Sx) / safe_sw, 0.0)
+    sj = np.maximum(sj, 0.0)
+    bad_tj = tj < 0.0
+    tj = np.where(bad_tj, 0.0, tj)
+    fallback = np.where(Sqq > 1.0e-30, Sqx / np.where(Sqq > 1.0e-30, Sqq, 1.0), 0.0)
+    sj = np.where(bad_tj, np.maximum(fallback, 0.0), sj)
+    return sj, tj
 
 
 def _solve(
@@ -317,6 +336,7 @@ def _solve(
     use_spectral: bool,
     max_iter: int,
     dc_carry: float,
+    legacy_1d: bool = False,
 ):
     x = np.asarray(x, dtype=np.float32)
     w = np.asarray(w, dtype=np.float32)
@@ -345,9 +365,7 @@ def _solve(
             else:
                 ls = np.zeros(NS, dtype=np.uint8)
         else:
-            d_new, ls_new = factor_1d(
-                s, ws, 15, cfg.factor_iters, cfg.factor_refine_1d
-            )
+            d_new, ls_new = _factor_1d_for(s, ws, 15, cfg, legacy_1d)
             d, ls = float(d_new), ls_new
 
         if pin & 2:
@@ -356,9 +374,7 @@ def _solve(
             else:
                 lm = np.zeros(NS, dtype=np.uint8)
         else:
-            m_new, lm_new = factor_1d(
-                t, wt, 15, cfg.factor_iters, cfg.factor_refine_1d
-            )
+            m_new, lm_new = _factor_1d_for(t, wt, 15, cfg, legacy_1d)
             m, lm = float(m_new), lm_new
 
         codes = codes_from_params(x, d, m, ls, lm)
@@ -432,48 +448,55 @@ def factor_q2k_block(
     max_iter = cfg.factor_standalone_iters if max_iter is None else max_iter
     starts = cfg.factor_starts if starts is None else starts
 
-    candidates = []
-    if init_s is not None and init_t is not None and d_seed > 1.0e-30 and m_seed > 1.0e-30:
-        d, m, ls, lm, codes = _seeded_init(x, init_s, init_t, d_seed, m_seed)
-        candidates.append((d, m, ls, lm, codes))
-    else:
-        d, m, ls, lm, codes = _range_init(x, w, cfg)
-        if d_seed > 1.0e-30:
-            d = float(d_seed)
-            if init_s is not None:
-                ls = np.clip(nearest_int(np.asarray(init_s) / d), 0, 15).astype(np.uint8)
-        if m_seed > 1.0e-30:
-            m = float(m_seed)
-            if init_t is not None:
-                lm = np.clip(nearest_int(np.asarray(init_t) / m), 0, 15).astype(np.uint8)
-        candidates.append((d, m, ls, lm, codes))
-        if starts > 1:
-            candidates.append(_ls_init(x, w, cfg))
+    def build_candidates(legacy_1d: bool):
+        out = []
+        if init_s is not None and init_t is not None and d_seed > 1.0e-30 and m_seed > 1.0e-30:
+            out.append(_seeded_init(x, init_s, init_t, d_seed, m_seed))
+        else:
+            d, m, ls, lm, codes = _range_init(x, w, cfg, legacy_1d)
+            if d_seed > 1.0e-30:
+                d = float(d_seed)
+                if init_s is not None:
+                    ls = np.clip(nearest_int(np.asarray(init_s) / d), 0, 15).astype(np.uint8)
+            if m_seed > 1.0e-30:
+                m = float(m_seed)
+                if init_t is not None:
+                    lm = np.clip(nearest_int(np.asarray(init_t) / m), 0, 15).astype(np.uint8)
+            out.append((d, m, ls, lm, codes))
+            if starts > 1:
+                out.append(_ls_init(x, w, cfg, legacy_1d))
+        return out
 
     best = None
     total_iters = 0
-    for d, m, ls, lm, codes in candidates:
-        d, m, ls, lm, codes, it = _solve(
-            x, w, cfg, d, m, ls, lm, codes, pin, use_spectral, max_iter, dc_carry
-        )
-        total_iters += it
-        d, m, d16, m16, ls, lm, codes, sse, obj = _emit_polish(
-            x, w, cfg, d, m, ls, lm, codes, dc_carry
-        )
-        if best is None or obj < best.objective - 1.0e-15:
-            best = FactorResult(
-                d=np.float32(d),
-                dmin=np.float32(m),
-                d16=np.uint16(d16),
-                m16=np.uint16(m16),
-                ls=ls.astype(np.uint8),
-                lm=lm.astype(np.uint8),
-                codes=codes.astype(np.uint8),
-                sse=float(sse),
-                objective=float(obj),
-                iterations=total_iters,
-                starts=len(candidates),
+    n_starts = 0
+    portfolio = cfg.factor_portfolio and (starts > 1 or max_iter > 1)
+    for legacy_1d in ((False, True) if portfolio else (False,)):
+        for d, m, ls, lm, codes in build_candidates(legacy_1d):
+            n_starts += 1
+            solve_iters = max(max_iter, 16) if legacy_1d else max_iter
+            d, m, ls, lm, codes, it = _solve(
+                x, w, cfg, d, m, ls, lm, codes, pin, use_spectral,
+                solve_iters, dc_carry, legacy_1d,
+            )
+            total_iters += it
+            d, m, d16, m16, ls, lm, codes, sse, obj = _emit_polish(
+                x, w, cfg, d, m, ls, lm, codes, dc_carry
             )
+            if best is None or obj < best.objective - 1.0e-15:
+                best = FactorResult(
+                    d=np.float32(d),
+                    dmin=np.float32(m),
+                    d16=np.uint16(d16),
+                    m16=np.uint16(m16),
+                    ls=ls.astype(np.uint8),
+                    lm=lm.astype(np.uint8),
+                    codes=codes.astype(np.uint8),
+                    sse=float(sse),
+                    objective=float(obj),
+                    iterations=total_iters,
+                    starts=n_starts,
+                )
     return best
 
 

pipeline.py

@@ -39,14 +39,12 @@ from .q2k import (
     codes_from_params,
     codes_from_params_batch,
     dequantize_params,
-    dequantize_params_batch,
     derive_subscales,
     make_block,
     pack_codes,
     pack_scales,
     unpack_codes,
     weighted_sse,
-    weighted_sse_batch,
 )
 from .seeds import build_seeds, seed_clamp_value
 from .sieve import block_chain_edges, sieve_measure_chain, subblock_chain_edges
@@ -144,17 +142,17 @@ def _sieve_marginals(candidate_errors, graph_blocks, stride, temperature, cfg):
         stop = min((gi + 1) * stride, n_blocks)
         group = candidate_errors[start:stop].astype(np.float64)
         finite = np.isfinite(group)
-        if np.any(finite):
-            agg = np.where(finite, group, np.nan).mean(axis=0)
-            valid = np.isfinite(agg)
-            mn = float(np.min(agg[valid])) if np.any(valid) else 0.0
-            coeff = np.where(
-                valid,
-                np.exp(-(agg - mn) / (2.0 * max(temperature, 1.0e-12))),
-                0.0,
-            )
-        else:
-            coeff = np.zeros(TOTAL_CANDIDATES, dtype=np.float64)
+        counts = finite.sum(axis=0)
+        valid = counts > 0
+        sums = np.where(finite, group, 0.0).sum(axis=0)
+        agg = np.where(valid, sums / np.maximum(counts, 1), 0.0)
+        mn = float(np.min(agg[valid])) if np.any(valid) else 0.0
+        expo = np.where(
+            valid,
+            (agg - mn) / (2.0 * max(temperature, 1.0e-12)),
+            np.inf,
+        )
+        coeff = np.exp(-expo)
         for v in range(6):
             coarse[gi, v] = float(np.sum(coeff[qi_d == v]))
             fine[gi, v] = float(np.sum(coeff[qi_m == v]))
@@ -229,28 +227,23 @@ def _per_sub_top_states(x, w, d, m, k=6):
         wb = np.ones((NS, 16), dtype=np.float32)
     else:
         wb = np.asarray(w, dtype=np.float32).reshape(NS, 16)
-    state_ls = np.zeros((NS, k), dtype=np.uint8)
-    state_lm = np.zeros((NS, k), dtype=np.uint8)
-    state_err = np.full((NS, k), np.inf, dtype=np.float64)
     ls_opts = np.arange(16, dtype=np.float32)[:, None]
     lm_opts = np.arange(16, dtype=np.float32)[None, :]
     d_sub = np.float32(d) * ls_opts
     m_sub = np.float32(m) * lm_opts
-    safe = d_sub[:, :, None] >= np.float32(1.0e-15)
-    safe_d = np.where(safe, d_sub[:, :, None], np.float32(1.0))
-    for j in range(NS):
-        sub = xb[j]
-        wj = wb[j]
-        raw = (sub[None, None, :] + m_sub[:, :, None]) / safe_d
-        q = np.clip(nearest_int(raw), 0, 3).astype(np.float32)
-        q = np.where(safe, q, 0.0)
-        deq = d_sub[:, :, None] * q - m_sub[:, :, None]
-        err = np.sum(wj[None, None, :] * (sub[None, None, :] - deq) ** 2, axis=2)
-        flat = err.ravel()
-        order = np.argsort(flat, kind="stable")[:k]
-        state_err[j] = flat[order]
-        state_ls[j] = (order // 16).astype(np.uint8)
-        state_lm[j] = (order % 16).astype(np.uint8)
+    safe = d_sub >= np.float32(1.0e-15)
+    safe_d = np.where(safe, d_sub, np.float32(1.0))
+    subs = xb[:, None, None, :]
+    raw = (subs + m_sub[None, :, :, None]) / safe_d[None, :, :, None]
+    q = np.clip(nearest_int(raw), 0, 3).astype(np.float32)
+    q = np.where(safe[None, :, :, None], q, 0.0)
+    deq = d_sub[None, :, :, None] * q - m_sub[None, :, :, None]
+    err = np.sum(wb[:, None, None, :] * (subs - deq) ** 2, axis=3)
+    flat = err.reshape(NS, 256)
+    order = np.argsort(flat, axis=1, kind="stable")[:, :k]
+    state_err = np.take_along_axis(flat, order, axis=1).astype(np.float64)
+    state_ls = (order // 16).astype(np.uint8)
+    state_lm = (order % 16).astype(np.uint8)
     return state_ls, state_lm, state_err
 
 
@@ -305,18 +298,28 @@ def _ulp_polish(x, w, d, m, ls, lm, radius):
     if radius <= 0:
         codes = codes_from_params(x, d, m, ls, lm)
         return float(d), float(m), codes, weighted_sse(x, w, d, m, ls, lm, codes)
-    d16 = np.uint16(to_fp16_bits(np.float32(d)))
-    m16 = np.uint16(to_fp16_bits(np.float32(m)))
-    best = None
-    for db in range(max(0, int(d16) - radius), min(0x7BFF, int(d16) + radius) + 1):
-        dd = float(from_fp16_bits(np.uint16(db)))
-        for mb in range(max(0, int(m16) - radius), min(0x7BFF, int(m16) + radius) + 1):
-            mm = float(from_fp16_bits(np.uint16(mb)))
-            c = codes_from_params(x, dd, mm, ls, lm)
-            s = weighted_sse(x, w, dd, mm, ls, lm, c)
-            if best is None or s < best[0] - 1.0e-15:
-                best = (s, dd, mm, c)
-    return best[1], best[2], best[3], best[0]
+    d16 = int(to_fp16_bits(np.float32(d)))
+    m16 = int(to_fp16_bits(np.float32(m)))
+    d_bits = np.arange(max(0, d16 - radius), min(0x7BFF, d16 + radius) + 1, dtype=np.uint16)
+    m_bits = np.arange(max(0, m16 - radius), min(0x7BFF, m16 + radius) + 1, dtype=np.uint16)
+    dd = from_fp16_bits(d_bits)
+    mm = from_fp16_bits(m_bits)
+    d_grid = np.repeat(dd, mm.size)
+    m_grid = np.tile(mm, dd.size)
+    ls_rep = np.repeat(np.asarray(ls, dtype=np.float32), 16)
+    lm_rep = np.repeat(np.asarray(lm, dtype=np.float32), 16)
+    d_sub = d_grid[:, None] * ls_rep[None, :]
+    m_sub = m_grid[:, None] * lm_rep[None, :]
+    safe = d_sub >= np.float32(1.0e-15)
+    safe_d = np.where(safe, d_sub, np.float32(1.0))
+    q = np.clip(nearest_int((x[None, :] + m_sub) / safe_d), 0, 3).astype(np.float32)
+    q = np.where(safe, q, 0.0)
+    deq = d_sub * q - m_sub
+    ww = np.ones(QK_K, dtype=np.float64) if w is None else np.asarray(w, dtype=np.float64)
+    err = np.sum(ww[None, :] * (np.asarray(x, dtype=np.float64)[None, :] - deq) ** 2, axis=1)
+    k = int(np.argmin(err))
+    codes = q[k].astype(np.uint8)
+    return float(d_grid[k]), float(m_grid[k]), codes, float(err[k])
 
 
 def _per_sub_scan(x, w, d, m):
@@ -341,19 +344,36 @@ def _assemble_block(x, w, seeds, base_d, base_m, cidx, dc_carry, cfg):
             break
         prev = cur
 
-    d, m, codes, _ = _ulp_polish(x, w, d, m, ls, lm, cfg.ulp_radius)
-    ls, lm = _per_sub_scan(x, w, d, m)
     codes = codes_from_params(x, d, m, ls, lm)
-    codes = refine_codes_hadamard(x, w, d, m, ls, lm, codes, cfg, dc_carry)
     obj, sse = objective(x, w, d, m, ls, lm, codes, cfg, dc_carry)
-    return d, m, ls, lm, codes, obj, sse
+    best = (obj, sse, d, m, ls.copy(), lm.copy(), codes.copy())
+
+    ud, um, uc, _ = _ulp_polish(x, w, d, m, ls, lm, cfg.ulp_radius)
+    uobj, usse = objective(x, w, ud, um, ls, lm, uc, cfg, dc_carry)
+    if uobj < best[0] - 1.0e-15:
+        best = (uobj, usse, ud, um, ls.copy(), lm.copy(), uc.copy())
 
+    sls, slm = _per_sub_scan(x, w, ud, um)
+    scodes = codes_from_params(x, ud, um, sls, slm)
+    sobj, ssse = objective(x, w, ud, um, sls, slm, scodes, cfg, dc_carry)
+    if sobj < best[0] - 1.0e-15:
+        best = (sobj, ssse, ud, um, sls.copy(), slm.copy(), scodes.copy())
 
-def _candidate_config(seeds, base_d, base_m, cidx):
+    _, _, bd, bm, bls, blm, bc = best
+    refined = refine_codes_hadamard(x, w, bd, bm, bls, blm, bc, cfg, dc_carry)
+    robj, rsse = objective(x, w, bd, bm, bls, blm, refined, cfg, dc_carry)
+    if robj < best[0] - 1.0e-15:
+        best = (robj, rsse, bd, bm, bls.copy(), blm.copy(), refined.copy())
+
+    _, sse, d, m, ls, lm, codes = best
+    return d, m, ls, lm, codes, best[0], sse
+
+
+def _candidate_config(scales, mins, base_d, base_m, cidx):
     d16, m16 = candidate_pair(base_d, base_m, cidx)
     d = float(from_fp16_bits(d16))
     m = float(from_fp16_bits(m16))
-    ls, lm = derive_subscales(seeds["scales"], seeds["mins"], d, m)
+    ls, lm = derive_subscales(scales, mins, d, m)
     return d, m, ls, lm
 
 
@@ -396,6 +416,8 @@ class Q2KQuantizer:
     ) -> QuantizedTensor:
         cfg = self.config
         x = np.asarray(weights, dtype=np.float32)
+        if x.size == 0:
+            raise ValueError("weight tensor is empty")
         if x.size % QK_K != 0:
             raise ValueError("weight count must be a multiple of 256")
         if x.ndim == 2 and row_width is None and x.shape[1] % QK_K == 0:
@@ -408,12 +430,14 @@ class Q2KQuantizer:
             imp = np.asarray(importance, dtype=np.float32).ravel()
             if imp.size != x.size:
                 raise ValueError("importance must match the weight count")
+            if np.any(imp < 0.0) or not np.all(np.isfinite(imp)):
+                raise ValueError("importance must be finite and non-negative")
             W = imp.reshape(n_blocks, QK_K)
         else:
             W = None
 
-        seeds = build_seeds(X, W, cfg)
         clamp_val = seed_clamp_value(x, cfg)
+        seeds = build_seeds(X, W, cfg, clamp_val)
         base_d, base_m = _phase2_wls(X, W, seeds, clamp_val)
 
         candidate_errors = np.full((n_blocks, TOTAL_CANDIDATES), np.inf, dtype=np.float32)
@@ -466,10 +490,7 @@ class Q2KQuantizer:
 
             if cfg.use_candidate_floor:
                 cd, cm, cls, clm = _candidate_config(
-                    {k: v[i] for k, v in seeds.items() if k in ("scales", "mins")},
-                    base_d[i],
-                    base_m[i],
-                    best[i],
+                    seeds["scales"][i], seeds["mins"][i], base_d[i], base_m[i], best[i]
                 )
                 ccodes = codes_from_params(X[i], cd, cm, cls, clm)
                 cobj, csse = objective(

seeds.py

@@ -157,18 +157,22 @@ def seed_weights(
 
 
 def build_seeds(
-    x: np.ndarray, importance: np.ndarray | None, cfg: HPCConfig
+    x: np.ndarray,
+    importance: np.ndarray | None,
+    cfg: HPCConfig,
+    clamp_val: float = float("inf"),
 ) -> dict[str, np.ndarray]:
     """Phase-1 seeds for all blocks: per-sub (scale,min) and shared (d,dmin)."""
     n_blocks = x.shape[0]
     blocks = np.asarray(x, dtype=np.float64).reshape(n_blocks, 16, 16)
+    clipped = np.clip(blocks, -clamp_val, clamp_val)
     if importance is None:
         sigma_sq = np.mean(blocks * blocks, axis=(1, 2), keepdims=True)
-        weights = np.sqrt(sigma_sq + blocks * blocks)
+        weights = np.sqrt(sigma_sq + clipped * clipped)
     else:
         weights = np.asarray(importance, dtype=np.float64).reshape(n_blocks, 16, 16)
 
-    flat_x = blocks.reshape(-1, 16)
+    flat_x = clipped.reshape(-1, 16)
     flat_w = weights.reshape(-1, 16)
     scales, mins, codes = make_qkx2_batch(flat_x, flat_w, nmax=3)
     scales = scales.reshape(n_blocks, 16)

spectral.py

@@ -13,27 +13,37 @@ D6/vesica greedy code descent.
 
 from __future__ import annotations
 
+from functools import lru_cache
+
 import numpy as np
 
 from .config import HPCConfig
-from .q2k import QK_K, dequantize_params, weighted_sse
+from .q2k import QK_K, dequantize_params
 
 
-def fold_weights(n: int, cfg: HPCConfig) -> np.ndarray:
+@lru_cache(maxsize=64)
+def _fold_weights_cached(n: int, dc_lambda: float, vw_lambda: float, depth: int, gamma: float) -> np.ndarray:
     levels = int(np.log2(n))
     if (1 << levels) != n:
         raise ValueError("fold pyramid requires a power-of-two length")
     weights = np.zeros(levels, dtype=np.float64)
-    gamma = 1.0
+    g = 1.0
     for k in range(1, levels + 1):
         if k == levels:
-            weights[k - 1] = cfg.dc_lambda / n
-        elif k <= cfg.fold_depth:
-            weights[k - 1] = (cfg.vw_lambda / n) * gamma
-        gamma *= cfg.fold_gamma
+            weights[k - 1] = dc_lambda / n
+        elif k <= depth:
+            weights[k - 1] = (vw_lambda / n) * g
+        g *= gamma
+    weights.setflags(write=False)
     return weights
 
 
+def fold_weights(n: int, cfg: HPCConfig) -> np.ndarray:
+    return _fold_weights_cached(
+        n, cfg.dc_lambda, cfg.vw_lambda, cfg.fold_depth, cfg.fold_gamma
+    )
+
+
 def fold_build(errors: np.ndarray) -> np.ndarray:
     e = np.asarray(errors, dtype=np.float64)
     n = e.size

viterbi.py

@@ -76,16 +76,17 @@ def viterbi_select(
         stop = min((gi + 1) * stride, n_blocks)
         group = candidate_errors[start:stop].astype(np.float64)
         flat = group.ravel()
+        bins_flat = np.tile(bins, group.shape[0])
         finite = np.isfinite(flat)
         group_min = np.full(N_STATES, np.inf)
         if np.any(finite):
-            np.minimum.at(group_min, bins[finite], flat[finite])
-        counts = np.bincount(bins[finite], minlength=N_STATES)
+            np.minimum.at(group_min, bins_flat[finite], flat[finite])
+        counts = np.bincount(bins_flat[finite], minlength=N_STATES)
         sums = np.zeros(N_STATES, dtype=np.float64)
         if np.any(finite):
-            shifted = flat[finite] - group_min[bins[finite]]
+            shifted = flat[finite] - group_min[bins_flat[finite]]
             amps = np.exp(-shifted / (2.0 * temp))
-            sums = np.bincount(bins[finite], weights=amps, minlength=N_STATES)
+            sums = np.bincount(bins_flat[finite], weights=amps, minlength=N_STATES)
         valid = (counts > 0) & np.isfinite(group_min)
         if np.any(valid):
             mean_exp = np.maximum(sums[valid] / counts[valid], 1.0e-300)
@@ -140,7 +141,7 @@ def viterbi_select(
         for gi in range(graph_blocks - 2, -1, -1):
             path[gi] = back[gi + 1, path[gi + 1]]
 
-    best_candidate = np.zeros(n_blocks, dtype=np.int64)
+    best_candidate = np.argmin(candidate_errors, axis=1).astype(np.int64)
     for gi in range(graph_blocks):
         s = int(path[gi])
         qd, qm = s // 6, s % 6