Upload q4_0.py

q4_0.py

@@ -27,6 +27,7 @@ per 18-byte block = 4.5 bpw):
 
 from __future__ import annotations
 
+import logging
 from math import ceil
 import numpy as np
 
@@ -38,7 +39,7 @@ except ImportError:  # pragma: no cover
 from .anneal import thermal_refine_q4_0, thermal_spectral_lams
 from .config import HPCConfig
 from .cuda_backend import cuda_is_available
-from .factored import CAND_TO_QUHIT, MULTS_D
+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
@@ -50,6 +51,16 @@ Q4_N_CAND = 24
 Q4_N_STATES = 6
 Q4_N_LEVELS = 16
 
+_log = logging.getLogger(__name__)
+
+# Chunk sizes (in 32-weight blocks).  The CUDA default keeps the historical
+# reference chunk (131072 blocks * 32 weights = 4,194,304 weights per chunk);
+# the CPU default is scaled back so the vectorized numpy temporaries
+# ((B, 24, 32) candidate-error fields) stay cache-resident instead of
+# allocating multi-hundred-MB arrays per chunk.
+Q4_0_CHUNK_BLOCKS_CUDA = 131072
+Q4_0_CHUNK_BLOCKS_CPU = 8192
+
 
 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)
@@ -354,6 +365,9 @@ def _quantize_q4_0_hpc_cuda(
 
         # 3. 6-State Sieve + 6-State Viterbi DP:
         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]
@@ -374,6 +388,7 @@ def _quantize_q4_0_hpc_cuda(
         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)
+        del e_u, q_u, qi_u, sh_u, sse_u
         best_d16 = all_b[rows, pick_u]
         best_d = all_d[rows, pick_u]
         best_q = q_u[rows, pick_u].to(torch.int64)
@@ -565,9 +580,14 @@ def quantize_tensor_q4_0(
     weights: np.ndarray,
     importance: np.ndarray | None = None,
     cfg: HPCConfig | None = None,
-    chunk_blocks: int = 131072,
+    chunk_blocks: int | None = None,
 ) -> tuple[bytes, int, float]:
-    """Quantize a tensor to GGML Q4_0 using the full HPC V2 pipeline."""
+    """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:
@@ -583,10 +603,19 @@ def quantize_tensor_q4_0(
     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 = Q4_0_CHUNK_BLOCKS_CUDA if use_cuda else 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