Upload 15 files

factored.py

@@ -32,7 +32,7 @@ from .analog import build_block_analogs
 from .config import HPCConfig
 from .fp16 import from_fp16_bits, fp16_ulp_neighbors, nearest_int, round_fp16, to_fp16_bits
 from .q2k import NS, QK_K, codes_from_params, derive_subscales
-from .spectral import objective, refine_codes_hadamard
+from .spectral import fold_weights, objective, refine_codes_hadamard
 
 N_CAND_D = 24
 N_CAND_M = 24
@@ -165,6 +165,10 @@ def _factor_1d_lloyd(v: np.ndarray, w: np.ndarray, nmax: int, iters: int):
     return best_d, best_codes
 
 
+_OPTS_15 = np.arange(16, dtype=np.float64)
+_OPTS_15_SQ = _OPTS_15 * _OPTS_15
+
+
 def _refine_integer_codes(v, w, d, codes, nmax, sweeps=2):
     v = np.asarray(v, dtype=np.float64)
     w = np.asarray(w, dtype=np.float64)
@@ -172,33 +176,46 @@ def _refine_integer_codes(v, w, d, codes, nmax, sweeps=2):
     d = float(d)
     if d <= 0.0:
         return d, codes.astype(np.uint8)
-    obj = float(np.sum(w * (v - d * codes) ** 2))
-    options = np.arange(nmax + 1, dtype=np.int64)
+    options = _OPTS_15 if nmax == 15 else np.arange(nmax + 1, dtype=np.float64)
+    opt_sq = _OPTS_15_SQ if nmax == 15 else options * options
+    wv = w * v
+    sum_lv = float(np.dot(wv, codes))
+    sum_l2 = float(np.dot(w, codes * codes))
+    if sum_l2 <= 1.0e-30:
+        return d, codes.astype(np.uint8)
+    best_gain = (sum_lv * sum_lv) / sum_l2
+
     for _ in range(sweeps):
+        c_f = codes[:, None].astype(np.float64)
+        slv_all = sum_lv + wv[:, None] * (options[None, :] - c_f)
+        sl2_all = sum_l2 + w[:, None] * (opt_sq[None, :] - c_f * c_f)
+        val_all = (sl2_all > 1.0e-30) & (slv_all > 0.0)
+        g_all = np.where(val_all, (slv_all * slv_all) / np.where(val_all, sl2_all, 1.0), -1.0)
+        if float(np.max(g_all)) <= best_gain + 1.0e-18:
+            break
         improved = False
         for i in range(v.size):
-            wi, vi, li = w[i], v[i], codes[i]
-            sum_lv = float(np.sum(w * v * codes))
-            sum_l2 = float(np.sum(w * codes * codes))
-            slv = sum_lv - wi * vi * li + wi * vi * options
-            sl2 = sum_l2 - wi * li * li + wi * options * options
-            valid = sl2 > 1.0e-30
-            if not np.any(valid):
+            wi = w[i]
+            if wi <= 0.0:
                 continue
-            dd = np.where(valid, slv / np.where(valid, sl2, 1.0), 0.0)
-            trial = codes[None, :].repeat(options.size, axis=0)
-            trial[:, i] = options
-            e = v[None, :] - dd[:, None] * trial.astype(np.float64)
-            scores = np.sum(w[None, :] * e * e, axis=1)
-            scores = np.where(valid & (dd > 0.0), scores, np.inf)
-            k = int(np.argmin(scores))
-            if scores[k] < obj - 1.0e-18 and options[k] != li:
-                obj = float(scores[k])
-                d = float(dd[k])
-                codes[i] = options[k]
+            wvi = wv[i]
+            li = int(codes[i])
+            base_lv = sum_lv - wvi * li
+            base_l2 = sum_l2 - wi * (li * li)
+            slv = base_lv + wvi * options
+            sl2 = base_l2 + wi * opt_sq
+            valid = (sl2 > 1.0e-30) & (slv > 0.0)
+            gains = np.where(valid, (slv * slv) / np.where(valid, sl2, 1.0), -1.0)
+            k = int(np.argmax(gains))
+            if k != li and gains[k] > best_gain + 1.0e-18:
+                best_gain = float(gains[k])
+                sum_lv = float(slv[k])
+                sum_l2 = float(sl2[k])
+                codes[i] = k
                 improved = True
         if not improved:
             break
+    d = sum_lv / sum_l2 if sum_l2 > 1.0e-30 else d
     return d, codes.astype(np.uint8)
 
 
@@ -242,28 +259,26 @@ def _factor_1d_for(v, w, nmax, cfg: HPCConfig, legacy: bool):
     )
 
 
+def _range_codes_and_st(x: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    blocks = np.asarray(x, dtype=np.float64).reshape(NS, 16)
+    xmin = np.minimum(np.min(blocks, axis=1), 0.0)
+    xmax = np.max(blocks, axis=1)
+    span = xmax - xmin
+    s = np.where(span > 1.0e-30, span / 3.0, 0.0)
+    t = -xmin
+    safe_s = np.where(s > 1.0e-30, s, 1.0)
+    q = np.where(
+        s[:, None] > 1.0e-30,
+        np.clip(nearest_int(((blocks + t[:, None]) / safe_s[:, None]).astype(np.float32)), 0, 3),
+        0,
+    ).astype(np.uint8)
+    return s, t, q.reshape(QK_K)
+
+
 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)
-    wblocks = w.reshape(NS, 16)
-    s = np.zeros(NS)
-    t = np.zeros(NS)
-    codes = np.zeros(QK_K, dtype=np.uint8)
-    for j in range(NS):
-        sub = blocks[j]
-        xmin = min(float(sub.min()), 0.0)
-        xmax = float(sub.max())
-        span = xmax - xmin
-        sj = span / 3.0 if span > 1.0e-30 else 0.0
-        tj = -xmin
-        s[j] = sj
-        t[j] = tj
-        if sj > 1.0e-30:
-            codes[16 * j : 16 * (j + 1)] = np.clip(
-                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)
+    s, t, codes = _range_codes_and_st(x)
+    sw = w.reshape(NS, 16).sum(axis=1)
     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
@@ -280,17 +295,7 @@ def _seeded_init(x, scales, mins, d_seed, m_seed):
 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)
-    codes = np.zeros(QK_K, dtype=np.uint8)
-    for j in range(NS):
-        sub = blocks[j]
-        xmin = min(float(sub.min()), 0.0)
-        xmax = float(sub.max())
-        span = xmax - xmin
-        if span > 1.0e-30:
-            codes[16 * j : 16 * (j + 1)] = np.clip(
-                nearest_int((sub - xmin) * 3.0 / span), 0, 3
-            ).astype(np.uint8)
+    _, _, codes = _range_codes_and_st(x)
     s, t = _sub_ls(x, w, codes)
     sw = w.reshape(NS, 16).sum(axis=1)
     d, ls = _factor_1d_for(s, sw, 15, cfg, legacy_1d)
@@ -302,7 +307,7 @@ def _ls_init(x, w, cfg: HPCConfig, legacy_1d: bool = False):
 def _analog_init(x, w, x_analog, cfg: HPCConfig):
     """Extract discrete Q2_K lattice coordinates from a geometry-preserving analog
     and project continuous scales onto the true weights ``x``."""
-    _, _, _, _, analog_codes = _range_init(x_analog, w, cfg, False)
+    _, _, analog_codes = _range_codes_and_st(x_analog)
     s, t = _sub_ls(x, w, analog_codes)
     q = analog_codes.astype(np.float64).reshape(NS, 16)
     wb = np.asarray(w, dtype=np.float64).reshape(NS, 16)
@@ -311,7 +316,7 @@ def _analog_init(x, w, x_analog, cfg: HPCConfig):
     Swq = np.sum(wb * q, axis=1)
     Sqq = np.sum(wb * q * q, axis=1)
     Sx = np.sum(wb * xb, axis=1)
-    q_bar = np.where(Sw > 1.0e-30, Swq / np.maximum(Sw, 1e-30), 0.0)
+    q_bar = np.where(Sw > 1.0e-30, Swq / np.maximum(Sw, 1.0e-30), 0.0)
     Vqq = np.maximum(Sqq - Swq * q_bar, 1.0e-12 * Sw)
     d, ls = _factor_1d_for(s, Vqq, 15, cfg, False)
     t_coupled = np.maximum(
@@ -474,22 +479,11 @@ def _solve(
             m, lm = float(m_new), lm_new
 
         codes = codes_from_params(x, d, m, ls, lm)
-        if coupled:
-            if pin == 0:
-                d, m = _global_bilinear_dm_solve(
-                    x, w, d, m, ls, lm, codes, cfg, dc_carry, use_spectral
-                )
-            if cfg.factor_joint_subscales:
-                d_f16 = float(d) if (pin & 1) else float(round_fp16(np.float32(d)))
-                m_f16 = float(m) if (pin & 2) else float(round_fp16(np.float32(m)))
-                ls_j, lm_j = _joint_subscale_scan(x, w, d_f16, m_f16)
-                codes_j = codes_from_params(x, d_f16, m_f16, ls_j, lm_j)
-                obj_j, _ = objective(x, w, d_f16, m_f16, ls_j, lm_j, codes_j, cfg, dc_carry)
-                if obj_j < best[0] - 1.0e-15:
-                    d, m, ls, lm, codes = d_f16, m_f16, ls_j, lm_j, codes_j
-                    best = (obj_j, d, m, ls.copy(), lm.copy(), codes.copy())
-                    stall = 0
-                    continue
+        if coupled and pin == 0:
+            d, m = _global_bilinear_dm_solve(
+                x, w, d, m, ls, lm, codes, cfg, dc_carry, use_spectral
+            )
+            codes = codes_from_params(x, d, m, ls, lm)
 
         cur, _ = objective(x, w, d, m, ls, lm, codes, cfg, dc_carry)
         if cur < best[0] - 1.0e-15:
@@ -501,6 +495,21 @@ def _solve(
             break
 
     obj, d, m, ls, lm, codes = best
+    if coupled and cfg.factor_joint_subscales:
+        d_f16 = float(d) if (pin & 1) else float(round_fp16(np.float32(d)))
+        m_f16 = float(m) if (pin & 2) else float(round_fp16(np.float32(m)))
+        ls_j, lm_j = _joint_subscale_scan(x, w, d_f16, m_f16)
+        codes_j = codes_from_params(x, d_f16, m_f16, ls_j, lm_j)
+        if pin == 0:
+            d_j, m_j = _global_bilinear_dm_solve(
+                x, w, d_f16, m_f16, ls_j, lm_j, codes_j, cfg, dc_carry, use_spectral
+            )
+            d_f16 = float(round_fp16(np.float32(d_j)))
+            m_f16 = float(round_fp16(np.float32(m_j)))
+            codes_j = codes_from_params(x, d_f16, m_f16, ls_j, lm_j)
+        obj_j, _ = objective(x, w, d_f16, m_f16, ls_j, lm_j, codes_j, cfg, dc_carry)
+        if obj_j < obj - 1.0e-15:
+            obj, d, m, ls, lm, codes = obj_j, d_f16, m_f16, ls_j, lm_j, codes_j
     return d, m, ls, lm, codes, iters
 
 
@@ -588,6 +597,7 @@ def factor_q2k_block(
     n_starts = 0
     portfolio = cfg.factor_portfolio and (starts > 1 or max_iter > 1)
     for legacy_1d in ((False, True) if portfolio else (False,)):
+        branch_best = None
         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
@@ -596,6 +606,15 @@ def factor_q2k_block(
                 solve_iters, dc_carry, legacy_1d,
             )
             total_iters += it
+            d_f16 = float(round_fp16(np.float32(d)))
+            m_f16 = float(round_fp16(np.float32(m)))
+            codes_f16 = codes_from_params(x, d_f16, m_f16, ls, lm)
+            raw_obj, _ = objective(x, w, d_f16, m_f16, ls, lm, codes_f16, cfg, dc_carry)
+            if branch_best is None or raw_obj < branch_best[0] - 1.0e-15:
+                branch_best = (raw_obj, d, m, ls, lm, codes)
+
+        if branch_best is not None:
+            _, d, m, ls, lm, codes = branch_best
             d, m, d16, m16, ls, lm, codes, sse, obj = _emit_polish(
                 x, w, cfg, d, m, ls, lm, codes, dc_carry
             )
@@ -626,40 +645,132 @@ def factor_candidates(
     cfg: HPCConfig,
     dc_carry: float = 0.0,
 ) -> np.ndarray:
-    """Factored candidate field indexed exactly like the C 24x24 lattice."""
-    out = np.full(TOTAL_CANDIDATES, np.inf, dtype=np.float32)
+    """Vectorized factored candidate field across the (nd x nm) lattice."""
+    x_f32 = np.asarray(x, dtype=np.float32).reshape(QK_K)
+    w_f64 = (
+        np.ones(QK_K, dtype=np.float64)
+        if w is None
+        else np.asarray(w, dtype=np.float64).reshape(QK_K)
+    )
+    scales_f32 = np.asarray(scales, dtype=np.float32).reshape(NS)
+    mins_f32 = np.asarray(mins, dtype=np.float32).reshape(NS)
+
     nd = int(cfg.n_cand_d)
     nm = int(cfg.n_cand_m)
-    for a in range(nd):
-        di = a if nd == N_CAND_D else (a * N_CAND_D) // nd
-        ds = float(round_fp16(np.float32(base_d) * MULTS_D[di]))
-        for b in range(nm):
-            mi = b if nm == N_CAND_M else (b * N_CAND_M) // nm
-            ms = float(round_fp16(np.float32(base_m) * MULTS_M[mi]))
-            res = factor_q2k_block(
-                x,
-                w,
-                cfg,
-                init_s=scales,
-                init_t=mins,
-                d_seed=ds,
-                m_seed=ms,
-                pin=cfg.factor_pin,
-                use_spectral=cfg.use_spectral,
-                max_iter=cfg.factor_iters,
-                starts=1,
-                dc_carry=dc_carry,
-            )
-            if nd == N_CAND_D and nm == N_CAND_M:
-                out[di * N_CAND_M + mi] = np.float32(res.objective)
-            else:
-                cb = di >> 2
-                fb = mi >> 2
+    di_idx = np.array(
+        [a if nd == N_CAND_D else (a * N_CAND_D) // nd for a in range(nd)],
+        dtype=np.int64,
+    )
+    mi_idx = np.array(
+        [b if nm == N_CAND_M else (b * N_CAND_M) // nm for b in range(nm)],
+        dtype=np.int64,
+    )
+
+    ds_vec = round_fp16(np.float32(base_d) * MULTS_D[di_idx]).astype(np.float64)
+    ms_vec = round_fp16(np.float32(base_m) * MULTS_M[mi_idx]).astype(np.float64)
+    ds = np.repeat(ds_vec, nm)
+    ms = np.tile(ms_vec, nd)
+    K = ds.size
+
+    safe_ds = np.where(ds[:, None] > 1.0e-30, ds[:, None], 1.0)
+    safe_ms = np.where(ms[:, None] > 1.0e-30, ms[:, None], 1.0)
+    ls0 = np.where(
+        ds[:, None] > 1.0e-30,
+        np.clip(nearest_int((scales_f32[None, :] / safe_ds).astype(np.float32)), 0, 15),
+        0,
+    ).astype(np.uint8)
+    lm0 = np.where(
+        ms[:, None] > 1.0e-30,
+        np.clip(nearest_int((mins_f32[None, :] / safe_ms).astype(np.float32)), 0, 15),
+        0,
+    ).astype(np.uint8)
+
+    def _eval_batch(ls_mat: np.ndarray, lm_mat: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+        d_sub = np.repeat(ds[:, None] * ls_mat.astype(np.float64), 16, axis=1)
+        m_sub = np.repeat(ms[:, None] * lm_mat.astype(np.float64), 16, axis=1)
+        safe = d_sub >= 1.0e-15
+        sd = np.where(safe, d_sub, 1.0)
+        q = np.where(
+            safe,
+            np.clip(nearest_int(((x_f32[None, :] + m_sub) / sd).astype(np.float32)), 0, 3),
+            0,
+        ).astype(np.float64)
+        e = x_f32.astype(np.float64)[None, :] - (d_sub * q - m_sub)
+        sse = np.sum(e * e * w_f64[None, :], axis=1)
+        if cfg.use_spectral and (cfg.dc_lambda != 0.0 or cfg.vw_lambda != 0.0):
+            fw = fold_weights(QK_K, cfg)
+            pen = np.zeros(K, dtype=np.float64)
+            src = e
+            m_len = QK_K // 2
+            lvl = 0
+            while m_len >= 1:
+                dst = src[:, :m_len] + src[:, m_len : 2 * m_len]
+                if m_len == 1 and dc_carry != 0.0:
+                    dst = dst + dc_carry
+                if fw[lvl] != 0.0:
+                    pen += fw[lvl] * np.sum(dst * dst, axis=1)
+                src = dst
+                m_len //= 2
+                lvl += 1
+            return sse + pen, q
+        return sse, q
+
+    obj0, q0 = _eval_batch(ls0, lm0)
+    best_obj = obj0
+
+    for _ in range(max(1, int(cfg.factor_iters))):
+        q_sub = q0.reshape(K, NS, 16)
+        xb = x_f32.astype(np.float64).reshape(1, NS, 16)
+        wb = w_f64.reshape(1, NS, 16)
+        Sqq = np.sum(wb * q_sub * q_sub, axis=2)
+        Swq = np.sum(wb * q_sub, axis=2)
+        Sw = np.sum(wb, axis=2)
+        Sqx = np.sum(wb * q_sub * xb, axis=2)
+        Sx = np.sum(wb * xb, axis=2)
+        det = Sqq * Sw - Swq * Swq
+        valid = det > 1.0e-30
+        s_det = np.where(valid, det, 1.0)
+        s_sw = np.where(Sw > 1.0e-30, Sw, 1.0)
+        sj = np.maximum(np.where(valid, (Sqx * Sw - Swq * Sx) / s_det, 0.0), 0.0)
+        tj = np.where(valid, (sj * Swq - Sx) / s_sw, 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)
+
+        ls1 = np.where(
+            ds[:, None] > 1.0e-30,
+            np.clip(nearest_int((sj / safe_ds).astype(np.float32)), 0, 15),
+            0,
+        ).astype(np.uint8)
+        t_tgt = (
+            np.maximum(0.0, (ds[:, None] * ls1.astype(np.float64) * Swq - Sx) / s_sw)
+            if cfg.factor_coupled_q2k
+            else tj
+        )
+        lm1 = np.where(
+            ms[:, None] > 1.0e-30,
+            np.clip(nearest_int((t_tgt / safe_ms).astype(np.float32)), 0, 15),
+            0,
+        ).astype(np.uint8)
+        obj1, q0 = _eval_batch(ls1, lm1)
+        best_obj = np.minimum(best_obj, obj1)
+
+    best_f32 = best_obj.astype(np.float32)
+    out = np.full(TOTAL_CANDIDATES, np.inf, dtype=np.float32)
+    if nd == N_CAND_D and nm == N_CAND_M:
+        out[:] = best_f32
+    else:
+        for a, di in enumerate(di_idx):
+            cb = int(di) >> 2
+            for b, mi in enumerate(mi_idx):
+                fb = int(mi) >> 2
+                val = best_f32[a * nm + b]
                 for aa in range(4):
                     for bb in range(4):
                         idx = (cb * 4 + aa) * N_CAND_M + fb * 4 + bb
-                        if res.objective < out[idx]:
-                            out[idx] = np.float32(res.objective)
+                        if val < out[idx]:
+                            out[idx] = val
     return out
 
 

pipeline.py

@@ -442,67 +442,55 @@ def _phase2b_one(i: int) -> tuple[np.ndarray, FactorResult]:
     return cand, fac_res
 
 
-def _phase4_one(i: int) -> tuple[Q2KBlock, int, int, float]:
+def _phase4_one(args: tuple[int, int, float, FactorResult]) -> tuple[Q2KBlock, int, int, float]:
+    i, best_i, carry_i, fac = args
     X = _MP_STATE["X"]
     W = _MP_STATE["W"]
     seeds = _MP_STATE["seeds"]
     base_d = _MP_STATE["base_d"]
     base_m = _MP_STATE["base_m"]
-    best = _MP_STATE["best"]
-    carry = _MP_STATE["carry"]
-    standalone = _MP_STATE["standalone"]
     cfg = _MP_STATE["cfg"]
 
     wi = None if W is None else W[i]
-    d16, m16 = candidate_pair(base_d[i], base_m[i], best[i])
-    at_slot = factor_q2k_block(
-        X[i],
-        wi,
-        cfg,
-        init_s=seeds["scales"][i],
-        init_t=seeds["mins"][i],
-        d_seed=float(from_fp16_bits(d16)),
-        m_seed=float(from_fp16_bits(m16)),
-        pin=cfg.factor_pin,
-        use_spectral=cfg.use_spectral,
-        max_iter=cfg.factor_iters,
-        starts=1,
-        dc_carry=0.0,
-    )
     d, m, ls, lm, codes, obj, sse = _assemble_block(
-        X[i], wi, seeds, base_d[i], base_m[i], best[i], carry[i], cfg
+        X[i], wi, seeds, base_d[i], base_m[i], best_i, carry_i, cfg
     )
     cand_hit = 0
     fac_hit = 0
 
     if cfg.use_candidate_floor:
         cd, cm, cls, clm = _candidate_config(
-            seeds["scales"][i], seeds["mins"][i], 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(
-            X[i], wi, cd, cm, cls, clm, ccodes, cfg, float(carry[i])
+            X[i], wi, cd, cm, cls, clm, ccodes, cfg, carry_i
         )
         if cobj < obj - 1.0e-15:
             d, m, ls, lm, codes, obj, sse = cd, cm, cls, clm, ccodes, cobj, csse
             cand_hit = 1
 
     if cfg.use_factor_floor:
-        for fac in (at_slot, standalone[i]):
-            fcodes = refine_codes_hadamard(
-                X[i], wi, fac.d, fac.dmin, fac.ls, fac.lm, fac.codes, cfg, float(carry[i])
+        fcodes = (
+            refine_codes_hadamard(
+                X[i], wi, fac.d, fac.dmin, fac.ls, fac.lm, fac.codes, cfg, carry_i
             )
-            fobj, fsse = objective(
-                X[i], wi, fac.d, fac.dmin, fac.ls, fac.lm, fcodes, cfg, float(carry[i])
-            )
-            if fobj < obj - 1.0e-15:
-                d = float(fac.d)
-                m = float(fac.dmin)
-                ls = fac.ls.copy()
-                lm = fac.lm.copy()
-                codes = fcodes.copy()
-                obj, sse = fobj, fsse
-                fac_hit = 1
+            if carry_i != 0.0
+            else fac.codes
+        )
+        fobj, fsse = (
+            objective(X[i], wi, fac.d, fac.dmin, fac.ls, fac.lm, fcodes, cfg, carry_i)
+            if carry_i != 0.0
+            else (fac.objective, fac.sse)
+        )
+        if fobj < obj - 1.0e-15:
+            d = float(fac.d)
+            m = float(fac.dmin)
+            ls = fac.ls.copy()
+            lm = fac.lm.copy()
+            codes = fcodes.copy()
+            obj, sse = fobj, fsse
+            fac_hit = 1
 
     blk_obj = make_block(d, m, ls, lm, codes)
     return blk_obj, cand_hit, fac_hit, blk_obj.error(X[i])
@@ -596,6 +584,11 @@ class Q2KQuantizer:
             }
         )
 
+        blocks: list[Q2KBlock] = []
+        factor_floor_used = 0
+        candidate_floor_used = 0
+        running_sse = 0.0
+
         if n_workers > 1 and "fork" in mp.get_all_start_methods():
             ctx = mp.get_context("fork")
             chunksize = max(1, min(16, n_blocks // (n_workers * 4) or 1))
@@ -624,77 +617,34 @@ class Q2KQuantizer:
                             eta=eta,
                             sigma=sigma,
                         )
-                pool.close()
-                pool.join()
-            except BaseException:
-                pool.terminate()
-                pool.join()
-                raise
-        else:
-            for i in range(n_blocks):
-                _check_stop()
-                cand_row, fac_res = _phase2b_one(i)
-                candidate_errors[i] = cand_row
-                standalone.append(fac_res)
-                now = time.time()
-                if (i + 1) % step_report == 0 or (i + 1) == n_blocks or (now - last_emit) >= 0.25:
-                    last_emit = now
-                    elapsed = max(now - t_cand, 1.0e-6)
-                    rate = (i + 1) / elapsed
-                    eta = (n_blocks - (i + 1)) / max(rate, 1.0e-9)
-                    _emit(
-                        "Phase 2b: Analog-Guided Factored Candidate Field",
-                        i + 1,
-                        n_blocks,
-                        f"blk {i+1}/{n_blocks} ({rate:.1f} blk/s, ETA {eta:.1f}s) "
-                        f"d={float(fac_res.d):.5f} dmin={float(fac_res.dmin):.5f} "
-                        f"obj={fac_res.objective:.5e} starts={fac_res.starts}",
-                        rate=rate,
-                        eta=eta,
-                        sigma=sigma,
-                    )
-
-        _check_stop()
-        _emit(
-            "Phase 3: 6-State Sieve & 36-State Viterbi DP",
-            0,
-            n_blocks,
-            f"Running Sieve + 36-state Viterbi sequence optimization ({cfg.optimizer.upper()})",
-            sigma=sigma,
-        )
-        best, temperature, _ = _phase3_select(candidate_errors, cfg)
-        carry = _compute_dc_carry(X, seeds, best, base_d, base_m, cfg, row_width)
-        _emit(
-            "Phase 3.9: Viterbi Path & Rolling DC Carry",
-            n_blocks,
-            n_blocks,
-            f"Boltzmann temp={temperature:.4e}, max |DC carry|={float(np.max(np.abs(carry))):.4e}",
-            temperature=temperature,
-            sigma=sigma,
-        )
 
-        _MP_STATE.update(
-            {
-                "best": best,
-                "carry": carry,
-                "standalone": standalone,
-            }
-        )
-
-        blocks: list[Q2KBlock] = []
-        factor_floor_used = 0
-        candidate_floor_used = 0
-        running_sse = 0.0
-        t_asm = time.time()
-        last_emit = 0.0
+                _check_stop()
+                _emit(
+                    "Phase 3: 6-State Sieve & 36-State Viterbi DP",
+                    0,
+                    n_blocks,
+                    f"Running Sieve + 36-state Viterbi sequence optimization ({cfg.optimizer.upper()})",
+                    sigma=sigma,
+                )
+                best, temperature, _ = _phase3_select(candidate_errors, cfg)
+                carry = _compute_dc_carry(X, seeds, best, base_d, base_m, cfg, row_width)
+                _emit(
+                    "Phase 3.9: Viterbi Path & Rolling DC Carry",
+                    n_blocks,
+                    n_blocks,
+                    f"Boltzmann temp={temperature:.4e}, max |DC carry|={float(np.max(np.abs(carry))):.4e}",
+                    temperature=temperature,
+                    sigma=sigma,
+                )
 
-        if n_workers > 1 and "fork" in mp.get_all_start_methods():
-            ctx = mp.get_context("fork")
-            chunksize = max(1, min(16, n_blocks // (n_workers * 4) or 1))
-            pool = ctx.Pool(processes=n_workers)
-            try:
+                t_asm = time.time()
+                last_emit = 0.0
+                p4_args = [
+                    (i, int(best[i]), float(carry[i]), standalone[i])
+                    for i in range(n_blocks)
+                ]
                 for i, (blk_obj, cand_hit, fac_hit, blk_err) in enumerate(
-                    pool.imap(_phase4_one, range(n_blocks), chunksize=chunksize)
+                    pool.imap(_phase4_one, p4_args, chunksize=chunksize)
                 ):
                     _check_stop()
                     blocks.append(blk_obj)
@@ -732,7 +682,53 @@ class Q2KQuantizer:
         else:
             for i in range(n_blocks):
                 _check_stop()
-                blk_obj, cand_hit, fac_hit, blk_err = _phase4_one(i)
+                cand_row, fac_res = _phase2b_one(i)
+                candidate_errors[i] = cand_row
+                standalone.append(fac_res)
+                now = time.time()
+                if (i + 1) % step_report == 0 or (i + 1) == n_blocks or (now - last_emit) >= 0.25:
+                    last_emit = now
+                    elapsed = max(now - t_cand, 1.0e-6)
+                    rate = (i + 1) / elapsed
+                    eta = (n_blocks - (i + 1)) / max(rate, 1.0e-9)
+                    _emit(
+                        "Phase 2b: Analog-Guided Factored Candidate Field",
+                        i + 1,
+                        n_blocks,
+                        f"blk {i+1}/{n_blocks} ({rate:.1f} blk/s, ETA {eta:.1f}s) "
+                        f"d={float(fac_res.d):.5f} dmin={float(fac_res.dmin):.5f} "
+                        f"obj={fac_res.objective:.5e} starts={fac_res.starts}",
+                        rate=rate,
+                        eta=eta,
+                        sigma=sigma,
+                    )
+
+            _check_stop()
+            _emit(
+                "Phase 3: 6-State Sieve & 36-State Viterbi DP",
+                0,
+                n_blocks,
+                f"Running Sieve + 36-state Viterbi sequence optimization ({cfg.optimizer.upper()})",
+                sigma=sigma,
+            )
+            best, temperature, _ = _phase3_select(candidate_errors, cfg)
+            carry = _compute_dc_carry(X, seeds, best, base_d, base_m, cfg, row_width)
+            _emit(
+                "Phase 3.9: Viterbi Path & Rolling DC Carry",
+                n_blocks,
+                n_blocks,
+                f"Boltzmann temp={temperature:.4e}, max |DC carry|={float(np.max(np.abs(carry))):.4e}",
+                temperature=temperature,
+                sigma=sigma,
+            )
+
+            t_asm = time.time()
+            last_emit = 0.0
+            for i in range(n_blocks):
+                _check_stop()
+                blk_obj, cand_hit, fac_hit, blk_err = _phase4_one(
+                    (i, int(best[i]), float(carry[i]), standalone[i])
+                )
                 blocks.append(blk_obj)
                 candidate_floor_used += cand_hit
                 factor_floor_used += fac_hit