| """Faithful Python port of the original C factored solver. |
| |
| This module is intentionally not on the production path. It exists so the |
| V2 improvements can be measured against the exact reference trajectory: |
| fixed six-iteration 1-D Lloyd, single range initialisation, no fp16-domain |
| polish and no spectral shaping inside the solver. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import numpy as np |
|
|
| from .config import HPCConfig |
| from .fp16 import from_fp16_bits, nearest_int, to_fp16_bits |
| from .q2k import QK_K, codes_from_params, dequantize_params, weighted_sse |
| from .spectral import objective, spectral_penalty |
|
|
|
|
| def reference_factor_1d(v, w, nmax: int = 15): |
| v = np.maximum(np.asarray(v, dtype=np.float64), 0.0) |
| w = np.maximum(np.asarray(w, dtype=np.float64), 0.0) |
| vmax = float(np.max(v)) |
| if vmax < 1.0e-30: |
| return 0.0, np.zeros(v.size, dtype=np.uint8) |
| d = vmax / nmax |
| for _ in range(6): |
| 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 reference_sub_ls(x, w, codes): |
| NS = 16 |
| q = np.asarray(codes, dtype=np.float64).reshape(NS, 16) |
| xb = np.asarray(x, dtype=np.float64).reshape(NS, 16) |
| wb = np.asarray(w, dtype=np.float64).reshape(NS, 16) |
| s = np.zeros(NS) |
| t = np.zeros(NS) |
| for j in range(NS): |
| ww, qq, xx = wb[j], q[j], 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 = 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], t[j] = sj, tj |
| return s, t |
|
|
|
|
| def reference_solve(x, w, cfg: HPCConfig, use_spectral: bool) -> tuple[float, float]: |
| NS = 16 |
| xb = np.asarray(x, dtype=np.float64).reshape(NS, 16) |
| wb = np.asarray(w, dtype=np.float64).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 = xb[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], t[j] = sj, 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 = wb.sum(axis=1) |
| d, ls = reference_factor_1d(s, sw) |
| m, lm = reference_factor_1d(t, sw) |
| L = codes_from_params(x, d, m, ls, lm) |
| obj, _ = objective(x, w, d, m, ls, lm, L, cfg) |
| best = (obj, d, m, ls.copy(), lm.copy(), L.copy()) |
| stall = 0 |
| for _ in range(16): |
| L = codes_from_params(x, d, m, ls, lm) |
| s, t = reference_sub_ls(x, w, L) |
| ws = np.sum(wb * L.astype(np.float64).reshape(NS, 16) ** 2, axis=1) |
| wt = wb.sum(axis=1) |
| d, ls = reference_factor_1d(s, ws) |
| m, lm = reference_factor_1d(t, wt) |
| L = codes_from_params(x, d, m, ls, lm) |
| cur, _ = objective(x, w, d, m, ls, lm, L, cfg) |
| if cur < best[0]: |
| best = (cur, d, m, ls.copy(), lm.copy(), L.copy()) |
| stall = 0 |
| else: |
| stall += 1 |
| if stall >= 2: |
| break |
|
|
| _, d, m, ls, lm, L = best |
| d = float(from_fp16_bits(to_fp16_bits(np.float32(d)))) |
| m = float(from_fp16_bits(to_fp16_bits(np.float32(m)))) |
| L = codes_from_params(x, d, m, ls, lm) |
| e = np.asarray(x, dtype=np.float64) - dequantize_params(d, m, ls, lm, L).astype(np.float64) |
| sse = weighted_sse(x, w, d, m, ls, lm, L) |
| obj = sse + (spectral_penalty(e, cfg) if use_spectral else 0.0) |
| return obj, sse |
|
|
|
|
| def synthetic_cases(seed: int = 123, per_kind: int = 12): |
| rng = np.random.default_rng(seed) |
| cases = [] |
| for _ in range(per_kind): |
| cases.append(("gaussian", rng.normal(0, 0.02, QK_K).astype(np.float32))) |
| for _ in range(per_kind): |
| cases.append(("student-t", (rng.standard_t(3, QK_K) * 0.01).astype(np.float32))) |
| for _ in range(per_kind): |
| mask = rng.random(QK_K) < 0.1 |
| cases.append(("sparse", (rng.normal(0, 0.02, QK_K) * mask).astype(np.float32))) |
| for _ in range(per_kind): |
| mask = rng.random(QK_K) < 0.5 |
| cases.append(("offset-sparse", (rng.normal(0, 0.02, QK_K) * mask + 0.01).astype(np.float32))) |
| return cases |
|
|
