Upload 4 files

gguf.py

@@ -1057,7 +1057,8 @@ def write_quantized_gguf(
                 t_q3_start = time.time()
                 q3_labels: list[np.ndarray] = []
 
-                for x_chunk, row, _r0, _r1 in reader.iter_tensor_row_chunks(ti.name, max_elements=8_388_608):
+                q3_max_elems = ti.n_elements if seg_active else 8_388_608
+                for x_chunk, row, _r0, _r1 in reader.iter_tensor_row_chunks(ti.name, max_elements=q3_max_elems):
                     if stop_event is not None and stop_event.is_set():
                         raise InterruptedError("Quantization cancelled by user")
                     imp_chunk = expand_importance(raw_imp, x_chunk.size, row) if has_imat else None

patch_llamacpp_hpcq_seg.py

@@ -78,11 +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 (128B/blk), 1 = 8-bit (256B/blk)
-    uint8_t * packed;      // compact packed label stream (128B or 256B per block)
+    uint8_t   enc;         // 0 = 4-bit nibbles, 1 = 8-bit, 2 = 256-entry static table + 1B/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;
+    if (enc == 0) return (size_t) n_blocks * 128u;
+    return (size_t) n_blocks * 256u;
+}
+
 GGML_API void ggml_hpcq_seg_bind_raw(
         const void * payload, int64_t n_bytes,
         const uint8_t * hqs1, size_t hqs1_len, int64_t block_bytes);
@@ -131,6 +137,18 @@ 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);
+    } else if (enc == 0) {
+        hpcq_seg_make_order_block(packed_base + (size_t) blk_idx * 128u, 0, order);
+    } else {
+        hpcq_seg_make_order_block(packed_base + (size_t) blk_idx * 256u, 1, order);
+    }
+}
+
 #ifdef __cplusplus
 }
 #endif

segmentation.py

@@ -113,17 +113,40 @@ def scatter_to_original(xh_slots: np.ndarray, order: np.ndarray) -> np.ndarray:
 # ---- label channel codec (in-format persistence inside a GGUF KV) --------
 #
 # 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),
-# enc=1 stores raw uint8 labels (k > 16 only).
+# | payload.
+#   enc=0: packs labels 0..15 as nibbles (2 labels per byte; used for n <= 65536)
+#   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).
 
 import struct as _struct
 
 SEG_MAGIC = b"HQS1"
+SEG_TABLE_SIZE = 256
+
+
+def _pack_nibbles_1d(flat: np.ndarray) -> bytes:
+    lo = flat[0::2].astype(np.uint8)
+    hi = flat[1::2].astype(np.uint8)
+    if hi.size < lo.size:
+        hi = np.concatenate([hi, np.zeros(1, dtype=np.uint8)])
+    return (lo | (hi << 4)).astype(np.uint8).tobytes()
+
+
+def _unpack_nibbles_1d(data: bytes, n: int) -> np.ndarray:
+    raw = np.frombuffer(data, dtype=np.uint8)
+    out = np.empty(raw.size * 2, dtype=np.uint8)
+    out[0::2] = raw & np.uint8(0x0F)
+    out[1::2] = raw >> np.uint8(4)
+    return out[:n]
 
 
 def label_channel_bytes(n: int, k: int) -> int:
     """Exact content size of the packed label channel."""
     if int(k) <= 16:
+        if n > SEG_TABLE_SIZE * 256 and n % 256 == 0:
+            return 10 + SEG_TABLE_SIZE * 128 + (n // 256)
         return 10 + (n + 1) // 2
     return 10 + n
 
@@ -135,12 +158,33 @@ def pack_label_channel(labels: np.ndarray, k: int) -> bytes:
     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})")
-        enc = 0
-        lo = flat[0::2].astype(np.uint8)
-        hi = flat[1::2].astype(np.uint8)
-        if hi.size < lo.size:
-            hi = np.concatenate([hi, np.zeros(1, dtype=np.uint8)])
-        data = (lo | (hi << 4)).astype(np.uint8).tobytes()
+        if n > SEG_TABLE_SIZE * 256 and n % 256 == 0:
+            enc = 2
+            n_blocks = n // 256
+            rows = flat.reshape(n_blocks, 256).astype(np.uint8)
+            lo = rows[:, 0::2]
+            hi = rows[:, 1::2]
+            packed_rows = (lo | (hi << 4)).astype(np.uint8)  # (n_blocks, 128)
+            row_dt = np.dtype((np.void, 128))
+            row_views = np.ascontiguousarray(packed_rows).view(row_dt).ravel()
+            uniq_views, inv_idx, counts = np.unique(
+                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[: uniq_rows.shape[0]] = uniq_rows
+                block_idx = inv_idx.astype(np.uint8)
+            else:
+                top_u = np.argsort(-counts)[:SEG_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)
+                block_idx = remap[inv_idx]
+            data = table.tobytes() + block_idx.tobytes()
+        else:
+            enc = 0
+            data = _pack_nibbles_1d(flat)
     else:
         enc = 1
         if flat.size and (int(flat.max()) > 255 or int(flat.min()) < 0):
@@ -157,15 +201,53 @@ def unpack_label_channel(content: bytes) -> tuple[int, np.ndarray]:
     n = _struct.unpack("<I", content[6:10])[0]
     data = content[10:]
     if enc == 0:
-        raw = np.frombuffer(data, dtype=np.uint8)
-        out = np.empty(raw.size * 2, dtype=np.uint8)
-        out[0::2] = raw & np.uint8(0x0F)
-        out[1::2] = raw >> np.uint8(4)
+        out = _unpack_nibbles_1d(data, n)
+    elif enc == 2:
+        n_blocks = n // 256
+        tbytes = SEG_TABLE_SIZE * 128
+        table = _unpack_nibbles_1d(data[:tbytes], SEG_TABLE_SIZE * 256).reshape(
+            SEG_TABLE_SIZE, 256
+        )
+        idx = np.frombuffer(data[tbytes : tbytes + n_blocks], dtype=np.uint8)
+        out = table[idx].reshape(-1)
     else:
         out = np.frombuffer(data, dtype=np.uint8).copy()
     return k, out[:n]
 
 
+def _order_to_canonical_labels(ord_row: np.ndarray, k: int = 16) -> np.ndarray:
+    lab = np.empty(256, dtype=np.int64)
+    step = 256 // k
+    for g in range(k):
+        s0 = g * step
+        s1 = 256 if g == k - 1 else (g + 1) * step
+        sub = np.sort(ord_row[s0:s1])
+        lab[sub] = g
+    return lab
+
+
+def _segment_stage_dense(
+    Xo: np.ndarray,
+    Wo: np.ndarray | None,
+    cfg: HPCConfig,
+    proxy,
+    k: int,
+) -> tuple[np.ndarray, np.ndarray, int]:
+    labels = band_labels(Xo, k)
+    order = make_order(labels)
+    rounds = 0
+    for _ in range(max(1, int(cfg.segment_iters))):
+        rounds += 1
+        xg = apply_order(Xo, order)
+        wgg = None if Wo is None else apply_order(Wo, order)
+        xh = proxy(xg, wgg)
+        labels_slot = band_labels(xh, k)
+        labels = np.empty_like(labels_slot)
+        np.put_along_axis(labels, order, labels_slot, axis=1)
+        order = make_order(labels)
+    return order, labels, rounds
+
+
 def segment_stage(
     X: np.ndarray,
     W: np.ndarray | None,
@@ -175,31 +257,67 @@ def segment_stage(
 ) -> tuple[np.ndarray, np.ndarray, int]:
     """Run the label fixed-point loop.
 
-    Returns ``(order, labels, rounds)`` where ``labels`` holds one band id
-    per ORIGINAL position (this is the label channel the decoder needs for
-    exact inversion) and ``order`` is the slot->original permutation.
-    ``proxy(xg, wg)`` must return a dequantized slot view from a cheap
-    quantize/dequantize of the permuted view; labels are refined from that
-    proxy reconstruction so the stored channel matches what a channelless
-    re-derivation would see.
+    For tensors with ``n_blocks <= 256`` (or ``k > 16``), refines a per-block
+    label vector directly.  For larger tensors (``n_blocks > 256, k <= 16``),
+    constructs a 256-entry generalized permutation table (Entry 0 = natural
+    order, Entries 1..2 = column RMS/abs profiles, Entries 3..255 = prototype
+    blocks refined by ``_segment_stage_dense``) and assigns each superblock its
+    optimal 1-byte table index (+0.03125 bpw static+index storage tax).
     """
     Xo = np.asarray(X, dtype=np.float64)
     Wo = None if W is None else np.asarray(W, dtype=np.float64)
     k = max(1, int(k))
-    labels = band_labels(Xo, k)          # label channel, ORIGINAL position order
-    order = make_order(labels)
-    wg = None if Wo is None else np.asarray(Wo, dtype=np.float64)
-    rounds = 0
-    for _ in range(max(1, int(cfg.segment_iters))):
-        rounds += 1
-        xg = apply_order(Xo, order)
-        wgg = None if wg is None else apply_order(wg, order)
-        xh = proxy(xg, wgg)              # slot-ordered reconstruction
-        labels_slot = band_labels(xh, k)  # labels in SLOT order
-        labels = np.empty_like(labels_slot)
-        np.put_along_axis(labels, order, labels_slot, axis=1)  # back to original order
-        order = make_order(labels)
-    return order, labels, rounds
+    n_blocks, B = Xo.shape
+    if n_blocks <= SEG_TABLE_SIZE 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)
+    # Entry 1 & 2: column-wide RMS and abs profiles
+    col_rms = np.sqrt(np.mean(Xo * Xo, axis=0))
+    table_labels[1] = _order_to_canonical_labels(np.argsort(col_rms, kind="stable"), k)
+    col_abs = np.mean(np.abs(Xo), axis=0)
+    table_labels[2] = _order_to_canonical_labels(np.argsort(col_abs, kind="stable"), k)
+
+    # Entries 3..M-1: prototype blocks refined with the proxy fixed-point loop
+    proto_idx = np.linspace(0, n_blocks - 1, M - 3, dtype=np.int64)
+    Wp = None if Wo is None else Wo[proto_idx]
+    ord_p, _, rounds = _segment_stage_dense(Xo[proto_idx], Wp, cfg, proxy, k)
+    for i in range(M - 3):
+        table_labels[3 + i] = _order_to_canonical_labels(ord_p[i], k)
+
+    table_orders = make_order(table_labels)  # (M, 256)
+
+    # Prescreen top-8 table entries per block via between-subblock variance GEMM
+    P = np.zeros((256, M * 16), dtype=np.float32)
+    for m in range(M):
+        for g in range(16):
+            P[table_orders[m, 16 * g : 16 * g + 16], m * 16 + g] = 0.25
+    Xf = Xo.astype(np.float32)
+    scores = ((Xf @ P) ** 2).reshape(n_blocks, M, 16).sum(axis=2) + 0.5 * (
+        (np.abs(Xf) @ P) ** 2
+    ).reshape(n_blocks, M, 16).sum(axis=2)
+    topK = 8
+    top_idx = np.argpartition(-scores, topK, axis=1)[:, :topK]
+    top_idx[:, 0] = 0  # always include identity (Entry 0) so no block ever regresses
+
+    best_sse = np.full(n_blocks, np.inf, dtype=np.float64)
+    best_m = np.zeros(n_blocks, dtype=np.int64)
+    for k_i in range(topK):
+        m_col = top_idx[:, k_i]
+        ords = table_orders[m_col]
+        Xk = apply_order(Xo, ords)
+        Wk = None if Wo is None else apply_order(Wo, ords)
+        xh = proxy(Xk, Wk)
+        err = Xk - xh
+        sse_k = np.sum(err * err if Wk is None else Wk * err * err, axis=1)
+        better = sse_k < best_sse
+        best_sse = np.where(better, sse_k, best_sse)
+        best_m = np.where(better, m_col, best_m)
+
+    return table_orders[best_m], table_labels[best_m], rounds
 
 
 def q2k_proxy_cfg(clamp_val: float, cfg: HPCConfig):