"""Full HPC V2 Q4_0 quantizer (CUDA GPU + Vectorized CPU). Implements the complete HPC-Quantize V2 decision stack for ``Q4_0`` (32 weights per 18-byte block = 4.5 bpw): 1. **Geometry-Preserving Analog + WLS Seeds**: Finds optimal ``d*`` from both the raw weights and an iso-moment companded analog of each 32-weight block. 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, including negative-scale sign-flip candidates ``-d`` (which mirror the asymmetric ``[-8, +7]`` grid to ``[-7, +8]`` when positive outliers dominate). 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 code 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 nibble codes are re-swept by the seeded pair heat-bath annealer (Anneal) on the 16-level Q4_0 lattice with extended-nonincreasing acceptance under the ``sse_budget`` cap. """ from __future__ import annotations import logging 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_q4_0, thermal_spectral_lams 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 .spectral import pin_dc_codes_uniform, pin_dc_codes_uniform_cuda from .viterbi import robust_temperature QK4_0 = 32 Q4_0_BLOCK_BYTES = 18 Q4_N_CAND = 24 Q4_N_STATES = 6 Q4_N_LEVELS = 16 _log = logging.getLogger(__name__) # Chunk sizes (in 32-weight blocks). The CUDA default keeps the historical # reference chunk (131072 blocks * 32 weights = 4,194,304 weights per chunk); # the CPU default is scaled back so the vectorized numpy temporaries # ((B, 24, 32) candidate-error fields) stay cache-resident instead of # allocating multi-hundred-MB arrays per chunk. Q4_0_CHUNK_BLOCKS_CUDA = 131072 Q4_0_CHUNK_BLOCKS_CPU = 8192 # Empirical peak VRAM per 32-weight block on the CUDA pipeline (~4 GiB # observed for the 131072-block reference chunk); used to auto-scale the GPU # chunk to the free memory actually available on the device. Q4_0_CUDA_BYTES_PER_BLOCK = 32768 Q4_0_CUDA_VRAM_FRACTION = 0.45 def dequantize_q4_0_bytes(raw: bytes | np.ndarray, n_blocks: int) -> np.ndarray: blocks = np.frombuffer(raw, dtype=np.uint8).reshape(n_blocks, Q4_0_BLOCK_BYTES) d = blocks[:, :2].copy().view("> 4).astype(np.float32) - 8.0 q = np.concatenate([low, high], axis=1) return (d * q).reshape(-1) def _q4_sieve_and_viterbi(cand_errors: np.ndarray, cfg: HPCConfig) -> np.ndarray: """6-state sequential Sieve + 6-state Viterbi DP over Q4_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, Q4_N_CAND), np.inf, dtype=np.float64)], axis=0) ce_grouped = ce.reshape(gb, stride, Q4_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) # Sieve local amplitudes from mean candidate errors per group: agg = np.mean(np.where(finite, ce_grouped, 0.0), axis=1).reshape(gb, Q4_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, Q4_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(Q4_N_STATES) trans = cfg.viterbi_alpha * scale_err * np.abs(states[None, :] - states[:, None]) dp = np.full((gb, Q4_N_STATES), np.inf, dtype=np.float64) back = np.zeros((gb, Q4_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, Q4_N_STATES, 4) bin_local_arg = np.argmin(ce_bins, axis=2) # (B, 6) 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_q4_0_hpc_cpu( X: np.ndarray, W: np.ndarray, cfg: HPCConfig ) -> tuple[bytes, float]: """Full HPC V2 Q4_0 pipeline on CPU for a chunk of (B, 32) blocks.""" B = X.shape[0] X_f32 = X.astype(np.float32) W_f32 = W.astype(np.float32) lam_dc = (float(cfg.dc_lambda) / QK4_0) if cfg.use_spectral else 0.0 lam_vw = (float(cfg.vw_lambda) / QK4_0) if cfg.use_spectral else 0.0 half = QK4_0 // 2 # 1. Analog + WLS seed for d* (testing both +d and asymmetric sign-flip -d): amax = np.max(np.abs(X_f32), axis=1) d_range = amax / 7.5 wsum = np.maximum(np.sum(W_f32, axis=1, keepdims=True), 1.0e-30) mu = np.sum(W_f32 * X_f32, axis=1, keepdims=True) / wsum var = np.sum(W_f32 * (X_f32 - mu) ** 2, axis=1, keepdims=True) / wsum sig = np.sqrt(np.maximum(var, 1.0e-30)) x_comp = mu + sig * 1.85 * np.tanh((X_f32 - mu) / (1.85 * sig)) d_analog = np.max(np.abs(x_comp), axis=1).astype(np.float32) / 7.5 wls_d = d_range.copy() best_seed_sse = np.full(B, np.inf, dtype=np.float64) for d_start in (d_range, -d_range, d_analog, -d_analog): d = d_start.astype(np.float64) for _ in range(4): act = np.abs(d) > 1.0e-15 sd = np.where(act, d, 1.0)[:, None] q = np.where(act[:, None], np.clip(nearest_int((X_f32 / sd).astype(np.float32)) + 8, 0, 15), 8) c = (q - 8).astype(np.float64) num = np.sum(W_f32 * X_f32 * c, axis=1, dtype=np.float64) den = np.sum(W_f32 * c * c, axis=1, dtype=np.float64) valid = den > 1.0e-30 d = np.where(valid, num / np.where(valid, den, 1.0), 0.0) df = from_fp16_bits(to_fp16_bits(d.astype(np.float32))) act = np.abs(df) > 1.0e-15 sdf = np.where(act, df, 1.0)[:, None] q = np.where(act[:, None], np.clip(nearest_int(X_f32 / sdf) + 8, 0, 15), 8) e = X_f32 - df[:, None] * (q.astype(np.float32) - 8.0) sse = np.sum(W_f32 * e * e, axis=1, dtype=np.float64) take = sse < best_seed_sse best_seed_sse = np.where(take, sse, best_seed_sse) wls_d = np.where(take, df, wls_d) # 2. 24-Candidate Field with Extended Spectral Objective: cand_d16 = to_fp16_bits(wls_d[:, None] * MULTS_D[None, :]) # (B, 24) cand_d = from_fp16_bits(cand_d16) # (B, 24) act_c = np.abs(cand_d) > 1.0e-15 sd_c = np.where(act_c, cand_d, 1.0)[:, :, None] q_c = np.where( act_c[:, :, None], np.clip(nearest_int(X_f32[:, None, :] / sd_c) + 8, 0, 15), 8, ).astype(np.float32) e_c = X_f32[:, None, :] - cand_d[:, :, None] * (q_c - 8.0) # (B, 24, 32) 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 = _q4_sieve_and_viterbi(cand_errors, cfg) rows = np.arange(B) d_sel = cand_d[rows, best_cand] # 4. Asymmetric Sign-Flip fp16-ULP Polish (+/-8 ULPs around +d_sel and -d_sel): rad = max(int(cfg.ulp_radius), 4) deltas = np.arange(-rad, rad + 1, dtype=np.int32) pos_b = np.clip(to_fp16_bits(np.abs(d_sel)).astype(np.int32)[:, None] + deltas[None, :], 0, 0x7BFF).astype(np.uint16) neg_b = (pos_b | np.uint16(0x8000)).astype(np.uint16) all_b = np.concatenate([pos_b, neg_b], axis=1) # (B, 2*(2*rad+1)) all_d = from_fp16_bits(all_b) act_u = np.abs(all_d) > 1.0e-15 sd_u = np.where(act_u, all_d, 1.0)[:, :, None] q_u = np.where( act_u[:, :, None], np.clip(nearest_int(X_f32[:, None, :] / sd_u) + 8, 0, 15), 8, ).astype(np.uint8) e_u = X_f32[:, None, :] - all_d[:, :, None] * (q_u.astype(np.float32) - 8.0) sse_u = np.sum(W_f32[:, None, :] * e_u * e_u, axis=2, dtype=np.float64) pick_u = np.argmin(sse_u, axis=1) best_d16 = all_b[rows, pick_u] best_d = all_d[rows, pick_u] best_q = q_u[rows, pick_u].astype(np.int64) # 5. Closed-Form 1-D Quadratic Vesica + DC Code Projection within sse_budget: if cfg.use_spectral and (lam_vw > 0.0 or lam_dc > 0.0): c_f = (best_q - 8).astype(np.float64) e = X_f32.astype(np.float64) - best_d[:, None].astype(np.float64) * 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) d_f64 = best_d.astype(np.float64)[:, None] denom = (W_f32 + lam_vw + lam_dc) * (d_f64 * d_f64) act_b = np.abs(d_f64) > 1.0e-15 s_den = np.where(act_b & (denom > 1.0e-30), denom, 1.0) v_full = np.concatenate([v, v], axis=1) grad = d_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, 15) e_new = X_f32.astype(np.float64) - d_f64 * (q_new - 8).astype(np.float64) 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_q4_0(X_f32, W_f32, best_d, best_q, cfg) best_q = thermal_q.astype(np.int64) # 5c. Zero-DC Pin (deterministic code-flip exchange within pin_budget): # Q4_0 codes are unsigned u in [0, 15] with dequant d*(u-8), so hand the # pin the zero-point-shifted target (X + 8d); its model residual # (X + 8d) - d*u then equals the true dequant error. if cfg.pin_dc: best_q, _ = pin_dc_codes_uniform( X_f32 + 8.0 * best_d[:, None].astype(np.float32), best_d, best_q, np.ones((B, QK4_0)), 0, 15, budget_frac=float(cfg.pin_budget), ) best_q_u8 = best_q.astype(np.uint8) out = np.empty((B, Q4_0_BLOCK_BYTES), dtype=np.uint8) out[:, :2] = best_d16.astype(" tuple[bytes, float]: """Full HPC V2 Q4_0 pipeline on CUDA GPU for a chunk of (B, 32) blocks.""" B = X_np.shape[0] lam_dc = (float(cfg.dc_lambda) / QK4_0) if cfg.use_spectral else 0.0 lam_vw = (float(cfg.vw_lambda) / QK4_0) if cfg.use_spectral else 0.0 half = QK4_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. Analog + WLS seed (including negative-scale sign-flip starts): amax = torch.max(torch.abs(X), dim=1).values d_range = amax / 7.5 wsum = torch.clamp(torch.sum(W, dim=1, keepdim=True), min=1.0e-30) mu = torch.sum(W * X, dim=1, keepdim=True) / wsum var = torch.sum(W * (X - mu) ** 2, dim=1, keepdim=True) / wsum sig = torch.sqrt(torch.clamp(var, min=1.0e-30)) x_comp = mu + sig * 1.85 * torch.tanh((X - mu) / (1.85 * sig)) d_analog = torch.max(torch.abs(x_comp), dim=1).values / 7.5 wls_d = d_range.clone() best_seed_sse = torch.full((B,), float("inf"), dtype=torch.float32, device=dev) for d_start in (d_range, -d_range, d_analog, -d_analog): d = d_start.clone() for _ in range(4): act = torch.abs(d) > 1.0e-15 sd = torch.where(act, d, torch.ones_like(d))[:, None] sh = (X / sd) + 12582912.0 qi = ((sh.view(torch.int32) & 0x007FFFFF) - 0x00400000 + 8).clamp(0, 15) q = torch.where(act[:, None], qi, torch.full_like(qi, 8)) c = (q - 8).to(torch.float32) num = torch.sum(W * X * c, dim=1) den = torch.sum(W * c * c, dim=1) valid = den > 1.0e-30 d = torch.where(valid, num / torch.where(valid, den, torch.ones_like(den)), torch.zeros_like(den)) df = d.to(torch.float16).to(torch.float32) act = torch.abs(df) > 1.0e-15 sdf = torch.where(act, df, torch.ones_like(df))[:, None] sh = (X / sdf) + 12582912.0 qi = ((sh.view(torch.int32) & 0x007FFFFF) - 0x00400000 + 8).clamp(0, 15) q = torch.where(act[:, None], qi, torch.full_like(qi, 8)) e = X - df[:, None] * (q.to(torch.float32) - 8.0) sse = torch.sum(W * e * e, dim=1) take = sse < best_seed_sse best_seed_sse = torch.where(take, sse, best_seed_sse) wls_d = torch.where(take, df, wls_d) # 2. 24-Candidate Field with Extended Spectral Objective on GPU: mults = torch.tensor(MULTS_D, dtype=torch.float32, device=dev) cand_d = (wls_d[:, None] * mults[None, :]).to(torch.float16).to(torch.float32) # (B, 24) act_c = torch.abs(cand_d) > 1.0e-15 sd_c = torch.where(act_c, cand_d, torch.ones_like(cand_d))[:, :, None] sh_c = (X[:, None, :] / sd_c) + 12582912.0 qi_c = ((sh_c.view(torch.int32) & 0x007FFFFF) - 0x00400000 + 8).clamp(0, 15) q_c = torch.where(act_c[:, :, None], qi_c, torch.full_like(qi_c, 8)).to(torch.float32) e_c = X[:, None, :] - cand_d[:, :, None] * (q_c - 8.0) 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 = _q4_sieve_and_viterbi(cand_err_t.cpu().numpy(), cfg) del cand_err_t, e_c, q_c, qi_c, sh_c, sse_c if cfg.use_spectral and (lam_dc > 0.0 or lam_vw > 0.0): del v_c, dc_c 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. Asymmetric Sign-Flip fp16-ULP Polish (+/-8 ULPs around +d_sel and -d_sel): rad = max(int(cfg.ulp_radius), 4) deltas = torch.arange(-rad, rad + 1, dtype=torch.int32, device=dev) abs_bits = torch.abs(d_sel).to(torch.float16).view(torch.int16).to(torch.int32) & 0x7FFF pos_b = torch.clamp(abs_bits[:, None] + deltas[None, :], 0, 0x7BFF) neg_b = pos_b | 0x8000 all_b = torch.cat([pos_b, neg_b], dim=1) # (B, 2*(2*rad+1)) all_d = all_b.to(torch.int16).view(torch.float16).to(torch.float32) act_u = torch.abs(all_d) > 1.0e-15 sd_u = torch.where(act_u, all_d, torch.ones_like(all_d))[:, :, None] sh_u = (X[:, None, :] / sd_u) + 12582912.0 qi_u = ((sh_u.view(torch.int32) & 0x007FFFFF) - 0x00400000 + 8).clamp(0, 15) q_u = torch.where(act_u[:, :, None], qi_u, torch.full_like(qi_u, 8)).to(torch.uint8) e_u = X[:, None, :] - all_d[:, :, None] * (q_u.to(torch.float32) - 8.0) sse_u = torch.sum(W[:, None, :] * e_u * e_u, dim=2) pick_u = torch.argmin(sse_u, dim=1) best_d16 = all_b[rows, pick_u] best_d = all_d[rows, pick_u] best_q = q_u[rows, pick_u].to(torch.int64) del e_u, q_u, qi_u, sh_u, sse_u # 5. Closed-Form 1-D Quadratic Vesica + DC Code Projection within sse_budget: if cfg.use_spectral and (lam_vw > 0.0 or lam_dc > 0.0): d_col = best_d[:, None] e = X - d_col * (best_q.to(torch.float32) - 8.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) * (d_col * d_col) act_b = torch.abs(d_col) > 1.0e-15 s_den = torch.where(act_b & (denom > 1.0e-30), denom, torch.ones_like(denom)) v_full = torch.cat([v, v], dim=1) grad = d_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, 15) e_new = X - d_col * (q_new.to(torch.float32) - 8.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, on-GPU): if cfg.therm_passes > 0: thermal_seed = cfg.therm_seed ^ ( (0x9E3779B97F4A7C15 * chunk_idx) & 0x3FFFFFFFFFFFFFFF ) best_q, _ = _q4_0_thermal_cuda(X, W, best_d, best_q, dev, cfg, thermal_seed) # 5c. Zero-DC Pin (deterministic code-flip exchange within pin_budget): # Unsigned codes u in [0, 15], zero-point 8 — pass (X + 8d) so the # pin's model residual equals the true dequant error. Runs fully on # the GPU. if cfg.pin_dc: best_q, _ = pin_dc_codes_uniform_cuda( X + 8.0 * best_d[:, None].float(), best_d, best_q.to(torch.int64), torch.ones_like(X), 0, 15, budget_frac=float(cfg.pin_budget), ) best_q_u8 = best_q.to(torch.uint8) deq_fin = best_d[:, None] * (best_q_u8.to(torch.float32) - 8.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, Q4_0_BLOCK_BYTES), dtype=np.uint8) out[:, :2] = d16_np.astype(" tuple["torch.Tensor", int]: """Seeded pair heat-bath annealing of Q4_0 nibble codes, fully on GPU. Torch port of ``anneal.thermal_refine_uniform_codes``: sequential Gibbs sweeps over the 16 vesica pairs of each 32-weight block with exact closed-form conditional energies (heat-bath sampling at the per-block temperature, greedy argmin for the closing sweeps), and elementwise-best end-of-pass selection (extended energy strictly decreasing and SSE under the ``sse_budget`` cap) so the anneal can never regress a block. At every block-level control point energies are recomputed exactly from the codes in fp64; the within-sweep pair math runs in fp32. RNG streams differ from the CPU annealer, but a fixed seed is reproducible on device. """ B = int(X.shape[0]) lam_vw, lam_dc = thermal_spectral_lams(cfg, QK4_0) if B == 0 or cfg.therm_passes == 0 or (lam_vw == 0.0 and lam_dc == 0.0): return codes, 0 lam_vw, lam_dc = float(lam_vw), float(lam_dc) n_passes, n_sweeps = int(cfg.therm_passes), int(cfg.therm_sweeps) half = QK4_0 // 2 ww = torch.ones_like(X) if W is None else W step = d.to(X.dtype) grid = torch.arange(Q4_N_LEVELS, device=dev, dtype=X.dtype) rows = torch.arange(B, device=dev) def exact_energies(qs: "torch.Tensor"): e64 = (X - step[:, None] * (qs - 8.0)).double() w64 = ww.double() sse = torch.sum(w64 * e64 * e64, dim=1) v = e64[:, :half] + e64[:, half:] dc = torch.sum(e64, dim=1) energy = sse + lam_vw * torch.sum(v * v, dim=1) + lam_dc * dc * dc return sse, energy _, energy_best = exact_energies(codes) energy0 = energy_best.clone() sse_cap = exact_energies(codes)[0] * (1.0 + float(cfg.sse_budget)) best = codes.clone() cur = codes.clone() gen = torch.Generator(device=dev) gen.manual_seed(int(seed) & 0x7FFFFFFFFFFFFFFF) for t in range(n_passes): temperature = ( float(cfg.therm_t0) * (cfg.therm_t_decay**t) * torch.clamp(energy_best.double(), min=1.0e-30) / QK4_0 ).float() for sweep in range(n_sweeps + 2): greedy = sweep >= n_sweeps _, energy_start = exact_energies(cur) e = X - step[:, None] * (cur - 8.0) v = e[:, :half] + e[:, half:] dc = torch.sum(e, dim=1) for p in range(half): e_lo, e_hi = e[:, p], e[:, half + p] w_lo, w_hi = ww[:, p], ww[:, half + p] v_p = v[:, p] de_lo = step[:, None] * (cur[:, p : p + 1] - grid[None, :]) de_hi = step[:, None] * (cur[:, half + p : half + p + 1] - grid[None, :]) de_sum = de_lo[:, :, None] + de_hi[:, None, :] d_e = ( w_lo[:, None, None] * (2.0 * e_lo[:, None, None] * de_lo[:, :, None] + de_lo[:, :, None] ** 2) + w_hi[:, None, None] * (2.0 * e_hi[:, None, None] * de_hi[:, None, :] + de_hi[:, None, :] ** 2) + lam_vw * (2.0 * v_p[:, None, None] * de_sum + de_sum * de_sum) + lam_dc * (2.0 * dc[:, None, None] * de_sum + de_sum * de_sum) ) cur_pair = cur[:, p] * Q4_N_LEVELS + cur[:, half + p] if greedy: flat = d_e.reshape(B, -1) pick = torch.argmin(flat, dim=1) best_val = flat[rows, pick] pick = torch.where(best_val < 0.0, pick, cur_pair) else: logits = -d_e / torch.clamp(temperature, min=1.0e-30)[:, None, None] logits -= logits.amax(dim=(1, 2), keepdim=True) probs = torch.exp(logits).reshape(B, -1) dead_idx = torch.where(step.abs() < 1.0e-15)[0] if bool(dead_idx.numel()): probs[dead_idx] = 0.0 probs[dead_idx, cur_pair[dead_idx]] = 1.0 cum = torch.cumsum(probs, dim=1) cum = cum / cum[:, -1:] u = torch.rand((B, 1), device=dev, generator=gen) pick = torch.sum(cum < u, dim=1).clamp(max=Q4_N_LEVELS * Q4_N_LEVELS - 1) ia, ib = pick // Q4_N_LEVELS, pick % Q4_N_LEVELS upd_lo = torch.gather(de_lo, 1, ia[:, None])[:, 0] upd_hi = torch.gather(de_hi, 1, ib[:, None])[:, 0] upd_sum = upd_lo + upd_hi e[:, p] = e_lo + upd_lo e[:, half + p] = e_hi + upd_hi v[:, p] = v_p + upd_sum dc = dc + upd_sum cur[:, p] = ia cur[:, half + p] = ib if greedy: _, energy_now = exact_energies(cur) if not bool(torch.any(energy_now < energy_start - 1.0e-10)): break sse_now, energy_now = exact_energies(cur) take = (energy_now < energy_best - 1.0e-18) & (sse_now <= sse_cap) best = torch.where(take[:, None], cur, best) energy_best = torch.where(take, energy_now, energy_best) cur = best.clone() n_improved = int(torch.sum(energy_best < energy0 - 1.0e-18).item()) return best, n_improved def quantize_tensor_q4_0( weights: np.ndarray, importance: np.ndarray | None = None, cfg: HPCConfig | None = None, chunk_blocks: int | None = None, ) -> tuple[bytes, int, float]: """Quantize a tensor to GGML Q4_0 using the full HPC V2 pipeline. ``chunk_blocks`` (32-weight blocks per chunk) defaults to ``Q4_0_CHUNK_BLOCKS_CUDA`` on the GPU path and the cache-resident ``Q4_0_CHUNK_BLOCKS_CPU`` on the CPU path. """ cfg = cfg or HPCConfig() x = np.asarray(weights, dtype=np.float32).ravel() if x.size == 0 or x.size % QK4_0 != 0: raise ValueError("Q4_0 requires a non-empty weight count divisible by 32") n_blocks = x.size // QK4_0 X = x.reshape(n_blocks, QK4_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, QK4_0) else: W_raw = None cuda_ok = cfg.device != "cpu" and cuda_is_available() use_cuda = cuda_ok 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 if chunk_blocks is None: if use_cuda and dev is not None: # VRAM-aware default: size the chunk from the free memory at # start-up, never below the reference chunk (which always fits # even on small cards). try: free_bytes, _ = torch.cuda.mem_get_info(dev) chunk_blocks = max( Q4_0_CHUNK_BLOCKS_CUDA, int(Q4_0_CUDA_VRAM_FRACTION * free_bytes) // Q4_0_CUDA_BYTES_PER_BLOCK, ) except Exception: chunk_blocks = Q4_0_CHUNK_BLOCKS_CUDA else: chunk_blocks = Q4_0_CHUNK_BLOCKS_CPU chunk_blocks = max(1, int(chunk_blocks)) _log.info( "Q4_0 routing: cfg.device=%r cuda_is_available=%s n_blocks=%d -> %s (chunk_blocks=%d)", cfg.device, cuda_ok if cfg.device != "cpu" else "n/a", n_blocks, f"CUDA[{dev}]" if use_cuda else "CPU", chunk_blocks, ) 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] sig2 = np.mean(Xb * Xb, axis=1, keepdims=True) if W_raw is None: Wb = np.sqrt(sig2 + Xb * Xb) else: wr = W_raw[start:stop] Wb = np.sqrt(np.maximum(wr, 1.0e-6)) * np.sqrt(sig2 + Xb * Xb) if cfg.temper_imatrix else wr wm = np.mean(Wb, axis=1, keepdims=True) Wb = np.where(wm > 1.0e-30, Wb / np.maximum(wm, 1.0e-30), 1.0).astype(np.float32) if use_cuda and dev is not None: raw_b, sse_b = _quantize_q4_0_hpc_cuda(Xb, Wb, dev, cfg, chunk_idx=start // chunk_blocks) else: raw_b, sse_b = _quantize_q4_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