Upload 6 files

__main__.py

@@ -59,6 +59,8 @@ def _config_from_args(args) -> HPCConfig:
         cfg.segment_k = int(args.segment_k)
     if getattr(args, "segment_iters", None) is not None:
         cfg.segment_iters = int(args.segment_iters)
+    if getattr(args, "segment_table", None) is not None:
+        cfg.segment_table_size = int(args.segment_table)
     if args.no_q4_fallback:
         cfg.enable_q4_fallback = False
     if args.keep_embd:
@@ -329,6 +331,13 @@ def main() -> None:
         metavar="N",
         help="proxy fixed-point rounds for superblock segmentation (default 4)",
     )
+    parser.add_argument(
+        "--segment-table",
+        type=int,
+        default=None,
+        metavar="M",
+        help="generalized permutation lookup table size per tensor (default 256 = +0.031 bpw; 0 = full per-block)",
+    )
     parser.add_argument(
         "--format",
         choices=("mixed", "mixed_q1", "custom", "q1_0", "q3_k", "q4_0", "q2_k"),

config.py

@@ -113,6 +113,7 @@ class HPCConfig:
     segment: bool = False
     segment_k: int = 16
     segment_iters: int = 4
+    segment_table_size: int = 256
 
     # ---- seed clamp -----------------------------------------------------
     clamp_kurtosis_off: float = 6.0
@@ -176,6 +177,8 @@ class HPCConfig:
                 raise ValueError("segment_k must lie in [2, 256]")
             if self.segment_iters < 1:
                 raise ValueError("segment_iters must be >= 1")
+            if self.segment_table_size < 0 or (0 < self.segment_table_size < 4) or self.segment_table_size > 65535:
+                raise ValueError("segment_table_size must be 0 (full per-block) or lie in [4, 65535]")
         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", "custom", "q1_0", "q3_k", "q4_0", "q2_k"):

gguf.py

@@ -798,6 +798,7 @@ def write_quantized_gguf(
     custom_map = getattr(cfg, "custom_quant_map", None) or None
     seg_active = bool(getattr(cfg, "segment", False))
     seg_k = max(2, int(getattr(cfg, "segment_k", 16)))
+    seg_table_size = int(getattr(cfg, "segment_table_size", 256))
 
     plan: list[str | None] = []
     quantized_count = 0
@@ -907,7 +908,7 @@ def write_quantized_gguf(
         for ti, fmt in zip(reader.tensors, plan):
             if fmt in ("Q2_K", "Q3_K"):
                 key = SEG_KV_PREFIX + ti.name
-                content_len = label_channel_bytes(ti.n_elements, seg_k)
+                content_len = label_channel_bytes(ti.n_elements, seg_k, seg_table_size)
                 seg_reserve[key] = content_len
                 # GGUF uint8 array: u32 element type + u64 count + zeros
                 kv_pairs.append(
@@ -936,7 +937,7 @@ def write_quantized_gguf(
             fout.write(raw)
 
         def _patch_seg_kv(fout, key: str, labels: np.ndarray) -> None:
-            packed = pack_label_channel(labels, seg_k)
+            packed = pack_label_channel(labels, seg_k, seg_table_size)
             if len(packed) != seg_reserve.get(key, -1):
                 raise RuntimeError(f"{key}: label channel size {len(packed)} != reserved")
             pos = fout.tell()

gui.py

@@ -70,10 +70,12 @@ from hpc_quantize_v2 import (
     read_imatrix,
 )
 from hpc_quantize_v2.gguf import (
+    SEG_KV_PREFIX,
     is_q1_0_eligible,
     is_q2k_eligible,
     is_q3_k_eligible,
     is_q4_0_eligible,
+    label_channel_bytes,
     target_quant_type,
     tensor_type_key,
     write_quantized_gguf,
@@ -187,6 +189,7 @@ class HPCQuantizeGUI:
         self.segment_var = tk.BooleanVar(value=False)
         self.segment_k_var = tk.IntVar(value=16)
         self.segment_iters_var = tk.IntVar(value=4)
+        self.segment_table_var = tk.IntVar(value=256)
 
         self.cli_preview_var = tk.StringVar()
         self.status_var = tk.StringVar(value="Ready.")
@@ -225,8 +228,11 @@ class HPCQuantizeGUI:
             self.segment_var,
             self.segment_k_var,
             self.segment_iters_var,
+            self.segment_table_var,
         ):
             v.trace_add("write", lambda *_: self._update_cli_preview())
+        for v in (self.segment_var, self.segment_k_var, self.segment_table_var):
+            v.trace_add("write", lambda *_: self._refresh_plan_view())
 
     def _apply_theme(self) -> None:
         style = ttk.Style(self.root)
@@ -504,10 +510,17 @@ class HPCQuantizeGUI:
         ttk.Spinbox(
             spec_f, from_=1, to=12, increment=1, textvariable=self.segment_iters_var, width=8
         ).grid(row=4, column=3, sticky=tk.EW, padx=4, pady=2)
+        ttk.Label(spec_f, text="Seg Lookup Table (0=full):").grid(row=5, column=0, columnspan=2, sticky=tk.W, padx=4, pady=2)
+        ttk.Combobox(
+            spec_f,
+            textvariable=self.segment_table_var,
+            values=(0, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096),
+            width=8,
+        ).grid(row=5, column=2, columnspan=2, sticky=tk.EW, padx=4, pady=2)
         ttk.Label(
             spec_f, text="(λ multiplier ≥1 during thermal anneal; ~1e2–1e3 ≈ zero-coherent-residual '0-error' anneal)",
             foreground="#93c5fd",
-        ).grid(row=5, column=0, columnspan=4, sticky=tk.W, padx=4, pady=(0, 2))
+        ).grid(row=6, column=0, columnspan=4, sticky=tk.W, padx=4, pady=(0, 2))
 
         # Right: GGUF Tensor Browser
         browser_frame = ttk.LabelFrame(right_pane, text="3. GGUF Tensor Inspector & Live Status", padding=8)
@@ -840,6 +853,7 @@ class HPCQuantizeGUI:
             segment=bool(self.segment_var.get()),
             segment_k=max(2, int(self.segment_k_var.get())),
             segment_iters=max(1, int(self.segment_iters_var.get())),
+            segment_table_size=max(0, int(self.segment_table_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() not in ("q2_k", "q3_k"),
@@ -899,6 +913,8 @@ class HPCQuantizeGUI:
                 cmd += ["--segment-k", str(int(self.segment_k_var.get()))]
             if int(self.segment_iters_var.get()) != 4:
                 cmd += ["--segment-iters", str(int(self.segment_iters_var.get()))]
+            if int(self.segment_table_var.get()) != 256:
+                cmd += ["--segment-table", str(int(self.segment_table_var.get()))]
         cmd.append("--verbose")
         return " ".join(shlex.quote(part) for part in cmd)
 
@@ -962,6 +978,43 @@ class HPCQuantizeGUI:
             self._refresh_plan_view()
         self._update_cli_preview()
 
+    def _estimate_seg_bytes(self, ti, tgt: str | None) -> int:
+        if not bool(self.segment_var.get()) or tgt not in ("Q2_K", "Q3_K"):
+            return 0
+        try:
+            seg_k = max(2, int(self.segment_k_var.get()))
+        except Exception:
+            seg_k = 16
+        try:
+            seg_tbl = max(0, int(self.segment_table_var.get()))
+        except Exception:
+            seg_tbl = 256
+        content_len = label_channel_bytes(ti.n_elements, seg_k, seg_tbl)
+        key_len = len((SEG_KV_PREFIX + ti.name).encode("utf-8"))
+        return 8 + key_len + 4 + 12 + content_len
+
+    def _format_size_summary(self, planned_bytes: int, total_seg_bytes: int) -> str:
+        if planned_bytes >= 1024 ** 3:
+            planned_text = f"{planned_bytes / (1024 ** 3):.2f} GiB"
+        elif planned_bytes >= 1024 ** 2:
+            planned_text = f"{planned_bytes / (1024 ** 2):.1f} MiB"
+        else:
+            planned_text = f"{planned_bytes / 1024:.1f} KiB"
+        if total_seg_bytes > 0:
+            try:
+                seg_tbl = max(0, int(self.segment_table_var.get()))
+            except Exception:
+                seg_tbl = 256
+            if total_seg_bytes >= 1024 ** 3:
+                seg_txt = f"+{total_seg_bytes / (1024 ** 3):.2f} GiB"
+            elif total_seg_bytes >= 1024 ** 2:
+                seg_txt = f"+{total_seg_bytes / (1024 ** 2):.1f} MiB"
+            else:
+                seg_txt = f"+{total_seg_bytes / 1024:.1f} KiB"
+            lut_lbl = f"LUT={seg_tbl}" if seg_tbl > 0 else "full"
+            planned_text = f"{planned_text} [seg {seg_txt}, {lut_lbl}]"
+        return planned_text
+
     def _refresh_plan_view(self) -> None:
         if self._reader is None:
             return
@@ -969,7 +1022,8 @@ class HPCQuantizeGUI:
         scope_selected = self.mode_var.get() == "selected"
         n_q2, n_q3, n_q4, n_copy = 0, 0, 0, 0
         n_q1 = 0
-        planned_bytes = 0
+        planned_bytes = int(getattr(self._reader, "data_start", 0))
+        total_seg_bytes = 0
         for ti in self._reader.tensors:
             tgt = self._plan_target(ti, chosen, scope_selected)
             if tgt == "Q2_K":
@@ -993,6 +1047,9 @@ class HPCQuantizeGUI:
                 planned_bytes += (ti.n_elements // QK1_0) * Q1_0_BLOCK_BYTES
             else:
                 planned_bytes += ti.n_bytes
+            seg_b = self._estimate_seg_bytes(ti, tgt)
+            planned_bytes += seg_b
+            total_seg_bytes += seg_b
 
             if not self.tree.exists(ti.name):
                 continue
@@ -1010,12 +1067,7 @@ class HPCQuantizeGUI:
             vals[6] = state_str
             self.tree.item(ti.name, tags=(tgt_label,), values=vals)
 
-        if planned_bytes >= 1024 ** 3:
-            planned_text = f"{planned_bytes / (1024 ** 3):.2f} GiB"
-        elif planned_bytes >= 1024 ** 2:
-            planned_text = f"{planned_bytes / (1024 ** 2):.1f} MiB"
-        else:
-            planned_text = f"{planned_bytes / 1024:.1f} KiB"
+        planned_text = self._format_size_summary(planned_bytes, total_seg_bytes)
         scope_text = "selected" if scope_selected else "all"
         self.status_var.set(
             f"Plan ({self.quant_format_var.get().upper()}, {scope_text}): {n_q1} Q1_0 | {n_q2} Q2_K | {n_q3} Q3_K | {n_q4} Q4_0 | {n_copy} Copy "
@@ -1032,7 +1084,8 @@ class HPCQuantizeGUI:
         plan_filt = self.plan_filter_var.get()
         n_q2, n_q3, n_q4, n_copy = 0, 0, 0, 0
         n_q1 = 0
-        planned_bytes = 0
+        planned_bytes = int(getattr(self._reader, "data_start", 0))
+        total_seg_bytes = 0
         for ti in self._reader.tensors:
             tgt = self._plan_target(ti, chosen, scope_selected)
             if tgt == "Q2_K":
@@ -1056,6 +1109,9 @@ class HPCQuantizeGUI:
                 planned_bytes += (ti.n_elements // QK1_0) * Q1_0_BLOCK_BYTES
             else:
                 planned_bytes += ti.n_bytes
+            seg_b = self._estimate_seg_bytes(ti, tgt)
+            planned_bytes += seg_b
+            total_seg_bytes += seg_b
 
             tgt_label = tgt if tgt else "Copy"
             if plan_filt != "ALL" and tgt_label.upper() != plan_filt:
@@ -1097,12 +1153,7 @@ class HPCQuantizeGUI:
         keep_selected = [n for n in chosen if self.tree.exists(n)]
         if keep_selected:
             self.tree.selection_set(keep_selected)
-        if planned_bytes >= 1024 ** 3:
-            planned_text = f"{planned_bytes / (1024 ** 3):.2f} GiB"
-        elif planned_bytes >= 1024 ** 2:
-            planned_text = f"{planned_bytes / (1024 ** 2):.1f} MiB"
-        else:
-            planned_text = f"{planned_bytes / 1024:.1f} KiB"
+        planned_text = self._format_size_summary(planned_bytes, total_seg_bytes)
         scope_text = "selected" if scope_selected else "all"
         self.status_var.set(
             f"Plan ({self.quant_format_var.get().upper()}, {scope_text}): {n_q1} Q1_0 | {n_q2} Q2_K | {n_q3} Q3_K | {n_q4} Q4_0 | {n_copy} Copy "

patch_llamacpp_hpcq_seg.py

@@ -78,13 +78,17 @@ struct hpcq_seg_entry {
     uintptr_t end;         // payload end (exclusive)
     int64_t   block_bytes; // 84 for Q2_K, 110 for Q3_K
     int64_t   n_blocks;    // number of 256-weight superblocks
-    uint8_t   enc;         // 0 = 4-bit nibbles, 1 = 8-bit, 2 = 256-entry static table + 1B/blk
+    uint8_t   enc;         // 0=nibbles, 1=bytes, 2=LUT+1B/blk, 3=LUT+2B/blk
     uint8_t * packed;      // compact packed label stream
     void    * dev_packed;  // lazily allocated device copy for CUDA (or NULL)
 };
 
-static inline HPCQ_HD size_t hpcq_seg_packed_bytes(int64_t n_blocks, uint8_t enc) {
-    if (enc == 2) return 32768u + (size_t) n_blocks;
+static inline HPCQ_HD size_t hpcq_seg_packed_bytes(int64_t n_blocks, uint8_t enc, const uint8_t * data) {
+    if (enc == 2 || enc == 3) {
+        const uint16_t tsize = (uint16_t) data[0] | ((uint16_t) data[1] << 8);
+        const size_t idx_bytes = (enc == 2) ? 1u : 2u;
+        return 2u + (size_t) tsize * 128u + (size_t) n_blocks * idx_bytes;
+    }
     if (enc == 0) return (size_t) n_blocks * 128u;
     return (size_t) n_blocks * 256u;
 }
@@ -139,9 +143,17 @@ static inline HPCQ_HD void hpcq_seg_make_order_block(
 
 static inline HPCQ_HD void hpcq_seg_make_order_at(
         const uint8_t * packed_base, uint8_t enc, int64_t blk_idx, uint16_t * order) {
-    if (enc == 2) {
-        const uint8_t m = packed_base[32768u + (size_t) blk_idx];
-        hpcq_seg_make_order_block(packed_base + (size_t) m * 128u, 0, order);
+    if (enc == 2 || enc == 3) {
+        const uint16_t tsize = (uint16_t) packed_base[0] | ((uint16_t) packed_base[1] << 8);
+        const size_t idx_off = 2u + (size_t) tsize * 128u;
+        uint16_t m = 0;
+        if (enc == 2) {
+            m = packed_base[idx_off + (size_t) blk_idx];
+        } else {
+            const uint8_t * p = packed_base + idx_off + (size_t) blk_idx * 2u;
+            m = (uint16_t) p[0] | ((uint16_t) p[1] << 8);
+        }
+        hpcq_seg_make_order_block(packed_base + 2u + (size_t) m * 128u, 0, order);
     } else if (enc == 0) {
         hpcq_seg_make_order_block(packed_base + (size_t) blk_idx * 128u, 0, order);
     } else {
@@ -180,7 +192,8 @@ void ggml_hpcq_seg_bind_raw(
     if (n_labels == 0 || (n_labels % 256) != 0) return;
     const int64_t n_blocks = (int64_t) (n_labels / 256);
     if (n_blocks * block_bytes != n_bytes) return;
-    const size_t need = hpcq_seg_packed_bytes(n_blocks, enc);
+    if ((enc == 2 || enc == 3) && hqs1_len < 12u) return;
+    const size_t need = hpcq_seg_packed_bytes(n_blocks, enc, hqs1 + 10);
     if (hqs1_len < 10u + need) return;
 
     const uintptr_t base = (uintptr_t) payload;
@@ -550,7 +563,7 @@ static inline const uint8_t * hpcq_cuda_get_dev_packed(
     if (!seg) return nullptr;
     *out_enc = seg->enc;
     if (seg->dev_packed == nullptr) {{
-        const size_t nbytes = hpcq_seg_packed_bytes(seg->n_blocks, seg->enc);
+        const size_t nbytes = hpcq_seg_packed_bytes(seg->n_blocks, seg->enc, seg->packed);
         void * d_ptr = nullptr;
         if (cudaMalloc(&d_ptr, nbytes) == cudaSuccess && d_ptr != nullptr) {{
             if (cudaMemcpy(d_ptr, seg->packed, nbytes, cudaMemcpyHostToDevice) == cudaSuccess) {{

segmentation.py

@@ -114,11 +114,12 @@ def scatter_to_original(xh_slots: np.ndarray, order: np.ndarray) -> np.ndarray:
 #
 # Layout: magic "HQS1" | uint8 k | uint8 enc | uint32 n_labels (little-endian)
 # | payload.
-#   enc=0: packs labels 0..15 as nibbles (2 labels per byte; used for n <= 65536)
+#   enc=0: packs labels 0..15 as nibbles (2 labels per byte; full per-block or small tensors)
 #   enc=1: stores raw uint8 labels (k > 16 only)
-#   enc=2: generalized 256-entry permutation table (32 KB nibble table + 1 byte
-#          index per 256-weight superblock = +0.03125 bpw; used for n > 65536,
-#          k <= 16, n % 256 == 0).
+#   enc=2: generalized permutation lookup table with 1-byte block indices (table_size in 4..256):
+#          uint16 table_size | uint8 table[table_size * 128] | uint8 block_idx[n_blocks]
+#   enc=3: generalized permutation lookup table with 2-byte block indices (table_size in 257..65535):
+#          uint16 table_size | uint8 table[table_size * 128] | uint16 block_idx[n_blocks]
 
 import struct as _struct
 
@@ -142,24 +143,29 @@ def _unpack_nibbles_1d(data: bytes, n: int) -> np.ndarray:
     return out[:n]
 
 
-def label_channel_bytes(n: int, k: int) -> int:
+def label_channel_bytes(n: int, k: int, table_size: int = SEG_TABLE_SIZE) -> int:
     """Exact content size of the packed label channel."""
+    table_size = int(table_size)
     if int(k) <= 16:
-        if n > SEG_TABLE_SIZE * 256 and n % 256 == 0:
-            return 10 + SEG_TABLE_SIZE * 128 + (n // 256)
+        if table_size > 0 and n > table_size * 256 and n % 256 == 0:
+            n_blocks = n // 256
+            idx_bytes = 1 if table_size <= 256 else 2
+            return 10 + 2 + table_size * 128 + n_blocks * idx_bytes
         return 10 + (n + 1) // 2
     return 10 + n
 
 
-def pack_label_channel(labels: np.ndarray, k: int) -> bytes:
+def pack_label_channel(labels: np.ndarray, k: int, table_size: int = SEG_TABLE_SIZE) -> bytes:
     flat = np.ascontiguousarray(labels, dtype=np.int64).ravel()
     n = int(flat.size)
     k = int(k)
+    table_size = int(table_size)
     if k <= 16:
         if flat.size and int(flat.max()) > 15:
             raise ValueError(f"label {int(flat.max())} does not fit nibble packing (k={k})")
-        if n > SEG_TABLE_SIZE * 256 and n % 256 == 0:
-            enc = 2
+        if table_size > 0 and n > table_size * 256 and n % 256 == 0:
+            enc = 2 if table_size <= 256 else 3
+            idx_dt = np.uint8 if enc == 2 else np.dtype("<u2")
             n_blocks = n // 256
             rows = flat.reshape(n_blocks, 256).astype(np.uint8)
             lo = rows[:, 0::2]
@@ -171,17 +177,17 @@ def pack_label_channel(labels: np.ndarray, k: int) -> bytes:
                 row_views, return_inverse=True, return_counts=True
             )
             uniq_rows = uniq_views.view(np.uint8).reshape(-1, 128)
-            table = np.zeros((SEG_TABLE_SIZE, 128), dtype=np.uint8)
-            if uniq_rows.shape[0] <= SEG_TABLE_SIZE:
+            table = np.zeros((table_size, 128), dtype=np.uint8)
+            if uniq_rows.shape[0] <= table_size:
                 table[: uniq_rows.shape[0]] = uniq_rows
-                block_idx = inv_idx.astype(np.uint8)
+                block_idx = inv_idx.astype(idx_dt)
             else:
-                top_u = np.argsort(-counts)[:SEG_TABLE_SIZE]
+                top_u = np.argsort(-counts)[:table_size]
                 table[:] = uniq_rows[top_u]
-                remap = np.zeros(uniq_rows.shape[0], dtype=np.uint8)
-                remap[top_u] = np.arange(SEG_TABLE_SIZE, dtype=np.uint8)
+                remap = np.zeros(uniq_rows.shape[0], dtype=idx_dt)
+                remap[top_u] = np.arange(table_size, dtype=idx_dt)
                 block_idx = remap[inv_idx]
-            data = table.tobytes() + block_idx.tobytes()
+            data = _struct.pack("<H", table_size) + table.tobytes() + block_idx.tobytes()
         else:
             enc = 0
             data = _pack_nibbles_1d(flat)
@@ -202,13 +208,19 @@ def unpack_label_channel(content: bytes) -> tuple[int, np.ndarray]:
     data = content[10:]
     if enc == 0:
         out = _unpack_nibbles_1d(data, n)
-    elif enc == 2:
+    elif enc in (2, 3):
         n_blocks = n // 256
-        tbytes = SEG_TABLE_SIZE * 128
-        table = _unpack_nibbles_1d(data[:tbytes], SEG_TABLE_SIZE * 256).reshape(
-            SEG_TABLE_SIZE, 256
+        tsize = _struct.unpack("<H", data[:2])[0]
+        tbytes = tsize * 128
+        table = _unpack_nibbles_1d(data[2 : 2 + tbytes], tsize * 256).reshape(
+            tsize, 256
         )
-        idx = np.frombuffer(data[tbytes : tbytes + n_blocks], dtype=np.uint8)
+        if enc == 2:
+            idx = np.frombuffer(data[2 + tbytes : 2 + tbytes + n_blocks], dtype=np.uint8)
+        else:
+            idx = np.frombuffer(
+                data[2 + tbytes : 2 + tbytes + 2 * n_blocks], dtype=np.dtype("<u2")
+            )
         out = table[idx].reshape(-1)
     else:
         out = np.frombuffer(data, dtype=np.uint8).copy()
@@ -268,10 +280,10 @@ def segment_stage(
     Wo = None if W is None else np.asarray(W, dtype=np.float64)
     k = max(1, int(k))
     n_blocks, B = Xo.shape
-    if n_blocks <= SEG_TABLE_SIZE or B != 256 or k > 16:
+    M = int(getattr(cfg, "segment_table_size", SEG_TABLE_SIZE))
+    if M <= 0 or n_blocks <= M or B != 256 or k > 16:
         return _segment_stage_dense(Xo, Wo, cfg, proxy, k)
 
-    M = SEG_TABLE_SIZE
     table_labels = np.empty((M, 256), dtype=np.int64)
     # Entry 0: natural order (monotone labels so make_order is identity [0..255])
     table_labels[0] = _order_to_canonical_labels(np.arange(256, dtype=np.int64), k)