Upload 13 files

__init__.py

@@ -6,6 +6,12 @@ Viterbi, spectral/DC shaping, local-safety floors) while replacing the
 original factored solver with a monotone, multi-start, fp16-exact version.
 """
 
+from .analog import (
+    build_block_analogs,
+    subblock_isomoment_analog,
+    subscale_lattice_analog,
+    vesica_gram_analog,
+)
 from .config import HPCConfig
 from .factored import (
     FactorResult,
@@ -26,6 +32,10 @@ __all__ = [
     "QuantizedTensor",
     "TensorStats",
     "QK_K",
+    "build_block_analogs",
+    "subblock_isomoment_analog",
+    "subscale_lattice_analog",
+    "vesica_gram_analog",
     "dequantize_params",
     "factor_1d",
     "factor_candidates",

analog.py

@@ -0,0 +1,213 @@
+"""Geometry-preserving weight analogs for Q2_K bilinear factorization.
+
+At 2 bits per weight (4 levels ``q in {0, 1, 2, 3}`` and 4-bit sub-block
+scales ``Ls, Lm in {0..15}``), raw weight blocks suffer from two structural
+bottlenecks:
+
+1. **Intra-sub-block outlier collapse**: a single heavy-tailed weight stretches
+   the sub-block span ``xmax - xmin``, forcing 13-15 of the 16 weights into
+   one or two quantization bins.
+2. **Inter-sub-block scale crushing**: a single high-variance sub-block inflates
+   the shared superblock scale ``d``, collapsing ``Ls_j`` of the remaining
+   15 sub-blocks toward small integers with severe rounding loss.
+
+Rather than factoring raw weights alone, V2 constructs a bank of
+**geometry-preserving weight analogs** ``x_tilde`` of the eligible weights:
+
+* :func:`subblock_isomoment_analog` maps each 16-weight sub-block through an
+  importance-weighted saturating transport blended with the optimal 4-level
+  quantile spacing, then orthogonally projects back onto the exact weighted
+  mean ``mu_j`` and weighted variance ``sigma_j^2`` of the original sub-block.
+* :func:`vesica_gram_analog` operates in the D6 / Vesica fold basis
+  ``(x_p + x_{p+128}, x_p - x_{p+128})``, regularizing outliers in the whitened
+  eigen-basis and coloring back so the exact 2x2 weighted Gram matrix and means
+  of the two half-blocks are preserved.
+* :func:`subscale_lattice_analog` tempers cross-sub-block dispersion toward the
+  4-bit ``{0..15}`` scale lattice while preserving every sub-block's unit
+  direction, mean, and total block weighted energy.
+
+The discrete Q2_K lattice coordinates ``(Ls, Lm, q)`` discovered on each analog
+are then projected back onto the **true weights** ``x`` via coupled bilinear
+least-squares, guaranteeing that the final emitted block is strictly optimal
+under the true objective.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+from .q2k import NS, QK_K
+
+_Q2_IDEAL_QUANTILES = np.linspace(-1.35, 1.35, 16, dtype=np.float64)
+
+
+def subblock_isomoment_analog(
+    x: np.ndarray,
+    w: np.ndarray | None = None,
+    tau: float = 1.65,
+    beta: float = 0.25,
+    quantile_blend: float = 0.30,
+    enforce_nonpositive_min: bool = False,
+) -> np.ndarray:
+    """Create a sub-block iso-moment analog preserving weighted mean & variance.
+
+    Invariants per 16-weight sub-block ``j`` (when ``enforce_nonpositive_min``
+    is ``False``):
+      * ``sum_k w_{j,k} x_tilde_{j,k} == sum_k w_{j,k} x_{j,k}``
+      * ``sum_k w_{j,k} (x_tilde_{j,k} - mu_j)^2 == sum_k w_{j,k} (x_{j,k} - mu_j)^2``
+      * strict rank-order preservation with respect to ``x_{j,k}`` (when ``beta == 0``)
+    """
+    xb = np.asarray(x, dtype=np.float64).reshape(NS, 16)
+    wb = (
+        np.ones((NS, 16), dtype=np.float64)
+        if w is None
+        else np.asarray(w, dtype=np.float64).reshape(NS, 16)
+    )
+    wsum = np.maximum(np.sum(wb, axis=1, keepdims=True), 1.0e-30)
+    mu = np.sum(wb * xb, axis=1, keepdims=True) / wsum
+    centered = xb - mu
+    var = np.sum(wb * centered * centered, axis=1, keepdims=True) / wsum
+    sigma = np.sqrt(np.maximum(var, 1.0e-30))
+    active = (sigma[:, 0] >= 1.0e-15)
+
+    z = centered / sigma
+    w_mean = wsum / 16.0
+    w_rel = np.power(np.maximum(wb / w_mean, 1.0e-12), beta) if beta != 0.0 else 1.0
+    z_comp = tau * np.tanh((z * w_rel) / max(tau, 1.0e-6))
+
+    if quantile_blend > 0.0:
+        order = np.argsort(z, axis=1, kind="stable")
+        z_quant = np.empty_like(z)
+        rows = np.arange(NS)[:, None]
+        z_quant[rows, order] = _Q2_IDEAL_QUANTILES[None, :]
+        z_analog = (1.0 - quantile_blend) * z_comp + quantile_blend * z_quant
+    else:
+        z_analog = z_comp
+
+    mu_a = np.sum(wb * z_analog, axis=1, keepdims=True) / wsum
+    z_analog = z_analog - mu_a
+    var_a = np.sum(wb * z_analog * z_analog, axis=1, keepdims=True) / wsum
+    scale = np.where(var_a > 1.0e-30, np.sqrt(var / np.maximum(var_a, 1.0e-30)), 0.0)
+    out = np.where(active[:, None], mu + z_analog * scale, xb)
+
+    if enforce_nonpositive_min:
+        orig_min = np.min(xb, axis=1)
+        new_min = np.min(out, axis=1)
+        shift = np.where((orig_min <= 0.0) & (new_min > 0.0), new_min, 0.0)
+        out = out - shift[:, None]
+
+    return out.reshape(QK_K).astype(np.float32)
+
+
+def vesica_gram_analog(
+    x: np.ndarray,
+    w: np.ndarray | None = None,
+    tau: float = 1.9,
+) -> np.ndarray:
+    """Create a D6 / Vesica fold analog preserving the 2x2 half-block Gram matrix.
+
+    Decomposes ``x`` into symmetric ``v = 0.5*(x[:128] + x[128:])`` and
+    antisymmetric ``u = 0.5*(x[:128] - x[128:])`` modes under pair weights
+    ``wp = 0.5*(w[:128] + w[128:])``.  Whitens ``(v, u)``, applies smooth
+    hyperbolic saturation in the principal axes, and colors back with the exact
+    2x2 covariance square-root so ``Gram(v_tilde, u_tilde) == Gram(v, u)``.
+    """
+    x_f64 = np.asarray(x, dtype=np.float64).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)
+    )
+    half = QK_K // 2
+    x0, x1 = x_f64[:half], x_f64[half:]
+    wp = np.maximum(0.5 * (w_f64[:half] + w_f64[half:]), 1.0e-30)
+    wsum = float(np.sum(wp))
+
+    v = 0.5 * (x0 + x1)
+    u = 0.5 * (x0 - x1)
+    mv = float(np.sum(wp * v) / wsum)
+    mu = float(np.sum(wp * u) / wsum)
+    Y = np.vstack([v - mv, u - mu])  # (2, 128)
+
+    G = (Y * wp[None, :]) @ Y.T / wsum  # 2x2 weighted covariance
+    evals, evecs = np.linalg.eigh(G)
+    if float(np.max(evals)) < 1.0e-28:
+        return x_f64.astype(np.float32)
+
+    pos = evals > 1.0e-28
+    inv_sqrt = np.where(pos, 1.0 / np.sqrt(np.maximum(evals, 1.0e-30)), 0.0)
+    W_white = (evecs * inv_sqrt[None, :]) @ evecs.T
+    Z = W_white @ Y  # whitened (2, 128)
+
+    Z_comp = tau * np.tanh(Z / max(tau, 1.0e-6))
+    Z_mean = np.sum(Z_comp * wp[None, :], axis=1, keepdims=True) / wsum
+    Z_comp = Z_comp - Z_mean
+
+    G_z = (Z_comp * wp[None, :]) @ Z_comp.T / wsum
+    ez, Vz = np.linalg.eigh(G_z)
+    pos_z = ez > 1.0e-28
+    if not np.all(pos_z):
+        return x_f64.astype(np.float32)
+    inv_sqrt_z = np.where(pos_z, 1.0 / np.sqrt(np.maximum(ez, 1.0e-30)), 0.0)
+    Z_orth = (Vz * inv_sqrt_z[None, :]) @ Vz.T @ Z_comp
+
+    sqrt_G = (evecs * np.sqrt(np.maximum(evals, 0.0))[None, :]) @ evecs.T
+    Y_tilde = sqrt_G @ Z_orth
+    v_new = mv + Y_tilde[0]
+    u_new = mu + Y_tilde[1]
+    return np.concatenate([v_new + u_new, v_new - u_new]).astype(np.float32)
+
+
+def subscale_lattice_analog(
+    x: np.ndarray,
+    w: np.ndarray | None = None,
+    gamma: float = 0.65,
+) -> np.ndarray:
+    """Create a cross-sub-block analog whose sub-block dispersions fit 4-bit scales.
+
+    Preserves each sub-block's mean ``mu_j`` and unit direction
+    ``(x_j - mu_j) / sigma_j``, while tempering extreme sub-block standard
+    deviations ``sigma_j`` toward the superblock root-mean-square dispersion
+    and rescaling to preserve total centered weighted energy.
+    """
+    xb = np.asarray(x, dtype=np.float64).reshape(NS, 16)
+    wb = (
+        np.ones((NS, 16), dtype=np.float64)
+        if w is None
+        else np.asarray(w, dtype=np.float64).reshape(NS, 16)
+    )
+    wsum = np.maximum(np.sum(wb, axis=1, keepdims=True), 1.0e-30)
+    mu = np.sum(wb * xb, axis=1, keepdims=True) / wsum
+    centered = xb - mu
+    energy_j = np.sum(wb * centered * centered, axis=1, keepdims=True)
+    total_energy = float(np.sum(energy_j))
+    if total_energy <= 1.0e-30:
+        return xb.reshape(QK_K).astype(np.float32)
+
+    sigma = np.sqrt(np.maximum(energy_j / wsum, 1.0e-30))
+    sigma_ref = float(np.sqrt(total_energy / max(float(np.sum(wsum)), 1.0e-30)))
+    target_sigma = np.power(np.maximum(sigma, 1.0e-30), gamma) * (
+        sigma_ref ** (1.0 - gamma)
+    )
+    scaled = np.where(sigma > 1.0e-15, centered * (target_sigma / sigma), 0.0)
+    new_energy = float(np.sum(wb * scaled * scaled))
+    if new_energy > 1.0e-30:
+        scaled = scaled * np.sqrt(total_energy / new_energy)
+    out = mu + scaled
+    return out.reshape(QK_K).astype(np.float32)
+
+
+def build_block_analogs(
+    x: np.ndarray,
+    w: np.ndarray | None = None,
+) -> list[np.ndarray]:
+    """Return the portfolio of geometry-preserving weight analogs for a block."""
+    a1 = subblock_isomoment_analog(
+        x, w, tau=1.60, beta=0.25, quantile_blend=0.30, enforce_nonpositive_min=True
+    )
+    a2 = subblock_isomoment_analog(
+        x, w, tau=2.15, beta=0.15, quantile_blend=0.15, enforce_nonpositive_min=True
+    )
+    a3 = vesica_gram_analog(x, w, tau=1.85)
+    a4 = subscale_lattice_analog(a1, w, gamma=0.70)
+    return [a1, a2, a3, a4]

config.py

@@ -33,6 +33,9 @@ class HPCConfig:
     factor_portfolio: bool = True
     factor_spectral_refine: bool = True
     factor_exact_emit: bool = True
+    use_weight_analogs: bool = True
+    factor_coupled_q2k: bool = True
+    factor_joint_subscales: bool = True
 
     # ---- objective ------------------------------------------------------
     use_spectral: bool = True

factored.py

@@ -28,6 +28,7 @@ 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
@@ -298,6 +299,83 @@ def _ls_init(x, w, cfg: HPCConfig, legacy_1d: bool = False):
     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_init(x_analog, w, cfg, False)
+    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, 1e-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)
@@ -350,14 +428,26 @@ def _solve(
     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)
-        wb = w.reshape(NS, 16)
-        ws = np.sum(wb * q * q, axis=1)
-        wt = np.sum(wb, axis=1)
+        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:
@@ -368,16 +458,39 @@ def _solve(
             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 / m), 0, 15).astype(np.uint8)
+                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, wt, 15, cfg, legacy_1d)
+            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:
+            if pin == 0:
+                d, m = _global_bilinear_dm_solve(
+                    x, w, d, m, ls, lm, codes, cfg, dc_carry, use_spectral
+                )
+            if 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)
+                obj_j, _ = objective(x, w, d_f16, m_f16, ls_j, lm_j, codes_j, cfg, dc_carry)
+                if obj_j < best[0] - 1.0e-15:
+                    d, m, ls, lm, codes = d_f16, m_f16, ls_j, lm_j, codes_j
+                    best = (obj_j, d, m, ls.copy(), lm.copy(), codes.copy())
+                    stall = 0
+                    continue
+
         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())
@@ -465,6 +578,9 @@ def factor_q2k_block(
             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

pipeline.py

@@ -501,18 +501,21 @@ class Q2KQuantizer:
                     candidate_floor_used += 1
 
             if cfg.use_factor_floor:
-                fac = at_slot[i]
-                fobj, fsse = objective(
-                    X[i], wi, fac.d, fac.dmin, fac.ls, fac.lm, fac.codes, cfg, float(carry[i])
-                )
-                if fobj < obj - 1.0e-15:
-                    d = float(fac.d)
-                    m = float(fac.dmin)
-                    ls = fac.ls.copy()
-                    lm = fac.lm.copy()
-                    codes = fac.codes.copy()
-                    obj, sse = fobj, fsse
-                    factor_floor_used += 1
+                for fac in (at_slot[i], standalone[i]):
+                    fcodes = refine_codes_hadamard(
+                        X[i], wi, fac.d, fac.dmin, fac.ls, fac.lm, fac.codes, cfg, float(carry[i])
+                    )
+                    fobj, fsse = objective(
+                        X[i], wi, fac.d, fac.dmin, fac.ls, fac.lm, fcodes, cfg, float(carry[i])
+                    )
+                    if fobj < obj - 1.0e-15:
+                        d = float(fac.d)
+                        m = float(fac.dmin)
+                        ls = fac.ls.copy()
+                        lm = fac.lm.copy()
+                        codes = fcodes.copy()
+                        obj, sse = fobj, fsse
+                        factor_floor_used += 1
 
             blocks.append(make_block(d, m, ls, lm, codes))