Upload 5 files

__main__.py

@@ -57,6 +57,29 @@ def _config_from_args(args) -> HPCConfig:
         cfg.quant_format = args.format
         if args.format == "q2_k":
             cfg.enable_q4_fallback = False
+    if getattr(args, "tensor_type", None):
+        custom: dict[str, str] = {}
+        for item in args.tensor_type:
+            if "=" not in item:
+                raise SystemExit(f"--tensor-type expects TYPE=QUANT, got {item!r}")
+            key, value = (part.strip() for part in item.split("=", 1))
+            if not key:
+                raise SystemExit(f"--tensor-type has an empty tensor type: {item!r}")
+            canonical = {
+                "q1_0": "Q1_0",
+                "q2_k": "Q2_K",
+                "q4_0": "Q4_0",
+                "copy": "Copy",
+                "auto": "Auto",
+            }
+            norm = canonical.get(value.lower())
+            if norm is None:
+                raise SystemExit(
+                    f"--tensor-type {key!r} has invalid quant {value!r}; "
+                    "use Q1_0, Q2_K, Q4_0, Copy, or Auto"
+                )
+            custom[key] = norm
+        cfg.custom_quant_map = custom
     return cfg.validate()
 
 
@@ -101,6 +124,7 @@ def _run_gguf(args, cfg: HPCConfig) -> None:
                 promote_attn_q4=cfg.promote_attn_q4,
                 quant_format=cfg.quant_format,
                 source_type=ti.type,
+                custom_quant_map=cfg.custom_quant_map or None,
             )
             tag = tgt if tgt else "COPY"
             print(
@@ -254,9 +278,17 @@ def main() -> None:
     parser.add_argument("--no-spectral", action="store_true", help="disable spectral/DC objective")
     parser.add_argument(
         "--format",
-        choices=("mixed", "mixed_q1", "q1_0", "q4_0", "q2_k"),
+        choices=("mixed", "mixed_q1", "custom", "q1_0", "q4_0", "q2_k"),
         default=None,
-        help="quantization format mode: mixed (Q2_K+Q4_0), mixed_q1 (Q1_0+Q4_0), q1_0 (pure Q1_0), q4_0 (pure Q4_0), or q2_k (pure Q2_K)",
+        help="quantization format mode: mixed (Q2_K+Q4_0), mixed_q1 (Q1_0+Q4_0), custom (per-tensor-type overrides), q1_0 (pure Q1_0), q4_0 (pure Q4_0), or q2_k (pure Q2_K)",
+    )
+    parser.add_argument(
+        "--tensor-type",
+        action="append",
+        default=[],
+        metavar="TYPE=QUANT",
+        help="per-tensor-type override, e.g. ffn_up=Q1_0 or attn_output=Q1_0 "
+        "(QUANT: Q1_0, Q2_K, Q4_0, Copy, Auto; repeatable)",
     )
     parser.add_argument("--no-q4-fallback", action="store_true", help="disable Q4_0 fallback for non-256-aligned tensors")
     parser.add_argument("--keep-embd", action="store_true", help="keep token_embd/output.weight at source precision")

config.py

@@ -8,7 +8,7 @@ The defaults are the C engine's production defaults unless marked
 
 from __future__ import annotations
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 
 
 @dataclass(slots=True)
@@ -69,6 +69,7 @@ class HPCConfig:
     use_candidate_floor: bool = True
     apply_dc_carry: bool = True
     quant_format: str = "mixed"
+    custom_quant_map: dict[str, str] = field(default_factory=dict)
     enable_q4_fallback: bool = True
     quantize_embd_q4: bool = True
     promote_attn_q4: bool = False
@@ -120,6 +121,15 @@ class HPCConfig:
             raise ValueError("workers must be >= 0 (0 = auto CPU count)")
         if self.device not in ("auto", "cuda", "cpu") and not self.device.startswith("cuda:"):
             raise ValueError("device must be 'auto', 'cuda', 'cuda:N', or 'cpu'")
-        if self.quant_format not in ("mixed", "mixed_q1", "q1_0", "q4_0", "q2_k"):
-            raise ValueError("quant_format must be 'mixed', 'mixed_q1', 'q1_0', 'q4_0', or 'q2_k'")
+        if self.quant_format not in ("mixed", "mixed_q1", "custom", "q1_0", "q4_0", "q2_k"):
+            raise ValueError(
+                "quant_format must be 'mixed', 'mixed_q1', 'custom', 'q1_0', 'q4_0', or 'q2_k'"
+            )
+        allowed_custom = {"q1_0", "q2_k", "q4_0", "copy", "auto"}
+        for tensor_type, quant in self.custom_quant_map.items():
+            if str(quant).strip().lower() not in allowed_custom:
+                raise ValueError(
+                    f"custom_quant_map[{tensor_type!r}] must be one of "
+                    f"Q1_0, Q2_K, Q4_0, Copy, Auto, got {quant!r}"
+                )
         return self

gguf.py

@@ -14,6 +14,7 @@ packed by an external GGUF writer.
 
 from __future__ import annotations
 
+import re
 import struct
 import time
 from dataclasses import dataclass
@@ -412,6 +413,18 @@ def expand_importance(
     )
 
 
+_BLOCK_PREFIX_RE = re.compile(r"^blk\.\d+\.")
+
+
+def tensor_type_key(name: str) -> str:
+    """Return the layer-independent tensor type (e.g. ``blk.4.ffn_up.weight`` -> ``ffn_up``)."""
+    key = _BLOCK_PREFIX_RE.sub("", name)
+    for suffix in (".weight", ".bias"):
+        if key.endswith(suffix):
+            return key[: -len(suffix)]
+    return key
+
+
 def _is_skipped_metadata_tensor(name: str, dims: tuple[int, ...]) -> bool:
     if len(dims) < 2:
         return True
@@ -539,6 +552,7 @@ def target_quant_type(
     promote_attn_q4: bool = False,
     quant_format: str = "mixed",
     source_type: int | None = None,
+    custom_quant_map: dict[str, str] | None = None,
 ) -> str | None:
     """Determine target quantization format ('Q1_0', 'Q2_K', 'Q4_0', or None for Copy).
 
@@ -547,15 +561,39 @@ def target_quant_type(
       tensors (and embeddings/head if quantize_embd_q4) -> Q4_0.
     * 'mixed_q1': FFN/expert 128-aligned tensors -> Q1_0; attention/SSM projections,
       embeddings/head, and non-128-aligned 32-aligned tensors -> Q4_0.
+    * 'custom': per-tensor-type overrides from ``custom_quant_map``; unlisted
+      tensor types fall back to the 'mixed' plan.
     * 'q4_0': all 32-aligned eligible weight tensors -> Q4_0.
     * 'q2_k': only Q2_K-eligible tensors -> Q2_K; everything else copied.
     * 'q1_0': all 128-aligned eligible weight tensors -> Q1_0.
 
+    ``custom_quant_map`` maps a layer-independent tensor type (see
+    :func:`tensor_type_key`) to 'Q1_0', 'Q2_K', 'Q4_0', 'Copy', or 'Auto'.
+    Non-Auto entries override the selected mode for that tensor type.  A target
+    is applied only when the tensor is actually eligible for it; explicit
+    'Copy' (and ineligible targets) leave the tensor at its source precision.
+
     Never promotes a lower-bit source tensor to a higher-bit target format
     (e.g., a Q2_K source tensor is never re-quantized to Q4_0).
     """
+    fmt = "mixed" if quant_format == "custom" else quant_format
     tgt: str | None = None
-    if quant_format == "q1_0":
+    custom_choice = None
+    if custom_quant_map:
+        custom_choice = custom_quant_map.get(tensor_type_key(name))
+    if custom_choice is not None and str(custom_choice).strip().lower() not in ("", "auto"):
+        choice = str(custom_choice).strip().upper()
+        if choice == "Q1_0":
+            if is_q1_0_eligible(name, dims, include_embd=True):
+                tgt = "Q1_0"
+        elif choice == "Q2_K":
+            if is_q2k_eligible(name, dims):
+                tgt = "Q2_K"
+        elif choice == "Q4_0":
+            if is_q4_0_eligible(name, dims, include_embd=True, allow_q2k_tensors=True):
+                tgt = "Q4_0"
+        # 'COPY', unknown values, and ineligible targets leave tgt as None.
+    elif fmt == "q1_0":
         if promote_attn_q4 and _is_attention_projection(name) and is_q4_0_eligible(
             name, dims, include_embd=quantize_embd_q4, promote_attn=True, allow_q2k_tensors=True
         ):
@@ -570,7 +608,7 @@ def target_quant_type(
             name, dims, include_embd=quantize_embd_q4, promote_attn=True, allow_q2k_tensors=True
         ):
             tgt = "Q4_0"
-    elif quant_format == "mixed_q1":
+    elif fmt == "mixed_q1":
         if (_is_attention_projection(name) or _is_tiny_ssm_gate(name, dims)) and is_q4_0_eligible(
             name, dims, include_embd=quantize_embd_q4, promote_attn=True, allow_q2k_tensors=True
         ):
@@ -584,12 +622,12 @@ def target_quant_type(
             name, dims, include_embd=quantize_embd_q4, promote_attn=True, allow_q2k_tensors=True
         ):
             tgt = "Q4_0"
-    elif quant_format == "q4_0":
+    elif fmt == "q4_0":
         if is_q4_0_eligible(
             name, dims, include_embd=quantize_embd_q4, allow_q2k_tensors=True
         ):
             tgt = "Q4_0"
-    elif quant_format == "q2_k":
+    elif fmt == "q2_k":
         if is_q2k_eligible(name, dims):
             tgt = "Q2_K"
     else:
@@ -649,6 +687,7 @@ def write_quantized_gguf(
     embd_q4 = bool(getattr(cfg, "quantize_embd_q4", True))
     attn_q4 = bool(getattr(cfg, "promote_attn_q4", False))
     q_format = str(getattr(cfg, "quant_format", "mixed"))
+    custom_map = getattr(cfg, "custom_quant_map", None) or None
 
     plan: list[str | None] = []
     quantized_count = 0
@@ -664,8 +703,16 @@ def write_quantized_gguf(
                     promote_attn_q4=attn_q4,
                     quant_format=q_format,
                     source_type=ti.type,
+                    custom_quant_map=custom_map,
+                )
+                custom_choice = None
+                if custom_map:
+                    custom_choice = custom_map.get(tensor_type_key(ti.name))
+                explicit_custom = (
+                    custom_choice is not None
+                    and str(custom_choice).strip().lower() not in ("", "auto")
                 )
-                if target_fmt is None:
+                if target_fmt is None and not explicit_custom:
                     src_bpw = source_type_bpw(ti.type)
                     if q_format == "q1_0" and len(ti.dims) >= 2 and ti.dims[0] % QK1_0 == 0 and src_bpw > 1.125:
                         target_fmt = "Q1_0"
@@ -686,6 +733,7 @@ def write_quantized_gguf(
                 promote_attn_q4=attn_q4,
                 quant_format=q_format,
                 source_type=ti.type,
+                custom_quant_map=custom_map,
             )
         if target_fmt and ti.type not in SUPPORTED_READ_TYPES:
             if verbose:

gui.py

@@ -61,10 +61,39 @@ from hpc_quantize_v2 import (
     quantize_tensor_q4_0,
     read_imatrix,
 )
-from hpc_quantize_v2.gguf import is_q2k_eligible, target_quant_type, write_quantized_gguf
+from hpc_quantize_v2.gguf import (
+    is_q1_0_eligible,
+    is_q2k_eligible,
+    is_q4_0_eligible,
+    target_quant_type,
+    tensor_type_key,
+    write_quantized_gguf,
+)
 from hpc_quantize_v2.reference import reference_solve, synthetic_cases
 
 
+DEFAULT_CUSTOM_RECIPE: dict[str, str] = {
+    # Q4_0: quality-critical embeddings, attention, and SSM projections
+    "token_embd": "Q4_0",
+    "output": "Q4_0",
+    "attn_gate": "Q4_0",
+    "ssm_out": "Q4_0",
+    "attn_q": "Q4_0",
+    "attn_k": "Q4_0",
+    "attn_v": "Q4_0",
+    # Q1_0: the one tensor type that tolerates it well
+    "attn_output": "Q1_0",
+    # Q2_K: the remaining compute-heavy tensors
+    "attn_qkv": "Q2_K",
+    "ffn_up": "Q2_K",
+    "ffn_gate": "Q2_K",
+    "ffn_down": "Q2_K",
+    "ssm_alpha": "Q2_K",
+    "ssm_beta": "Q2_K",
+    "nextn.eh_proj": "Q2_K",
+}
+
+
 def _ts() -> str:
     return dt.datetime.now().strftime("%H:%M:%S.%f")[:-3]
 
@@ -100,6 +129,8 @@ class HPCQuantizeGUI:
         self._init_vars()
         self._apply_theme()
         self._build_layout()
+        if self.quant_format_var.get() == "custom":
+            self._on_format_changed()
         self._auto_populate_defaults()
         self._update_cli_preview()
         self.root.after(80, self._poll_queue)
@@ -133,8 +164,9 @@ class HPCQuantizeGUI:
         self.q4_fallback_var = tk.BooleanVar(value=True)
         self.embd_q4_var = tk.BooleanVar(value=True)
         self.promote_attn_q4_var = tk.BooleanVar(value=False)
-        self.quant_format_var = tk.StringVar(value="mixed")
+        self.quant_format_var = tk.StringVar(value="custom")
         self.plan_filter_var = tk.StringVar(value="ALL")
+        self.custom_type_vars: dict[str, tk.StringVar] = {}
 
         self.dc_lambda_var = tk.DoubleVar(value=1.0)
         self.vw_lambda_var = tk.DoubleVar(value=1.0)
@@ -310,8 +342,31 @@ class HPCQuantizeGUI:
                 command=self._on_format_changed,
             ).pack(side=tk.LEFT, padx=3)
 
+        custom_row = ttk.Frame(opt_frame)
+        custom_row.pack(fill=tk.X, pady=(0, 4))
+        ttk.Radiobutton(
+            custom_row,
+            text="Custom (per tensor type)",
+            variable=self.quant_format_var,
+            value="custom",
+            command=self._on_format_changed,
+        ).pack(side=tk.LEFT, padx=4)
+        ttk.Label(
+            custom_row,
+            text="Override the quant type per tensor category (Auto = mixed plan).",
+        ).pack(side=tk.LEFT, padx=6)
+
+        self.custom_frame = ttk.LabelFrame(opt_frame, text="Custom Tensor-Type Targets", padding=6)
+        self.custom_body = ttk.Frame(self.custom_frame)
+        self.custom_body.pack(fill=tk.X)
+        self.custom_hint_var = tk.StringVar(value="Load a source GGUF to list tensor types.")
+        ttk.Label(
+            self.custom_frame, textvariable=self.custom_hint_var, style="TeleKey.TLabel"
+        ).pack(anchor=tk.W, pady=(4, 0))
+
         opt_row = ttk.Frame(opt_frame)
         opt_row.pack(fill=tk.X, pady=(0, 4))
+        self._custom_anchor = opt_row
         ttk.Label(opt_row, text="Global Optimizer:").pack(side=tk.LEFT, padx=(0, 8))
         ttk.Radiobutton(
             opt_row, text="HPC (Sieve + 36-State Viterbi)", variable=self.optimizer_var, value="hpc"
@@ -566,6 +621,8 @@ class HPCQuantizeGUI:
             return "-Q1_0.gguf"
         if fmt == "mixed_q1":
             return "-Q1_0-Q4_0.gguf"
+        if fmt == "custom":
+            return "-Custom.gguf"
         if fmt == "q4_0":
             return "-Q4_0.gguf"
         return "-Q2K.gguf"
@@ -575,11 +632,96 @@ class HPCQuantizeGUI:
         gguf_p = self.gguf_var.get().strip()
         if gguf_p and (
             not cur_out
-            or any(cur_out.endswith(sfx) for sfx in ("-Q2K.gguf", "-Q1_0.gguf", "-Q1_0-Q4_0.gguf", "-Q4_0.gguf"))
+            or any(
+                cur_out.endswith(sfx)
+                for sfx in ("-Q2K.gguf", "-Q1_0.gguf", "-Q1_0-Q4_0.gguf", "-Custom.gguf", "-Q4_0.gguf")
+            )
         ):
             p = Path(gguf_p)
             self.output_var.set(str(p.with_name(f"{p.stem}{self._default_out_suffix()}")))
+        if self.quant_format_var.get() == "custom":
+            self.custom_frame.pack(fill=tk.X, pady=(0, 4), before=self._custom_anchor)
+            self._populate_custom_tensor_rows()
+        else:
+            self.custom_frame.pack_forget()
+        self._populate_tensor_tree()
+
+    def _custom_quant_map(self) -> dict[str, str]:
+        return {
+            key: var.get()
+            for key, var in self.custom_type_vars.items()
+            if var.get() not in ("", "Auto")
+        }
+
+    def _custom_map_for_plan(self) -> dict[str, str] | None:
+        if self.quant_format_var.get() != "custom":
+            return None
+        return self._custom_quant_map() or None
+
+    def _populate_custom_tensor_rows(self) -> None:
+        if not hasattr(self, "custom_body"):
+            return
+        for child in self.custom_body.winfo_children():
+            child.destroy()
+
+        if self._reader is None:
+            self.custom_hint_var.set("Load a source GGUF to list tensor types.")
+            self.custom_type_vars = {}
+            return
+
+        groups: dict[str, list] = {}
+        for ti in self._reader.tensors:
+            groups.setdefault(tensor_type_key(ti.name), []).append(ti)
+
+        previous = {key: var.get() for key, var in self.custom_type_vars.items()}
+        self.custom_type_vars = {}
+        rows = 0
+        for key in sorted(groups):
+            tensors = groups[key]
+            options = ["Auto", "Copy"]
+            if any(is_q1_0_eligible(t.name, t.dims) for t in tensors):
+                options.append("Q1_0")
+            if any(is_q2k_eligible(t.name, t.dims) for t in tensors):
+                options.append("Q2_K")
+            if any(
+                is_q4_0_eligible(t.name, t.dims, include_embd=True, allow_q2k_tensors=True)
+                for t in tensors
+            ):
+                options.append("Q4_0")
+            if len(options) == 2:
+                continue
+            var = tk.StringVar(
+                value=previous.get(key, DEFAULT_CUSTOM_RECIPE.get(key, "Auto"))
+            )
+            if var.get() not in options:
+                var.set("Auto")
+            var.trace_add(
+                "write",
+                lambda *_a: self._on_custom_type_changed(),
+            )
+            self.custom_type_vars[key] = var
+            r, c = divmod(rows, 2)
+            ttk.Label(self.custom_body, text=key, width=18, anchor=tk.W).grid(
+                row=r, column=c * 2, sticky=tk.W, padx=(0, 4), pady=1
+            )
+            ttk.Combobox(
+                self.custom_body,
+                textvariable=var,
+                values=options,
+                state="readonly",
+                width=8,
+            ).grid(row=r, column=c * 2 + 1, sticky=tk.W, padx=(0, 12), pady=1)
+            rows += 1
+
+        self.custom_hint_var.set(
+            "Quantizable tensor types only; norms/1D tensors and other non-eligible "
+            "weights are copied unchanged."
+        )
+        self._update_cli_preview()
+
+    def _on_custom_type_changed(self) -> None:
         self._populate_tensor_tree()
+        self._update_cli_preview()
 
     def _browse_gguf(self) -> None:
         path = filedialog.askopenfilename(
@@ -636,6 +778,7 @@ class HPCQuantizeGUI:
             dc_decay=float(self.dc_decay_var.get()),
             sse_budget=float(self.sse_budget_var.get()),
             quant_format=self.quant_format_var.get(),
+            custom_quant_map=self._custom_map_for_plan() or {},
             enable_q4_fallback=bool(self.q4_fallback_var.get()) and self.quant_format_var.get() != "q2_k",
             quantize_embd_q4=bool(self.embd_q4_var.get()),
             promote_attn_q4=bool(self.promote_attn_q4_var.get()),
@@ -664,6 +807,10 @@ class HPCQuantizeGUI:
         if self.limit_var.get() > 0:
             cmd += ["--limit", str(self.limit_var.get())]
         cmd += ["--format", self.quant_format_var.get()]
+        if self.quant_format_var.get() == "custom":
+            custom = self._custom_quant_map()
+            for key in sorted(custom):
+                cmd += ["--tensor-type", f"{key}={custom[key]}"]
         cmd += ["--device", self.device_var.get()]
         cmd += ["--workers", str(self.workers_var.get())]
         cmd += ["--cand-d", str(self.cand_d_var.get()), "--cand-m", str(self.cand_m_var.get())]
@@ -709,6 +856,8 @@ class HPCQuantizeGUI:
             return
         try:
             self._reader = GGUFReader(gguf_path)
+            if self.quant_format_var.get() == "custom":
+                self._populate_custom_tensor_rows()
             self._populate_tensor_tree()
         except Exception as exc:
             self._reader = None
@@ -731,6 +880,7 @@ class HPCQuantizeGUI:
                 promote_attn_q4=bool(self.promote_attn_q4_var.get()),
                 quant_format=self.quant_format_var.get(),
                 source_type=ti.type,
+                custom_quant_map=self._custom_map_for_plan(),
             )
             if tgt == "Q2_K":
                 n_q2 += 1
@@ -793,6 +943,7 @@ class HPCQuantizeGUI:
                 promote_attn_q4=bool(self.promote_attn_q4_var.get()),
                 quant_format=self.quant_format_var.get(),
                 source_type=ti.type,
+                custom_quant_map=self._custom_map_for_plan(),
             )
             is not None
             and self.tree.exists(ti.name)
@@ -884,7 +1035,17 @@ class HPCQuantizeGUI:
                     raw = lookup_imatrix(imat, tname)
                     if raw is not None:
                         imp = expand_importance(raw, x.size, row)
-                if cfg.quant_format in ("q1_0", "mixed_q1"):
+                tgt = target_quant_type(
+                    tname,
+                    ti.dims,
+                    enable_q4_fallback=bool(getattr(cfg, "enable_q4_fallback", True)),
+                    quantize_embd_q4=bool(getattr(cfg, "quantize_embd_q4", True)),
+                    promote_attn_q4=bool(getattr(cfg, "promote_attn_q4", False)),
+                    quant_format=str(getattr(cfg, "quant_format", "mixed")),
+                    source_type=ti.type,
+                    custom_quant_map=getattr(cfg, "custom_quant_map", None) or None,
+                )
+                if tgt == "Q1_0":
                     raw_q1, nb_q1, rmse_q1 = quantize_tensor_q1_0(x, importance=imp, cfg=cfg)
                     np.savez_compressed(
                         out_path,
@@ -896,7 +1057,7 @@ class HPCQuantizeGUI:
                         ("done", f"Saved Q1_0 {out_path} (RMSE={rmse_q1:.6e}, {time.time() - t0:.1f}s)")
                     )
                     return
-                if cfg.quant_format == "q4_0":
+                if tgt == "Q4_0":
                     raw_q4, nb_q4, rmse_q4 = quantize_tensor_q4_0(x, importance=imp, cfg=cfg)
                     np.savez_compressed(
                         out_path,
@@ -908,6 +1069,8 @@ class HPCQuantizeGUI:
                         ("done", f"Saved Q4_0 {out_path} (RMSE={rmse_q4:.6e}, {time.time() - t0:.1f}s)")
                     )
                     return
+                if tgt != "Q2_K":
+                    raise ValueError(f"{tname}: target {tgt!r} is not quantizable to .npz")
                 res = quantizer.quantize(
                     x,
                     importance=imp,