| """End-to-end Q2_K pipeline preserving the HPC machinery. |
| |
| Pipeline stages (mirroring the C engine, with V2 fixes): |
| |
| 1. kurtosis-gated outlier clamp and batched GGML seeds; |
| 2. vectorised WLS superblock solve; |
| 3. factored candidate field + standalone factored floor; |
| 4. sequential Sieve over coarse/fine symbolic states; |
| 5. 36-state Viterbi sequence optimisation; |
| 6. candidate-aware factor consistency pass; |
| 7. rolling DC residual carry; |
| 8. sub-block Sieve assembly, LS refit, fp16-ULP polish, spectral shaping; |
| 9. objective-consistent local-safety floors and DC nudge. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| from math import ceil |
| import multiprocessing as mp |
| import os |
| import time |
|
|
| import numpy as np |
|
|
| from .anneal import thermal_refine_tensor |
| from .config import HPCConfig |
| from .factored import ( |
| CAND_TO_QUHIT, |
| MULTS_D, |
| MULTS_M, |
| N_CAND_M, |
| TOTAL_CANDIDATES, |
| FactorResult, |
| candidate_pair, |
| factor_candidates, |
| factor_only_solve, |
| factor_q2k_block, |
| ) |
| from .fp16 import from_fp16_bits, nearest_int, round_fp16, to_fp16_bits |
| from .q2k import ( |
| NS, |
| QK_K, |
| Q2KBlock, |
| codes_from_params, |
| codes_from_params_batch, |
| dequantize_params, |
| dequantize_params_batch, |
| derive_subscales, |
| make_block, |
| pack_codes, |
| pack_scales, |
| pin_q2k_dc_residual, |
| unpack_codes, |
| weighted_sse, |
| ) |
| from .seeds import block_clamp_values, build_seeds, seed_clamp_value |
| from .sieve import block_chain_edges, sieve_measure_chain, subblock_chain_edges |
| from .spectral import dc_nudge_codes, fold_weights, objective, refine_codes_hadamard |
| from .viterbi import robust_temperature, viterbi_select |
|
|
|
|
| @dataclass(slots=True) |
| class TensorStats: |
| n_blocks: int |
| total_sse: float |
| weighted_sse: float |
| rmse: float |
| sigma: float |
| rmse_over_sigma: float |
| temperature: float |
| factor_floor_used: int |
| candidate_floor_used: int |
| therm_improved: int = 0 |
|
|
|
|
| @dataclass(slots=True) |
| class QuantizedTensor: |
| blocks: list[Q2KBlock] |
| scales: np.ndarray |
| qs: np.ndarray |
| dm: np.ndarray |
| best_candidate: np.ndarray |
| factor_objectives: np.ndarray |
| config: HPCConfig |
| stats: TensorStats |
|
|
| def dequantize(self) -> np.ndarray: |
| return np.concatenate([b.dequantize() for b in self.blocks]) |
|
|
|
|
| def _phase2_wls(X, W, seeds, cfg: HPCConfig): |
| n = X.shape[0] |
| per_blk_clamp = block_clamp_values(X, cfg).reshape(n, 1) |
| clipped = np.clip(X, -per_blk_clamp, per_blk_clamp).astype(np.float32) |
| base_d = seeds["d"].astype(np.float64).copy() |
| base_m = seeds["dmin"].astype(np.float64).copy() |
| ls = seeds["ls"].copy() |
| lm = seeds["lm"].copy() |
| scales = seeds["scales"].astype(np.float64) |
| mins = seeds["mins"].astype(np.float64) |
| ww = np.ones_like(X, dtype=np.float64) if W is None else W.astype(np.float64) |
|
|
| for _ in range(5): |
| codes = codes_from_params_batch(clipped, base_d, base_m, ls, lm).astype(np.float64) |
| A = np.repeat(ls.astype(np.float64), 16, axis=1) * codes |
| B = np.repeat(lm.astype(np.float64), 16, axis=1) |
| Saa = np.sum(ww * A * A, axis=1) |
| Sab = np.sum(ww * A * B, axis=1) |
| Sbb = np.sum(ww * B * B, axis=1) |
| Sxa = np.sum(ww * clipped.astype(np.float64) * A, axis=1) |
| Sxb = np.sum(ww * clipped.astype(np.float64) * B, axis=1) |
| det = Saa * Sbb - Sab * Sab |
| valid = det > 1.0e-30 |
| safe_det = np.where(valid, det, 1.0) |
| d_new = (Sbb * Sxa - Sab * Sxb) / safe_det |
| m_new = (Sab * Sxa - Saa * Sxb) / safe_det |
| upd_d = valid & (d_new > 0.0) & np.isfinite(d_new) |
| upd_m = valid & (m_new > 0.0) & np.isfinite(m_new) |
| base_d = np.where(upd_d, round_fp16(d_new.astype(np.float32)).astype(np.float64), base_d) |
| base_m = np.where(upd_m, round_fp16(m_new.astype(np.float32)).astype(np.float64), base_m) |
|
|
| safe_d = np.where(base_d > 1.0e-15, base_d, 1.0) |
| safe_m = np.where(base_m > 1.0e-15, base_m, 1.0) |
| ls = np.where( |
| base_d[:, None] > 1.0e-15, |
| np.clip(nearest_int(scales / safe_d[:, None]), 0, 15), |
| 0, |
| ).astype(np.uint8) |
| lm = np.where( |
| base_m[:, None] > 1.0e-15, |
| np.clip(nearest_int(mins / safe_m[:, None]), 0, 15), |
| 0, |
| ).astype(np.uint8) |
|
|
| return base_d.astype(np.float32), base_m.astype(np.float32) |
|
|
|
|
| def _bin_geometry(): |
| di = np.arange(TOTAL_CANDIDATES) // N_CAND_M |
| mi = np.arange(TOTAL_CANDIDATES) % N_CAND_M |
| return CAND_TO_QUHIT[di], CAND_TO_QUHIT[mi] |
|
|
|
|
| def _sieve_marginals(candidate_errors, graph_blocks, stride, temperature, cfg): |
| n_blocks = candidate_errors.shape[0] |
| pad = graph_blocks * stride - n_blocks |
| ce = candidate_errors.astype(np.float64) |
| if pad > 0: |
| ce = np.concatenate([ce, np.full((pad, TOTAL_CANDIDATES), np.inf, dtype=np.float64)], axis=0) |
| grouped = ce.reshape(graph_blocks, stride, TOTAL_CANDIDATES) |
| finite = np.isfinite(grouped) |
| counts = np.sum(finite, axis=1) |
| valid = counts > 0 |
| sums = np.sum(np.where(finite, grouped, 0.0), axis=1) |
| agg = np.where(valid, sums / np.maximum(counts, 1), np.inf) |
| mn = np.min(agg, axis=1, keepdims=True) |
| mn = np.where(np.isfinite(mn), mn, 0.0) |
| expo = np.where(valid, (agg - mn) / (2.0 * max(temperature, 1.0e-12)), np.inf) |
| coeff = np.exp(-expo).reshape(graph_blocks, 6, 4, 6, 4) |
|
|
| coarse = np.sum(coeff, axis=(2, 3, 4)) |
| fine = np.sum(coeff, axis=(1, 2, 4)) |
| cn = np.linalg.norm(coarse, axis=1, keepdims=True) |
| fn = np.linalg.norm(fine, axis=1, keepdims=True) |
| coarse = np.where(cn > 1.0e-30, coarse / np.maximum(cn, 1.0e-30), 0.0) |
| fine = np.where(fn > 1.0e-30, fine / np.maximum(fn, 1.0e-30), 0.0) |
|
|
| amps = np.zeros((2 * graph_blocks, 6), dtype=np.float64) |
| amps[0::2] = coarse |
| amps[1::2] = fine |
| marg, _ = sieve_measure_chain( |
| amps, |
| block_chain_edges(graph_blocks), |
| same_penalty=cfg.same_penalty, |
| log_slack=cfg.log_slack, |
| deterministic=True, |
| ) |
| return marg[0::2], marg[1::2] |
|
|
|
|
| def _phase3_select(candidate_errors, cfg): |
| n_blocks = candidate_errors.shape[0] |
| if cfg.optimizer != "hpc" or n_blocks < 2: |
| best = np.argmin(candidate_errors, axis=1).astype(np.int64) |
| return best, 0.0, np.zeros(n_blocks, dtype=np.int64) |
|
|
| stride = max(1, ceil(n_blocks / cfg.graph_max_blocks)) |
| graph_blocks = ceil(n_blocks / stride) |
| temperature = robust_temperature(candidate_errors, cfg) |
| coarse_marg, fine_marg = _sieve_marginals( |
| candidate_errors, graph_blocks, stride, temperature, cfg |
| ) |
| best, path = viterbi_select( |
| candidate_errors, |
| coarse_marg, |
| fine_marg, |
| temperature, |
| stride, |
| graph_blocks, |
| cfg, |
| ) |
| return best, temperature, path |
|
|
|
|
| def _compute_dc_carry(X, seeds, best, base_d, base_m, cfg, row_width): |
| n = X.shape[0] |
| carry = np.zeros(n, dtype=np.float32) |
| blocks_per_row = row_width // QK_K if row_width and row_width % QK_K == 0 else 0 |
| rolling = 0.0 |
| for i in range(n): |
| if blocks_per_row > 0 and i % blocks_per_row == 0: |
| rolling = 0.0 |
| d16, m16 = candidate_pair(base_d[i], base_m[i], best[i]) |
| d = float(from_fp16_bits(d16)) |
| m = float(from_fp16_bits(m16)) |
| ls, lm = derive_subscales(seeds["scales"][i], seeds["mins"][i], d, m) |
| codes = codes_from_params(X[i], d, m, ls, lm) |
| xh = dequantize_params(d, m, ls, lm, codes) |
| carry[i] = np.float32(cfg.dc_decay * rolling) |
| dc_res = float(np.sum(X[i].astype(np.float64) - xh.astype(np.float64))) |
| rolling = carry[i] + dc_res if cfg.carry_cumulative else dc_res |
| return carry |
|
|
|
|
| def _per_sub_top_states(x, w, d, m, k=6): |
| k = max(1, min(int(k), 6)) |
| xb = np.asarray(x, dtype=np.float32).reshape(NS, 16) |
| if w is None: |
| wb = np.ones((NS, 16), dtype=np.float32) |
| else: |
| 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) |
| flat = err.reshape(NS, 256) |
| order = np.argsort(flat, axis=1, kind="stable")[:, :k] |
| state_err = np.take_along_axis(flat, order, axis=1).astype(np.float64) |
| state_ls = (order // 16).astype(np.uint8) |
| state_lm = (order % 16).astype(np.uint8) |
| return state_ls, state_lm, state_err |
|
|
|
|
| def _select_sub_states(state_ls, state_lm, state_err, cfg): |
| if not cfg.subblock_sieve: |
| return state_ls[:, 0].copy(), state_lm[:, 0].copy() |
| amps = np.zeros((NS, 6), dtype=np.float64) |
| for j in range(NS): |
| err0 = state_err[j, 0] |
| spread = float(state_err[j, 5] - state_err[j, 0]) |
| temp = spread * 0.3 if spread > 1.0e-15 else 0.1 |
| temp = max(temp, 1.0e-12) |
| amp = np.exp(-(state_err[j] - err0) / temp) |
| norm = float(np.linalg.norm(amp)) |
| amps[j] = amp / norm if norm > 1.0e-30 else 1.0 / np.sqrt(6.0) |
| marg, _ = sieve_measure_chain( |
| amps, |
| subblock_chain_edges(NS), |
| same_penalty=cfg.same_penalty, |
| log_slack=cfg.log_slack, |
| deterministic=True, |
| ) |
| pick = np.argmax(marg, axis=1) |
| ls = np.array([state_ls[j, pick[j]] for j in range(NS)], dtype=np.uint8) |
| lm = np.array([state_lm[j, pick[j]] for j in range(NS)], dtype=np.uint8) |
| return ls, lm |
|
|
|
|
| def _ls_refit(x, w, d, m, ls, lm, codes, max_d, max_m): |
| x = np.asarray(x, dtype=np.float64) |
| w = np.ones(QK_K, dtype=np.float64) if w is None else np.asarray(w, dtype=np.float64) |
| 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 * A * A)) |
| Sab = float(np.sum(w * A * B)) |
| Sbb = float(np.sum(w * B * B)) |
| Sxa = float(np.sum(w * x * A)) |
| Sxb = float(np.sum(w * x * B)) |
| det = Saa * Sbb - Sab * Sab |
| if det <= 1.0e-30: |
| return float(d), float(m) |
| d_new = (Sbb * Sxa - Sab * Sxb) / det |
| m_new = (Sab * Sxa - Saa * Sxb) / det |
| if d_new <= 0.0 or d_new >= 4.0 * (max_d + 1.0e-10): |
| d_new = d |
| if m_new <= 0.0 or m_new >= 4.0 * (max_m + 1.0e-10): |
| m_new = m |
| return float(round_fp16(np.float32(d_new))), float(round_fp16(np.float32(m_new))) |
|
|
|
|
| def _ulp_polish(x, w, d, m, ls, lm, radius): |
| if radius <= 0: |
| codes = codes_from_params(x, d, m, ls, lm) |
| return float(d), float(m), codes, weighted_sse(x, w, d, m, ls, lm, codes) |
| d16 = int(to_fp16_bits(np.float32(d))) |
| m16 = int(to_fp16_bits(np.float32(m))) |
| d_bits = np.arange(max(0, d16 - radius), min(0x7BFF, d16 + radius) + 1, dtype=np.uint16) |
| m_bits = np.arange(max(0, m16 - radius), min(0x7BFF, m16 + radius) + 1, dtype=np.uint16) |
| dd = from_fp16_bits(d_bits) |
| mm = from_fp16_bits(m_bits) |
| d_grid = np.repeat(dd, mm.size) |
| m_grid = np.tile(mm, dd.size) |
| ls_rep = np.repeat(np.asarray(ls, dtype=np.float32), 16) |
| lm_rep = np.repeat(np.asarray(lm, dtype=np.float32), 16) |
| d_sub = d_grid[:, None] * ls_rep[None, :] |
| m_sub = m_grid[:, None] * lm_rep[None, :] |
| safe = d_sub >= np.float32(1.0e-15) |
| safe_d = np.where(safe, d_sub, np.float32(1.0)) |
| q = np.clip(nearest_int((x[None, :] + m_sub) / safe_d), 0, 3).astype(np.float32) |
| q = np.where(safe, q, 0.0) |
| deq = d_sub * q - m_sub |
| ww = np.ones(QK_K, dtype=np.float64) if w is None else np.asarray(w, dtype=np.float64) |
| err = np.sum(ww[None, :] * (np.asarray(x, dtype=np.float64)[None, :] - deq) ** 2, axis=1) |
| k = int(np.argmin(err)) |
| codes = q[k].astype(np.uint8) |
| return float(d_grid[k]), float(m_grid[k]), codes, float(err[k]) |
|
|
|
|
| def _per_sub_scan(x, w, d, m): |
| state_ls, state_lm, _ = _per_sub_top_states(x, w, d, m, k=1) |
| return state_ls[:, 0], state_lm[:, 0] |
|
|
|
|
| def _assemble_block(x, w, seeds, base_d, base_m, cidx, dc_carry, cfg): |
| d16, m16 = candidate_pair(base_d, base_m, cidx) |
| d = float(from_fp16_bits(d16)) |
| m = float(from_fp16_bits(m16)) |
| ls, lm = derive_subscales(seeds["scales"], seeds["mins"], d, m) |
| max_d, max_m = d, m |
| prev = None |
| for _ in range(5): |
| s_ls, s_lm, s_err = _per_sub_top_states(x, w, d, m, cfg.subblock_states) |
| ls, lm = _select_sub_states(s_ls, s_lm, s_err, cfg) |
| codes = codes_from_params(x, d, m, ls, lm) |
| d, m = _ls_refit(x, w, d, m, ls, lm, codes, max_d, max_m) |
| cur = (int(to_fp16_bits(np.float32(d))), int(to_fp16_bits(np.float32(m)))) |
| if cur == prev: |
| break |
| prev = cur |
|
|
| codes = codes_from_params(x, d, m, ls, lm) |
| obj, sse = objective(x, w, d, m, ls, lm, codes, cfg, dc_carry) |
| best = (obj, sse, d, m, ls.copy(), lm.copy(), codes.copy()) |
|
|
| ud, um, uc, _ = _ulp_polish(x, w, d, m, ls, lm, cfg.ulp_radius) |
| uobj, usse = objective(x, w, ud, um, ls, lm, uc, cfg, dc_carry) |
| if uobj < best[0] - 1.0e-15: |
| best = (uobj, usse, ud, um, ls.copy(), lm.copy(), uc.copy()) |
|
|
| sls, slm = _per_sub_scan(x, w, ud, um) |
| scodes = codes_from_params(x, ud, um, sls, slm) |
| sobj, ssse = objective(x, w, ud, um, sls, slm, scodes, cfg, dc_carry) |
| if sobj < best[0] - 1.0e-15: |
| best = (sobj, ssse, ud, um, sls.copy(), slm.copy(), scodes.copy()) |
|
|
| _, _, bd, bm, bls, blm, bc = best |
| refined = refine_codes_hadamard(x, w, bd, bm, bls, blm, bc, cfg, dc_carry) |
| robj, rsse = objective(x, w, bd, bm, bls, blm, refined, cfg, dc_carry) |
| if robj < best[0] - 1.0e-15: |
| best = (robj, rsse, bd, bm, bls.copy(), blm.copy(), refined.copy()) |
|
|
| _, sse, d, m, ls, lm, codes = best |
| return d, m, ls, lm, codes, best[0], sse |
|
|
|
|
| def _batch_objective( |
| X: np.ndarray, |
| W_f64: np.ndarray, |
| d: np.ndarray, |
| m: np.ndarray, |
| ls: np.ndarray, |
| lm: np.ndarray, |
| codes: np.ndarray, |
| carry: np.ndarray, |
| cfg: HPCConfig, |
| ) -> tuple[np.ndarray, np.ndarray]: |
| xh = dequantize_params_batch(d, m, ls, lm, codes).astype(np.float64) |
| e = X.astype(np.float64) - xh |
| sse = np.sum(W_f64 * e * e, axis=1) |
| if not cfg.use_spectral or (cfg.dc_lambda == 0.0 and cfg.vw_lambda == 0.0): |
| return sse, sse |
| fw = fold_weights(QK_K, cfg) |
| pen = np.zeros(X.shape[0], dtype=np.float64) |
| src = e |
| m_len = QK_K // 2 |
| lvl = 0 |
| while m_len >= 1: |
| dst = src[:, :m_len] + src[:, m_len : 2 * m_len] |
| if m_len == 1: |
| dst = dst + carry[:, None].astype(np.float64) |
| if fw[lvl] != 0.0: |
| pen += fw[lvl] * np.sum(dst * dst, axis=1) |
| src = dst |
| m_len //= 2 |
| lvl += 1 |
| return sse + pen, sse |
|
|
|
|
| def _batch_extended_dm_solve_and_ulp( |
| X_f64: np.ndarray, |
| W_f64: np.ndarray, |
| ls: np.ndarray, |
| lm: np.ndarray, |
| codes: np.ndarray, |
| d_cur: np.ndarray, |
| m_cur: np.ndarray, |
| carry: np.ndarray, |
| cfg: HPCConfig, |
| ) -> tuple[np.ndarray, np.ndarray]: |
| """Closed-form 2x2 Cramer solve + O(1) quadratic ULP polish on true weighted SSE.""" |
| N = X_f64.shape[0] |
| A = np.repeat(ls.astype(np.float64), 16, axis=1) * codes.astype(np.float64) |
| B = np.repeat(lm.astype(np.float64), 16, axis=1) |
| Saa = np.sum(W_f64 * A * A, axis=1) |
| Sab = np.sum(W_f64 * A * B, axis=1) |
| Sbb = np.sum(W_f64 * B * B, axis=1) |
| Sxa = np.sum(W_f64 * X_f64 * A, axis=1) |
| Sxb = np.sum(W_f64 * X_f64 * B, axis=1) |
|
|
| det = Saa * Sbb - Sab * Sab |
| valid = det > 1.0e-30 |
| s_det = np.where(valid, det, 1.0) |
| d_new = (Sbb * Sxa - Sab * Sxb) / s_det |
| m_new = (Sab * Sxa - Saa * Sxb) / s_det |
| d_opt = np.where(valid & (d_new > 0.0) & np.isfinite(d_new), d_new, d_cur) |
| m_opt = np.where(valid & (m_new >= 0.0) & np.isfinite(m_new), m_new, m_cur) |
|
|
| rad = max(int(cfg.ulp_radius), 1) |
| deltas = np.arange(-rad, rad + 1, dtype=np.int32) |
| d16_c = to_fp16_bits(d_opt.astype(np.float32)).astype(np.int32) |
| m16_c = to_fp16_bits(m_opt.astype(np.float32)).astype(np.int32) |
| d_bits = np.clip(d16_c[:, None] + deltas[None, :], 0, 0x7BFF).astype(np.uint16) |
| m_bits = np.clip(m16_c[:, None] + deltas[None, :], 0, 0x7BFF).astype(np.uint16) |
| dd = from_fp16_bits(d_bits).astype(np.float64)[:, :, None] |
| mm = from_fp16_bits(m_bits).astype(np.float64)[:, None, :] |
|
|
| q_form = ( |
| Saa[:, None, None] * (dd * dd) |
| - 2.0 * Sab[:, None, None] * (dd * mm) |
| + Sbb[:, None, None] * (mm * mm) |
| - 2.0 * Sxa[:, None, None] * dd |
| + 2.0 * Sxb[:, None, None] * mm |
| ).reshape(N, -1) |
| best_idx = np.argmin(q_form, axis=1) |
| K = deltas.size |
| rows = np.arange(N) |
| return dd[rows, best_idx // K, 0].astype(np.float32), mm[rows, 0, best_idx % K].astype(np.float32) |
|
|
|
|
| def _batch_algebraic_refine_codes( |
| X_f64: np.ndarray, |
| W_f64: np.ndarray, |
| d: np.ndarray, |
| m: np.ndarray, |
| ls: np.ndarray, |
| lm: np.ndarray, |
| codes: np.ndarray, |
| carry: np.ndarray, |
| cfg: HPCConfig, |
| ) -> np.ndarray: |
| """Closed-form 1D quadratic coordinate projection for Vesica + DC + weighted SSE.""" |
| N = X_f64.shape[0] |
| out = codes.astype(np.int64).copy() |
| lam_vw = (float(cfg.vw_lambda) / QK_K) if cfg.use_spectral else 0.0 |
| lam_dc = (float(cfg.dc_lambda) / QK_K) if cfg.use_spectral else 0.0 |
| if lam_vw == 0.0 and lam_dc == 0.0: |
| return out.astype(np.uint8) |
|
|
| d_sub = np.repeat(d.astype(np.float64)[:, None] * ls.astype(np.float64), 16, axis=1) |
| m_sub = np.repeat(m.astype(np.float64)[:, None] * lm.astype(np.float64), 16, axis=1) |
| active = d_sub >= 1.0e-15 |
| e = X_f64 - (d_sub * out - m_sub) |
| sse_live = np.sum(W_f64 * e * e, axis=1) |
| sse_cap = sse_live * (1.0 + cfg.sse_budget) |
| half = QK_K // 2 |
| v = e[:, :half] + e[:, half:] |
| dc = np.sum(e, axis=1) + carry.astype(np.float64) |
| denom = (W_f64 + lam_vw + lam_dc) * (d_sub * d_sub) |
| safe_den = np.where(active & (denom > 1.0e-30), denom, 1.0) |
|
|
| v_full = np.concatenate([v, v], axis=1) |
| grad = d_sub * (W_f64 * e + lam_vw * v_full + lam_dc * dc[:, None]) |
| dq = np.where(active, nearest_int((grad / safe_den).astype(np.float32)), 0) |
| q_new = np.clip(out + dq, 0, 3) |
|
|
| for b in range(N): |
| changed = np.where(q_new[b] != out[b])[0] |
| if changed.size == 0: |
| continue |
| xb, wb, dsb, msb = X_f64[b], W_f64[b], d_sub[b], m_sub[b] |
| eb, vb, dcb, sseb, capb = e[b], v[b], float(dc[b]), float(sse_live[b]), float(sse_cap[b]) |
| e_cand = xb[changed] - (dsb[changed] * q_new[b, changed] - msb[changed]) |
| de_cand = e_cand - eb[changed] |
| dsse_cand = wb[changed] * (e_cand * e_cand - eb[changed] * eb[changed]) |
| pi_cand = np.where(changed < half, changed, changed - half) |
| v_old = vb[pi_cand] |
| v_new = v_old + de_cand |
| d_spec = lam_vw * (v_old * v_old - v_new * v_new) + lam_dc * (dcb * dcb - (dcb + de_cand) ** 2) |
| order = changed[np.argsort(-(d_spec - dsse_cand))] |
| for k in order: |
| q_try = int(q_new[b, k]) |
| if q_try == int(out[b, k]): |
| continue |
| e_new = xb[k] - (dsb[k] * q_try - msb[k]) |
| de = e_new - eb[k] |
| dsse = wb[k] * (e_new * e_new - eb[k] * eb[k]) |
| if sseb + dsse > capb: |
| continue |
| p = k if k < half else k - half |
| v_np = vb[p] + de |
| dc_n = dcb + de |
| spec_gain = lam_vw * (vb[p] * vb[p] - v_np * v_np) + lam_dc * (dcb * dcb - dc_n * dc_n) |
| if spec_gain > dsse + 1.0e-18: |
| out[b, k] = q_try |
| eb[k] = e_new |
| vb[p] = v_np |
| dcb = dc_n |
| sseb += dsse |
| return out.astype(np.uint8) |
|
|
|
|
| def _refine_trajectory_algebraic( |
| X: np.ndarray, |
| X_f64: np.ndarray, |
| W_f64: np.ndarray, |
| Xb: np.ndarray, |
| Wb: np.ndarray, |
| Sw: np.ndarray, |
| Sx: np.ndarray, |
| safe_Sw: np.ndarray, |
| d0: np.ndarray, |
| m0: np.ndarray, |
| ls0: np.ndarray, |
| lm0: np.ndarray, |
| codes0: np.ndarray, |
| carry: np.ndarray, |
| cfg: HPCConfig, |
| ): |
| N = X.shape[0] |
| best_d, best_m = d0.copy(), m0.copy() |
| best_ls, best_lm = ls0.copy(), lm0.copy() |
| best_codes = codes0.copy() |
| best_obj, best_sse = _batch_objective(X, W_f64, best_d, best_m, best_ls, best_lm, best_codes, carry, cfg) |
| d, m, ls, lm, codes = best_d.copy(), best_m.copy(), best_ls.copy(), best_lm.copy(), best_codes.copy() |
| ls_deltas = np.array([-3, -2, -1, 0, 1, 2, 3], dtype=np.int64) |
| lm_deltas = np.array([-1, 0, 1], dtype=np.int64) |
| n_idx = np.arange(N)[:, None] |
| s_idx = np.arange(NS)[None, :] |
|
|
| for _ in range(cfg.refine_passes): |
| d_new, m_new = _batch_extended_dm_solve_and_ulp(X_f64, W_f64, ls, lm, codes, d, m, carry, cfg) |
| obj_dm, sse_dm = _batch_objective(X, W_f64, d_new, m_new, ls, lm, codes, carry, cfg) |
| imp_dm = sse_dm < best_sse - 1.0e-15 |
| if np.any(imp_dm): |
| best_obj = np.where(imp_dm, obj_dm, best_obj) |
| best_sse = np.where(imp_dm, sse_dm, best_sse) |
| best_d = np.where(imp_dm, d_new, best_d) |
| best_m = np.where(imp_dm, m_new, best_m) |
| best_ls = np.where(imp_dm[:, None], ls, best_ls) |
| best_lm = np.where(imp_dm[:, None], lm, best_lm) |
| best_codes = np.where(imp_dm[:, None], codes, best_codes) |
| d, m = best_d.copy(), best_m.copy() |
|
|
| codes_re = codes_from_params_batch(X, d, m, ls, lm) |
| obj_re, sse_re = _batch_objective(X, W_f64, d, m, ls, lm, codes_re, carry, cfg) |
| imp_re = sse_re < best_sse - 1.0e-15 |
| if np.any(imp_re): |
| best_obj = np.where(imp_re, obj_re, best_obj) |
| best_sse = np.where(imp_re, sse_re, best_sse) |
| best_codes = np.where(imp_re[:, None], codes_re, best_codes) |
| codes = best_codes.copy() |
|
|
| q = codes.astype(np.float64).reshape(N, NS, 16) |
| Sqq = np.sum(Wb * q * q, axis=2) |
| Swq = np.sum(Wb * q, axis=2) |
| Sqx = np.sum(Wb * q * Xb, axis=2) |
| q_bar = Swq / safe_Sw |
| Vqq = np.maximum(Sqq - Swq * q_bar, 0.0) |
| Cqx = Sqx - Sx * q_bar |
| s_star = np.where(Vqq > 1.0e-30, np.maximum(Cqx / np.maximum(Vqq, 1.0e-30), 0.0), d[:, None] * ls) |
| sd = np.where(d[:, None] > 1.0e-15, d[:, None], 1.0) |
| sm = np.where(m[:, None] > 1.0e-15, m[:, None], 1.0) |
| ls_center = np.where(d[:, None] > 1.0e-15, np.clip(nearest_int((s_star / sd).astype(np.float32)), 0, 15), 0) |
|
|
| c_ls_base = np.clip(ls_center[:, :, None] + ls_deltas[None, None, :], 0, 15).astype(np.float64) |
| d_sub_base = d.astype(np.float64)[:, None, None] * c_ls_base |
| t_opt_base = np.maximum(0.0, (d_sub_base * Swq[:, :, None] - Sx[:, :, None]) / safe_Sw[:, :, None]) |
| lm_center = np.where( |
| m[:, None, None] > 1.0e-15, |
| np.clip(nearest_int((t_opt_base / sm[:, :, None]).astype(np.float32)), 0, 15), |
| 0, |
| ) |
|
|
| cand_ls = np.repeat(c_ls_base, 3, axis=2) |
| cand_lm = np.clip( |
| np.repeat(lm_center, 3, axis=2) + np.tile(lm_deltas, 7)[None, None, :], 0, 15 |
| ).astype(np.float64) |
|
|
| d_sub_c = d.astype(np.float64)[:, None, None] * cand_ls |
| m_sub_c = m.astype(np.float64)[:, None, None] * cand_lm |
| safe_dc = d_sub_c >= 1.0e-15 |
| sd_c = np.where(safe_dc, d_sub_c, 1.0) |
| q_c = np.where( |
| safe_dc[:, :, :, None], |
| np.clip( |
| nearest_int(((Xb[:, :, None, :] + m_sub_c[:, :, :, None]) / sd_c[:, :, :, None]).astype(np.float32)), |
| 0, |
| 3, |
| ), |
| 0, |
| ).astype(np.float64) |
| e_c = Xb[:, :, None, :] - (d_sub_c[:, :, :, None] * q_c - m_sub_c[:, :, :, None]) |
| err_c = np.sum(Wb[:, :, None, :] * e_c * e_c, axis=3) |
| pick = np.argmin(err_c, axis=2) |
| ls_try = cand_ls[n_idx, s_idx, pick].astype(np.uint8) |
| lm_try = cand_lm[n_idx, s_idx, pick].astype(np.uint8) |
| codes_try = codes_from_params_batch(X, d, m, ls_try, lm_try) |
| d_try, m_try = _batch_extended_dm_solve_and_ulp(X_f64, W_f64, ls_try, lm_try, codes_try, d, m, carry, cfg) |
| obj_try, sse_try = _batch_objective(X, W_f64, d_try, m_try, ls_try, lm_try, codes_try, carry, cfg) |
| imp_sub = sse_try < best_sse - 1.0e-15 |
| if np.any(imp_sub): |
| best_obj = np.where(imp_sub, obj_try, best_obj) |
| best_sse = np.where(imp_sub, sse_try, best_sse) |
| best_d = np.where(imp_sub, d_try, best_d) |
| best_m = np.where(imp_sub, m_try, best_m) |
| best_ls = np.where(imp_sub[:, None], ls_try, best_ls) |
| best_lm = np.where(imp_sub[:, None], lm_try, best_lm) |
| best_codes = np.where(imp_sub[:, None], codes_try, best_codes) |
| d, m, ls, lm, codes = best_d.copy(), best_m.copy(), best_ls.copy(), best_lm.copy(), best_codes.copy() |
|
|
| refined = _batch_algebraic_refine_codes(X_f64, W_f64, best_d, best_m, best_ls, best_lm, best_codes, carry, cfg) |
| d_r, m_r = _batch_extended_dm_solve_and_ulp(X_f64, W_f64, best_ls, best_lm, refined, best_d, best_m, carry, cfg) |
| obj_r, sse_r = _batch_objective(X, W_f64, d_r, m_r, best_ls, best_lm, refined, carry, cfg) |
| imp_r = (sse_r <= best_sse * (1.0 + cfg.sse_budget)) & (obj_r < best_obj - 1.0e-15) |
| if np.any(imp_r): |
| best_obj = np.where(imp_r, obj_r, best_obj) |
| best_sse = np.where(imp_r, sse_r, best_sse) |
| best_d = np.where(imp_r, d_r, best_d) |
| best_m = np.where(imp_r, m_r, best_m) |
| best_codes = np.where(imp_r[:, None], refined, best_codes) |
|
|
| return best_d, best_m, best_ls, best_lm, best_codes, best_obj, best_sse |
|
|
|
|
| def _phase4_algebraic_batch( |
| X: np.ndarray, |
| W: np.ndarray | None, |
| scales_seed: np.ndarray, |
| mins_seed: np.ndarray, |
| base_d: np.ndarray, |
| base_m: np.ndarray, |
| best: np.ndarray, |
| carry: np.ndarray, |
| standalone: list[FactorResult], |
| cfg: HPCConfig, |
| ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, int, int]: |
| """Batched algebraic Phase-4 solver across N blocks (replaces brute-force 16x16 & ULP grids).""" |
| N = X.shape[0] |
| X_f64 = X.astype(np.float64) |
| W_f64 = np.ones_like(X_f64) if W is None else W.astype(np.float64) |
| Xb = X_f64.reshape(N, NS, 16) |
| Wb = W_f64.reshape(N, NS, 16) |
| Sw = np.sum(Wb, axis=2) |
| Sx = np.sum(Wb * Xb, axis=2) |
| safe_Sw = np.maximum(Sw, 1.0e-30) |
|
|
| di = best // N_CAND_M |
| mi = best % N_CAND_M |
| d_v = round_fp16(base_d * MULTS_D[di]) |
| m_v = round_fp16(base_m * MULTS_M[mi]) |
| sd_v = np.where(d_v[:, None] > 1.0e-15, d_v[:, None], 1.0) |
| sm_v = np.where(m_v[:, None] > 1.0e-15, m_v[:, None], 1.0) |
| ls_v = np.where(d_v[:, None] > 1.0e-15, np.clip(nearest_int(scales_seed / sd_v), 0, 15), 0).astype(np.uint8) |
| lm_v = np.where(m_v[:, None] > 1.0e-15, np.clip(nearest_int(mins_seed / sm_v), 0, 15), 0).astype(np.uint8) |
| codes_v = codes_from_params_batch(X, d_v, m_v, ls_v, lm_v) |
|
|
| d1, m1, ls1, lm1, c1, _, sse1 = _refine_trajectory_algebraic( |
| X, X_f64, W_f64, Xb, Wb, Sw, Sx, safe_Sw, d_v, m_v, ls_v, lm_v, codes_v, carry, cfg |
| ) |
| cand_hits = 0 |
| fac_hits = 0 |
|
|
| if cfg.use_factor_floor: |
| d_f = np.array([f.d for f in standalone], dtype=np.float32) |
| m_f = np.array([f.dmin for f in standalone], dtype=np.float32) |
| ls_f = np.stack([f.ls for f in standalone], axis=0) |
| lm_f = np.stack([f.lm for f in standalone], axis=0) |
| codes_f = np.stack([f.codes for f in standalone], axis=0) |
| d2, m2, ls2, lm2, c2, _, sse2 = _refine_trajectory_algebraic( |
| X, X_f64, W_f64, Xb, Wb, Sw, Sx, safe_Sw, d_f, m_f, ls_f, lm_f, codes_f, carry, cfg |
| ) |
| take2 = sse2 < sse1 - 1.0e-15 |
| fac_hits = int(np.sum(take2)) |
| d1 = np.where(take2, d2, d1) |
| m1 = np.where(take2, m2, m1) |
| ls1 = np.where(take2[:, None], ls2, ls1) |
| lm1 = np.where(take2[:, None], lm2, lm1) |
| c1 = np.where(take2[:, None], c2, c1) |
|
|
| return d1, m1, ls1, lm1, c1, cand_hits, fac_hits |
|
|
|
|
| def _apply_dc_nudge(X, W, blocks, cfg, row_width): |
| n = X.shape[0] |
| blocks_per_row = row_width // QK_K if row_width and row_width % QK_K == 0 else 0 |
| rolling = 0.0 |
| for i in range(n): |
| if blocks_per_row > 0 and i % blocks_per_row == 0: |
| rolling = 0.0 |
| block = blocks[i] |
| carry = np.float32(cfg.dc_decay * rolling) |
| d = float(block.d_float) |
| m = float(block.dmin_float) |
| ls, lm = block.unpack_scales() |
| codes = unpack_codes(block.qs) |
| w = None if W is None else W[i] |
| codes = dc_nudge_codes(X[i], d, m, ls, lm, codes, w, float(carry), cfg) |
| blocks[i] = Q2KBlock( |
| scales=pack_scales(ls, lm), |
| qs=pack_codes(codes), |
| d=block.d, |
| dmin=block.dmin, |
| ) |
| xh = blocks[i].dequantize() |
| dc_res = float(np.sum(X[i].astype(np.float64) - xh.astype(np.float64))) |
| rolling = float(carry) + dc_res if cfg.carry_cumulative else dc_res |
| return blocks |
|
|
|
|
| _MP_STATE: dict[str, object] = {} |
|
|
|
|
| def _resolve_workers(cfg: HPCConfig, n_blocks: int) -> int: |
| if cfg.workers == 1 or n_blocks < 4: |
| return 1 |
| if cfg.workers > 1: |
| return min(int(cfg.workers), n_blocks) |
| if n_blocks < 16: |
| return 1 |
| cpu = os.cpu_count() or 1 |
| return min(cpu, n_blocks) |
|
|
|
|
| def _phase2b_one(i: int) -> tuple[np.ndarray, FactorResult]: |
| X = _MP_STATE["X"] |
| W = _MP_STATE["W"] |
| seeds = _MP_STATE["seeds"] |
| base_d = _MP_STATE["base_d"] |
| base_m = _MP_STATE["base_m"] |
| cfg = _MP_STATE["cfg"] |
| wi = None if W is None else W[i] |
| cand = factor_candidates( |
| X[i], |
| wi, |
| seeds["scales"][i], |
| seeds["mins"][i], |
| base_d[i], |
| base_m[i], |
| cfg, |
| ) |
| fac_res = factor_only_solve(X[i], wi, cfg) |
| return cand, fac_res |
|
|
|
|
| def _phase4_one(args: tuple[int, int, float, FactorResult]) -> tuple[Q2KBlock, int, int, float]: |
| i, best_i, carry_i, fac = args |
| X = _MP_STATE["X"] |
| W = _MP_STATE["W"] |
| seeds = _MP_STATE["seeds"] |
| base_d = _MP_STATE["base_d"] |
| base_m = _MP_STATE["base_m"] |
| cfg = _MP_STATE["cfg"] |
|
|
| wi = None if W is None else W[i] |
| d, m, ls, lm, codes, obj, sse = _assemble_block( |
| X[i], wi, seeds, base_d[i], base_m[i], best_i, carry_i, cfg |
| ) |
| cand_hit = 0 |
| fac_hit = 0 |
|
|
| if cfg.use_candidate_floor: |
| cd, cm, cls, clm = _candidate_config( |
| seeds["scales"][i], seeds["mins"][i], base_d[i], base_m[i], best_i |
| ) |
| ccodes = codes_from_params(X[i], cd, cm, cls, clm) |
| cobj, csse = objective( |
| X[i], wi, cd, cm, cls, clm, ccodes, cfg, carry_i |
| ) |
| if cobj < obj - 1.0e-15: |
| d, m, ls, lm, codes, obj, sse = cd, cm, cls, clm, ccodes, cobj, csse |
| cand_hit = 1 |
|
|
| if cfg.use_factor_floor: |
| fcodes = ( |
| refine_codes_hadamard( |
| X[i], wi, fac.d, fac.dmin, fac.ls, fac.lm, fac.codes, cfg, carry_i |
| ) |
| if carry_i != 0.0 |
| else fac.codes |
| ) |
| fobj, fsse = ( |
| objective(X[i], wi, fac.d, fac.dmin, fac.ls, fac.lm, fcodes, cfg, carry_i) |
| if carry_i != 0.0 |
| else (fac.objective, fac.sse) |
| ) |
| 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 |
| fac_hit = 1 |
|
|
| blk_obj = make_block(d, m, ls, lm, codes) |
| return blk_obj, cand_hit, fac_hit, blk_obj.error(X[i]) |
|
|
|
|
| class Q2KQuantizer: |
| def __init__(self, config: HPCConfig | None = None): |
| self.config = (config or HPCConfig()).validate() |
|
|
| def quantize( |
| self, |
| weights: np.ndarray, |
| importance: np.ndarray | None = None, |
| row_width: int | None = None, |
| status_callback=None, |
| stop_event=None, |
| ) -> QuantizedTensor: |
| def _emit(phase: str, done: int, total: int, detail: str, **extra) -> None: |
| if status_callback is not None: |
| status_callback( |
| { |
| "phase": phase, |
| "done": done, |
| "total": total, |
| "detail": detail, |
| **extra, |
| } |
| ) |
|
|
| def _check_stop() -> None: |
| if stop_event is not None and stop_event.is_set(): |
| raise InterruptedError("Quantization cancelled by user") |
|
|
| cfg = self.config |
| x = np.asarray(weights, dtype=np.float32) |
| if x.size == 0: |
| raise ValueError("weight tensor is empty") |
| if x.size % QK_K != 0: |
| raise ValueError("weight count must be a multiple of 256") |
| if x.ndim == 2 and row_width is None and x.shape[1] % QK_K == 0: |
| row_width = int(x.shape[1]) |
| x = x.ravel() |
| n_blocks = x.size // QK_K |
|
|
| if cfg.device != "cpu": |
| from .cuda_backend import cuda_is_available, quantize_cuda |
|
|
| if cfg.device.startswith("cuda") or ( |
| cfg.device == "auto" and cfg.workers == 0 and n_blocks >= 32 and cuda_is_available() |
| ): |
| dev_str = "cuda" if cfg.device == "auto" else cfg.device |
| res = quantize_cuda( |
| weights, |
| importance, |
| row_width, |
| cfg, |
| status_callback=status_callback, |
| stop_event=stop_event, |
| device=dev_str, |
| ) |
| if cfg.verbose: |
| _report(res.stats) |
| return res |
|
|
| X = x.reshape(n_blocks, QK_K) |
|
|
| if importance is not None: |
| imp = np.asarray(importance, dtype=np.float32).ravel() |
| if imp.size != x.size: |
| raise ValueError("importance must match the weight count") |
| if np.any(imp < 0.0) or not np.all(np.isfinite(imp)): |
| raise ValueError("importance must be finite and non-negative") |
| W_raw = imp.reshape(n_blocks, QK_K) |
| if cfg.temper_imatrix: |
| sig2 = np.mean(X.astype(np.float64) ** 2, axis=1, keepdims=True) |
| W = ( |
| np.sqrt(np.maximum(W_raw.astype(np.float64), 1.0e-6)) |
| * np.sqrt(sig2 + X.astype(np.float64) ** 2) |
| ) |
| w_mean = np.mean(W, axis=1, keepdims=True) |
| W = np.where(w_mean > 1.0e-30, W / np.maximum(w_mean, 1.0e-30), 1.0).astype(np.float32) |
| else: |
| W = W_raw |
| else: |
| W = None |
|
|
| sigma = float(np.sqrt(np.mean(x.astype(np.float64) ** 2))) |
| _check_stop() |
| clamp_val = seed_clamp_value(x, cfg) |
| _emit( |
| "Phase 1: Kurtosis Clamp & Batched GGML Seeds", |
| 0, |
| n_blocks, |
| f"blocks={n_blocks} sigma={sigma:.5f} clamp={clamp_val:.5f} imatrix={'yes' if W is not None else 'no'}", |
| sigma=sigma, |
| ) |
| seeds = build_seeds(X, W, cfg, clamp_val) |
| _check_stop() |
| _emit( |
| "Phase 2a: Vectorized Superblock WLS", |
| 0, |
| n_blocks, |
| f"Solving 5-iteration WLS across {n_blocks} superblocks", |
| sigma=sigma, |
| ) |
| base_d, base_m = _phase2_wls(X, W, seeds, cfg) |
|
|
| candidate_errors = np.full((n_blocks, TOTAL_CANDIDATES), np.inf, dtype=np.float32) |
| standalone: list[FactorResult] = [] |
| t_cand = time.time() |
| last_emit = 0.0 |
| step_report = max(1, min(32, n_blocks // 20 or 1)) |
| n_workers = _resolve_workers(cfg, n_blocks) |
| _MP_STATE.update( |
| { |
| "X": X, |
| "W": W, |
| "seeds": seeds, |
| "base_d": base_d, |
| "base_m": base_m, |
| "cfg": cfg, |
| } |
| ) |
|
|
| blocks: list[Q2KBlock] = [] |
| factor_floor_used = 0 |
| candidate_floor_used = 0 |
| running_sse = 0.0 |
|
|
| if n_workers > 1 and "fork" in mp.get_all_start_methods(): |
| ctx = mp.get_context("fork") |
| chunksize = max(1, min(16, n_blocks // (n_workers * 4) or 1)) |
| pool = ctx.Pool(processes=n_workers) |
| try: |
| for i, (cand_row, fac_res) in enumerate( |
| pool.imap(_phase2b_one, range(n_blocks), chunksize=chunksize) |
| ): |
| _check_stop() |
| candidate_errors[i] = cand_row |
| standalone.append(fac_res) |
| now = time.time() |
| if (i + 1) % step_report == 0 or (i + 1) == n_blocks or (now - last_emit) >= 0.25: |
| last_emit = now |
| elapsed = max(now - t_cand, 1.0e-6) |
| rate = (i + 1) / elapsed |
| eta = (n_blocks - (i + 1)) / max(rate, 1.0e-9) |
| _emit( |
| f"Phase 2b: Analog-Guided Factored Field ({n_workers} workers)", |
| i + 1, |
| n_blocks, |
| f"blk {i+1}/{n_blocks} ({rate:.1f} blk/s, ETA {eta:.1f}s) " |
| f"d={float(fac_res.d):.5f} dmin={float(fac_res.dmin):.5f} " |
| f"obj={fac_res.objective:.5e} starts={fac_res.starts}", |
| rate=rate, |
| eta=eta, |
| sigma=sigma, |
| ) |
| pool.close() |
| pool.join() |
| except BaseException: |
| pool.terminate() |
| pool.join() |
| raise |
| else: |
| for i in range(n_blocks): |
| _check_stop() |
| cand_row, fac_res = _phase2b_one(i) |
| candidate_errors[i] = cand_row |
| standalone.append(fac_res) |
| now = time.time() |
| if (i + 1) % step_report == 0 or (i + 1) == n_blocks or (now - last_emit) >= 0.25: |
| last_emit = now |
| elapsed = max(now - t_cand, 1.0e-6) |
| rate = (i + 1) / elapsed |
| eta = (n_blocks - (i + 1)) / max(rate, 1.0e-9) |
| _emit( |
| "Phase 2b: Analog-Guided Factored Candidate Field", |
| i + 1, |
| n_blocks, |
| f"blk {i+1}/{n_blocks} ({rate:.1f} blk/s, ETA {eta:.1f}s) " |
| f"d={float(fac_res.d):.5f} dmin={float(fac_res.dmin):.5f} " |
| f"obj={fac_res.objective:.5e} starts={fac_res.starts}", |
| rate=rate, |
| eta=eta, |
| sigma=sigma, |
| ) |
|
|
| _MP_STATE.clear() |
| _check_stop() |
| _emit( |
| "Phase 3: 6-State Sieve & 36-State Viterbi DP", |
| 0, |
| n_blocks, |
| f"Running Sieve + 36-state Viterbi sequence optimization ({cfg.optimizer.upper()})", |
| sigma=sigma, |
| ) |
| best, temperature, _ = _phase3_select(candidate_errors, cfg) |
| carry = _compute_dc_carry(X, seeds, best, base_d, base_m, cfg, row_width) |
| _emit( |
| "Phase 3.9: Viterbi Path & Rolling DC Carry", |
| n_blocks, |
| n_blocks, |
| f"Boltzmann temp={temperature:.4e}, max |DC carry|={float(np.max(np.abs(carry))):.4e}", |
| temperature=temperature, |
| sigma=sigma, |
| ) |
|
|
| t_asm = time.time() |
| batch_sz = 256 |
| for start in range(0, n_blocks, batch_sz): |
| _check_stop() |
| stop = min(start + batch_sz, n_blocks) |
| W_chunk = None if W is None else W[start:stop] |
| d_b, m_b, ls_b, lm_b, codes_b, c_hits, f_hits = _phase4_algebraic_batch( |
| X[start:stop], |
| W_chunk, |
| seeds["scales"][start:stop], |
| seeds["mins"][start:stop], |
| base_d[start:stop], |
| base_m[start:stop], |
| best[start:stop], |
| carry[start:stop], |
| standalone[start:stop], |
| cfg, |
| ) |
| candidate_floor_used += c_hits |
| factor_floor_used += f_hits |
| for k in range(stop - start): |
| blk_obj = make_block(d_b[k], m_b[k], ls_b[k], lm_b[k], codes_b[k]) |
| blocks.append(blk_obj) |
| running_sse += blk_obj.error(X[start + k]) |
|
|
| now = time.time() |
| elapsed = max(now - t_asm, 1.0e-6) |
| rate = stop / elapsed |
| eta = (n_blocks - stop) / max(rate, 1.0e-9) |
| cur_rmse = float(np.sqrt(running_sse / (stop * QK_K))) |
| _emit( |
| "Phase 4: Algebraic Coupled Bilinear Assembly", |
| stop, |
| n_blocks, |
| f"blk {stop}/{n_blocks} ({rate:.1f} blk/s, ETA {eta:.1f}s) " |
| f"rmse={cur_rmse:.6e} rmse/sigma={cur_rmse/max(sigma, 1e-15):.4f} " |
| f"floors(cand={candidate_floor_used}, fac={factor_floor_used})", |
| rmse=cur_rmse, |
| sigma=sigma, |
| temperature=temperature, |
| cand_floors=candidate_floor_used, |
| fac_floors=factor_floor_used, |
| rate=rate, |
| eta=eta, |
| ) |
|
|
| therm_improved = 0 |
| if cfg.therm_passes > 0 and n_blocks > 0: |
| _check_stop() |
| t_therm = time.time() |
| _emit( |
| "Phase 4.5: Thermal Refinery (Pair Heat-Bath)", |
| 0, |
| n_blocks, |
| f"annealing {n_blocks} blocks: passes={cfg.therm_passes} " |
| f"sweeps={cfg.therm_sweeps} T0={cfg.therm_t0} decay={cfg.therm_t_decay} " |
| f"seed={cfg.therm_seed}", |
| sigma=sigma, |
| temperature=temperature, |
| ) |
| therm_d = np.empty(n_blocks, dtype=np.float64) |
| therm_m = np.empty(n_blocks, dtype=np.float64) |
| therm_ls = np.empty((n_blocks, NS), dtype=np.uint8) |
| therm_lm = np.empty((n_blocks, NS), dtype=np.uint8) |
| therm_codes = np.empty((n_blocks, QK_K), dtype=np.uint8) |
| for i, b in enumerate(blocks): |
| therm_d[i] = b.d_float |
| therm_m[i] = b.dmin_float |
| therm_ls[i], therm_lm[i] = b.unpack_scales() |
| therm_codes[i] = unpack_codes(b.qs) |
| last_therm = [0.0, 0] |
|
|
| def _therm_progress(done, total, improved): |
| now = time.time() |
| if (now - last_therm[0]) >= 0.25 or done == total: |
| last_therm[0] = now |
| _emit( |
| "Phase 4.5: Thermal Refinery (Pair Heat-Bath)", |
| done, |
| total, |
| f"blk {done}/{total} ({max(done, 1)/max(now - t_therm, 1.0e-6):.1f} blk/s) " |
| f"improved={improved}", |
| sigma=sigma, |
| temperature=temperature, |
| ) |
|
|
| therm_d, therm_m, therm_ls, therm_lm, therm_codes, therm_improved = thermal_refine_tensor( |
| X, |
| None if W is None else W, |
| therm_d, |
| therm_m, |
| therm_ls, |
| therm_lm, |
| therm_codes, |
| carry, |
| cfg, |
| dm_solve=_batch_extended_dm_solve_and_ulp, |
| progress=_therm_progress, |
| ) |
| blocks = [ |
| make_block(therm_d[i], therm_m[i], therm_ls[i], therm_lm[i], therm_codes[i]) |
| for i in range(n_blocks) |
| ] |
| _emit( |
| "Phase 4.5: Thermal Refinery (Pair Heat-Bath)", |
| n_blocks, |
| n_blocks, |
| f"done in {time.time() - t_therm:.1f}s, blocks improved={therm_improved}/{n_blocks}", |
| sigma=sigma, |
| temperature=temperature, |
| ) |
|
|
| if cfg.pin_dc and n_blocks > 0: |
| _check_stop() |
| t_pin = time.time() |
| _emit( |
| "Phase 4.7: Zero-DC Pin (Q2_K Free Lunch)", |
| 0, |
| n_blocks, |
| "re-solving (d, m) on the exact dc=0 constraint line " |
| "(exact-JS zero-coherent-residual pin, SSE-nondecreasing)", |
| sigma=sigma, |
| temperature=temperature, |
| ) |
| pin_d = np.empty(n_blocks, dtype=np.float64) |
| pin_m = np.empty(n_blocks, dtype=np.float64) |
| pin_ls = np.empty((n_blocks, NS), dtype=np.uint8) |
| pin_lm = np.empty((n_blocks, NS), dtype=np.uint8) |
| pin_codes = np.empty((n_blocks, QK_K), dtype=np.uint8) |
| for i, b in enumerate(blocks): |
| pin_d[i] = b.d_float |
| pin_m[i] = b.dmin_float |
| pin_ls[i], pin_lm[i] = b.unpack_scales() |
| pin_codes[i] = unpack_codes(b.qs) |
| pin_d, pin_m = pin_q2k_dc_residual( |
| pin_d, pin_m, pin_ls, pin_lm, pin_codes, X.astype(np.float64) |
| ) |
| blocks = [ |
| make_block(pin_d[i], pin_m[i], pin_ls[i], pin_lm[i], pin_codes[i]) |
| for i in range(n_blocks) |
| ] |
| _emit( |
| "Phase 4.7: Zero-DC Pin (Q2_K Free Lunch)", |
| n_blocks, |
| n_blocks, |
| f"done in {time.time() - t_pin:.1f}s", |
| sigma=sigma, |
| temperature=temperature, |
| ) |
|
|
| if cfg.apply_dc_carry and (cfg.dc_lambda > 0.0 or cfg.dc_decay > 0.0): |
| _check_stop() |
| _emit( |
| "Phase 4.8: Sequential Row-Aware DC Residual Nudge", |
| n_blocks, |
| n_blocks, |
| f"Cancelling cumulative row DC residuals (row_width={row_width or QK_K})", |
| sigma=sigma, |
| temperature=temperature, |
| ) |
| blocks = _apply_dc_nudge(X, W, blocks, cfg, row_width) |
|
|
| total_sse = 0.0 |
| weighted = 0.0 |
| for i, b in enumerate(blocks): |
| e = X[i].astype(np.float64) - b.dequantize().astype(np.float64) |
| total_sse += float(np.dot(e, e)) |
| if W is None: |
| weighted += float(np.dot(e, e)) |
| else: |
| weighted += float(np.dot(W[i].astype(np.float64), e * e)) |
|
|
| rmse = float(np.sqrt(total_sse / x.size)) |
| stats = TensorStats( |
| n_blocks=n_blocks, |
| total_sse=total_sse, |
| weighted_sse=weighted, |
| rmse=rmse, |
| sigma=sigma, |
| rmse_over_sigma=rmse / sigma if sigma > 1.0e-15 else 0.0, |
| temperature=temperature, |
| factor_floor_used=factor_floor_used, |
| candidate_floor_used=candidate_floor_used, |
| therm_improved=therm_improved, |
| ) |
| _emit( |
| "Phase 5: Tensor Complete", |
| n_blocks, |
| n_blocks, |
| f"Done: blocks={n_blocks} rmse={rmse:.6e} rmse/sigma={stats.rmse_over_sigma:.5f} " |
| f"wSSE={weighted:.5e} temp={temperature:.3e} " |
| f"floors(cand={candidate_floor_used}, fac={factor_floor_used})", |
| rmse=rmse, |
| sigma=sigma, |
| temperature=temperature, |
| cand_floors=candidate_floor_used, |
| fac_floors=factor_floor_used, |
| therm_improved=therm_improved, |
| ) |
| scales, qs, dm = _stack_blocks(blocks) |
| factor_objectives = np.array([r.objective for r in standalone], dtype=np.float64) |
| result = QuantizedTensor( |
| blocks=blocks, |
| scales=scales, |
| qs=qs, |
| dm=dm, |
| best_candidate=best, |
| factor_objectives=factor_objectives, |
| config=cfg, |
| stats=stats, |
| ) |
| if cfg.verbose: |
| _report(stats) |
| return result |
|
|
|
|
| def _stack_blocks(blocks): |
| n = len(blocks) |
| scales = np.empty((n, NS), dtype=np.uint8) |
| qs = np.empty((n, QK_K // 4), dtype=np.uint8) |
| dm = np.empty((n, 2), dtype=np.uint16) |
| for i, b in enumerate(blocks): |
| scales[i] = b.scales |
| qs[i] = b.qs |
| dm[i, 0] = b.d |
| dm[i, 1] = b.dmin |
| return scales, qs, dm |
|
|
|
|
| def _report(stats: TensorStats) -> None: |
| extra = f" therm={stats.therm_improved}" if stats.therm_improved else "" |
| print( |
| f"Q2_K V2: blocks={stats.n_blocks} rmse={stats.rmse:.6e} " |
| f"rmse/sigma={stats.rmse_over_sigma:.5f} temp={stats.temperature:.3e} " |
| f"floors(cand={stats.candidate_floor_used}, factor={stats.factor_floor_used}){extra}" |
| ) |
|
|
