| """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: |
| torch = None |
|
|
| 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__) |
|
|
| |
| |
| |
| |
| |
| Q4_0_CHUNK_BLOCKS_CUDA = 131072 |
| Q4_0_CHUNK_BLOCKS_CPU = 8192 |
| |
| |
| |
| 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("<f2").reshape(n_blocks, 1).astype(np.float32) |
| qs = blocks[:, 2:] |
| low = (qs & 0x0F).astype(np.float32) - 8.0 |
| high = (qs >> 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)) |
| 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, 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) |
| 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 |
|
|
| |
| 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) |
|
|
| |
| cand_d16 = to_fp16_bits(wls_d[:, None] * MULTS_D[None, :]) |
| cand_d = from_fp16_bits(cand_d16) |
| 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) |
| 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) |
|
|
| |
| best_cand = _q4_sieve_and_viterbi(cand_errors, cfg) |
| rows = np.arange(B) |
| d_sel = cand_d[rows, best_cand] |
|
|
| |
| 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) |
| 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) |
|
|
| |
| 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) |
|
|
| |
| 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) |
|
|
| |
| |
| |
| |
| 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("<u2").view(np.uint8).reshape(B, 2) |
| out[:, 2:] = (best_q_u8[:, :16] & 0x0F) | ((best_q_u8[:, 16:] & 0x0F) << np.uint8(4)) |
|
|
| deq_final = best_d[:, None] * (best_q_u8.astype(np.float32) - 8.0) |
| raw_sse = float(np.sum((X_f32 - deq_final) ** 2, dtype=np.float64)) |
| return out.tobytes(), raw_sse |
|
|
|
|
| def _quantize_q4_0_hpc_cuda( |
| X_np: np.ndarray, W_np: np.ndarray, dev: "torch.device", cfg: HPCConfig, chunk_idx: int = 0 |
| ) -> 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) |
|
|
| |
| 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) |
|
|
| |
| mults = torch.tensor(MULTS_D, dtype=torch.float32, device=dev) |
| cand_d = (wls_d[:, None] * mults[None, :]).to(torch.float16).to(torch.float32) |
| 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 |
|
|
| |
| 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] |
|
|
| |
| 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) |
| 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 |
|
|
| |
| 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) |
|
|
| |
| 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) |
|
|
| |
| |
| |
| |
| 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("<u2").view(np.uint8).reshape(B, 2) |
| out[:, 2:] = (q_np[:, :16] & 0x0F) | ((q_np[:, 16:] & 0x0F) << np.uint8(4)) |
| return out.tobytes(), raw_sse |
|
|
|
|
| def _q4_0_thermal_cuda( |
| X: "torch.Tensor", |
| W: "torch.Tensor | None", |
| d: "torch.Tensor", |
| codes: "torch.Tensor", |
| dev: "torch.device", |
| cfg: HPCConfig, |
| seed: int, |
| ) -> 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: |
| |
| |
| |
| 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 |
|
|
