Upload 3 files

q3_k.py

@@ -50,11 +50,11 @@ except ImportError:  # pragma: no cover
 
 from .anneal import thermal_refine_q3_k, thermal_refine_q3_k_cuda
 from .config import HPCConfig
-from .cuda_backend import cuda_is_available
+from .cuda_backend import _from_fp16_bits_t, 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
+from .spectral import pin_dc_codes_uniform, pin_dc_codes_uniform_cuda
 from .viterbi import robust_temperature
 
 QK3_K = 256
@@ -851,17 +851,22 @@ def _quantize_q3_k_hpc_cuda(
     d16, sc, code = _fit_q3_k_blocks_cuda(X_f32, W_f32, sub_scales, dev, cfg)
 
     # Zero-DC Pin (deterministic code-flip exchange within pin_budget):
+    # Runs fully on the GPU — d from the fp16 bits, unit-grid = sc repeated.
     if cfg.pin_dc:
-        code_pin, _ = pin_dc_codes_uniform(
-            X_f32,
-            from_fp16_bits(d16),
-            code.reshape(B, QK3_K).astype(np.int64),
-            np.repeat(sc, Q3_K_SUBBLOCK, axis=1),
+        d_pin = _from_fp16_bits_t(torch.from_numpy(d16).to(dev))
+        sc_gpu = torch.from_numpy(sc.astype(np.float32)).to(dev)
+        codes_gpu, _ = pin_dc_codes_uniform_cuda(
+            torch.from_numpy(X_f32).to(dev),
+            d_pin,
+            torch.from_numpy(code.reshape(B, QK3_K).astype(np.int64)).to(dev),
+            sc_gpu.repeat_interleave(Q3_K_SUBBLOCK, dim=1),
             Q3_CODE_MIN,
             Q3_CODE_MAX,
             budget_frac=float(cfg.pin_budget),
         )
-        code = code_pin.reshape(B, Q3_K_N_SUBBLOCKS, Q3_K_SUBBLOCK).astype(np.int16)
+        code = codes_gpu.cpu().numpy().reshape(
+            B, Q3_K_N_SUBBLOCKS, Q3_K_SUBBLOCK
+        ).astype(np.int16)
 
     hmask, qs = _pack_q3_k_codes(code)
     scales = _pack_q3_k_scales(sc)

q4_0.py

@@ -41,7 +41,7 @@ from .cuda_backend import cuda_is_available
 from .factored import CAND_TO_QUHIT, 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
+from .spectral import pin_dc_codes_uniform, pin_dc_codes_uniform_cuda
 from .viterbi import robust_temperature
 
 QK4_0 = 32
@@ -416,19 +416,18 @@ def _quantize_q4_0_hpc_cuda(
 
         # 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.
+        # pin's model residual equals the true dequant error.  Runs fully on
+        # the GPU.
         if cfg.pin_dc:
-            d_np = best_d.cpu().numpy()
-            q_np = pin_dc_codes_uniform(
-                X_np + 8.0 * d_np[:, None].astype(np.float32),
-                d_np,
-                best_q.cpu().numpy().astype(np.int64),
-                np.ones((X_np.shape[0], QK4_0)),
+            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),
-            )[0]
-            best_q = torch.from_numpy(q_np).to(dev)
+            )
 
         best_q_u8 = best_q.to(torch.uint8)
         deq_fin = best_d[:, None] * (best_q_u8.to(torch.float32) - 8.0)

spectral.py

@@ -315,3 +315,85 @@ def pin_dc_codes_uniform(
         touched[bi] = True
 
     return codes, int(np.sum(touched))
+
+
+def pin_dc_codes_uniform_cuda(
+    x: "torch.Tensor",
+    d: "torch.Tensor",
+    codes: "torch.Tensor",
+    unit: "torch.Tensor",
+    lo: int,
+    hi: int,
+    budget_frac: float = 0.05,
+    max_rounds: int = 96,
+) -> tuple["torch.Tensor", int]:
+    """GPU port of :func:`pin_dc_codes_uniform` (same greedy exchange).
+
+    ``x``/``d``/``codes``/``unit`` all live on one CUDA device; ``codes`` is
+    int64 and is updated in place (returned tensor is the same object).
+    Float32 throughout — matches the CUDA pipeline's precision, and the
+    exchange is self-correcting (any round that slightly mis-ranks a flip
+    simply loses one candidate, the residual still never regresses SSE
+    beyond ``budget_frac`` of the block's starting SSE).
+    """
+    import torch
+
+    if budget_frac <= 0.0 or codes.shape[0] == 0:
+        return codes, 0
+    X = x.float()
+    B, n = X.shape
+    dev = X.device
+    if codes.dtype != torch.int64:
+        codes = codes.to(torch.int64)
+    d = d.to(dev).float()
+    if unit.dim() == 1 and unit.numel() == n:
+        unit = unit[None, :]
+    unit = unit.to(dev).float()
+    if unit.shape[0] == 1 and B != 1:
+        unit = unit.expand(B, n)
+
+    d_col = d[:, None]
+    e = X - d_col * codes.float() * unit
+    dc = e.sum(dim=1)
+    sse0 = (e * e).sum(dim=1)
+    cap = sse0 * float(budget_frac)
+    spent = torch.zeros(B, device=dev)
+    rows = torch.arange(B, device=dev)
+    tol = d.abs() * 0.5
+    touched = 0
+
+    for _ in range(max_rounds):
+        want_cost = torch.full((B,), float("inf"), device=dev)
+        want_pos = torch.full((B,), -1, dtype=torch.int64, device=dev)
+        want_dir = torch.zeros(B, dtype=torch.int64, device=dev)
+        for direction in (1, -1):
+            sel = (codes < hi) if direction == 1 else (codes > lo)
+            de = sel.float() * (float(direction) * unit * -d_col)
+            dE = 2.0 * e * de + de * de
+            valid = sel & (dc[:, None] * de < 0.0)
+            ratio = torch.where(valid, dE.abs() / de.abs().clamp_min(1.0e-30), torch.inf)
+            pmin = ratio.argmin(dim=1)
+            rmin = ratio[rows, pmin]
+            cmin = dE[rows, pmin]
+            take = torch.isfinite(rmin) & (rmin < want_cost)
+            want_cost = torch.where(take, cmin, want_cost)
+            want_pos = torch.where(take, pmin, want_pos)
+            want_dir = torch.where(take, torch.full_like(want_dir, direction), want_dir)
+        want = (
+            (dc.abs() > tol)
+            & (want_pos >= 0)
+            & (spent + want_cost <= cap)
+        )
+        if not bool(want.any()):
+            break
+        bi = want.nonzero(as_tuple=False).squeeze(1)
+        pos = want_pos[bi]
+        dirv = want_dir[bi]
+        de = -d[bi] * dirv.float() * unit[bi, pos]
+        spent[bi] += 2.0 * e[bi, pos] * de + de * de
+        e[bi, pos] += de
+        codes[bi, pos] += dirv
+        dc[bi] += de
+        touched += int(bi.numel())
+
+    return codes, touched