"""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) # (graph_blocks, 576) 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)) # (graph_blocks, 6) fine = np.sum(coeff, axis=(1, 2, 4)) # (graph_blocks, 6) 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}" )