| """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 |
|
|
