Wire CUDA pinning

q3_k.py

@@ -40,6 +40,7 @@ The final ``block_q3_K`` payload is emitted in the exact GGML layout
 
 from __future__ import annotations
 
+import logging
 from math import ceil
 import numpy as np
 
@@ -50,7 +51,7 @@ 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 _from_fp16_bits_t, cuda_is_available
+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
@@ -68,6 +69,24 @@ Q3_CODE_MAX = 3
 Q3_N_CAND = 24
 Q3_N_STATES = 6
 
+_log = logging.getLogger(__name__)
+
+# Chunk sizes (in 256-weight super-blocks).  The CUDA default matches the
+# Q4_0 reference chunk (131072 blocks * 32 weights = 4,194,304 weights per
+# chunk): a Q3_K block is 8x wider, so 131072 / 8 = 16384 blocks keeps the same
+# per-chunk GPU footprint while amortising the per-chunk CPU work (two
+# Sieve/Viterbi passes, ~512 Python-loop steps each) and kernel-launch overhead.
+# The CPU default stays small (numpy temporaries stay cache-resident).
+Q3_K_CHUNK_BLOCKS_CUDA = 16384
+Q3_K_CHUNK_BLOCKS_CPU = 1024
+
+
+def _to_dev(x, dev: "torch.device") -> "torch.Tensor":
+    """float32 tensor on ``dev``; no copy if ``x`` is already a device tensor."""
+    if isinstance(x, torch.Tensor):
+        return x.to(dev, dtype=torch.float32)
+    return torch.from_numpy(np.ascontiguousarray(x)).to(dev, dtype=torch.float32)
+
 
 def _nearest_int_t(value):
     """GGML round-to-nearest-even for torch tensors."""
@@ -349,17 +368,21 @@ def _subblock_scales_hpc(
 
 
 def _subblock_scales_hpc_cuda(
-    X_np: np.ndarray, W_np: np.ndarray, dev: "torch.device", cfg: HPCConfig
-) -> np.ndarray:
-    """Run the HPC 3-bit sub-block pipeline on CUDA and return signed scales."""
-    M = X_np.shape[0]
+    X_src, W_src, dev: "torch.device", cfg: HPCConfig, *, as_tensor: bool = False
+):
+    """Run the HPC 3-bit sub-block pipeline on CUDA and return signed scales.
+
+    ``X_src`` / ``W_src`` may be numpy arrays or device tensors.  With
+    ``as_tensor=True`` the scales stay on ``dev`` (no D2H copy).
+    """
+    M = X_src.shape[0]
     lam_dc = (float(cfg.dc_lambda) / Q3_K_SUBBLOCK) if cfg.use_spectral else 0.0
     lam_vw = (float(cfg.vw_lambda) / Q3_K_SUBBLOCK) if cfg.use_spectral else 0.0
     half = Q3_K_SUBBLOCK // 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)
+        X = _to_dev(X_src, dev)
+        W = _to_dev(W_src, dev)
 
         amax = torch.max(torch.abs(X), dim=1).values
         d_range = amax / 3.5
@@ -464,6 +487,8 @@ def _subblock_scales_hpc_cuda(
             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 as_tensor:
+            return best_d
         return best_d.cpu().numpy().astype(np.float32)
 
 
@@ -658,23 +683,31 @@ def _fit_q3_k_blocks(
 
 
 def _fit_q3_k_blocks_cuda(
-    X_np: np.ndarray,
-    W_np: np.ndarray,
-    sub_scales: np.ndarray,
+    X_src,
+    W_src,
+    sub_scales,
     dev: "torch.device",
     cfg: HPCConfig,
-) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
-    """Full HPC V2 super-block fit for (B, 256) blocks on CUDA."""
-    B = X_np.shape[0]
+    *,
+    as_tensor: bool = False,
+):
+    """Full HPC V2 super-block fit for (B, 256) blocks on CUDA.
+
+    Inputs may be numpy arrays or device tensors.  Default return is the numpy
+    triple ``(d16 uint16, sc int16, code int16)``.  With ``as_tensor=True`` the
+    results stay on ``dev`` as ``(d16, d, sc, code)`` where ``d`` is the fp16-
+    decoded float32 super-scale, ``sc`` is (B, 16) and ``code`` is (B, 16, 16).
+    """
+    B = X_src.shape[0]
     lam_dc = (float(cfg.dc_lambda) / QK3_K) if cfg.use_spectral else 0.0
     lam_vw = (float(cfg.vw_lambda) / QK3_K) if cfg.use_spectral else 0.0
 
     with torch.no_grad():
-        X3 = torch.from_numpy(X_np).to(dev, dtype=torch.float32).reshape(B, Q3_K_N_SUBBLOCKS, Q3_K_SUBBLOCK)
-        W3 = torch.from_numpy(W_np).to(dev, dtype=torch.float32).reshape(B, Q3_K_N_SUBBLOCKS, Q3_K_SUBBLOCK)
+        X3 = _to_dev(X_src, dev).reshape(B, Q3_K_N_SUBBLOCKS, Q3_K_SUBBLOCK)
+        W3 = _to_dev(W_src, dev).reshape(B, Q3_K_N_SUBBLOCKS, Q3_K_SUBBLOCK)
         Xf = X3.reshape(B, QK3_K)
         Wf = W3.reshape(B, QK3_K)
-        s = torch.from_numpy(np.ascontiguousarray(sub_scales.reshape(B, Q3_K_N_SUBBLOCKS))).to(dev, dtype=torch.float32)
+        s = _to_dev(sub_scales, dev).reshape(B, Q3_K_N_SUBBLOCKS)
         rows = torch.arange(B, device=dev)
         Wj = W3.sum(dim=2)  # (B, 16) per-sub-block importance
 
@@ -827,6 +860,8 @@ def _fit_q3_k_blocks_cuda(
             # thermal_refine_q3_k_cuda already returns source-signed codes [-4, 3]
             code = thermal_c.reshape(B, Q3_K_N_SUBBLOCKS, Q3_K_SUBBLOCK).to(torch.int16)
 
+        if as_tensor:
+            return best_d16, best_d, sc, code.reshape(B, Q3_K_N_SUBBLOCKS, Q3_K_SUBBLOCK)
         return (
             best_d16.cpu().numpy().astype(np.uint16),
             sc.cpu().numpy().astype(np.int16),
@@ -837,36 +872,53 @@ def _fit_q3_k_blocks_cuda(
 def _quantize_q3_k_hpc_cuda(
     X: np.ndarray, W: np.ndarray, dev: "torch.device", cfg: HPCConfig
 ) -> tuple[bytes, float]:
-    """Full HPC V2 Q3_K pipeline on CUDA for a chunk of (B, 256) blocks."""
-    B = X.shape[0]
-    X_f32 = X.astype(np.float32)
-    W_f32 = W.astype(np.float32)
+    """Full HPC V2 Q3_K pipeline on CUDA for a chunk of (B, 256) blocks.
 
-    sub_scales = _subblock_scales_hpc_cuda(
-        X_f32.reshape(B * Q3_K_N_SUBBLOCKS, Q3_K_SUBBLOCK),
-        W_f32.reshape(B * Q3_K_N_SUBBLOCKS, Q3_K_SUBBLOCK),
-        dev,
-        cfg,
-    )
-    d16, sc, code = _fit_q3_k_blocks_cuda(X_f32, W_f32, sub_scales, dev, cfg)
+    X/W are uploaded once.  Level A scales, Level B fit, the zero-DC pin and the
+    raw-SSE reduction all stay on the device; the only D2H copies are the
+    Sieve/Viterbi candidate-error matrices (sequential DP, runs in numpy) and
+    the final (d16, sc, code) needed for byte packing.
+    """
+    B = X.shape[0]
 
-    # 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:
-        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),
+    with torch.no_grad():
+        X_t = torch.from_numpy(np.ascontiguousarray(X, dtype=np.float32)).to(dev)
+        W_t = torch.from_numpy(np.ascontiguousarray(W, dtype=np.float32)).to(dev)
+
+        sub_scales_t = _subblock_scales_hpc_cuda(
+            X_t.reshape(B * Q3_K_N_SUBBLOCKS, Q3_K_SUBBLOCK),
+            W_t.reshape(B * Q3_K_N_SUBBLOCKS, Q3_K_SUBBLOCK),
+            dev,
+            cfg,
+            as_tensor=True,
+        )
+        d16_t, d_t, sc_t, code_t = _fit_q3_k_blocks_cuda(
+            X_t, W_t, sub_scales_t, dev, cfg, as_tensor=True
         )
-        code = codes_gpu.cpu().numpy().reshape(
-            B, Q3_K_N_SUBBLOCKS, Q3_K_SUBBLOCK
-        ).astype(np.int16)
+        del sub_scales_t, W_t
+
+        # Zero-DC Pin (deterministic code-flip exchange within pin_budget), on
+        # the device.  ``d_t`` is already the fp16-decoded super-scale, so no
+        # uint16 numpy -> torch round-trip is needed (torch's uint16 support is
+        # version-dependent and fails on older releases).
+        if cfg.pin_dc:
+            codes_pin, _ = pin_dc_codes_uniform_cuda(
+                X_t,
+                d_t,
+                code_t.reshape(B, QK3_K).to(torch.int64),
+                sc_t.to(torch.float32).repeat_interleave(Q3_K_SUBBLOCK, dim=1),
+                Q3_CODE_MIN,
+                Q3_CODE_MAX,
+                budget_frac=float(cfg.pin_budget),
+            )
+            code_t = codes_pin.reshape(B, Q3_K_N_SUBBLOCKS, Q3_K_SUBBLOCK)
+
+        deq = d_t[:, None, None] * sc_t[:, :, None].to(torch.float32) * code_t.to(torch.float32)
+        raw_sse = float(torch.sum((X_t - deq.reshape(B, QK3_K)) ** 2, dtype=torch.float64).item())
+
+        d16 = d16_t.cpu().numpy().astype(np.uint16)
+        sc = sc_t.cpu().numpy().astype(np.int16)
+        code = code_t.cpu().numpy().astype(np.int16)
 
     hmask, qs = _pack_q3_k_codes(code)
     scales = _pack_q3_k_scales(sc)
@@ -875,10 +927,6 @@ def _quantize_q3_k_hpc_cuda(
     out[:, QK3_K // 8 : QK3_K // 8 + QK3_K // 4] = qs
     out[:, QK3_K // 8 + QK3_K // 4 : Q3_K_BLOCK_BYTES - 2] = scales
     out[:, Q3_K_BLOCK_BYTES - 2 :] = d16.astype("<u2").view(np.uint8).reshape(B, 2)
-
-    d = from_fp16_bits(d16).astype(np.float32)
-    deq = (d[:, None, None] * sc[:, :, None].astype(np.float32) * code.astype(np.float32)).reshape(B, QK3_K)
-    raw_sse = float(np.sum((X_f32 - deq) ** 2, dtype=np.float64))
     return out.tobytes(), raw_sse
 
 
@@ -928,9 +976,14 @@ def quantize_tensor_q3_k(
     weights: np.ndarray,
     importance: np.ndarray | None = None,
     cfg: HPCConfig | None = None,
-    chunk_blocks: int = 1024,
+    chunk_blocks: int | None = None,
 ) -> tuple[bytes, int, float]:
-    """Quantize a tensor to GGML Q3_K using the full HPC V2 pipeline."""
+    """Quantize a tensor to GGML Q3_K using the full HPC V2 pipeline.
+
+    ``chunk_blocks`` (256-weight super-blocks per chunk) defaults to
+    ``Q3_K_CHUNK_BLOCKS_CUDA`` on the GPU path (same weights-per-chunk as the
+    Q4_0 reference) and ``Q3_K_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 % QK3_K != 0:
@@ -946,10 +999,19 @@ def quantize_tensor_q3_k(
     else:
         W_raw = None
 
-    use_cuda = cfg.device != "cpu" and cuda_is_available() and (
+    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:
+        chunk_blocks = Q3_K_CHUNK_BLOCKS_CUDA if use_cuda else Q3_K_CHUNK_BLOCKS_CPU
+    chunk_blocks = max(1, int(chunk_blocks))
+    _log.info(
+        "Q3_K 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