Upload 3 files

gguf.py

@@ -26,8 +26,16 @@ from .q2k import QK_K, blocks_to_bytes, dequantize_params_batch, unpack_codes
 from .q3_k import Q3_K_BLOCK_BYTES, QK3_K, dequantize_q3_k_bytes, quantize_tensor_q3_k
 from .q4_0 import Q4_0_BLOCK_BYTES, QK4_0, quantize_tensor_q4_0
 from .q1_0 import QK1_0, Q1_0_BLOCK_BYTES, quantize_tensor_q1_0
+from .segmentation import (
+    label_channel_bytes,
+    pack_label_channel,
+    recover_order,
+    scatter_to_original,
+    unpack_label_channel,
+)
 
 GGUF_MAGIC = 0x46554747
+SEG_KV_PREFIX = "HPCQ.SEG.V1."
 DEFAULT_ALIGNMENT = 32
 
 GGML_TYPE_F32 = 0
@@ -274,10 +282,15 @@ class GGUFReader:
 
     def read_tensor(self, name: str) -> tuple[np.ndarray, int]:
         raw = self.read_tensor_raw(name)
-        return _dequantize(raw, self.info(name)), self.info(name).row_width
+        x = _dequantize(raw, self.info(name))
+        return self._unsegment(x, name), self.info(name).row_width
 
     def iter_tensor_row_chunks(self, name: str, max_elements: int = 8_388_608):
-        """Stream a tensor in row-aligned float32 chunks to bound peak RAM."""
+        """Stream a tensor in row-aligned float32 chunks to bound peak RAM.
+
+        Chunks are unsegmented: if the tensor carries an HPCQ label channel,
+        each yielded chunk is scattered back to original order.
+        """
         ti = self.info(name)
         row = ti.row_width
         blk = BLOCK_SIZE[ti.type]
@@ -285,6 +298,10 @@ class GGUFReader:
             x, r = self.read_tensor(name)
             yield x, r, 0, ti.n_elements // max(r, 1)
             return
+        seg_2d = None
+        seg = self.seg_labels(name)
+        if seg is not None:
+            seg_2d = seg[0]
         n_rows = ti.n_elements // row
         rows_per_chunk = max(1, max_elements // row)
         bytes_per_row = (row // blk) * BLOCK_BYTES[ti.type]
@@ -297,8 +314,57 @@ class GGUFReader:
                 raw = f.read(chunk_bytes)
                 if len(raw) != chunk_bytes:
                     raise ValueError(f"{name}: truncated tensor data")
-                sub_ti = TensorInfo(ti.name, (chunk_elems,), ti.type, 0)
-                yield _dequantize(raw, sub_ti), row, r0, r1
+                sub = _dequantize(raw, TensorInfo(ti.name, (chunk_elems,), ti.type, 0))
+                if seg_2d is not None:
+                    b0 = r0 * (row // blk)
+                    b1 = b0 + (r1 - r0) * (row // blk)
+                    sub = _scatter_rows(sub, seg_2d[b0:b1])
+                yield sub, row, r0, r1
+
+    def _seg_content(self, name: str) -> bytes | None:
+        """Raw packed label channel bytes for ``name`` (undecoded KV value)."""
+        for key, vtype, raw in self.kv_pairs:
+            if key == SEG_KV_PREFIX + name:
+                if vtype != 8:
+                    raise ValueError(f"{name}: label channel KV has non-string type {vtype}")
+                if len(raw) < 8:
+                    raise ValueError(f"{name}: corrupt label channel KV")
+                return raw[8:]
+        return None
+
+    def seg_labels(self, name: str) -> tuple[np.ndarray, int] | None:
+        """Decode a tensor's label channel (labels per ORIGINAL position).
+
+        Returns ``(labels_2d, k)`` where ``labels_2d`` is shaped
+        ``(n_blocks, block_qk)``, or ``None`` when no channel is present.
+        """
+        content = self._seg_content(name)
+        if content is None:
+            return None
+        ti = self.info(name)
+        if ti.type not in (GGML_TYPE_Q2_K, GGML_TYPE_Q3_K):
+            raise ValueError(f"{name}: label channel on non-segmentable type {ti.type_name}")
+        blk = QK_K if ti.type == GGML_TYPE_Q2_K else QK3_K
+        k, labels = unpack_label_channel(content)
+        if labels.size != ti.n_elements:
+            raise ValueError(
+                f"{name}: label channel size {labels.size} != tensor elements {ti.n_elements}"
+            )
+        return labels.reshape(ti.n_elements // blk, blk), k
+
+    def _unsegment(self, x: np.ndarray, name: str) -> np.ndarray:
+        seg = self.seg_labels(name)
+        if seg is None:
+            return x
+        return _scatter_rows(x, seg[0])
+
+
+def _scatter_rows(x: np.ndarray, labels_2d: np.ndarray) -> np.ndarray:
+    """Scatter slot-order values back to original positions via the channel."""
+    out = scatter_to_original(
+        x.reshape(labels_2d.shape).astype(np.float32), recover_order(labels_2d)
+    )
+    return out.reshape(-1)
 
 
 def _dequantize(raw: bytes, ti: TensorInfo) -> np.ndarray:
@@ -718,6 +784,8 @@ def write_quantized_gguf(
     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
+    seg_active = bool(getattr(cfg, "segment", False))
+    seg_k = max(2, int(getattr(cfg, "segment_k", 16)))
 
     plan: list[str | None] = []
     quantized_count = 0
@@ -812,11 +880,24 @@ def write_quantized_gguf(
     ftype_code = 11 if has_q3k else (10 if has_q2k else (40 if has_q1_0 else 2))
     kv_pairs = []
     for key, vtype, raw in reader.kv_pairs:
+        if key.startswith(SEG_KV_PREFIX):
+            continue
         if quantized_count and key == "general.file_type" and vtype == 4:
             raw = struct.pack("<I", ftype_code)
         elif quantized_count and key == "general.quantization_version" and vtype == 4:
             raw = struct.pack("<I", 2)
         kv_pairs.append((key, vtype, raw))
+    # Reserve label channels for every planned segmented tensor so the KV
+    # section layout is final before data streaming; contents are patched
+    # in place (same length) after each tensor is quantized.
+    seg_reserve: dict[str, int] = {}
+    if seg_active:
+        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)
+                seg_reserve[key] = content_len
+                kv_pairs.append((key, 8, struct.pack("<Q", content_len) + bytes(content_len)))
 
     output_path = Path(output_path)
     output_path.parent.mkdir(parents=True, exist_ok=True)
@@ -831,10 +912,22 @@ def write_quantized_gguf(
     model_t0 = time.time()
     with output_path.open("wb") as fout:
         fout.write(struct.pack("<IIQQ", GGUF_MAGIC, reader.version, len(out_infos), len(kv_pairs)))
+        seg_patch_offsets: dict[str, int] = {}
         for key, vtype, raw in kv_pairs:
             _write_string(fout, key)
             fout.write(struct.pack("<I", vtype))
+            if key in seg_reserve:
+                seg_patch_offsets[key] = fout.tell()
             fout.write(raw)
+
+        def _patch_seg_kv(fout, key: str, labels: np.ndarray) -> None:
+            packed = pack_label_channel(labels, seg_k)
+            if len(packed) != seg_reserve.get(key, -1):
+                raise RuntimeError(f"{key}: label channel size {len(packed)} != reserved")
+            pos = fout.tell()
+            fout.seek(seg_patch_offsets[key] + 8)  # skip the u64 value length prefix
+            fout.write(packed)
+            fout.seek(pos)
         for ti, out_type, out_size, offset in out_infos:
             _write_string(fout, ti.name)
             fout.write(struct.pack("<I", len(ti.dims)))
@@ -918,6 +1011,8 @@ def write_quantized_gguf(
                 if len(payload) != out_size:
                     raise RuntimeError(f"{ti.name}: Q2_K payload size mismatch")
                 fout.write(payload)
+                if seg_active and res.segment_labels is not None:
+                    _patch_seg_kv(fout, SEG_KV_PREFIX + ti.name, res.segment_labels)
                 total_rmse += res.stats.rmse * res.stats.rmse * x.size
                 total_elements += x.size
                 total_quantized += 1
@@ -945,14 +1040,25 @@ def write_quantized_gguf(
                 sse_q3 = 0.0
                 done_q3_blks = 0
                 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):
                     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
-                    q3_bytes, nb_chunk, chunk_rmse = quantize_tensor_q3_k(
-                        x_chunk, importance=imp_chunk, cfg=cfg, chunk_blocks=8192
-                    )
+                    if seg_active:
+                        q3_bytes, nb_chunk, chunk_rmse, chunk_labels = quantize_tensor_q3_k(
+                            x_chunk,
+                            importance=imp_chunk,
+                            cfg=cfg,
+                            chunk_blocks=8192,
+                            return_labels=True,
+                        )
+                        q3_labels.append(np.asarray(chunk_labels, dtype=np.int64).ravel())
+                    else:
+                        q3_bytes, nb_chunk, chunk_rmse = quantize_tensor_q3_k(
+                            x_chunk, importance=imp_chunk, cfg=cfg, chunk_blocks=8192
+                        )
                     fout.write(q3_bytes)
                     written_bytes += len(q3_bytes)
                     sse_q3 += (chunk_rmse * chunk_rmse) * x_chunk.size
@@ -988,6 +1094,8 @@ def write_quantized_gguf(
 
                 if written_bytes != out_size:
                     raise RuntimeError(f"{ti.name}: Q3_K payload size mismatch")
+                if seg_active:
+                    _patch_seg_kv(fout, SEG_KV_PREFIX + ti.name, np.concatenate(q3_labels))
                 completed_planned_blocks += eff_tb
                 q3_rmse = float(np.sqrt(sse_q3 / ti.n_elements))
                 total_rmse += sse_q3

q3_k.py

@@ -983,12 +983,17 @@ def quantize_tensor_q3_k(
     importance: np.ndarray | None = None,
     cfg: HPCConfig | None = None,
     chunk_blocks: int | None = None,
-) -> tuple[bytes, int, float]:
+    return_labels: bool = False,
+):
     """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.
+
+    With ``return_labels=True`` returns ``(payload, n_blocks, rmse, labels)``
+    where ``labels`` is the segmentation label channel (one band id per
+    original position, or None when segmentation is off).
     """
     cfg = cfg or HPCConfig()
     x = np.asarray(weights, dtype=np.float32).ravel()
@@ -1016,6 +1021,10 @@ def quantize_tensor_q3_k(
             W_raw = np.ascontiguousarray(
                 apply_order(W_raw.astype(np.float64), seg_order)
             ).astype(np.float32)
+    else:
+        seg_order = None
+        seg_labels = None
+        seg_rounds = 0
 
     cuda_ok = cfg.device != "cpu" and cuda_is_available()
     use_cuda = cuda_ok and (
@@ -1066,4 +1075,7 @@ def quantize_tensor_q3_k(
             seg_rounds, seg_flips, rmse,
         )
 
-    return b"".join(chunks), n_blocks, rmse
+    payload = b"".join(chunks)
+    if return_labels:
+        return payload, n_blocks, rmse, seg_labels
+    return payload, n_blocks, rmse

segmentation.py

@@ -110,6 +110,62 @@ def scatter_to_original(xh_slots: np.ndarray, order: np.ndarray) -> np.ndarray:
     return out
 
 
+# ---- 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).
+
+import struct as _struct
+
+SEG_MAGIC = b"HQS1"
+
+
+def label_channel_bytes(n: int, k: int) -> int:
+    """Exact content size of the packed label channel."""
+    if int(k) <= 16:
+        return 10 + (n + 1) // 2
+    return 10 + n
+
+
+def pack_label_channel(labels: np.ndarray, k: int) -> bytes:
+    flat = np.ascontiguousarray(labels, dtype=np.int64).ravel()
+    n = int(flat.size)
+    k = int(k)
+    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()
+    else:
+        enc = 1
+        if flat.size and (int(flat.max()) > 255 or int(flat.min()) < 0):
+            raise ValueError("labels must be in [0, 255] for 8-bit packing")
+        data = flat.astype(np.uint8).tobytes()
+    return SEG_MAGIC + bytes([k, enc]) + _struct.pack("<I", n) + data
+
+
+def unpack_label_channel(content: bytes) -> tuple[int, np.ndarray]:
+    if content[:4] != SEG_MAGIC:
+        raise ValueError(f"bad label channel magic {content[:4]!r}")
+    k = int(content[4])
+    enc = int(content[5])
+    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)
+    else:
+        out = np.frombuffer(data, dtype=np.uint8).copy()
+    return k, out[:n]
+
+
 def segment_stage(
     X: np.ndarray,
     W: np.ndarray | None,