| """HPC spectral objective: folded error pyramid, DC/vesica penalty, shaping. |
| |
| The C engine scores a block with:: |
| |
| E = sum_i w_i e_i^2 |
| + (dc_lambda / n) * (sum_i e_i)^2 |
| + (vw_lambda / n) * sum_{pairs} (e_i + e_{i+n/2})^2 |
| |
| and the fold tree generalises the vesica term to every periodic level. V2 |
| keeps the exact same objective and adds a vectorised implementation of the |
| D6/vesica greedy code descent. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from functools import lru_cache |
|
|
| import numpy as np |
|
|
| from .config import HPCConfig |
| from .q2k import QK_K, dequantize_params |
|
|
|
|
| @lru_cache(maxsize=64) |
| def _fold_weights_cached(n: int, dc_lambda: float, vw_lambda: float, depth: int, gamma: float) -> np.ndarray: |
| levels = int(np.log2(n)) |
| if (1 << levels) != n: |
| raise ValueError("fold pyramid requires a power-of-two length") |
| weights = np.zeros(levels, dtype=np.float64) |
| g = 1.0 |
| for k in range(1, levels + 1): |
| if k == levels: |
| weights[k - 1] = dc_lambda / n |
| elif k <= depth: |
| weights[k - 1] = (vw_lambda / n) * g |
| g *= gamma |
| weights.setflags(write=False) |
| return weights |
|
|
|
|
| def fold_weights(n: int, cfg: HPCConfig) -> np.ndarray: |
| return _fold_weights_cached( |
| n, cfg.dc_lambda, cfg.vw_lambda, cfg.fold_depth, cfg.fold_gamma |
| ) |
|
|
|
|
| def fold_build(errors: np.ndarray) -> np.ndarray: |
| e = np.asarray(errors, dtype=np.float64) |
| n = e.size |
| py = np.empty(n - 1, dtype=np.float64) |
| src = e |
| offset = 0 |
| m = n // 2 |
| while m >= 1: |
| dst = src[:m] + src[m : 2 * m] |
| py[offset : offset + m] = dst |
| src = dst |
| offset += m |
| m //= 2 |
| return py |
|
|
|
|
| def fold_energy(pyramid: np.ndarray, weights: np.ndarray) -> float: |
| acc = 0.0 |
| offset = 0 |
| m = pyramid.size // 2 + 1 |
| level = 0 |
| while m >= 1: |
| if weights[level] != 0.0: |
| v = pyramid[offset : offset + m] |
| acc += float(weights[level] * np.dot(v, v)) |
| offset += m |
| m //= 2 |
| level += 1 |
| return acc |
|
|
|
|
| def spectral_penalty( |
| errors: np.ndarray, cfg: HPCConfig, dc_carry: float = 0.0 |
| ) -> float: |
| if not cfg.use_spectral or (cfg.dc_lambda == 0.0 and cfg.vw_lambda == 0.0): |
| return 0.0 |
| pyramid = fold_build(errors) |
| pyramid[-1] += dc_carry |
| return fold_energy(pyramid, fold_weights(errors.size, cfg)) |
|
|
|
|
| def objective( |
| x: np.ndarray, |
| w: np.ndarray | None, |
| d, |
| dmin, |
| ls: np.ndarray, |
| lm: np.ndarray, |
| codes: np.ndarray, |
| cfg: HPCConfig, |
| dc_carry: float = 0.0, |
| ) -> tuple[float, float]: |
| xh = dequantize_params(d, dmin, ls, lm, codes) |
| e = np.asarray(x, dtype=np.float64) - xh.astype(np.float64) |
| sse = float(np.dot(np.asarray(w, dtype=np.float64), e * e)) if w is not None else float(np.dot(e, e)) |
| return sse + spectral_penalty(e, cfg, dc_carry), sse |
|
|
|
|
| def refine_codes_hadamard( |
| x: np.ndarray, |
| w: np.ndarray | None, |
| d, |
| dmin, |
| ls: np.ndarray, |
| lm: np.ndarray, |
| codes: np.ndarray, |
| cfg: HPCConfig, |
| dc_carry: float = 0.0, |
| ) -> np.ndarray: |
| """Deterministic greedy descent of ``4*vesica^2 + (DC + carry)^2``. |
| |
| Same acceptance rule as the C Phase-4 shaper, but the whole 256-element |
| search is evaluated as four vectorised candidate matrices per pass. |
| """ |
| x = np.asarray(x, dtype=np.float64) |
| out = np.asarray(codes, dtype=np.uint8).copy() |
| ww = np.ones(QK_K, dtype=np.float64) if w is None else np.asarray(w, dtype=np.float64) |
| d_sub = np.repeat(np.float32(d) * np.asarray(ls, dtype=np.float32), 16).astype(np.float64) |
| m_sub = np.repeat(np.float32(dmin) * np.asarray(lm, dtype=np.float32), 16).astype(np.float64) |
| active = d_sub >= 1.0e-15 |
| if not np.any(active): |
| return out |
|
|
| xh = d_sub * out.astype(np.float64) - m_sub |
| e_live = x - xh |
| sse_live = float(np.dot(ww, e_live * e_live)) |
| sse_cap = sse_live * (1.0 + cfg.sse_budget) |
| v_live = e_live[: QK_K // 2] + e_live[QK_K // 2 :] |
| vesica = float(np.dot(v_live, v_live)) |
| dc = float(np.sum(e_live)) + dc_carry |
| metric = 4.0 * vesica + dc * dc |
|
|
| idx = np.arange(QK_K) |
| pi = np.where(idx < QK_K // 2, idx, idx - QK_K // 2) |
| for _ in range(QK_K): |
| best_delta = 0.0 |
| best_k = -1 |
| best_q = 0 |
| for q_try in range(4): |
| e_new = x - (d_sub * q_try - m_sub) |
| delta_e = e_new - e_live |
| sse_alt = sse_live + ww * (e_new * e_new - e_live * e_live) |
| mask = (sse_alt <= sse_cap) & active & (out != q_try) |
| if not np.any(mask): |
| continue |
| v_new = v_live[pi] + delta_e |
| vesica_alt = vesica - v_live[pi] * v_live[pi] + v_new * v_new |
| dc_alt = dc + delta_e |
| delta = metric - (4.0 * vesica_alt + dc_alt * dc_alt) |
| delta = np.where(mask, delta, -np.inf) |
| k = int(np.argmax(delta)) |
| if delta[k] > best_delta: |
| best_delta = float(delta[k]) |
| best_k = k |
| best_q = q_try |
| if best_k < 0: |
| break |
| old_e = e_live[best_k] |
| new_e = x[best_k] - (d_sub[best_k] * best_q - m_sub[best_k]) |
| de = new_e - old_e |
| p = pi[best_k] |
| v_live[p] += de |
| vesica += v_live[p] * v_live[p] - (v_live[p] - de) * (v_live[p] - de) |
| dc += de |
| metric = 4.0 * vesica + dc * dc |
| sse_live += ww[best_k] * (new_e * new_e - old_e * old_e) |
| e_live[best_k] = new_e |
| out[best_k] = np.uint8(best_q) |
| return out |
|
|
|
|
| def dc_nudge_codes( |
| x: np.ndarray, |
| d, |
| dmin, |
| ls: np.ndarray, |
| lm: np.ndarray, |
| codes: np.ndarray, |
| w: np.ndarray | None, |
| carry: float, |
| cfg: HPCConfig, |
| max_passes: int = 4, |
| ) -> np.ndarray: |
| """Spend a bounded SSE budget to cancel the rolling DC residual. |
| |
| Equivalent in spirit to Phase 4.8 of the C engine: the reconstruction |
| target is always the true weight, only the codeword choice changes. |
| """ |
| x = np.asarray(x, dtype=np.float64) |
| out = np.asarray(codes, dtype=np.uint8).copy() |
| ww = np.ones(QK_K, dtype=np.float64) if w is None else np.asarray(w, dtype=np.float64) |
| d_sub = np.repeat(np.float32(d) * np.asarray(ls, dtype=np.float32), 16).astype(np.float64) |
| m_sub = np.repeat(np.float32(dmin) * np.asarray(lm, dtype=np.float32), 16).astype(np.float64) |
| active = d_sub >= 1.0e-15 |
| if not np.any(active): |
| return out |
| e = x - (d_sub * out.astype(np.float64) - m_sub) |
| sse = float(np.dot(ww, e * e)) |
| cap = sse * (1.0 + cfg.sse_budget) |
| residual = float(np.sum(e)) + carry |
| for _ in range(max_passes): |
| best_improvement = 0.0 |
| best_k = -1 |
| best_q = 0 |
| for q_try in range(4): |
| e_new = x - (d_sub * q_try - m_sub) |
| mask = active & (out != q_try) |
| sse_alt = sse + ww * (e_new * e_new - e * e) |
| mask &= sse_alt <= cap |
| if not np.any(mask): |
| continue |
| improvement = np.abs(residual) - np.abs(residual + (e_new - e)) |
| improvement = np.where(mask, improvement, -np.inf) |
| k = int(np.argmax(improvement)) |
| if improvement[k] > best_improvement: |
| best_improvement = float(improvement[k]) |
| best_k = k |
| best_q = q_try |
| if best_k < 0: |
| break |
| old_e = e[best_k] |
| new_e = x[best_k] - (d_sub[best_k] * best_q - m_sub[best_k]) |
| de = new_e - old_e |
| residual += de |
| sse += ww[best_k] * (new_e * new_e - old_e * old_e) |
| e[best_k] = new_e |
| out[best_k] = np.uint8(best_q) |
| if abs(residual) < 1.0e-12: |
| break |
| return out |
|
|
|
|
| def pin_dc_codes_uniform( |
| x: np.ndarray, |
| d: np.ndarray, |
| codes: np.ndarray, |
| unit: np.ndarray, |
| lo: int, |
| hi: int, |
| budget_frac: float = 0.05, |
| max_rounds: int = 96, |
| ) -> tuple[np.ndarray, int]: |
| """Zero-DC pin for single-scale lattices (``Q4_0`` nibble, ``Q3_K`` signed grid). |
| |
| These formats carry only one free fp16 parameter per block, so the exact |
| ``sum(e) == 0`` scale re-solve that gives ``Q2_K`` its free lunch is |
| structurally degenerate (``sum(a) ~= 0`` for symmetric code vectors, and |
| the ``fp16`` scale granularity turns any forced ``d`` move into a large |
| SSE penalty). The pin therefore spends ``budget_frac`` of each block's |
| SSE (an exchange, never a free move) flipping codes toward the DC |
| residual: each round picks, per block, the cheapest-to-DC flip direction |
| and position by the ``SSE cost / |delta DC|`` ratio (cheap "free" |
| strict-improvement flips first), applies it until the budget is spent or |
| the residual is below half a flip quantum. Measured on real tensors: |
| ``Q3_K`` DC RMS −82% at ~0.9% extra SSE; ``Q4_0`` −43% at ~1.3% (the |
| symmetric ``[-8, 7]`` grid has half the DC slack of ``Q3_K``'s signed |
| ``sc`` units, so its residual lands at the flip-quantum floor rather |
| than ~0). Deterministic; safe on any input state. |
| """ |
| X = np.asarray(x, dtype=np.float64) |
| B, n = X.shape |
| codes = np.asarray(codes, dtype=np.int64).copy() |
| unit = np.broadcast_to(np.asarray(unit, dtype=np.float64), codes.shape) |
| if B == 0 or budget_frac <= 0.0: |
| return codes, 0 |
| d = np.asarray(d, dtype=np.float64) |
| d_sub = np.repeat(d[:, None], n, axis=1) |
| rows = np.arange(B) |
| e = X - d_sub * codes * unit |
| dc = e.sum(axis=1) |
| sse0 = np.sum(e * e, axis=1) |
| cap = sse0 * float(budget_frac) |
| spent = np.zeros(B) |
| touched = np.zeros(B, dtype=bool) |
| tol = np.abs(d) * 0.5 |
|
|
| for _ in range(max_rounds): |
| want_cost = np.full(B, np.inf) |
| want_pos = np.full(B, -1, dtype=np.int64) |
| want_dir = np.zeros(B, dtype=np.int64) |
| for direction in (1, -1): |
| if direction == 1: |
| sel = codes < hi |
| else: |
| sel = codes > lo |
| de = np.where(sel, float(direction), 0.0) * unit * -d_sub |
| dE = 2.0 * e * de + de * de |
| push = dc[:, None] * de < 0.0 |
| valid = sel & push |
| ratio = np.where(valid & (np.abs(de) > 1.0e-300), dE / np.abs(np.where(valid, de, 1.0)), np.inf) |
| pmin = np.argmin(ratio, axis=1) |
| rmin = ratio[rows, pmin] |
| cmin = dE[rows, pmin] |
| take = np.isfinite(rmin) & (rmin < want_cost) |
| want_cost = np.where(take, cmin, want_cost) |
| want_pos = np.where(take, pmin, want_pos) |
| want_dir = np.where(take, direction, want_dir) |
| want = (np.abs(dc) > tol) & (want_pos >= 0) & (spent + want_cost <= cap) |
| if not np.any(want): |
| break |
| bi = np.where(want)[0] |
| pos = want_pos[bi] |
| dirv = want_dir[bi] |
| de = -d_sub[bi, pos] * dirv * unit[bi, pos] |
| spent[bi] += 2.0 * e[bi, pos] * de + de * de |
| e[bi, pos] += de |
| codes[bi, pos] += dirv |
| dc[bi] += de |
| touched[bi] = True |
|
|
| return codes, int(np.sum(touched)) |
|
|
|
|
| def pin_dc_codes_uniform_cuda( |
| x: "torch.Tensor", |
| d: "torch.Tensor", |
| codes: "torch.Tensor", |
| unit: "torch.Tensor", |
| lo: int, |
| hi: int, |
| budget_frac: float = 0.05, |
| max_rounds: int = 96, |
| ) -> tuple["torch.Tensor", int]: |
| """GPU port of :func:`pin_dc_codes_uniform` (same greedy exchange). |
| |
| ``x``/``d``/``codes``/``unit`` all live on one CUDA device; ``codes`` is |
| int64 and is updated in place (returned tensor is the same object). |
| Float32 throughout — matches the CUDA pipeline's precision, and the |
| exchange is self-correcting (any round that slightly mis-ranks a flip |
| simply loses one candidate, the residual still never regresses SSE |
| beyond ``budget_frac`` of the block's starting SSE). |
| """ |
| import torch |
|
|
| if budget_frac <= 0.0 or codes.shape[0] == 0: |
| return codes, 0 |
| X = x.float() |
| B, n = X.shape |
| dev = X.device |
| if codes.dtype != torch.int64: |
| codes = codes.to(torch.int64) |
| d = d.to(dev).float() |
| if unit.dim() == 1 and unit.numel() == n: |
| unit = unit[None, :] |
| unit = unit.to(dev).float() |
| if unit.shape[0] == 1 and B != 1: |
| unit = unit.expand(B, n) |
|
|
| d_col = d[:, None] |
| e = X - d_col * codes.float() * unit |
| dc = e.sum(dim=1) |
| sse0 = (e * e).sum(dim=1) |
| cap = sse0 * float(budget_frac) |
| spent = torch.zeros(B, device=dev) |
| rows = torch.arange(B, device=dev) |
| tol = d.abs() * 0.5 |
| touched = 0 |
|
|
| for _ in range(max_rounds): |
| want_cost = torch.full((B,), float("inf"), device=dev) |
| want_pos = torch.full((B,), -1, dtype=torch.int64, device=dev) |
| want_dir = torch.zeros(B, dtype=torch.int64, device=dev) |
| for direction in (1, -1): |
| sel = (codes < hi) if direction == 1 else (codes > lo) |
| de = sel.float() * (float(direction) * unit * -d_col) |
| dE = 2.0 * e * de + de * de |
| valid = sel & (dc[:, None] * de < 0.0) |
| ratio = torch.where(valid, dE.abs() / de.abs().clamp_min(1.0e-30), torch.inf) |
| pmin = ratio.argmin(dim=1) |
| rmin = ratio[rows, pmin] |
| cmin = dE[rows, pmin] |
| take = torch.isfinite(rmin) & (rmin < want_cost) |
| want_cost = torch.where(take, cmin, want_cost) |
| want_pos = torch.where(take, pmin, want_pos) |
| want_dir = torch.where(take, torch.full_like(want_dir, direction), want_dir) |
| want = ( |
| (dc.abs() > tol) |
| & (want_pos >= 0) |
| & (spent + want_cost <= cap) |
| ) |
| if not bool(want.any()): |
| break |
| bi = want.nonzero(as_tuple=False).squeeze(1) |
| pos = want_pos[bi] |
| dirv = want_dir[bi] |
| de = -d[bi] * dirv.float() * unit[bi, pos] |
| spent[bi] += 2.0 * e[bi, pos] * de + de * de |
| e[bi, pos] += de |
| codes[bi, pos] += dirv |
| dc[bi] += de |
| touched += int(bi.numel()) |
|
|
| return codes, touched |
|
|
