Download q1_0.py from CompressedGemma/HPC-Quantize-V2: direct link, hf CLI and curl.
- Browser
- Download file 22.6 kB
-
https://huggingface.co/CompressedGemma/HPC-Quantize-V2/resolve/main/q1_0.py
- Command line
-
hf download hf://CompressedGemma/HPC-Quantize-V2/q1_0.py
-
curl -L -o q1_0.py https://huggingface.co/CompressedGemma/HPC-Quantize-V2/resolve/main/q1_0.py
22.6 kB
| """Full HPC V2 Q1_0 quantizer (CUDA GPU + Vectorized CPU). | |
| Implements the complete HPC-Quantize V2 decision stack for ``Q1_0`` (128 weights | |
| per 18-byte block = 1.125 bpw, matching ``block_q1_0`` in ``llama.cpp``): | |
| 1. **Geometry-Preserving Analog + WLS Seeds**: | |
| Finds optimal ``d*`` from the llama.cpp reference ``mean(|x|)``, importance- | |
| weighted WLS ``sum(w*|x|)/sum(w)``, and an iso-moment companded analog of | |
| each 128-weight block (including negative-scale sign-flip starts ``-d``). | |
| 2. **24-Candidate Scale Lattice & Extended Spectral Objective**: | |
| Evaluates 24 scale multipliers (mapped to 6 symbolic quhit states) on | |
| importance-weighted SSE + D6/Vesica fold pyramid + DC penalty. | |
| 3. **6-State Sequential Sieve & 6-State Viterbi DP**: | |
| Runs the 6-state sequential Sieve with parity back-action and a 6-state | |
| Viterbi dynamic program with free-energy bin reduction and Manhattan | |
| transition penalty across the block chain. | |
| 4. **Asymmetric Sign-Flip fp16-ULP Polish & Closed-Form Vesica/DC Shaping**: | |
| Polishes ``d`` over ``+/-8`` fp16 ULPs of both ``+d`` and ``-d``, and | |
| applies closed-form 1-D quadratic Vesica + DC bit-flip projection bounded by | |
| ``sse_budget`` plus a strict monotone candidate floor. | |
| 5. **Optional Thermal Refinery (Pair Heat-Bath, opt-in)**: | |
| With ``cfg.therm_passes > 0`` the finished bits are re-swept by the seeded | |
| pair heat-bath annealer (Anneal) on the 2-level Q1_0 lattice in the same | |
| companded target space, non-increasing under the ``sse_budget`` cap. | |
| """ | |
| from __future__ import annotations | |
| from math import ceil | |
| import numpy as np | |
| try: | |
| import torch | |
| except ImportError: # pragma: no cover | |
| torch = None # type: ignore[assignment] | |
| from .anneal import thermal_refine_q1_0 | |
| from .config import HPCConfig | |
| from .cuda_backend import cuda_is_available | |
| from .factored import MULTS_D | |
| from .fp16 import from_fp16_bits, nearest_int, to_fp16_bits | |
| from .sieve import sieve_measure_chain, subblock_chain_edges | |
| from .viterbi import robust_temperature | |
| QK1_0 = 128 | |
| Q1_0_BLOCK_BYTES = 18 | |
| Q1_N_CAND = 24 | |
| Q1_N_STATES = 6 | |
| _BIT_WEIGHTS = (1 << np.arange(8, dtype=np.uint8))[None, None, :] | |
| def _pack_q1_bits_np(q_u8: np.ndarray) -> np.ndarray: | |
| """Pack (B, 128) uint8 bits {0, 1} into (B, 16) LSB-first bytes.""" | |
| B = q_u8.shape[0] | |
| bits = (q_u8.reshape(B, 16, 8) & 1).astype(np.uint8) | |
| return np.sum(bits * _BIT_WEIGHTS, axis=2, dtype=np.uint8) | |
| def dequantize_q1_0_bytes(raw: bytes | np.ndarray, n_blocks: int) -> np.ndarray: | |
| """Decode GGML Q1_0 blocks into float32 weights.""" | |
| blocks = np.frombuffer(raw, dtype=np.uint8).reshape(n_blocks, Q1_0_BLOCK_BYTES) | |
| d = blocks[:, :2].copy().view("<f2").reshape(n_blocks, 1).astype(np.float32) | |
| qs = blocks[:, 2:] | |
| bits = np.unpackbits(qs, axis=1, bitorder="little")[:, :QK1_0].astype(np.float32) | |
| signs = bits * 2.0 - 1.0 | |
| return (d * signs).reshape(-1) | |
| def _q1_sieve_and_viterbi(cand_errors: np.ndarray, cfg: HPCConfig) -> np.ndarray: | |
| """6-state sequential Sieve + 6-state Viterbi DP over Q1_0 candidate lattice.""" | |
| B = cand_errors.shape[0] | |
| if cfg.optimizer != "hpc" or B < 2: | |
| return np.argmin(cand_errors, axis=1).astype(np.int64) | |
| max_gb = min(int(cfg.graph_max_blocks), 512) | |
| stride = max(1, ceil(B / max_gb)) | |
| gb = ceil(B / stride) | |
| temp = max(float(robust_temperature(cand_errors, cfg)), 1.0e-12) | |
| pad = gb * stride - B | |
| ce = cand_errors.astype(np.float64) | |
| if pad > 0: | |
| ce = np.concatenate([ce, np.full((pad, Q1_N_CAND), np.inf, dtype=np.float64)], axis=0) | |
| ce_grouped = ce.reshape(gb, stride, Q1_N_STATES, 4) | |
| group_min = np.min(ce_grouped, axis=(1, 3)) # (gb, 6) | |
| finite = np.isfinite(ce_grouped) | |
| counts = np.sum(finite, axis=(1, 3)) | |
| valid = (counts > 0) & np.isfinite(group_min) | |
| gm_b = np.where(np.isfinite(group_min), group_min, 0.0)[:, None, :, None] | |
| shifted = np.where(finite, (ce_grouped - gm_b) / (2.0 * temp), np.inf) | |
| sums = np.sum(np.exp(-shifted), axis=(1, 3)) | |
| mean_exp = np.maximum(sums / np.maximum(counts, 1), 1.0e-300) | |
| bin_err = np.where(valid, group_min - 2.0 * temp * np.log(mean_exp), np.inf) | |
| agg = np.mean(np.where(finite, ce_grouped, 0.0), axis=1).reshape(gb, Q1_N_CAND) | |
| mn = np.min(agg, axis=1, keepdims=True) | |
| coeff = np.exp(-np.clip((agg - mn) / (2.0 * temp), 0.0, 80.0)).reshape(gb, Q1_N_STATES, 4) | |
| amps = np.sum(coeff, axis=2) | |
| norms = np.linalg.norm(amps, axis=1, keepdims=True) | |
| amps = np.where(norms > 1.0e-30, amps / np.maximum(norms, 1.0e-30), 1.0 / np.sqrt(6.0)) | |
| marg, _ = sieve_measure_chain( | |
| amps, | |
| subblock_chain_edges(gb), | |
| same_penalty=cfg.same_penalty, | |
| log_slack=cfg.log_slack, | |
| deterministic=True, | |
| ) | |
| m_tot = np.sum(marg, axis=1, keepdims=True) | |
| pm = np.where(m_tot > 1.0e-30, marg / np.maximum(m_tot, 1.0e-30), 1.0 / 6.0) | |
| log_prior = np.log(pm + 1.0e-30) | |
| fin_err = bin_err[np.isfinite(bin_err)] | |
| scale_err = max(float(fin_err.mean()) if fin_err.size else 1.0e-10, 1.0e-20) | |
| local = np.where(np.isfinite(bin_err), bin_err - cfg.viterbi_beta * scale_err * log_prior, np.inf) | |
| states = np.arange(Q1_N_STATES) | |
| trans = cfg.viterbi_alpha * scale_err * np.abs(states[None, :] - states[:, None]) | |
| dp = np.full((gb, Q1_N_STATES), np.inf, dtype=np.float64) | |
| back = np.zeros((gb, Q1_N_STATES), dtype=np.int64) | |
| dp[0] = local[0] | |
| for gi in range(1, gb): | |
| totals = dp[gi - 1][:, None] + trans | |
| best_sp = np.argmin(totals, axis=0) | |
| best_val = totals[best_sp, states] | |
| dp[gi] = np.where(np.isfinite(best_val), best_val + local[gi], np.inf) | |
| back[gi] = best_sp | |
| path = np.zeros(gb, dtype=np.int64) | |
| if np.any(np.isfinite(dp[-1])): | |
| path[-1] = int(np.argmin(dp[-1])) | |
| for gi in range(gb - 2, -1, -1): | |
| path[gi] = back[gi + 1, path[gi + 1]] | |
| ce_bins = cand_errors.reshape(B, Q1_N_STATES, 4) | |
| bin_local_arg = np.argmin(ce_bins, axis=2) | |
| blk_gi = np.arange(B) // stride | |
| chosen_bin = path[blk_gi] | |
| chosen_cand = chosen_bin * 4 + bin_local_arg[np.arange(B), chosen_bin] | |
| global_best = np.argmin(cand_errors, axis=1).astype(np.int64) | |
| cur_e = cand_errors[np.arange(B), chosen_cand] | |
| g_e = cand_errors[np.arange(B), global_best] | |
| return np.where(g_e < cur_e * cfg.greedy_override_ratio, global_best, chosen_cand) | |
| def _quantize_q1_0_hpc_cpu( | |
| X: np.ndarray, W: np.ndarray, cfg: HPCConfig | |
| ) -> tuple[bytes, float]: | |
| """Full HPC V2 Q1_0 pipeline on CPU for a chunk of (B, 128) blocks.""" | |
| B = X.shape[0] | |
| X_f32 = X.astype(np.float32) | |
| W_f32 = W.astype(np.float32) | |
| lam_dc = (float(cfg.dc_lambda) / QK1_0) if cfg.use_spectral else 0.0 | |
| lam_vw = (float(cfg.vw_lambda) / QK1_0) if cfg.use_spectral else 0.0 | |
| half = QK1_0 // 2 | |
| # 1. Iso-Energy / Geometric-Projection Analog + WLS seed for d*: | |
| # Standard d_mean = mean(|x|) attenuates each linear projection by 2/pi (~0.637), | |
| # shrinking SwiGLU FFN residual updates to 0.30x across layers. | |
| # Projecting onto the geometric mean of the Iso-Energy RMS scale d_rms and the | |
| # unbiased linear-projection scale d_proj = d_rms^2 / d_mean restores multi-layer | |
| # residual norm (~0.82x == cos_theta) and improves multi-layer cosine similarity. | |
| d_mean = np.mean(np.abs(X_f32), axis=1).astype(np.float32) | |
| wsum = np.maximum(np.sum(W_f32, axis=1, keepdims=True), 1.0e-30) | |
| d_wls_l1 = (np.sum(W_f32 * np.abs(X_f32), axis=1, keepdims=True) / wsum)[:, 0].astype(np.float32) | |
| d_wls_rms = np.sqrt(np.maximum((np.sum(W_f32 * X_f32 * X_f32, axis=1, keepdims=True) / wsum)[:, 0], 0.0)).astype(np.float32) | |
| if cfg.use_weight_analogs: | |
| ratio = np.clip(d_wls_rms / np.maximum(d_wls_l1, 1.0e-30), 1.0, 1.45).astype(np.float32) | |
| d_center = (d_wls_rms * np.sqrt(ratio)).astype(np.float32) | |
| d_lo = (0.85 * d_wls_rms).astype(np.float32) | |
| d_hi = (1.35 * d_wls_rms).astype(np.float32) | |
| else: | |
| d_center = d_wls_l1 | |
| d_lo = (0.80 * d_mean).astype(np.float32) | |
| d_hi = (1.25 * d_mean).astype(np.float32) | |
| d_center = np.clip(d_center, d_lo, d_hi).astype(np.float32) | |
| gain = (d_center / np.maximum(d_wls_l1, 1.0e-30))[:, None].astype(np.float32) | |
| X_target = (X_f32 * gain).astype(np.float32) | |
| q_sign = (X_f32 >= 0.0).astype(np.float32) | |
| c_sign = 2.0 * q_sign - 1.0 | |
| wls_d = from_fp16_bits(to_fp16_bits(d_center)) | |
| # 2. 24-Candidate Field with Extended Spectral Objective (clamped to [d_lo, d_hi]): | |
| cand_raw = np.clip(wls_d[:, None] * MULTS_D[None, :], d_lo[:, None], d_hi[:, None]).astype(np.float32) | |
| cand_d16 = to_fp16_bits(cand_raw) # (B, 24) | |
| cand_d = from_fp16_bits(cand_d16) # (B, 24) | |
| e_c = X_target[:, None, :] - cand_d[:, :, None] * c_sign[:, None, :] # (B, 24, 128) | |
| sse_c = np.sum(W_f32[:, None, :] * e_c * e_c, axis=2) | |
| if cfg.use_spectral and (lam_dc > 0.0 or lam_vw > 0.0): | |
| v_c = e_c[:, :, :half] + e_c[:, :, half:] | |
| dc_c = np.sum(e_c, axis=2) | |
| cand_errors = (sse_c + lam_vw * np.sum(v_c * v_c, axis=2) + lam_dc * dc_c * dc_c).astype(np.float32) | |
| else: | |
| cand_errors = sse_c.astype(np.float32) | |
| # 3. 6-State Sieve + 6-State Viterbi DP: | |
| best_cand = _q1_sieve_and_viterbi(cand_errors, cfg) | |
| rows = np.arange(B) | |
| d_sel = cand_d[rows, best_cand] | |
| # 4. FP16-ULP Polish (+/-rad ULPs around d_sel, wls_d, and d_wls_rms within [d_lo, d_hi]): | |
| rad = max(int(cfg.ulp_radius), 4) | |
| deltas = np.arange(-rad, rad + 1, dtype=np.int32) | |
| lo_b = to_fp16_bits(d_lo).astype(np.int32)[:, None] | |
| hi_b = np.maximum(to_fp16_bits(d_hi).astype(np.int32)[:, None], lo_b) | |
| sel_b = np.clip(to_fp16_bits(np.abs(d_sel)).astype(np.int32)[:, None] + deltas[None, :], lo_b, hi_b).astype(np.uint16) | |
| wls_b = np.clip(to_fp16_bits(np.abs(wls_d)).astype(np.int32)[:, None] + deltas[None, :], lo_b, hi_b).astype(np.uint16) | |
| ref_b = np.clip(to_fp16_bits(np.abs(d_wls_rms if cfg.use_weight_analogs else d_mean)).astype(np.int32)[:, None] + deltas[None, :], lo_b, hi_b).astype(np.uint16) | |
| all_b = np.concatenate([sel_b, wls_b, ref_b], axis=1) | |
| all_d = from_fp16_bits(all_b) | |
| e_u = X_target[:, None, :] - all_d[:, :, None] * c_sign[:, None, :] | |
| sse_u = np.sum(W_f32[:, None, :] * e_u * e_u, axis=2, dtype=np.float64) | |
| if cfg.use_spectral and (lam_dc > 0.0 or lam_vw > 0.0): | |
| v_u = e_u[:, :, :half] + e_u[:, :, half:] | |
| dc_u = np.sum(e_u, axis=2, dtype=np.float64) | |
| sse_min_u = np.min(sse_u, axis=1, keepdims=True) | |
| obj_u = sse_u + lam_vw * np.sum(v_u * v_u, axis=2, dtype=np.float64) + lam_dc * dc_u * dc_u | |
| obj_u = np.where(sse_u <= sse_min_u * (1.0 + cfg.sse_budget), obj_u, np.inf) | |
| pick_u = np.argmin(obj_u, axis=1) | |
| else: | |
| pick_u = np.argmin(sse_u, axis=1) | |
| best_d16 = all_b[rows, pick_u] | |
| best_d = all_d[rows, pick_u] | |
| best_q = (X_f32 >= 0.0).astype(np.int64) | |
| # 5. Boundary-Guarded 1-D Quadratic Vesica + DC Bit-Flip Projection (only near-zero |x| <= 0.10*d): | |
| if cfg.use_spectral and (lam_vw > 0.0 or lam_dc > 0.0): | |
| c_f = (2 * best_q - 1).astype(np.float64) | |
| d_f64 = best_d.astype(np.float64)[:, None] | |
| step_f64 = 2.0 * d_f64 | |
| e = X_target.astype(np.float64) - d_f64 * c_f | |
| sse_live = np.sum(W_f32 * e * e, axis=1, dtype=np.float64) | |
| sse_cap = sse_live * (1.0 + cfg.sse_budget) | |
| v = e[:, :half] + e[:, half:] | |
| dc = np.sum(e, axis=1) | |
| denom = (W_f32 + lam_vw + lam_dc) * (step_f64 * step_f64) | |
| near_zero = np.abs(X_target) <= 0.10 * d_f64 | |
| act_b = (d_f64 > 1.0e-15) & near_zero | |
| s_den = np.where(act_b & (denom > 1.0e-30), denom, 1.0) | |
| v_full = np.concatenate([v, v], axis=1) | |
| grad = step_f64 * (W_f32 * e + lam_vw * v_full + lam_dc * dc[:, None]) | |
| dq = np.where(act_b, nearest_int((grad / s_den).astype(np.float32)), 0) | |
| q_new = np.clip(best_q + dq, 0, 1) | |
| c_new = (2 * q_new - 1).astype(np.float64) | |
| e_new = X_target.astype(np.float64) - d_f64 * c_new | |
| sse_new = np.sum(W_f32 * e_new * e_new, axis=1, dtype=np.float64) | |
| v_new = e_new[:, :half] + e_new[:, half:] | |
| dc_new = np.sum(e_new, axis=1) | |
| spec_old = lam_vw * np.sum(v * v, axis=1) + lam_dc * dc * dc | |
| spec_new = lam_vw * np.sum(v_new * v_new, axis=1) + lam_dc * dc_new * dc_new | |
| accept = (sse_new <= sse_cap) & ((sse_new + spec_new) < (sse_live + spec_old) - 1.0e-15) | |
| best_q = np.where(accept[:, None], q_new, best_q) | |
| # 5. Optional Thermal Refinery (Pair Heat-Bath, non-regressing): | |
| if cfg.therm_passes > 0: | |
| thermal_q, _ = thermal_refine_q1_0( | |
| X_target.astype(np.float32), W_f32, best_d, best_q, cfg | |
| ) | |
| best_q = thermal_q.astype(np.int64) | |
| best_q_u8 = best_q.astype(np.uint8) | |
| out = np.empty((B, Q1_0_BLOCK_BYTES), dtype=np.uint8) | |
| out[:, :2] = best_d16.astype("<u2").view(np.uint8).reshape(B, 2) | |
| out[:, 2:] = _pack_q1_bits_np(best_q_u8) | |
| deq_final = best_d[:, None] * (2.0 * best_q_u8.astype(np.float32) - 1.0) | |
| raw_sse = float(np.sum((X_f32 - deq_final) ** 2, dtype=np.float64)) | |
| return out.tobytes(), raw_sse | |
| def _quantize_q1_0_hpc_cuda( | |
| X_np: np.ndarray, W_np: np.ndarray, dev: "torch.device", cfg: HPCConfig | |
| ) -> tuple[bytes, float]: | |
| """Full HPC V2 Q1_0 pipeline on CUDA GPU for a chunk of (B, 128) blocks.""" | |
| B = X_np.shape[0] | |
| lam_dc = (float(cfg.dc_lambda) / QK1_0) if cfg.use_spectral else 0.0 | |
| lam_vw = (float(cfg.vw_lambda) / QK1_0) if cfg.use_spectral else 0.0 | |
| half = QK1_0 // 2 | |
| with torch.no_grad(): | |
| X = torch.from_numpy(X_np).to(dev, dtype=torch.float32) | |
| W = torch.from_numpy(W_np).to(dev, dtype=torch.float32) | |
| # 1. Iso-Energy / Geometric-Projection Analog + WLS seed on GPU: | |
| d_mean = torch.mean(torch.abs(X), dim=1) | |
| wsum = torch.clamp(torch.sum(W, dim=1, keepdim=True), min=1.0e-30) | |
| d_wls_l1 = (torch.sum(W * torch.abs(X), dim=1, keepdim=True) / wsum).squeeze(1) | |
| d_wls_rms = torch.sqrt(torch.clamp((torch.sum(W * X * X, dim=1, keepdim=True) / wsum).squeeze(1), min=0.0)) | |
| if cfg.use_weight_analogs: | |
| ratio = torch.clamp(d_wls_rms / torch.clamp(d_wls_l1, min=1.0e-30), min=1.0, max=1.45) | |
| d_center = d_wls_rms * torch.sqrt(ratio) | |
| d_lo = 0.85 * d_wls_rms | |
| d_hi = 1.35 * d_wls_rms | |
| else: | |
| d_center = d_wls_l1 | |
| d_lo = 0.80 * d_mean | |
| d_hi = 1.25 * d_mean | |
| d_center = torch.clamp(d_center, min=d_lo, max=d_hi) | |
| gain = (d_center / torch.clamp(d_wls_l1, min=1.0e-30)).unsqueeze(1) | |
| X_target = X * gain | |
| c_sign = torch.where(X >= 0.0, 1.0, -1.0) | |
| wls_d = d_center.to(torch.float16).to(torch.float32) | |
| # 2. 24-Candidate Field with Extended Spectral Objective on GPU (clamped to [d_lo, d_hi]): | |
| mults = torch.tensor(MULTS_D, dtype=torch.float32, device=dev) | |
| cand_raw = torch.clamp(wls_d[:, None] * mults[None, :], min=d_lo[:, None], max=d_hi[:, None]) | |
| cand_d = cand_raw.to(torch.float16).to(torch.float32) # (B, 24) | |
| e_c = X_target[:, None, :] - cand_d[:, :, None] * c_sign[:, None, :] | |
| sse_c = torch.sum(W[:, None, :] * e_c * e_c, dim=2) | |
| if cfg.use_spectral and (lam_dc > 0.0 or lam_vw > 0.0): | |
| v_c = e_c[:, :, :half] + e_c[:, :, half:] | |
| dc_c = torch.sum(e_c, dim=2) | |
| cand_err_t = sse_c + lam_vw * torch.sum(v_c * v_c, dim=2) + lam_dc * dc_c * dc_c | |
| else: | |
| cand_err_t = sse_c | |
| # 3. 6-State Sieve + 6-State Viterbi DP: | |
| best_cand_np = _q1_sieve_and_viterbi(cand_err_t.cpu().numpy(), cfg) | |
| best_cand_t = torch.from_numpy(best_cand_np).to(dev) | |
| rows = torch.arange(B, device=dev) | |
| d_sel = cand_d[rows, best_cand_t] | |
| # 4. FP16-ULP Polish (+/-rad ULPs around d_sel, wls_d, and d_wls_rms within [d_lo, d_hi]): | |
| rad = max(int(cfg.ulp_radius), 4) | |
| deltas = torch.arange(-rad, rad + 1, dtype=torch.int32, device=dev) | |
| lo_bits = (d_lo.to(torch.float16).view(torch.int16).to(torch.int32) & 0x7FFF)[:, None] | |
| hi_bits = torch.maximum((d_hi.to(torch.float16).view(torch.int16).to(torch.int32) & 0x7FFF)[:, None], lo_bits) | |
| abs_bits = torch.abs(d_sel).to(torch.float16).view(torch.int16).to(torch.int32) & 0x7FFF | |
| wls_bits = torch.abs(wls_d).to(torch.float16).view(torch.int16).to(torch.int32) & 0x7FFF | |
| ref_src = d_wls_rms if cfg.use_weight_analogs else d_mean | |
| ref_bits = torch.abs(ref_src).to(torch.float16).view(torch.int16).to(torch.int32) & 0x7FFF | |
| pos_b = torch.clamp(abs_bits[:, None] + deltas[None, :], lo_bits, hi_bits) | |
| wls_b = torch.clamp(wls_bits[:, None] + deltas[None, :], lo_bits, hi_bits) | |
| mean_b = torch.clamp(ref_bits[:, None] + deltas[None, :], lo_bits, hi_bits) | |
| all_b = torch.cat([pos_b, wls_b, mean_b], dim=1) | |
| all_d = all_b.to(torch.int16).view(torch.float16).to(torch.float32) | |
| e_u = X_target[:, None, :] - all_d[:, :, None] * c_sign[:, None, :] | |
| sse_u = torch.sum(W[:, None, :] * e_u * e_u, dim=2) | |
| if cfg.use_spectral and (lam_dc > 0.0 or lam_vw > 0.0): | |
| v_u = e_u[:, :, :half] + e_u[:, :, half:] | |
| dc_u = torch.sum(e_u, dim=2) | |
| sse_min_u = torch.min(sse_u, dim=1, keepdim=True).values | |
| obj_u = sse_u + lam_vw * torch.sum(v_u * v_u, dim=2) + lam_dc * dc_u * dc_u | |
| obj_u = torch.where(sse_u <= sse_min_u * (1.0 + float(cfg.sse_budget)), obj_u, torch.full_like(obj_u, float("inf"))) | |
| pick_u = torch.argmin(obj_u, dim=1) | |
| else: | |
| pick_u = torch.argmin(sse_u, dim=1) | |
| best_d16 = all_b[rows, pick_u] | |
| best_d = all_d[rows, pick_u] | |
| best_q = (X >= 0.0).to(torch.int64) | |
| # 5. Boundary-Guarded 1-D Quadratic Vesica + DC Bit-Flip Projection (only near-zero |x| <= 0.10*d): | |
| if cfg.use_spectral and (lam_vw > 0.0 or lam_dc > 0.0): | |
| d_col = best_d[:, None] | |
| step_col = 2.0 * d_col | |
| e = X_target - d_col * (2.0 * best_q.to(torch.float32) - 1.0) | |
| sse_live = torch.sum(W * e * e, dim=1) | |
| sse_cap = sse_live * (1.0 + float(cfg.sse_budget)) | |
| v = e[:, :half] + e[:, half:] | |
| dc = torch.sum(e, dim=1) | |
| denom = (W + lam_vw + lam_dc) * (step_col * step_col) | |
| near_zero = torch.abs(X_target) <= 0.10 * d_col | |
| act_b = (d_col > 1.0e-15) & near_zero | |
| s_den = torch.where(act_b & (denom > 1.0e-30), denom, torch.ones_like(denom)) | |
| v_full = torch.cat([v, v], dim=1) | |
| grad = step_col * (W * e + lam_vw * v_full + lam_dc * dc[:, None]) | |
| sh_dq = (grad / s_den) + 12582912.0 | |
| dq = torch.where( | |
| act_b, | |
| ((sh_dq.view(torch.int32) & 0x007FFFFF) - 0x00400000).to(torch.int64), | |
| torch.zeros_like(best_q), | |
| ) | |
| q_new = torch.clamp(best_q + dq, 0, 1) | |
| e_new = X_target - d_col * (2.0 * q_new.to(torch.float32) - 1.0) | |
| sse_new = torch.sum(W * e_new * e_new, dim=1) | |
| v_new = e_new[:, :half] + e_new[:, half:] | |
| dc_new = torch.sum(e_new, dim=1) | |
| spec_old = lam_vw * torch.sum(v * v, dim=1) + lam_dc * dc * dc | |
| spec_new = lam_vw * torch.sum(v_new * v_new, dim=1) + lam_dc * dc_new * dc_new | |
| accept = (sse_new <= sse_cap) & ((sse_new + spec_new) < (sse_live + spec_old) - 1.0e-12) | |
| best_q = torch.where(accept[:, None], q_new, best_q) | |
| # 5. Optional Thermal Refinery (Pair Heat-Bath, non-regressing): | |
| if cfg.therm_passes > 0: | |
| thermal_q, _ = thermal_refine_q1_0( | |
| X_target.cpu().numpy(), | |
| None if W_np is None else W.cpu().numpy(), | |
| best_d.cpu().numpy(), | |
| best_q.cpu().numpy(), | |
| cfg, | |
| ) | |
| best_q = torch.from_numpy(thermal_q.astype(np.int64)).to(dev) | |
| best_q_u8 = best_q.to(torch.uint8) | |
| deq_fin = best_d[:, None] * (2.0 * best_q_u8.to(torch.float32) - 1.0) | |
| raw_sse = float(torch.sum((X - deq_fin) ** 2, dtype=torch.float64).item()) | |
| d16_np = best_d16.to(torch.int16).cpu().numpy().view(np.uint16) | |
| q_np = best_q_u8.cpu().numpy() | |
| out = np.empty((B, Q1_0_BLOCK_BYTES), dtype=np.uint8) | |
| out[:, :2] = d16_np.astype("<u2").view(np.uint8).reshape(B, 2) | |
| out[:, 2:] = _pack_q1_bits_np(q_np) | |
| return out.tobytes(), raw_sse | |
| def quantize_tensor_q1_0( | |
| weights: np.ndarray, | |
| importance: np.ndarray | None = None, | |
| cfg: HPCConfig | None = None, | |
| chunk_blocks: int = 65536, | |
| ) -> tuple[bytes, int, float]: | |
| """Quantize a tensor to GGML Q1_0 using the full HPC V2 pipeline.""" | |
| cfg = cfg or HPCConfig() | |
| x = np.asarray(weights, dtype=np.float32).ravel() | |
| if x.size == 0 or x.size % QK1_0 != 0: | |
| raise ValueError("Q1_0 requires a non-empty weight count divisible by 128") | |
| n_blocks = x.size // QK1_0 | |
| X = x.reshape(n_blocks, QK1_0) | |
| if importance is not None: | |
| imp = np.asarray(importance, dtype=np.float32).ravel() | |
| if imp.size != x.size: | |
| raise ValueError("importance length must match weight count") | |
| W_raw = imp.reshape(n_blocks, QK1_0) | |
| else: | |
| W_raw = None | |
| use_cuda = cfg.device != "cpu" and cuda_is_available() and ( | |
| cfg.device.startswith("cuda") or (cfg.device == "auto" and n_blocks >= 64) | |
| ) | |
| dev = torch.device("cuda" if cfg.device == "auto" else cfg.device) if use_cuda else None | |
| chunks: list[bytes] = [] | |
| total_sse = 0.0 | |
| for start in range(0, n_blocks, chunk_blocks): | |
| stop = min(start + chunk_blocks, n_blocks) | |
| Xb = X[start:stop] | |
| if W_raw is None: | |
| Wb = np.ones_like(Xb, dtype=np.float32) | |
| else: | |
| wr = np.maximum(W_raw[start:stop], 1.0e-6) | |
| med = np.maximum(np.median(wr, axis=1, keepdims=True), 1.0e-6) | |
| wt = np.clip((wr / med) ** 0.25 if cfg.temper_imatrix else wr / med, 0.5, 2.5) | |
| wm = np.mean(wt, axis=1, keepdims=True) | |
| Wb = np.where(wm > 1.0e-30, wt / np.maximum(wm, 1.0e-30), 1.0).astype(np.float32) | |
| if use_cuda and dev is not None: | |
| raw_b, sse_b = _quantize_q1_0_hpc_cuda(Xb, Wb, dev, cfg) | |
| else: | |
| raw_b, sse_b = _quantize_q1_0_hpc_cpu(Xb, Wb, cfg) | |
| chunks.append(raw_b) | |
| total_sse += sse_b | |
| rmse = float(np.sqrt(total_sse / x.size)) | |
| return b"".join(chunks), n_blocks, rmse | |