File size: 4,909 Bytes
3324725 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 | """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
|