"""CUDA GPU-accelerated backend for HPC-Quantize V2 (PyTorch CUDA). Executes the entire HPC-Quantize V2 pipeline in batched FP32 GPU tensor kernels: * Batched GGML seed solvers (``make_qkx2``, ``make_qp``) and 5-step WLS; * Batched Geometry-Preserving Weight Analogs (Sub-block Iso-Moment, Vesica Gram-Preserving, and Cross-Sub-block Scale-Lattice); * Batched Coupled Orthogonal-Basis Bilinear Q2_K Factorization with rank-1 1-D integer refinement and 36-seed candidate lattice evaluation; * Sequential 6-state Sieve & 36-state Viterbi DP + vectorized row DC carry; * Batched Algebraic Phase-4 Coupled Assembly, Vesica+DC augmented 2x2 Cramer (d, dmin) solve, O(1) quadratic fp16-ULP polish, closed-form 1-D quadratic Vesica/DC code projection, and row-parallel GPU DC residual nudge. """ from __future__ import annotations import time import numpy as np try: import torch except ImportError: # pragma: no cover torch = None # type: ignore[assignment] from .config import HPCConfig from .factored import MULTS_D, MULTS_M, N_CAND_D, N_CAND_M, TOTAL_CANDIDATES from .pipeline import QuantizedTensor, TensorStats, _phase3_select from .q2k import NS, QK_K, Q2KBlock, pin_q2k_dc_residual from .seeds import seed_clamp_value from .spectral import fold_weights def cuda_is_available() -> bool: return bool(torch is not None and torch.cuda.is_available()) def cuda_device_name(device: str = "cuda") -> str: if not cuda_is_available(): return "CPU" idx = torch.device(device).index or 0 return str(torch.cuda.get_device_name(idx)) def _nearest_int_t(val: "torch.Tensor") -> "torch.Tensor": """Bit-exact GGML round-to-nearest-even on GPU.""" shifted = val.to(torch.float32) + 12582912.0 bits = shifted.view(torch.int32) return ((bits & 0x007FFFFF) - 0x00400000).to(torch.int64) def _round_fp16_t(val: "torch.Tensor") -> "torch.Tensor": return val.to(torch.float16).to(torch.float32) def _to_fp16_bits_t(val: "torch.Tensor") -> "torch.Tensor": return val.to(torch.float16).view(torch.int16).to(torch.int32) & 0xFFFF def _from_fp16_bits_t(bits: "torch.Tensor") -> "torch.Tensor": return bits.to(torch.int16).view(torch.float16).to(torch.float32) def _codes_from_params_t( X: "torch.Tensor", d: "torch.Tensor", m: "torch.Tensor", ls: "torch.Tensor", lm: "torch.Tensor", ) -> "torch.Tensor": d_sub = (d[:, None].to(torch.float32) * ls.to(torch.float32)).repeat_interleave(16, dim=1) m_sub = (m[:, None].to(torch.float32) * lm.to(torch.float32)).repeat_interleave(16, dim=1) safe = d_sub >= 1.0e-15 safe_d = torch.where(safe, d_sub, torch.ones_like(d_sub)) q = _nearest_int_t((X.to(torch.float32) + m_sub) / safe_d).clamp(0, 3) return torch.where(safe, q, torch.zeros_like(q)).to(torch.uint8) def _dequantize_t( d: "torch.Tensor", m: "torch.Tensor", ls: "torch.Tensor", lm: "torch.Tensor", codes: "torch.Tensor", ) -> "torch.Tensor": d_sub = (d[:, None].to(torch.float32) * ls.to(torch.float32)).repeat_interleave(16, dim=1) m_sub = (m[:, None].to(torch.float32) * lm.to(torch.float32)).repeat_interleave(16, dim=1) return d_sub * codes.to(torch.float32) - m_sub def _objective_t( X: "torch.Tensor", W: "torch.Tensor", d: "torch.Tensor", m: "torch.Tensor", ls: "torch.Tensor", lm: "torch.Tensor", codes: "torch.Tensor", carry: "torch.Tensor | None", cfg: HPCConfig, ) -> tuple["torch.Tensor", "torch.Tensor"]: xh = _dequantize_t(d, m, ls, lm, codes) e = X - xh sse = torch.sum(W * e * e, dim=1, dtype=torch.float32) 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 = torch.zeros_like(sse) 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 and carry is not None: dst = dst + carry[:, None].to(torch.float32) if fw[lvl] != 0.0: pen = pen + float(fw[lvl]) * torch.sum(dst * dst, dim=1) src = dst m_len //= 2 lvl += 1 return sse + pen, sse def _make_qkx2_cuda(x: "torch.Tensor", w: "torch.Tensor", nmax: int = 3): x = x.to(torch.float32) w = w.to(torch.float32) xmin = torch.clamp(torch.min(x, dim=1).values, max=0.0) xmax = torch.max(x, dim=1).values span = xmax - xmin degenerate = (xmax == xmin) | (span <= 1.0e-30) sum_w = torch.sum(w, dim=1) sum_x = torch.sum(w * x, dim=1) safe_span = torch.where(span > 1.0e-30, span, torch.ones_like(span)) iscale = nmax / safe_span scale = 1.0 / iscale codes = _nearest_int_t(iscale[:, None] * (x - xmin[:, None])).clamp(0, nmax) best_mad = torch.sum(w * torch.abs(scale[:, None] * codes.to(torch.float32) + xmin[:, None] - x), dim=1) best_scale = scale.clone() best_min = xmin.clone() best_codes = codes.clone() steps = torch.arange(16, dtype=torch.float32, device=x.device)[None, :] # (1, 16) try_iscale = (-0.5 + 0.1 * steps + nmax) / safe_span[:, None] # (N, 16) aux = _nearest_int_t(try_iscale[:, :, None] * (x - xmin[:, None])[:, None, :]).clamp(0, nmax) aux_f = aux.to(torch.float32) w_e = w[:, None, :] x_e = x[:, None, :] sl = torch.sum(w_e * aux_f, dim=2) sl2 = torch.sum(w_e * aux_f * aux_f, dim=2) sxl = torch.sum(w_e * aux_f * x_e, dim=2) det = sum_w[:, None] * sl2 - sl * sl valid = det > 1.0e-30 s_det = torch.where(valid, det, torch.ones_like(det)) this_scale = (sum_w[:, None] * sxl - sum_x[:, None] * sl) / s_det this_min = (sl2 * sum_x[:, None] - sl * sxl) / s_det pos_min = this_min > 0.0 this_min = torch.where(pos_min, torch.zeros_like(this_min), this_min) s_sl2 = torch.where(sl2 > 1.0e-30, sl2, torch.ones_like(sl2)) this_scale = torch.where(pos_min & (sl2 > 1.0e-30), sxl / s_sl2, this_scale) mad = torch.sum(w_e * torch.abs(this_scale[:, :, None] * aux_f + this_min[:, :, None] - x_e), dim=2) mad = torch.where(valid, mad, torch.full_like(mad, float("inf"))) min_mad, best_k = torch.min(mad, dim=1) take = min_mad < best_mad rows = torch.arange(x.shape[0], device=x.device) best_scale = torch.where(take, this_scale[rows, best_k], best_scale) best_min = torch.where(take, this_min[rows, best_k], best_min) best_codes = torch.where(take[:, None], aux[rows, best_k], best_codes) best_scale = torch.where(degenerate, torch.zeros_like(best_scale), best_scale) best_codes = torch.where(degenerate[:, None], torch.zeros_like(best_codes), best_codes) best_min = torch.where(degenerate, xmin, best_min) return best_scale, -best_min, best_codes.to(torch.uint8) def _make_qp_cuda(values: "torch.Tensor", weights: "torch.Tensor", nmax: int = 15): x = values.to(torch.float32) w = weights.to(torch.float32) xmax = torch.max(x, dim=1).values dead = xmax < 1.0e-15 safe_max = torch.where(dead, torch.ones_like(xmax), xmax) steps = torch.arange(-4, 5, dtype=torch.float32, device=x.device)[None, :] # (1, 9) t_is = (0.1 * steps + nmax) / safe_max[:, None] # (N, 9) t_sc = 1.0 / t_is t_cd = _nearest_int_t(t_is[:, :, None] * x[:, None, :]).clamp(0, nmax).to(torch.float32) mse = torch.sum(w[:, None, :] * (x[:, None, :] - t_sc[:, :, None] * t_cd) ** 2, dim=2) best_k = torch.argmin(mse, dim=1) rows = torch.arange(x.shape[0], device=x.device) best_iscale = t_is[rows, best_k] codes = _nearest_int_t(best_iscale[:, None] * x).clamp(0, nmax) cf = codes.to(torch.float32) sumlx = torch.sum(w * x * cf, dim=1) suml2 = torch.sum(w * cf * cf, dim=1) for _ in range(3): changed = False for i in range(x.shape[1]): wi = w[:, i] xi = x[:, i] li = codes[:, i].to(torch.float32) slx = sumlx - wi * xi * li sl2 = suml2 - wi * li * li valid = (slx > 0.0) & (sl2 > 0.0) s_slx = torch.where(valid, slx, torch.ones_like(slx)) new_l = _nearest_int_t(xi * sl2 / s_slx).clamp(0, nmax) nl_f = new_l.to(torch.float32) trial = valid & (new_l != codes[:, i]) slx2 = slx + wi * xi * nl_f sl2_2 = sl2 + wi * nl_f * nl_f better = trial & (slx2 * slx2 * suml2 > sumlx * sumlx * sl2_2) if torch.any(better): codes[:, i] = torch.where(better, new_l, codes[:, i]) sumlx = torch.where(better, slx2, sumlx) suml2 = torch.where(better, sl2_2, suml2) changed = True if not changed: break s_l2 = torch.where(suml2 > 1.0e-30, suml2, torch.ones_like(suml2)) d = torch.where((suml2 > 1.0e-30) & (~dead), sumlx / s_l2, torch.zeros_like(sumlx)) codes = torch.where(dead[:, None], torch.zeros_like(codes), codes) return d, codes.to(torch.uint8) def _block_clamp_cuda(X: "torch.Tensor", cfg: HPCConfig) -> "torch.Tensor": """Per-block kurtosis-gated outlier clamp thresholds on GPU (shape (B, 1, 1)).""" b2 = X * X mean_sq = torch.mean(b2, dim=1) mean_4 = torch.mean(b2 * b2, dim=1) kurt = torch.where(mean_sq > 1.0e-30, mean_4 / torch.clamp(mean_sq * mean_sq, min=1.0e-30), torch.full_like(mean_sq, 3.0)) sigma = torch.sqrt(torch.clamp(mean_sq, min=0.0)) mult = torch.where( kurt <= cfg.clamp_kurtosis_off, torch.full_like(kurt, 1.0e9), torch.where( kurt <= cfg.clamp_kurtosis_soft, torch.full_like(kurt, float(cfg.clamp_sigma_soft)), torch.full_like(kurt, float(cfg.clamp_sigma_hard)), ), ) clamp = torch.where(mean_sq > 1.0e-30, sigma * mult, torch.full_like(sigma, float("inf"))) return clamp[:, None, None] def _build_seeds_cuda(X: "torch.Tensor", W: "torch.Tensor | None", cfg: HPCConfig): B = X.shape[0] blocks = X.reshape(B, 16, 16) per_blk_clamp = _block_clamp_cuda(X, cfg) clipped = torch.clamp(blocks, min=-per_blk_clamp, max=per_blk_clamp) if W is None: sigma_sq = torch.mean(blocks * blocks, dim=(1, 2), keepdim=True) weights = torch.sqrt(sigma_sq + clipped * clipped) else: weights = W.reshape(B, 16, 16) scales, mins, codes = _make_qkx2_cuda(clipped.reshape(-1, 16), weights.reshape(-1, 16), 3) scales = scales.reshape(B, 16) mins = mins.reshape(B, 16) codes = codes.reshape(B, 256) sw = torch.sum(weights, dim=2) d, ls = _make_qp_cuda(scales, sw, 15) dmin, lm = _make_qp_cuda(mins, sw, 15) return { "scales": scales, "mins": mins, "codes": codes, "sw": sw, "d": d, "dmin": dmin, "ls": ls, "lm": lm, "clamp": per_blk_clamp.reshape(B, 1), } def _phase2_wls_cuda(X: "torch.Tensor", W: "torch.Tensor", seeds: dict): per_blk_clamp = seeds["clamp"] clipped = torch.clamp(X, min=-per_blk_clamp, max=per_blk_clamp) base_d = seeds["d"].clone() base_m = seeds["dmin"].clone() ls = seeds["ls"].clone() lm = seeds["lm"].clone() scales = seeds["scales"] mins = seeds["mins"] for _ in range(5): codes = _codes_from_params_t(clipped, base_d, base_m, ls, lm).to(torch.float32) A = ls.to(torch.float32).repeat_interleave(16, dim=1) * codes Bm = lm.to(torch.float32).repeat_interleave(16, dim=1) Saa = torch.sum(W * A * A, dim=1) Sab = torch.sum(W * A * Bm, dim=1) Sbb = torch.sum(W * Bm * Bm, dim=1) Sxa = torch.sum(W * clipped * A, dim=1) Sxb = torch.sum(W * clipped * Bm, dim=1) det = Saa * Sbb - Sab * Sab valid = det > 1.0e-30 s_det = torch.where(valid, det, torch.ones_like(det)) d_new = (Sbb * Sxa - Sab * Sxb) / s_det m_new = (Sab * Sxa - Saa * Sxb) / s_det upd_d = valid & (d_new > 0.0) & torch.isfinite(d_new) upd_m = valid & (m_new > 0.0) & torch.isfinite(m_new) base_d = torch.where(upd_d, _round_fp16_t(d_new), base_d) base_m = torch.where(upd_m, _round_fp16_t(m_new), base_m) sd = torch.where(base_d > 1.0e-15, base_d, torch.ones_like(base_d)) sm = torch.where(base_m > 1.0e-15, base_m, torch.ones_like(base_m)) ls = torch.where( base_d[:, None] > 1.0e-15, _nearest_int_t(scales / sd[:, None]).clamp(0, 15), torch.zeros_like(ls, dtype=torch.int64), ).to(torch.uint8) lm = torch.where( base_m[:, None] > 1.0e-15, _nearest_int_t(mins / sm[:, None]).clamp(0, 15), torch.zeros_like(lm, dtype=torch.int64), ).to(torch.uint8) return base_d, base_m def _factor_1d_cuda( V: "torch.Tensor", W: "torch.Tensor", nmax: int = 15, iters: int = 5, refine: bool = True, ) -> tuple["torch.Tensor", "torch.Tensor"]: """Batched FP32 1-D factorization across B blocks on GPU with rank-1 sweep.""" V = torch.clamp(V.to(torch.float32), min=0.0) W = torch.clamp(W.to(torch.float32), min=0.0) B, n = V.shape vmax = torch.max(V, dim=1).values active = vmax >= 1.0e-30 d = torch.where(active, vmax / nmax, torch.zeros_like(vmax)) sd = torch.where(d > 1.0e-30, d, torch.ones_like(d)) codes = torch.where( active[:, None] & (d[:, None] > 1.0e-30), _nearest_int_t(V / sd[:, None]).clamp(0, nmax), torch.zeros((B, n), dtype=torch.int64, device=V.device), ) for _ in range(iters): cf = codes.to(torch.float32) den = torch.sum(W * cf * cf, dim=1) valid = active & (den > 1.0e-30) s_den = torch.where(valid, den, torch.ones_like(den)) dn = torch.sum(W * V * cf, dim=1) / s_den d = torch.where(valid & (dn > 0.0), dn, d) sd = torch.where(d > 1.0e-30, d, torch.ones_like(d)) codes = torch.where( active[:, None] & (d[:, None] > 1.0e-30), _nearest_int_t(V / sd[:, None]).clamp(0, nmax), torch.zeros_like(codes), ) if refine: opts = torch.arange(nmax + 1, dtype=torch.float32, device=V.device)[None, :] opts_sq = opts * opts wv = W * V cf = codes.to(torch.float32) sum_lv = torch.sum(wv * cf, dim=1) sum_l2 = torch.sum(W * cf * cf, dim=1) s_l2 = torch.where(sum_l2 > 1.0e-30, sum_l2, torch.ones_like(sum_l2)) best_gain = torch.where(sum_l2 > 1.0e-30, (sum_lv * sum_lv) / s_l2, torch.zeros_like(sum_lv)) rows = torch.arange(B, device=V.device) for i in range(n): wi = W[:, i : i + 1] wvi = wv[:, i : i + 1] li = codes[:, i : i + 1].to(torch.float32) slv = (sum_lv[:, None] - wvi * li) + wvi * opts sl2 = (sum_l2[:, None] - wi * (li * li)) + wi * opts_sq valid = (wi > 0.0) & (sl2 > 1.0e-30) & (slv > 0.0) gains = torch.where(valid, (slv * slv) / torch.where(valid, sl2, torch.ones_like(sl2)), -torch.ones_like(slv)) max_g, k_idx = torch.max(gains, dim=1) improve = active & (k_idx != codes[:, i]) & (max_g > best_gain + 1.0e-12) if torch.any(improve): best_gain = torch.where(improve, max_g, best_gain) sum_lv = torch.where(improve, slv[rows, k_idx], sum_lv) sum_l2 = torch.where(improve, sl2[rows, k_idx], sum_l2) codes[:, i] = torch.where(improve, k_idx, codes[:, i]) s_l2 = torch.where(sum_l2 > 1.0e-30, sum_l2, torch.ones_like(sum_l2)) d = torch.where(active & (sum_l2 > 1.0e-30), sum_lv / s_l2, d) return d, codes.to(torch.uint8) def _analogs_cuda(X: "torch.Tensor", W: "torch.Tensor") -> list["torch.Tensor"]: B = X.shape[0] xb = X.reshape(B, NS, 16) wb = W.reshape(B, NS, 16) wsum = torch.clamp(torch.sum(wb, dim=2, keepdim=True), min=1.0e-30) mu = torch.sum(wb * xb, dim=2, keepdim=True) / wsum centered = xb - mu var = torch.sum(wb * centered * centered, dim=2, keepdim=True) / wsum sigma = torch.sqrt(torch.clamp(var, min=1.0e-30)) active = sigma >= 1.0e-15 z = centered / sigma q_ideal = torch.linspace(-1.35, 1.35, 16, dtype=torch.float32, device=X.device) order = torch.argsort(z, dim=2, stable=True) z_quant = torch.empty_like(z) z_quant.scatter_(2, order, q_ideal[None, None, :].expand_as(z)) analogs = [] for tau, beta, q_blend in ((1.60, 0.25, 0.30), (2.15, 0.15, 0.15)): w_rel = torch.clamp(wb / (wsum / 16.0), min=1.0e-12) ** beta z_comp = tau * torch.tanh((z * w_rel) / tau) z_a = (1.0 - q_blend) * z_comp + q_blend * z_quant mu_a = torch.sum(wb * z_a, dim=2, keepdim=True) / wsum z_a = z_a - mu_a var_a = torch.sum(wb * z_a * z_a, dim=2, keepdim=True) / wsum scale = torch.where(var_a > 1.0e-30, torch.sqrt(var / torch.clamp(var_a, min=1.0e-30)), torch.zeros_like(var_a)) sub_a = torch.where(active, mu + z_a * scale, xb) orig_min = torch.min(xb, dim=2, keepdim=True).values new_min = torch.min(sub_a, dim=2, keepdim=True).values shift = torch.where((orig_min <= 0.0) & (new_min > 0.0), new_min, torch.zeros_like(new_min)) analogs.append((sub_a - shift).reshape(B, QK_K)) return analogs def _sub_ls_cuda(X: "torch.Tensor", W: "torch.Tensor", codes: "torch.Tensor"): B = X.shape[0] q = codes.to(torch.float32).reshape(B, NS, 16) xb = X.reshape(B, NS, 16) wb = W.reshape(B, NS, 16) Sqq = torch.sum(wb * q * q, dim=2) Swq = torch.sum(wb * q, dim=2) Sw = torch.sum(wb, dim=2) Sqx = torch.sum(wb * q * xb, dim=2) Sx = torch.sum(wb * xb, dim=2) det = Sqq * Sw - Swq * Swq valid = det > 1.0e-30 s_det = torch.where(valid, det, torch.ones_like(det)) s_sw = torch.where(Sw > 1.0e-30, Sw, torch.ones_like(Sw)) sj = torch.clamp(torch.where(valid, (Sqx * Sw - Swq * Sx) / s_det, torch.zeros_like(det)), min=0.0) tj = torch.where(valid, (sj * Swq - Sx) / s_sw, torch.zeros_like(det)) bad_tj = tj < 0.0 tj = torch.where(bad_tj, torch.zeros_like(tj), tj) s_sqq = torch.where(Sqq > 1.0e-30, Sqq, torch.ones_like(Sqq)) fallback = torch.clamp(torch.where(Sqq > 1.0e-30, Sqx / s_sqq, torch.zeros_like(Sqq)), min=0.0) sj = torch.where(bad_tj, fallback, sj) return sj, tj, Sqq, Swq, Sw, Sx def _factor_candidates_cuda( X: "torch.Tensor", W: "torch.Tensor", scales: "torch.Tensor", mins: "torch.Tensor", base_d: "torch.Tensor", base_m: "torch.Tensor", cfg: HPCConfig, ) -> np.ndarray: """Vectorized FP32 36-seed candidate lattice across B blocks on GPU.""" B = X.shape[0] device = X.device nd, nm = int(cfg.n_cand_d), int(cfg.n_cand_m) di_idx = [a if nd == N_CAND_D else (a * N_CAND_D) // nd for a in range(nd)] mi_idx = [b if nm == N_CAND_M else (b * N_CAND_M) // nm for b in range(nm)] md = torch.tensor(MULTS_D[di_idx], dtype=torch.float32, device=device) mm = torch.tensor(MULTS_M[mi_idx], dtype=torch.float32, device=device) ds_mat = _round_fp16_t(base_d[:, None] * md[None, :]) ms_mat = _round_fp16_t(base_m[:, None] * mm[None, :]) ds = ds_mat.repeat_interleave(nm, dim=1) ms = ms_mat.repeat(1, nd) K = ds.shape[1] sc = scales[:, None, :] mn = mins[:, None, :] s_ds = torch.where(ds[:, :, None] > 1.0e-30, ds[:, :, None], torch.ones_like(ds[:, :, None])) s_ms = torch.where(ms[:, :, None] > 1.0e-30, ms[:, :, None], torch.ones_like(ms[:, :, None])) ls0 = torch.where( ds[:, :, None] > 1.0e-30, _nearest_int_t(sc / s_ds).clamp(0, 15), torch.zeros((B, K, NS), dtype=torch.int64, device=device), ).to(torch.float32) lm0 = torch.where( ms[:, :, None] > 1.0e-30, _nearest_int_t(mn / s_ms).clamp(0, 15), torch.zeros((B, K, NS), dtype=torch.int64, device=device), ).to(torch.float32) X_e = X[:, None, :] W_e = W[:, None, :] fw = fold_weights(QK_K, cfg) def _eval_k(ls_k: "torch.Tensor", lm_k: "torch.Tensor"): d_sub = (ds[:, :, None] * ls_k).repeat_interleave(16, dim=2) m_sub = (ms[:, :, None] * lm_k).repeat_interleave(16, dim=2) safe = d_sub >= 1.0e-15 sd = torch.where(safe, d_sub, torch.ones_like(d_sub)) q = torch.where( safe, _nearest_int_t((X_e + m_sub) / sd).clamp(0, 3), torch.zeros((B, K, QK_K), dtype=torch.int64, device=device), ).to(torch.float32) e = X_e - (d_sub * q - m_sub) sse = torch.sum(W_e * e * e, dim=2) if cfg.use_spectral and (cfg.dc_lambda != 0.0 or cfg.vw_lambda != 0.0): pen = torch.zeros_like(sse) src = e m_len = QK_K // 2 lvl = 0 while m_len >= 1: dst = src[:, :, :m_len] + src[:, :, m_len : 2 * m_len] if fw[lvl] != 0.0: pen = pen + float(fw[lvl]) * torch.sum(dst * dst, dim=2) src = dst m_len //= 2 lvl += 1 return sse + pen, q return sse, q obj0, q0 = _eval_k(ls0, lm0) q_sub = q0.reshape(B, K, NS, 16) xb = X.reshape(B, 1, NS, 16) wb = W.reshape(B, 1, NS, 16) Sqq = torch.sum(wb * q_sub * q_sub, dim=3) Swq = torch.sum(wb * q_sub, dim=3) Sw = torch.sum(wb, dim=3) Sqx = torch.sum(wb * q_sub * xb, dim=3) Sx = torch.sum(wb * xb, dim=3) det = Sqq * Sw - Swq * Swq valid = det > 1.0e-30 s_det = torch.where(valid, det, torch.ones_like(det)) s_sw = torch.where(Sw > 1.0e-30, Sw, torch.ones_like(Sw)) sj = torch.clamp(torch.where(valid, (Sqx * Sw - Swq * Sx) / s_det, torch.zeros_like(det)), min=0.0) tj = torch.where(valid, (sj * Swq - Sx) / s_sw, torch.zeros_like(det)) bad_tj = tj < 0.0 tj = torch.where(bad_tj, torch.zeros_like(tj), tj) s_sqq = torch.where(Sqq > 1.0e-30, Sqq, torch.ones_like(Sqq)) sj = torch.where(bad_tj, torch.clamp(torch.where(Sqq > 1.0e-30, Sqx / s_sqq, torch.zeros_like(Sqq)), min=0.0), sj) ls1 = torch.where( ds[:, :, None] > 1.0e-30, _nearest_int_t(sj / s_ds).clamp(0, 15), torch.zeros((B, K, NS), dtype=torch.int64, device=device), ).to(torch.float32) t_tgt = torch.clamp((ds[:, :, None] * ls1 * Swq - Sx) / s_sw, min=0.0) if cfg.factor_coupled_q2k else tj lm1 = torch.where( ms[:, :, None] > 1.0e-30, _nearest_int_t(t_tgt / s_ms).clamp(0, 15), torch.zeros((B, K, NS), dtype=torch.int64, device=device), ).to(torch.float32) obj1, _ = _eval_k(ls1, lm1) best_k = torch.minimum(obj0, obj1).cpu().numpy() out = np.full((B, TOTAL_CANDIDATES), np.inf, dtype=np.float32) if nd == N_CAND_D and nm == N_CAND_M: out[:, :] = best_k else: for a, di in enumerate(di_idx): cb = int(di) >> 2 for b, mi in enumerate(mi_idx): fb = int(mi) >> 2 col = best_k[:, a * nm + b] for aa in range(4): for bb in range(4): idx = (cb * 4 + aa) * N_CAND_M + fb * 4 + bb out[:, idx] = np.minimum(out[:, idx], col) return out def _standalone_coupled_solve_cuda( X: "torch.Tensor", W: "torch.Tensor", cfg: HPCConfig ) -> tuple["torch.Tensor", "torch.Tensor", "torch.Tensor", "torch.Tensor", "torch.Tensor", "torch.Tensor"]: B = X.shape[0] Sw = torch.sum(W.reshape(B, NS, 16), dim=2) s_Sw = torch.clamp(Sw, min=1.0e-30) sources = [X] if cfg.use_weight_analogs: sources.extend(_analogs_cuda(X, W)) best_obj = torch.full((B,), float("inf"), dtype=torch.float32, device=X.device) best_d = torch.zeros((B,), dtype=torch.float32, device=X.device) best_m = torch.zeros((B,), dtype=torch.float32, device=X.device) best_ls = torch.zeros((B, NS), dtype=torch.uint8, device=X.device) best_lm = torch.zeros((B, NS), dtype=torch.uint8, device=X.device) best_codes = torch.zeros((B, QK_K), dtype=torch.uint8, device=X.device) for src in sources: sb = src.reshape(B, NS, 16) xmin = torch.clamp(torch.min(sb, dim=2).values, max=0.0) xmax = torch.max(sb, dim=2).values span = xmax - xmin s0 = torch.where(span > 1.0e-30, span / 3.0, torch.zeros_like(span)) t0 = -xmin s_s0 = torch.where(s0 > 1.0e-30, s0, torch.ones_like(s0)) q0 = torch.where( s0[:, :, None] > 1.0e-30, _nearest_int_t((sb + t0[:, :, None]) / s_s0[:, :, None]).clamp(0, 3), torch.zeros((B, NS, 16), dtype=torch.int64, device=X.device), ).reshape(B, QK_K).to(torch.uint8) codes = q0 for _ in range(3): s, t, Sqq, Swq, _, Sx = _sub_ls_cuda(X, W, codes) q_bar = Swq / s_Sw Vqq = torch.clamp(Sqq - Swq * q_bar, min=1.0e-12 * Sw) d, ls = _factor_1d_cuda(s, Vqq, 15, 4, cfg.factor_refine_1d) t_c = torch.clamp((d[:, None] * ls.to(torch.float32) * Swq - Sx) / s_Sw, min=0.0) m, lm = _factor_1d_cuda(t_c, Sw, 15, 4, cfg.factor_refine_1d) d = _round_fp16_t(d) m = _round_fp16_t(m) codes = _codes_from_params_t(X, d, m, ls, lm) _, sse_cur = _objective_t(X, W, d, m, ls, lm, codes, None, cfg) take = sse_cur < best_obj - 1.0e-12 if torch.any(take): best_obj = torch.where(take, sse_cur, best_obj) best_d = torch.where(take, d, best_d) best_m = torch.where(take, m, best_m) best_ls = torch.where(take[:, None], ls, best_ls) best_lm = torch.where(take[:, None], lm, best_lm) best_codes = torch.where(take[:, None], codes, best_codes) return best_d, best_m, best_ls, best_lm, best_codes, best_obj def _extended_dm_solve_and_ulp_cuda( X: "torch.Tensor", W: "torch.Tensor", ls: "torch.Tensor", lm: "torch.Tensor", codes: "torch.Tensor", d_cur: "torch.Tensor", m_cur: "torch.Tensor", carry: "torch.Tensor", cfg: HPCConfig, ) -> tuple["torch.Tensor", "torch.Tensor"]: B = X.shape[0] A = ls.to(torch.float32).repeat_interleave(16, dim=1) * codes.to(torch.float32) Bm = lm.to(torch.float32).repeat_interleave(16, dim=1) Saa = torch.sum(W * A * A, dim=1) Sab = torch.sum(W * A * Bm, dim=1) Sbb = torch.sum(W * Bm * Bm, dim=1) Sxa = torch.sum(W * X * A, dim=1) Sxb = torch.sum(W * X * Bm, dim=1) det = Saa * Sbb - Sab * Sab valid = det > 1.0e-30 s_det = torch.where(valid, det, torch.ones_like(det)) d_new = (Sbb * Sxa - Sab * Sxb) / s_det m_new = (Sab * Sxa - Saa * Sxb) / s_det d_opt = torch.where(valid & (d_new > 0.0) & torch.isfinite(d_new), d_new, d_cur) m_opt = torch.where(valid & (m_new >= 0.0) & torch.isfinite(m_new), m_new, m_cur) rad = max(int(cfg.ulp_radius), 1) deltas = torch.arange(-rad, rad + 1, dtype=torch.int32, device=X.device) d16_c = _to_fp16_bits_t(d_opt) m16_c = _to_fp16_bits_t(m_opt) d_bits = torch.clamp(d16_c[:, None] + deltas[None, :], 0, 0x7BFF) m_bits = torch.clamp(m16_c[:, None] + deltas[None, :], 0, 0x7BFF) dd = _from_fp16_bits_t(d_bits)[:, :, None] mm = _from_fp16_bits_t(m_bits)[:, 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(B, -1) best_idx = torch.argmin(q_form, dim=1) K = deltas.numel() rows = torch.arange(B, device=X.device) return dd[rows, best_idx // K, 0], mm[rows, 0, best_idx % K] def _refine_trajectory_cuda( X: "torch.Tensor", W: "torch.Tensor", Xb: "torch.Tensor", Wb: "torch.Tensor", Sw: "torch.Tensor", Sx: "torch.Tensor", safe_Sw: "torch.Tensor", d0: "torch.Tensor", m0: "torch.Tensor", ls0: "torch.Tensor", lm0: "torch.Tensor", codes0: "torch.Tensor", carry: "torch.Tensor", cfg: HPCConfig, ): B = X.shape[0] device = X.device best_d, best_m = d0.clone(), m0.clone() best_ls, best_lm = ls0.clone(), lm0.clone() best_codes = codes0.clone() _, best_sse = _objective_t(X, W, best_d, best_m, best_ls, best_lm, best_codes, carry, cfg) d, m, ls, lm, codes = best_d.clone(), best_m.clone(), best_ls.clone(), best_lm.clone(), best_codes.clone() ls_deltas = torch.tensor([-3, -2, -1, 0, 1, 2, 3], dtype=torch.int64, device=device) lm_deltas = torch.tensor([-1, 0, 1], dtype=torch.int64, device=device).repeat(7) b_idx = torch.arange(B, device=device)[:, None] s_idx = torch.arange(NS, device=device)[None, :] for _ in range(cfg.refine_passes): d_new, m_new = _extended_dm_solve_and_ulp_cuda(X, W, ls, lm, codes, d, m, carry, cfg) _, sse_dm = _objective_t(X, W, d_new, m_new, ls, lm, codes, carry, cfg) imp_dm = sse_dm < best_sse - 1.0e-12 if torch.any(imp_dm): best_sse = torch.where(imp_dm, sse_dm, best_sse) best_d = torch.where(imp_dm, d_new, best_d) best_m = torch.where(imp_dm, m_new, best_m) best_ls = torch.where(imp_dm[:, None], ls, best_ls) best_lm = torch.where(imp_dm[:, None], lm, best_lm) best_codes = torch.where(imp_dm[:, None], codes, best_codes) d, m = best_d.clone(), best_m.clone() codes_re = _codes_from_params_t(X, d, m, ls, lm) _, sse_re = _objective_t(X, W, d, m, ls, lm, codes_re, carry, cfg) imp_re = sse_re < best_sse - 1.0e-12 if torch.any(imp_re): best_sse = torch.where(imp_re, sse_re, best_sse) best_codes = torch.where(imp_re[:, None], codes_re, best_codes) codes = best_codes.clone() q = codes.to(torch.float32).reshape(B, NS, 16) Sqq = torch.sum(Wb * q * q, dim=2) Swq = torch.sum(Wb * q, dim=2) Sqx = torch.sum(Wb * q * Xb, dim=2) q_bar = Swq / safe_Sw Vqq = torch.clamp(Sqq - Swq * q_bar, min=0.0) Cqx = Sqx - Sx * q_bar s_star = torch.where( Vqq > 1.0e-30, torch.clamp(Cqx / torch.clamp(Vqq, min=1.0e-30), min=0.0), d[:, None] * ls.to(torch.float32), ) sd = torch.where(d[:, None] > 1.0e-15, d[:, None], torch.ones((B, 1), dtype=torch.float32, device=device)) sm = torch.where(m[:, None] > 1.0e-15, m[:, None], torch.ones((B, 1), dtype=torch.float32, device=device)) ls_center = torch.where( d[:, None] > 1.0e-15, _nearest_int_t(s_star / sd).clamp(0, 15), torch.zeros((B, NS), dtype=torch.int64, device=device), ) c_ls_base = torch.clamp(ls_center[:, :, None] + ls_deltas[None, None, :], 0, 15).to(torch.float32) d_sub_base = d[:, None, None] * c_ls_base t_opt_base = torch.clamp((d_sub_base * Swq[:, :, None] - Sx[:, :, None]) / safe_Sw[:, :, None], min=0.0) lm_center = torch.where( m[:, None, None] > 1.0e-15, _nearest_int_t(t_opt_base / sm[:, :, None]).clamp(0, 15), torch.zeros_like(c_ls_base, dtype=torch.int64), ) cand_ls = c_ls_base.repeat_interleave(3, dim=2) cand_lm = torch.clamp(lm_center.repeat_interleave(3, dim=2) + lm_deltas[None, None, :], 0, 15).to(torch.float32) d_sub_c = d[:, None, None] * cand_ls m_sub_c = m[:, None, None] * cand_lm safe_dc = d_sub_c >= 1.0e-15 sd_c = torch.where(safe_dc, d_sub_c, torch.ones_like(d_sub_c)) q_c = torch.where( safe_dc[:, :, :, None], _nearest_int_t((Xb[:, :, None, :] + m_sub_c[:, :, :, None]) / sd_c[:, :, :, None]).clamp(0, 3), torch.zeros((B, NS, 21, 16), dtype=torch.int64, device=device), ).to(torch.float32) e_c = Xb[:, :, None, :] - (d_sub_c[:, :, :, None] * q_c - m_sub_c[:, :, :, None]) err_c = torch.sum(Wb[:, :, None, :] * e_c * e_c, dim=3) pick = torch.argmin(err_c, dim=2) ls_try = cand_ls[b_idx, s_idx, pick].to(torch.uint8) lm_try = cand_lm[b_idx, s_idx, pick].to(torch.uint8) codes_try = _codes_from_params_t(X, d, m, ls_try, lm_try) d_try, m_try = _extended_dm_solve_and_ulp_cuda(X, W, ls_try, lm_try, codes_try, d, m, carry, cfg) _, sse_try = _objective_t(X, W, d_try, m_try, ls_try, lm_try, codes_try, carry, cfg) imp_sub = sse_try < best_sse - 1.0e-12 if torch.any(imp_sub): best_sse = torch.where(imp_sub, sse_try, best_sse) best_d = torch.where(imp_sub, d_try, best_d) best_m = torch.where(imp_sub, m_try, best_m) best_ls = torch.where(imp_sub[:, None], ls_try, best_ls) best_lm = torch.where(imp_sub[:, None], lm_try, best_lm) best_codes = torch.where(imp_sub[:, None], codes_try, best_codes) d, m, ls, lm, codes = best_d.clone(), best_m.clone(), best_ls.clone(), best_lm.clone(), best_codes.clone() return best_d, best_m, best_ls, best_lm, best_codes, best_sse def _compute_dc_carry_fast( X_np: np.ndarray, scales_seed: np.ndarray, mins_seed: np.ndarray, best: np.ndarray, base_d: np.ndarray, base_m: np.ndarray, cfg: HPCConfig, row_width: int | None, ) -> np.ndarray: """Vectorized residual DC computation + fast row recurrence.""" n = X_np.shape[0] di = best // N_CAND_M mi = best % N_CAND_M d = (base_d * MULTS_D[di]).astype(np.float16).astype(np.float32) m = (base_m * MULTS_M[mi]).astype(np.float16).astype(np.float32) 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 = np.where(d[:, None] > 1.0e-15, np.clip(np.rint(scales_seed / sd), 0, 15), 0).astype(np.float32) lm = np.where(m[:, None] > 1.0e-15, np.clip(np.rint(mins_seed / sm), 0, 15), 0).astype(np.float32) d_sub = np.repeat(d[:, None] * ls, 16, axis=1) m_sub = np.repeat(m[:, None] * lm, 16, axis=1) safe = d_sub >= 1.0e-15 q = np.where(safe, np.clip(np.rint((X_np + m_sub) / np.where(safe, d_sub, 1.0)), 0, 3), 0.0) dc_res = np.sum(X_np - (d_sub * q - m_sub), axis=1, dtype=np.float64) bpr = row_width // QK_K if row_width and row_width % QK_K == 0 else 0 carry = np.zeros(n, dtype=np.float32) decay = float(cfg.dc_decay) cum = bool(cfg.carry_cumulative) if bpr > 0 and n % bpr == 0: n_rows = n // bpr res_2d = dc_res.reshape(n_rows, bpr) car_2d = np.zeros((n_rows, bpr), dtype=np.float64) rolling = np.zeros(n_rows, dtype=np.float64) for col in range(bpr): car_2d[:, col] = decay * rolling rolling = (car_2d[:, col] + res_2d[:, col]) if cum else res_2d[:, col] carry = car_2d.reshape(n).astype(np.float32) else: rolling_s = 0.0 for i in range(n): if bpr > 0 and i % bpr == 0: rolling_s = 0.0 carry[i] = np.float32(decay * rolling_s) rolling_s = (carry[i] + dc_res[i]) if cum else dc_res[i] return carry def _apply_dc_nudge_cuda( X_np: np.ndarray, W_np: np.ndarray | None, final_d: np.ndarray, final_m: np.ndarray, final_ls: np.ndarray, final_lm: np.ndarray, final_codes: np.ndarray, cfg: HPCConfig, row_width: int | None, dev: "torch.device", ) -> np.ndarray: """Row-parallel GPU DC residual nudge across all rows simultaneously.""" n = X_np.shape[0] bpr = row_width // QK_K if row_width and row_width % QK_K == 0 else 0 if bpr <= 0 or n % bpr != 0: return final_codes n_rows = n // bpr codes_out = final_codes.copy().reshape(n_rows, bpr, QK_K) X_r = X_np.reshape(n_rows, bpr, QK_K) W_r = None if W_np is None else W_np.reshape(n_rows, bpr, QK_K) d_r = final_d.reshape(n_rows, bpr) m_r = final_m.reshape(n_rows, bpr) ls_r = final_ls.reshape(n_rows, bpr, NS) lm_r = final_lm.reshape(n_rows, bpr, NS) decay = float(cfg.dc_decay) cum = bool(cfg.carry_cumulative) row_batch = 1024 q_opts = torch.arange(4, dtype=torch.float32, device=dev)[None, :, None] # (1, 4, 1) for r0 in range(0, n_rows, row_batch): r1 = min(r0 + row_batch, n_rows) R = r1 - r0 rolling = torch.zeros(R, dtype=torch.float32, device=dev) X_batch = torch.from_numpy(X_r[r0:r1]).to(dev) W_batch = ( torch.ones_like(X_batch) if W_r is None else torch.from_numpy(W_r[r0:r1]).to(dev) ) d_batch = torch.from_numpy(d_r[r0:r1]).to(dev) m_batch = torch.from_numpy(m_r[r0:r1]).to(dev) ls_batch = torch.from_numpy(ls_r[r0:r1]).to(dev, dtype=torch.float32) lm_batch = torch.from_numpy(lm_r[r0:r1]).to(dev, dtype=torch.float32) cd_batch = torch.from_numpy(codes_out[r0:r1]).to(dev, dtype=torch.int64) rows = torch.arange(R, device=dev) for col in range(bpr): carry = decay * rolling x_c = X_batch[:, col] w_c = W_batch[:, col] d_c = d_batch[:, col] m_c = m_batch[:, col] ls_c = ls_batch[:, col] lm_c = lm_batch[:, col] cd_c = cd_batch[:, col] d_sub = (d_c[:, None] * ls_c).repeat_interleave(16, dim=1) m_sub = (m_c[:, None] * lm_c).repeat_interleave(16, dim=1) active = d_sub >= 1.0e-15 e = x_c - (d_sub * cd_c.to(torch.float32) - m_sub) sse = torch.sum(w_c * e * e, dim=1) cap = sse * (1.0 + float(cfg.sse_budget)) res = torch.sum(e, dim=1) + carry for _ in range(2): e_new = x_c[:, None, :] - (d_sub[:, None, :] * q_opts - m_sub[:, None, :]) dsse = w_c[:, None, :] * (e_new * e_new - e[:, None, :] * e[:, None, :]) mask = ( active[:, None, :] & (q_opts.to(torch.int64) != cd_c[:, None, :]) & ((sse[:, None, None] + dsse) <= cap[:, None, None]) ) new_res = res[:, None, None] + (e_new - e[:, None, :]) impr = torch.where( mask, torch.abs(res)[:, None, None] - torch.abs(new_res), torch.full_like(new_res, -float("inf")), ).reshape(R, -1) best_imp, flat_idx = torch.max(impr, dim=1) do_upd = best_imp > 1.0e-12 if not torch.any(do_upd): break q_pick = flat_idx // QK_K k_pick = flat_idx % QK_K old_e = e[rows, k_pick] chosen_e = e_new[rows, q_pick, k_pick] de = chosen_e - old_e res = torch.where(do_upd, res + de, res) sse = torch.where( do_upd, sse + w_c[rows, k_pick] * (chosen_e * chosen_e - old_e * old_e), sse, ) e[rows, k_pick] = torch.where(do_upd, chosen_e, old_e) cd_c[rows, k_pick] = torch.where(do_upd, q_pick, cd_c[rows, k_pick]) cd_batch[:, col] = cd_c dc_res = torch.sum(e, dim=1) rolling = (carry + dc_res) if cum else dc_res codes_out[r0:r1] = cd_batch.cpu().numpy().astype(np.uint8) return codes_out.reshape(n, QK_K) # VRAM-aware CUDA batching for the Q2_K factored pipeline. The candidate # kernels materialize several live (B, K, 256) fp32 tensors (K = n_cand_d * # n_cand_m), so the per-block footprint scales with K (~6 fp32 temporaries of # 256 weights each). The floor is the historical 4096-block batch (known to # fit even on small cards); on large GPUs the batch grows with the free VRAM. Q2_K_CUDA_FLOOR_BLOCKS = 4096 Q2_K_CUDA_VRAM_FRACTION = 0.45 Q2_K_CUDA_FP32_TEMPS_PER_CANDIDATE = 6 def quantize_cuda( weights: np.ndarray, importance: np.ndarray | None, row_width: int | None, cfg: HPCConfig, status_callback=None, stop_event=None, device: str = "cuda", ) -> QuantizedTensor: """Run the full HPC-Quantize V2 pipeline on a CUDA GPU.""" if not cuda_is_available(): raise RuntimeError("PyTorch CUDA is not available on this system") 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") 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 X_np = 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_np = imp.reshape(n_blocks, QK_K) else: W_np = None sigma = float(np.sqrt(np.mean(x.astype(np.float64) ** 2))) clamp_val = seed_clamp_value(x, cfg) dev = torch.device(device) gpu_name = cuda_device_name(device) batch_blocks = Q2_K_CUDA_FLOOR_BLOCKS try: # VRAM-aware default: scale the batch to free memory at start-up, # never below the floor. Footprint per block: K live (B, K, 256) # fp32 temporaries drive the peak (see _factor_candidates_cuda). k_eff = int(cfg.n_cand_d) * int(cfg.n_cand_m) bytes_per_block = ( Q2_K_CUDA_FP32_TEMPS_PER_CANDIDATE * k_eff * QK_K * 4 ) free_bytes, _ = torch.cuda.mem_get_info(dev) batch_blocks = max( Q2_K_CUDA_FLOOR_BLOCKS, int(Q2_K_CUDA_VRAM_FRACTION * free_bytes) // max(bytes_per_block, 1), ) except Exception: batch_blocks = Q2_K_CUDA_FLOOR_BLOCKS candidate_errors = np.empty((n_blocks, TOTAL_CANDIDATES), dtype=np.float32) scales_seed = np.empty((n_blocks, NS), dtype=np.float32) mins_seed = np.empty((n_blocks, NS), dtype=np.float32) base_d = np.empty((n_blocks,), dtype=np.float32) base_m = np.empty((n_blocks,), dtype=np.float32) solo_d = np.empty((n_blocks,), dtype=np.float32) solo_m = np.empty((n_blocks,), dtype=np.float32) solo_ls = np.empty((n_blocks, NS), dtype=np.uint8) solo_lm = np.empty((n_blocks, NS), dtype=np.uint8) solo_codes = np.empty((n_blocks, QK_K), dtype=np.uint8) solo_obj = np.empty((n_blocks,), dtype=np.float64) t0 = time.time() with torch.no_grad(): for start in range(0, n_blocks, batch_blocks): _check_stop() stop = min(start + batch_blocks, n_blocks) X_t = torch.from_numpy(X_np[start:stop]).to(dev) if W_np is None: W_t = torch.ones_like(X_t) else: W_raw_t = torch.from_numpy(W_np[start:stop]).to(dev) if cfg.temper_imatrix: sig2_t = torch.mean(X_t * X_t, dim=1, keepdim=True) W_t = torch.sqrt(torch.clamp(W_raw_t, min=1.0e-6)) * torch.sqrt(sig2_t + X_t * X_t) wm_t = torch.mean(W_t, dim=1, keepdim=True) W_t = torch.where(wm_t > 1.0e-30, W_t / torch.clamp(wm_t, min=1.0e-30), torch.ones_like(W_t)) else: W_t = W_raw_t seeds_t = _build_seeds_cuda(X_t, None if W_np is None else W_t, cfg) bd_t, bm_t = _phase2_wls_cuda(X_t, W_t, seeds_t) scales_seed[start:stop] = seeds_t["scales"].cpu().numpy() mins_seed[start:stop] = seeds_t["mins"].cpu().numpy() base_d[start:stop] = bd_t.cpu().numpy() base_m[start:stop] = bm_t.cpu().numpy() candidate_errors[start:stop] = _factor_candidates_cuda( X_t, W_t, seeds_t["scales"], seeds_t["mins"], bd_t, bm_t, cfg ) sd_t, sm_t, sls_t, slm_t, sc_t, sobj_t = _standalone_coupled_solve_cuda(X_t, W_t, cfg) _, sse_Solo = _objective_t(X_t, W_t, sd_t, sm_t, sls_t, slm_t, sc_t, None, cfg) solo_d[start:stop] = sd_t.cpu().numpy() solo_m[start:stop] = sm_t.cpu().numpy() solo_ls[start:stop] = sls_t.cpu().numpy() solo_lm[start:stop] = slm_t.cpu().numpy() solo_codes[start:stop] = sc_t.cpu().numpy() solo_obj[start:stop] = sse_Solo.cpu().numpy() elapsed = max(time.time() - t0, 1.0e-6) rate = stop / elapsed eta = (n_blocks - stop) / max(rate, 1.0e-9) _emit( f"CUDA Phase 1–2b ({gpu_name}): Seeds + Analogs + Factored Field", stop, n_blocks, f"blk {stop}/{n_blocks} ({rate:.1f} blk/s, ETA {eta:.1f}s) " f"mean_obj={float(np.mean(solo_obj[:stop])):.5e}", rate=rate, eta=eta, sigma=sigma, ) _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_fast( X_np, scales_seed, mins_seed, best, base_d, base_m, cfg, row_width ) final_d = np.empty((n_blocks,), dtype=np.float32) final_m = np.empty((n_blocks,), dtype=np.float32) final_ls = np.empty((n_blocks, NS), dtype=np.uint8) final_lm = np.empty((n_blocks, NS), dtype=np.uint8) final_codes = np.empty((n_blocks, QK_K), dtype=np.uint8) factor_floor_used = 0 t_p4 = time.time() for start in range(0, n_blocks, batch_blocks): _check_stop() stop = min(start + batch_blocks, n_blocks) B_c = stop - start X_t = torch.from_numpy(X_np[start:stop]).to(dev) if W_np is None: W_t = torch.ones_like(X_t) else: W_raw_t = torch.from_numpy(W_np[start:stop]).to(dev) if cfg.temper_imatrix: sig2_t = torch.mean(X_t * X_t, dim=1, keepdim=True) W_t = torch.sqrt(torch.clamp(W_raw_t, min=1.0e-6)) * torch.sqrt(sig2_t + X_t * X_t) wm_t = torch.mean(W_t, dim=1, keepdim=True) W_t = torch.where(wm_t > 1.0e-30, W_t / torch.clamp(wm_t, min=1.0e-30), torch.ones_like(W_t)) else: W_t = W_raw_t Xb = X_t.reshape(B_c, NS, 16) Wb = W_t.reshape(B_c, NS, 16) Sw = torch.sum(Wb, dim=2) Sx = torch.sum(Wb * Xb, dim=2) safe_Sw = torch.clamp(Sw, min=1.0e-30) carry_t = torch.from_numpy(carry[start:stop]).to(dev) b_slice = best[start:stop] di = b_slice // N_CAND_M mi = b_slice % N_CAND_M d_v = torch.from_numpy(base_d[start:stop] * MULTS_D[di]).to(dev) m_v = torch.from_numpy(base_m[start:stop] * MULTS_M[mi]).to(dev) d_v = _round_fp16_t(d_v) m_v = _round_fp16_t(m_v) sc_s = torch.from_numpy(scales_seed[start:stop]).to(dev) mn_s = torch.from_numpy(mins_seed[start:stop]).to(dev) sd_v = torch.where(d_v[:, None] > 1.0e-15, d_v[:, None], torch.ones_like(d_v[:, None])) sm_v = torch.where(m_v[:, None] > 1.0e-15, m_v[:, None], torch.ones_like(m_v[:, None])) ls_v = torch.where( d_v[:, None] > 1.0e-15, _nearest_int_t(sc_s / sd_v).clamp(0, 15), torch.zeros((B_c, NS), dtype=torch.int64, device=dev), ).to(torch.uint8) lm_v = torch.where( m_v[:, None] > 1.0e-15, _nearest_int_t(mn_s / sm_v).clamp(0, 15), torch.zeros((B_c, NS), dtype=torch.int64, device=dev), ).to(torch.uint8) codes_v = _codes_from_params_t(X_t, d_v, m_v, ls_v, lm_v) d1, m1, ls1, lm1, c1, obj1 = _refine_trajectory_cuda( X_t, W_t, Xb, Wb, Sw, Sx, safe_Sw, d_v, m_v, ls_v, lm_v, codes_v, carry_t, cfg ) if cfg.use_factor_floor: d_f = torch.from_numpy(solo_d[start:stop]).to(dev) m_f = torch.from_numpy(solo_m[start:stop]).to(dev) ls_f = torch.from_numpy(solo_ls[start:stop]).to(dev) lm_f = torch.from_numpy(solo_lm[start:stop]).to(dev) c_f = torch.from_numpy(solo_codes[start:stop]).to(dev) d2, m2, ls2, lm2, c2, obj2 = _refine_trajectory_cuda( X_t, W_t, Xb, Wb, Sw, Sx, safe_Sw, d_f, m_f, ls_f, lm_f, c_f, carry_t, cfg ) take2 = obj2 < obj1 - 1.0e-12 factor_floor_used += int(torch.sum(take2).item()) d1 = torch.where(take2, d2, d1) m1 = torch.where(take2, m2, m1) ls1 = torch.where(take2[:, None], ls2, ls1) lm1 = torch.where(take2[:, None], lm2, lm1) c1 = torch.where(take2[:, None], c2, c1) final_d[start:stop] = d1.cpu().numpy() final_m[start:stop] = m1.cpu().numpy() final_ls[start:stop] = ls1.cpu().numpy() final_lm[start:stop] = lm1.cpu().numpy() final_codes[start:stop] = c1.cpu().numpy() elapsed = max(time.time() - t_p4, 1.0e-6) rate = stop / elapsed eta = (n_blocks - stop) / max(rate, 1.0e-9) _emit( f"CUDA Phase 4 ({gpu_name}): Algebraic Coupled Bilinear Assembly", stop, n_blocks, f"blk {stop}/{n_blocks} ({rate:.1f} blk/s, ETA {eta:.1f}s) floors(fac={factor_floor_used})", rate=rate, eta=eta, sigma=sigma, temperature=temperature, cand_floors=0, fac_floors=factor_floor_used, ) if cfg.apply_dc_carry and (cfg.dc_lambda > 0.0 or cfg.dc_decay > 0.0): _check_stop() final_codes = _apply_dc_nudge_cuda( X_np, W_np, final_d, final_m, final_ls, final_lm, final_codes, cfg, row_width, dev ) if cfg.pin_dc: _check_stop() t_pin = time.time() _emit( f"Phase 4.7: Zero-DC Pin ({gpu_name})", 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, ) final_d, final_m = pin_q2k_dc_residual( final_d.copy(), final_m.copy(), final_ls, final_lm, final_codes, X_np.astype(np.float64) ) _emit( f"Phase 4.7: Zero-DC Pin ({gpu_name})", n_blocks, n_blocks, f"done in {time.time() - t_pin:.1f}s", sigma=sigma, temperature=temperature, ) scales = (final_ls & 0x0F) | ((final_lm & 0x0F) << np.uint8(4)) qs = np.empty((n_blocks, QK_K // 4), dtype=np.uint8) for base in (0, 128): qs[:, base // 4 : base // 4 + 32] = ( final_codes[:, base : base + 32] | (final_codes[:, base + 32 : base + 64] << np.uint8(2)) | (final_codes[:, base + 64 : base + 96] << np.uint8(4)) | (final_codes[:, base + 96 : base + 128] << np.uint8(6)) ) dm = np.empty((n_blocks, 2), dtype=np.uint16) dm[:, 0] = final_d.astype(np.float16).view(np.uint16) dm[:, 1] = final_m.astype(np.float16).view(np.uint16) d_f32 = dm[:, 0].view(np.float16).astype(np.float32) m_f32 = dm[:, 1].view(np.float16).astype(np.float32) d_sub = np.repeat(d_f32[:, None] * final_ls.astype(np.float32), 16, axis=1) m_sub = np.repeat(m_f32[:, None] * final_lm.astype(np.float32), 16, axis=1) xh_all = d_sub * final_codes.astype(np.float32) - m_sub err_mat = X_np.astype(np.float64) - xh_all.astype(np.float64) total_sse = float(np.sum(err_mat * err_mat)) weighted = total_sse if W_np is None else float(np.sum(W_np.astype(np.float64) * err_mat * err_mat)) rmse = float(np.sqrt(total_sse / x.size)) blocks = [ Q2KBlock(scales=scales[i], qs=qs[i], d=np.uint16(dm[i, 0]), dmin=np.uint16(dm[i, 1])) for i in range(n_blocks) ] 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=0, ) _emit( f"Phase 5: Tensor Complete ({gpu_name})", 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} floors(fac={factor_floor_used})", rmse=rmse, sigma=sigma, temperature=temperature, cand_floors=0, fac_floors=factor_floor_used, ) return QuantizedTensor( blocks=blocks, scales=scales, qs=qs, dm=dm, best_candidate=best, factor_objectives=solo_obj, config=cfg, stats=stats, )