Upload 6 files

gui.py

@@ -30,6 +30,22 @@ from tkinter import filedialog, messagebox, ttk
 
 import numpy as np
 
+_HERE = Path(__file__).resolve().parent
+if not (_HERE / "hpc_quantize_v2").is_dir() and (_HERE / "__init__.py").is_file():
+    import importlib.util
+    import sys
+
+    if "hpc_quantize_v2" not in sys.modules:
+        _spec = importlib.util.spec_from_file_location(
+            "hpc_quantize_v2",
+            _HERE / "__init__.py",
+            submodule_search_locations=[str(_HERE)],
+        )
+        if _spec is not None and _spec.loader is not None:
+            _mod = importlib.util.module_from_spec(_spec)
+            sys.modules["hpc_quantize_v2"] = _mod
+            _spec.loader.exec_module(_mod)
+
 from hpc_quantize_v2 import (
     GGUFReader,
     HPCConfig,

q1_0.py

@@ -201,37 +201,48 @@ def _quantize_q1_0_hpc_cpu(
     rows = np.arange(B)
     d_sel = cand_d[rows, best_cand]
 
-    # 4. Asymmetric Sign-Flip fp16-ULP Polish (+/-8 ULPs around +d_sel, -d_sel, and d_mean):
+    # 4. Asymmetric Sign-Flip fp16-ULP Polish (+/-8 ULPs around +d_sel, -d_sel, wls_d, and d_mean):
     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)
+    wls_b = np.clip(to_fp16_bits(np.abs(wls_d)).astype(np.int32)[:, None] + deltas[None, :], 0, 0x7BFF).astype(np.uint16)
+    wls_neg_b = (wls_b | np.uint16(0x8000)).astype(np.uint16)
     ref_b = np.clip(to_fp16_bits(np.abs(d_mean)).astype(np.int32)[:, None] + deltas[None, :], 0, 0x7BFF).astype(np.uint16)
-    all_b = np.concatenate([pos_b, neg_b, ref_b], axis=1)
+    all_b = np.concatenate([pos_b, neg_b, wls_b, wls_neg_b, ref_b], axis=1)
     all_d = from_fp16_bits(all_b)
     q_u = np.where(all_d[:, :, None] >= 0.0, X_f32[:, None, :] >= 0.0, X_f32[:, None, :] < 0.0).astype(np.uint8)
     c_u = 2.0 * q_u.astype(np.float32) - 1.0
     e_u = X_f32[:, None, :] - all_d[:, :, None] * c_u
     sse_u = np.sum(W_f32[:, None, :] * e_u * e_u, axis=2, dtype=np.float64)
-    pick_u = np.argmin(sse_u, axis=1)
+    if cfg.use_spectral and (lam_dc > 0.0 or lam_vw > 0.0):
+        v_u = e_u[:, :, :half] + e_u[:, :, half:]
+        dc_u = np.sum(e_u, axis=2, dtype=np.float64)
+        sse_min_u = np.min(sse_u, axis=1, keepdims=True)
+        obj_u = sse_u + lam_vw * np.sum(v_u * v_u, axis=2, dtype=np.float64) + lam_dc * dc_u * dc_u
+        obj_u = np.where(sse_u <= sse_min_u * (1.0 + cfg.sse_budget), obj_u, np.inf)
+        pick_u = np.argmin(obj_u, axis=1)
+    else:
+        pick_u = np.argmin(sse_u, axis=1)
     best_d16 = all_b[rows, pick_u]
     best_d = all_d[rows, pick_u]
     best_q = q_u[rows, pick_u].astype(np.int64)
 
-    # 5. Closed-Form 1-D Quadratic Vesica + DC Bit-Flip Projection within sse_budget:
+    # 5. Closed-Form 1-D Quadratic Vesica + DC Bit-Flip Projection (step = 2*d) within sse_budget:
     if cfg.use_spectral and (lam_vw > 0.0 or lam_dc > 0.0):
         c_f = (2 * best_q - 1).astype(np.float64)
         d_f64 = best_d.astype(np.float64)[:, None]
+        step_f64 = 2.0 * d_f64
         e = X_f32.astype(np.float64) - d_f64 * c_f
         sse_live = np.sum(W_f32 * e * e, axis=1, dtype=np.float64)
         sse_cap = sse_live * (1.0 + cfg.sse_budget)
         v = e[:, :half] + e[:, half:]
         dc = np.sum(e, axis=1)
-        denom = 2.0 * (W_f32 + lam_vw + lam_dc) * (d_f64 * d_f64)
+        denom = (W_f32 + lam_vw + lam_dc) * (step_f64 * step_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])
+        grad = step_f64 * (W_f32 * e + lam_vw * v_full + lam_dc * dc[:, None])
         dq = np.where(act_b, nearest_int((grad / s_den).astype(np.float32)), 0)
         q_new = np.clip(best_q + dq, 0, 1)
         c_new = (2 * q_new - 1).astype(np.float64)
@@ -317,38 +328,50 @@ def _quantize_q1_0_hpc_cuda(
         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, -d_sel, and d_mean):
+        # 4. Asymmetric Sign-Flip fp16-ULP Polish (+/-8 ULPs around +d_sel, -d_sel, wls_d, and d_mean):
         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
+        wls_bits = torch.abs(wls_d).to(torch.float16).view(torch.int16).to(torch.int32) & 0x7FFF
         ref_bits = torch.abs(d_mean).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
+        wls_b = torch.clamp(wls_bits[:, None] + deltas[None, :], 0, 0x7BFF)
+        wls_neg_b = wls_b | 0x8000
         mean_b = torch.clamp(ref_bits[:, None] + deltas[None, :], 0, 0x7BFF)
-        all_b = torch.cat([pos_b, neg_b, mean_b], dim=1)
+        all_b = torch.cat([pos_b, neg_b, wls_b, wls_neg_b, mean_b], dim=1)
         all_d = all_b.to(torch.int16).view(torch.float16).to(torch.float32)
         q_u = torch.where(all_d[:, :, None] >= 0.0, X[:, None, :] >= 0.0, X[:, None, :] < 0.0).to(torch.uint8)
         c_u = 2.0 * q_u.to(torch.float32) - 1.0
         e_u = X[:, None, :] - all_d[:, :, None] * c_u
         sse_u = torch.sum(W[:, None, :] * e_u * e_u, dim=2)
-        pick_u = torch.argmin(sse_u, dim=1)
+        if cfg.use_spectral and (lam_dc > 0.0 or lam_vw > 0.0):
+            v_u = e_u[:, :, :half] + e_u[:, :, half:]
+            dc_u = torch.sum(e_u, dim=2)
+            sse_min_u = torch.min(sse_u, dim=1, keepdim=True).values
+            obj_u = sse_u + lam_vw * torch.sum(v_u * v_u, dim=2) + lam_dc * dc_u * dc_u
+            obj_u = torch.where(sse_u <= sse_min_u * (1.0 + float(cfg.sse_budget)), obj_u, torch.full_like(obj_u, float("inf")))
+            pick_u = torch.argmin(obj_u, dim=1)
+        else:
+            pick_u = torch.argmin(sse_u, dim=1)
         best_d16 = all_b[rows, pick_u]
         best_d = all_d[rows, pick_u]
         best_q = q_u[rows, pick_u].to(torch.int64)
 
-        # 5. Closed-Form 1-D Quadratic Vesica + DC Bit-Flip Projection within sse_budget:
+        # 5. Closed-Form 1-D Quadratic Vesica + DC Bit-Flip Projection (step = 2*d) within sse_budget:
         if cfg.use_spectral and (lam_vw > 0.0 or lam_dc > 0.0):
             d_col = best_d[:, None]
+            step_col = 2.0 * d_col
             e = X - d_col * (2.0 * best_q.to(torch.float32) - 1.0)
             sse_live = torch.sum(W * e * e, dim=1)
             sse_cap = sse_live * (1.0 + float(cfg.sse_budget))
             v = e[:, :half] + e[:, half:]
             dc = torch.sum(e, dim=1)
-            denom = 2.0 * (W + lam_vw + lam_dc) * (d_col * d_col)
+            denom = (W + lam_vw + lam_dc) * (step_col * step_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])
+            grad = step_col * (W * e + lam_vw * v_full + lam_dc * dc[:, None])
             sh_dq = (grad / s_den) + 12582912.0
             dq = torch.where(
                 act_b,
@@ -409,14 +432,12 @@ def quantize_tensor_q1_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)
+            Wb = np.ones_like(Xb, dtype=np.float32)
         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)
+            wr = np.maximum(W_raw[start:stop], 1.0e-6)
+            wm = np.mean(wr, axis=1, keepdims=True)
+            Wb = np.where(wm > 1.0e-30, wr / np.maximum(wm, 1.0e-30), 1.0).astype(np.float32)
 
         if use_cuda and dev is not None:
             raw_b, sse_b = _quantize_q1_0_hpc_cuda(Xb, Wb, dev, cfg)