Upload 14 files

gguf.py

@@ -376,6 +376,8 @@ def write_quantized_gguf(
     tensor_names: set[str] | None = None,
     limit: int = 0,
     verbose: bool = True,
+    progress_callback=None,
+    stop_event=None,
 ) -> dict[str, float]:
     """Quantize a source GGUF with the V2 engine and emit a valid GGUF.
 
@@ -461,6 +463,8 @@ def write_quantized_gguf(
         fout.write(b"\x00" * (data_start - fout.tell()))
 
         for idx, (ti, out_type, out_size, offset) in enumerate(out_infos):
+            if stop_event is not None and stop_event.is_set():
+                raise InterruptedError("Quantization cancelled by user")
             target = data_start + offset
             if target > fout.tell():
                 fout.write(b"\x00" * (target - fout.tell()))
@@ -484,10 +488,21 @@ def write_quantized_gguf(
                         f"  Q2_K {ti.name:48s} row={row:6d} "
                         f"rmse={res.stats.rmse:.6e} imatrix={'yes' if imp is not None else 'no'}"
                     )
+                if progress_callback is not None:
+                    progress_callback(
+                        idx + 1,
+                        len(out_infos),
+                        ti.name,
+                        "Q2_K",
+                        float(res.stats.rmse),
+                        imp is not None,
+                    )
             else:
                 reader.copy_tensor_raw_to(ti.name, fout)
                 if verbose and idx % 25 == 0:
                     print(f"  copy {ti.name}")
+                if progress_callback is not None:
+                    progress_callback(idx + 1, len(out_infos), ti.name, "copy", None, False)
 
     rmse = float(np.sqrt(total_rmse / total_elements)) if total_elements else 0.0
     return {

reference.py

@@ -0,0 +1,138 @@
+"""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