"""Grid-free factored Q2_K solver (V2). The reference solver alternates four exact coordinate solves:: q | (d, Ls, dmin, Lm) s_j,t_j | q d,Ls | s dmin,Lm | t V2 keeps that decomposition but fixes three weaknesses: 1. ``factor_1d`` is no longer a fixed six-iteration Lloyd heuristic. It starts from several deterministic candidates and finishes with an exact integer-code coordinate refinement, so its weighted error is monotone. 2. The block solve is multi-start, accepts only objective-decreasing states, and returns the best snapshot rather than the last iterate. 3. The emitted fp16 ``(d, dmin)`` pair is polished in its own domain, and the spectral/DC shaping stage is applied inside the solver. The returned objective therefore equals the objective of the bytes that are written. The result object is immutable and self-describing, replacing the C engine's many out-parameters. """ from __future__ import annotations from dataclasses import dataclass import numpy as np 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 fold_weights, objective, refine_codes_hadamard N_CAND_D = 24 N_CAND_M = 24 TOTAL_CANDIDATES = N_CAND_D * N_CAND_M MULTS_D = np.array( [ 0.780, 0.835, 0.880, 0.915, 0.943, 0.963, 0.978, 0.988, 0.994, 0.997, 0.999, 1.000, 1.002, 1.005, 1.011, 1.021, 1.035, 1.054, 1.080, 1.115, 1.160, 1.215, 1.275, 1.340, ], dtype=np.float32, ) MULTS_M = np.array( [ 0.750, 0.800, 0.840, 0.870, 0.900, 0.920, 0.940, 0.955, 0.970, 0.985, 0.995, 1.000, 1.005, 1.015, 1.030, 1.045, 1.060, 1.080, 1.100, 1.130, 1.160, 1.200, 1.250, 1.300, ], dtype=np.float32, ) CAND_TO_QUHIT = np.array( [0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4, 5, 5, 5, 5], dtype=np.int64, ) @dataclass(slots=True) class FactorResult: d: np.float32 dmin: np.float32 d16: np.uint16 m16: np.uint16 ls: np.ndarray lm: np.ndarray codes: np.ndarray sse: float objective: float iterations: int starts: int def subscales(self) -> np.ndarray: return (self.ls & 0x0F) | ((self.lm & 0x0F) << 4) def _factor_1d_lloyd(v: np.ndarray, w: np.ndarray, nmax: int, iters: int): vmax = float(v.max()) if vmax < 1.0e-30: return 0.0, np.zeros(v.size, dtype=np.uint8) best_d = vmax / nmax best_codes = np.clip(nearest_int(v / best_d), 0, nmax) best_obj = float(np.sum(w * (v - best_d * best_codes) ** 2)) starts = [vmax / nmax] init_codes = np.clip(nearest_int(v / starts[0]), 0, nmax) den0 = float(np.sum(w * init_codes * init_codes)) if den0 > 1.0e-30: starts.append(float(np.sum(w * v * init_codes) / den0)) seen: set[float] = set() for d0 in starts: if not np.isfinite(d0) or d0 <= 0.0: continue key = float(np.float32(d0)) if key in seen: continue seen.add(key) d = float(d0) codes = np.clip(nearest_int(v / d), 0, nmax) for _ in range(iters): den = float(np.sum(w * codes * codes)) if den < 1.0e-30: break dn = float(np.sum(w * v * codes) / den) if dn <= 0.0: break conv = abs(dn - d) <= 1.0e-7 * abs(d) + 1.0e-30 d = dn codes = np.clip(nearest_int(v / d), 0, nmax) if conv: break obj = float(np.sum(w * (v - d * codes) ** 2)) if obj < best_obj: best_obj = obj best_d = d best_codes = codes.copy() 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) codes = np.asarray(codes, dtype=np.int64).copy() d = float(d) if d <= 0.0: return d, codes.astype(np.uint8) 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 = w[i] if wi <= 0.0: continue 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) 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, sweeps) return d, codes 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_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): w = np.asarray(w, dtype=np.float64) 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 def _seeded_init(x, scales, mins, d_seed, m_seed): d = float(d_seed) m = float(m_seed) ls, lm = derive_subscales(scales, mins, d, m) codes = codes_from_params(x, d, m, ls, lm) return d, m, ls, lm, codes 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) _, _, 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) 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 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_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) xb = np.asarray(x, dtype=np.float64).reshape(NS, 16) Sw = np.sum(wb, axis=1) 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, 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( 0.0, (d * ls.astype(np.float64) * Swq - Sx) / np.maximum(Sw, 1e-30) ) m, lm = _factor_1d_for(t_coupled, Sw, 15, cfg, False) codes = codes_from_params(x, d, m, ls, lm) return d, m, ls, lm, codes def _joint_subscale_scan(x, w, d: float, m: float) -> tuple[np.ndarray, np.ndarray]: """Vectorized exact (Ls_j, Lm_j) in {0..15}^2 minimizer per sub-block.""" xb = np.asarray(x, dtype=np.float32).reshape(NS, 16) wb = np.asarray(w, dtype=np.float32).reshape(NS, 16) 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 >= 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).reshape(NS, 256) best_idx = np.argmin(err, axis=1) return (best_idx // 16).astype(np.uint8), (best_idx % 16).astype(np.uint8) def _global_bilinear_dm_solve( x, w, d: float, m: float, ls, lm, codes, cfg: HPCConfig, dc_carry: float, use_spectral: bool ) -> tuple[float, float]: """Exact 2x2 weighted least-squares solve for superblock (d, dmin).""" x_f64 = np.asarray(x, dtype=np.float64).ravel() w_f64 = np.asarray(w, dtype=np.float64).ravel() A = np.repeat(np.asarray(ls, dtype=np.float64), 16) * np.asarray(codes, dtype=np.float64) B = np.repeat(np.asarray(lm, dtype=np.float64), 16) Saa = float(np.sum(w_f64 * A * A)) Sab = float(np.sum(w_f64 * A * B)) Sbb = float(np.sum(w_f64 * B * B)) Sxa = float(np.sum(w_f64 * x_f64 * A)) Sxb = float(np.sum(w_f64 * x_f64 * B)) if use_spectral and cfg.dc_lambda > 0.0: rw = float(cfg.dc_lambda) / QK_K rA = float(np.sum(A)) rB = float(np.sum(B)) rS = float(np.sum(x_f64)) + float(dc_carry) Saa += rw * rA * rA Sab += rw * rA * rB Sbb += rw * rB * rB Sxa += rw * rS * rA Sxb += rw * rS * rB det = Saa * Sbb - Sab * Sab if det > 1.0e-30: d_opt = (Sbb * Sxa - Sab * Sxb) / det m_opt = (Sab * Sxa - Saa * Sxb) / det if d_opt > 0.0 and np.isfinite(d_opt): d = float(d_opt) if m_opt >= 0.0 and np.isfinite(m_opt): m = float(m_opt) return d, m def _sub_ls(x, w, codes): x = np.asarray(x, dtype=np.float64) w = np.asarray(w, dtype=np.float64) q = np.asarray(codes, dtype=np.float64).reshape(NS, 16) xb = x.reshape(NS, 16) wb = w.reshape(NS, 16) 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( x, w, cfg: HPCConfig, d, m, ls, lm, codes, pin: int, 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) ls = np.asarray(ls, dtype=np.uint8).copy() lm = np.asarray(lm, dtype=np.uint8).copy() codes = np.asarray(codes, dtype=np.uint8).copy() d = float(d) m = float(m) obj, _ = objective(x, w, d, m, ls, lm, codes, cfg, dc_carry) best = (obj, d, m, ls.copy(), lm.copy(), codes.copy()) stall = 0 iters = 0 xb_f64 = x.astype(np.float64).reshape(NS, 16) wb_f64 = w.astype(np.float64).reshape(NS, 16) Sw = np.sum(wb_f64, axis=1) Sx = np.sum(wb_f64 * xb_f64, axis=1) coupled = (not legacy_1d) and cfg.factor_coupled_q2k for it in range(max(1, max_iter)): iters = it + 1 codes = codes_from_params(x, d, m, ls, lm) s, t = _sub_ls(x, w, codes) q = codes.astype(np.float64).reshape(NS, 16) Sqq = np.sum(wb_f64 * q * q, axis=1) Swq = np.sum(wb_f64 * q, axis=1) if coupled: q_bar = np.where(Sw > 1.0e-30, Swq / np.maximum(Sw, 1.0e-30), 0.0) ws = np.maximum(Sqq - Swq * q_bar, 1.0e-12 * Sw) else: ws = Sqq wt = Sw if pin & 1: if d > 1.0e-30: ls = np.clip(nearest_int(s / d), 0, 15).astype(np.uint8) else: ls = np.zeros(NS, dtype=np.uint8) else: d_new, ls_new = _factor_1d_for(s, ws, 15, cfg, legacy_1d) d, ls = float(d_new), ls_new t_target = ( np.maximum(0.0, (d * ls.astype(np.float64) * Swq - Sx) / np.maximum(Sw, 1.0e-30)) if coupled else t ) if pin & 2: if m > 1.0e-30: lm = np.clip(nearest_int(t_target / m), 0, 15).astype(np.uint8) else: lm = np.zeros(NS, dtype=np.uint8) else: m_new, lm_new = _factor_1d_for(t_target, wt, 15, cfg, legacy_1d) m, lm = float(m_new), lm_new codes = codes_from_params(x, d, m, ls, lm) 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: best = (cur, d, m, ls.copy(), lm.copy(), codes.copy()) stall = 0 else: stall += 1 if stall >= 2: 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 def _emit_polish(x, w, cfg, d, m, ls, lm, codes, dc_carry): d16 = np.uint16(to_fp16_bits(np.float32(d))) m16 = np.uint16(to_fp16_bits(np.float32(m))) d = float(from_fp16_bits(d16)) m = float(from_fp16_bits(m16)) codes = codes_from_params(x, d, m, ls, lm) obj, sse = objective(x, w, d, m, ls, lm, codes, cfg, dc_carry) best = (obj, sse, d, m, d16, m16, ls.copy(), lm.copy(), codes.copy()) if cfg.factor_exact_emit and cfg.factor_ulp_radius > 0: d_choices = fp16_ulp_neighbors(d16, cfg.factor_ulp_radius) m_choices = fp16_ulp_neighbors(m16, cfg.factor_ulp_radius) for db in d_choices: dd = float(from_fp16_bits(db)) for mb in m_choices: mm = float(from_fp16_bits(mb)) c = codes_from_params(x, dd, mm, ls, lm) o, s = objective(x, w, dd, mm, ls, lm, c, cfg, dc_carry) if o < best[0] - 1.0e-15: best = (o, s, dd, mm, np.uint16(db), np.uint16(mb), ls.copy(), lm.copy(), c) obj, sse, d, m, d16, m16, ls, lm, codes = best if cfg.factor_spectral_refine: refined = refine_codes_hadamard( x, w, d, m, ls, lm, codes, cfg, dc_carry ) o, s = objective(x, w, d, m, ls, lm, refined, cfg, dc_carry) if o < obj - 1.0e-15: obj, sse, codes = o, s, refined return d, m, d16, m16, ls, lm, codes, sse, obj def factor_q2k_block( x, w=None, cfg: HPCConfig | None = None, init_s=None, init_t=None, d_seed: float = 0.0, m_seed: float = 0.0, pin: int = 0, use_spectral: bool | None = None, max_iter: int | None = None, starts: int | None = None, dc_carry: float = 0.0, ) -> FactorResult: cfg = (cfg or HPCConfig()).validate() x = np.asarray(x, dtype=np.float32).reshape(QK_K) w = ( np.ones(QK_K, dtype=np.float32) if w is None else np.asarray(w, dtype=np.float32).reshape(QK_K) ) use_spectral = cfg.use_spectral if use_spectral is None else use_spectral max_iter = cfg.factor_standalone_iters if max_iter is None else max_iter starts = cfg.factor_starts if starts is None else starts 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)) if (not legacy_1d) and cfg.use_weight_analogs and d_seed <= 1.0e-30 and m_seed <= 1.0e-30: for x_analog in build_block_analogs(x, w): out.append(_analog_init(x, w, x_analog, cfg)) return out best = None total_iters = 0 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 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_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 ) 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 def factor_candidates( x, w, scales, mins, base_d, base_m, cfg: HPCConfig, dc_carry: float = 0.0, ) -> np.ndarray: """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) 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 val < out[idx]: out[idx] = val return out def factor_only_solve( x, w, cfg: HPCConfig, dc_carry: float = 0.0 ) -> FactorResult: return factor_q2k_block( x, w, cfg, pin=0, use_spectral=cfg.use_spectral, max_iter=cfg.factor_standalone_iters, starts=cfg.factor_starts, dc_carry=dc_carry, ) def candidate_pair(base_d, base_m, cidx: int) -> tuple[np.uint16, np.uint16]: di = int(cidx) // N_CAND_M mi = int(cidx) % N_CAND_M d16 = np.uint16(to_fp16_bits(np.float32(base_d) * MULTS_D[di])) m16 = np.uint16(to_fp16_bits(np.float32(base_m) * MULTS_M[mi])) return d16, m16