Highly experimental

config.py

@@ -103,6 +103,17 @@ class HPCConfig:
     # (the Q2_K closed-form pin never spends any).  0 disables the expense.
     pin_budget: float = 0.05
 
+    # ---- self-describing segmentation (V2-only, opt-in) -----------------
+    # Per-superblock deterministic band clustering.  Labels are a pure
+    # function of the block's own dequantized values, so the decoder
+    # recovers the slot->position permutation by recomputation -- no
+    # sidecar, GGUF layout untouched.  Q2_K / Q3_K use segment_k bands
+    # over the 256-weight superblock; Q4_0's 32-wide unit block uses k=1
+    # (ascending value order).
+    segment: bool = False
+    segment_k: int = 16
+    segment_iters: int = 4
+
     # ---- seed clamp -----------------------------------------------------
     clamp_kurtosis_off: float = 6.0
     clamp_kurtosis_soft: float = 20.0
@@ -160,6 +171,11 @@ class HPCConfig:
             raise ValueError("therm_t_decay must lie in (0, 1)")
         if self.therm_lambda_tilt < 1.0:
             raise ValueError("therm_lambda_tilt must be >= 1.0 (1.0 = off)")
+        if self.segment:
+            if not (2 <= self.segment_k <= 256):
+                raise ValueError("segment_k must lie in [2, 256]")
+            if self.segment_iters < 1:
+                raise ValueError("segment_iters must be >= 1")
         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"):

pipeline.py

@@ -2,6 +2,7 @@
 
 Pipeline stages (mirroring the C engine, with V2 fixes):
 
+0. self-describing per-superblock segmentation (opt-in, ``cfg.segment``);
 1. kurtosis-gated outlier clamp and batched GGML seeds;
 2. vectorised WLS superblock solve;
 3. factored candidate field + standalone factored floor;
@@ -72,6 +73,8 @@ class TensorStats:
     factor_floor_used: int
     candidate_floor_used: int
     therm_improved: int = 0
+    segment_rounds: int = 0
+    segment_flips: int = 0
 
 
 @dataclass(slots=True)
@@ -84,10 +87,22 @@ class QuantizedTensor:
     factor_objectives: np.ndarray
     config: HPCConfig
     stats: TensorStats
+    segment_perm: np.ndarray | None = None
+    segment_labels: np.ndarray | None = None
 
     def dequantize(self) -> np.ndarray:
         return np.concatenate([b.dequantize() for b in self.blocks])
 
+    def dequantize_original(self) -> np.ndarray:
+        """Consumer-side reconstruction: scatter dequantized slots back to
+        original positions using the stored label channel (exact)."""
+        if self.segment_perm is None:
+            return self.dequantize()
+        Xh = np.stack([b.dequantize() for b in self.blocks])
+        out = np.empty_like(Xh)
+        np.put_along_axis(out, self.segment_perm, Xh, axis=1)
+        return out.ravel()
+
 
 def _phase2_wls(X, W, seeds, cfg: HPCConfig):
     n = X.shape[0]
@@ -864,6 +879,33 @@ class Q2KQuantizer:
         x = x.ravel()
         n_blocks = x.size // QK_K
 
+        seg_perm = None
+        seg_labels = None
+        seg_rounds = 0
+        seg_flips = 0
+
+        if cfg.segment:
+            from .segmentation import apply_order, q2k_proxy_cfg, segment_stage
+
+            X_p = x.reshape(n_blocks, QK_K)
+            imp2d = (
+                None
+                if importance is None
+                else np.asarray(importance, dtype=np.float32).ravel().reshape(n_blocks, QK_K)
+            )
+            clamp_val0 = float(seed_clamp_value(x, cfg))
+            proxy = q2k_proxy_cfg(clamp_val0, cfg)
+            seg_perm, seg_labels, seg_rounds = segment_stage(
+                X_p, imp2d, cfg, proxy, max(2, int(cfg.segment_k))
+            )
+            X_p = apply_order(X_p, seg_perm).astype(np.float32)
+            x = np.ascontiguousarray(X_p.ravel())
+            if imp2d is not None:
+                importance = np.ascontiguousarray(
+                    apply_order(imp2d, seg_perm).astype(np.float32).ravel()
+                )
+            X_p = None
+
         if cfg.device != "cpu":
             from .cuda_backend import cuda_is_available, quantize_cuda
 
@@ -872,7 +914,7 @@ class Q2KQuantizer:
             ):
                 dev_str = "cuda" if cfg.device == "auto" else cfg.device
                 res = quantize_cuda(
-                    weights,
+                    x if seg_labels is not None else weights,
                     importance,
                     row_width,
                     cfg,
@@ -1199,6 +1241,16 @@ class Q2KQuantizer:
             else:
                 weighted += float(np.dot(W[i].astype(np.float64), e * e))
 
+        if seg_perm is not None:
+            # informational: how many labels a channelless decoder would
+            # mis-derive from the stored slot values
+            from .segmentation import band_labels
+
+            Xh = np.stack([b.dequantize() for b in blocks]).astype(np.float64)
+            labels_slot = np.take_along_axis(seg_labels, seg_perm, axis=1)
+            labels_rec = band_labels(Xh, max(2, int(cfg.segment_k)))
+            seg_flips = int(np.count_nonzero(labels_rec != labels_slot))
+
         rmse = float(np.sqrt(total_sse / x.size))
         stats = TensorStats(
             n_blocks=n_blocks,
@@ -1211,6 +1263,8 @@ class Q2KQuantizer:
             factor_floor_used=factor_floor_used,
             candidate_floor_used=candidate_floor_used,
             therm_improved=therm_improved,
+            segment_rounds=seg_rounds,
+            segment_flips=seg_flips,
         )
         _emit(
             "Phase 5: Tensor Complete",
@@ -1237,6 +1291,8 @@ class Q2KQuantizer:
             factor_objectives=factor_objectives,
             config=cfg,
             stats=stats,
+            segment_perm=seg_perm,
+            segment_labels=seg_labels,
         )
         if cfg.verbose:
             _report(stats)
@@ -1258,6 +1314,8 @@ def _stack_blocks(blocks):
 
 def _report(stats: TensorStats) -> None:
     extra = f" therm={stats.therm_improved}" if stats.therm_improved else ""
+    if stats.segment_rounds:
+        extra += f" seg(rounds={stats.segment_rounds}, flips={stats.segment_flips})"
     print(
         f"Q2_K V2: blocks={stats.n_blocks} rmse={stats.rmse:.6e} "
         f"rmse/sigma={stats.rmse_over_sigma:.5f} temp={stats.temperature:.3e} "

q3_k.py

@@ -1005,6 +1005,18 @@ def quantize_tensor_q3_k(
     else:
         W_raw = None
 
+    if cfg.segment:
+        from .segmentation import apply_order, band_labels, q3k_proxy, segment_stage
+
+        seg_order, seg_labels, seg_rounds = segment_stage(
+            X, W_raw, cfg, q3k_proxy, max(2, int(cfg.segment_k))
+        )
+        X = np.ascontiguousarray(apply_order(X.astype(np.float64), seg_order)).astype(np.float32)
+        if W_raw is not None:
+            W_raw = np.ascontiguousarray(
+                apply_order(W_raw.astype(np.float64), seg_order)
+            ).astype(np.float32)
+
     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)
@@ -1041,4 +1053,17 @@ def quantize_tensor_q3_k(
         total_sse += sse_b
 
     rmse = float(np.sqrt(total_sse / x.size))
+
+    if cfg.segment:
+        xh = np.asarray(dequantize_q3_k_bytes(b"".join(chunks), n_blocks), dtype=np.float64).reshape(
+            n_blocks, QK3_K
+        )
+        labels_slot = np.take_along_axis(seg_labels, seg_order, axis=1)
+        labels_rec = band_labels(xh, max(2, int(cfg.segment_k)))
+        seg_flips = int(np.count_nonzero(labels_rec != labels_slot))
+        _log.info(
+            "Q3_K segmentation: rounds=%d channelless flips=%d rmse=%.6e",
+            seg_rounds, seg_flips, rmse,
+        )
+
     return b"".join(chunks), n_blocks, rmse

q4_0.py

@@ -28,6 +28,8 @@ per 18-byte block = 4.5 bpw):
 from __future__ import annotations
 
 from math import ceil
+import logging
+
 import numpy as np
 
 try:
@@ -48,6 +50,8 @@ QK4_0 = 32
 Q4_0_BLOCK_BYTES = 18
 Q4_N_CAND = 24
 Q4_N_STATES = 6
+
+_log = logging.getLogger(__name__)
 Q4_N_LEVELS = 16
 
 
@@ -583,6 +587,16 @@ def quantize_tensor_q4_0(
     else:
         W_raw = None
 
+    if cfg.segment:
+        # Q4_0's unit block (32 weights, one scale+min) has no per-band
+        # sub-scale dof to align, position-ascending k=1 layout is the
+        # identity permutation, and the permutation channel would pay 8
+        # bits per weight for nothing.  Explicit no-op for now.
+        _log.info(
+            "Q4_0 segmentation disabled: unit blocks (d, min) have no band dof "
+            "to align (k=1 layout is the identity)"
+        )
+
     use_cuda = cfg.device != "cpu" and cuda_is_available() and (
         cfg.device.startswith("cuda") or (cfg.device == "auto" and n_blocks >= 64)
     )

segmentation.py

@@ -0,0 +1,248 @@
+"""Self-describing per-block segmentation (V2-only stage).
+
+Deterministic cluster labels are recomputed from a block's own DEQUANTIZED
+values with an unweighted K-means (Lloyd, quantile init, centers sorted
+ascending, stable value tie-breaks).  Because the label rule is a pure
+function of the stored slot values, the decoder recovers the identical
+slot->position permutation by recomputation -- no sidecar storage, the
+GGUF block layout is untouched.
+
+Encode loop:
+
+    order_0 = S(x)                      (labels from the originals)
+    x_slot  = x[order_0]
+    order_{n+1} = S(x̂_slot)             (labels from the dequantized view)
+    ... repeat until the permutation is stable; store quantize(x[order]).
+
+Decode:
+
+    x̂_slot = dequantize(payload)
+    order  = S(x̂_slot)
+    x̂_orig = scatter(x̂_slot, order)
+
+Q2_K / Q3_K segment the 256-weight superblock into ``k`` bands aligned with
+the 16-weight sub-scale groups; Q4_0's unit block is 32 wide with a single
+scale, so it uses ``k=1`` (pure ascending value order, which the spectral
+pair solvers can exploit).
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+from .config import HPCConfig
+from .fp16 import nearest_int
+
+_ROW_CHUNK = 4096
+
+
+def _rowwise(fn, v: np.ndarray, chunk: int = _ROW_CHUNK) -> np.ndarray:
+    n = v.shape[0]
+    if n <= chunk:
+        return fn(v)
+    return np.concatenate([fn(v[s : s + chunk]) for s in range(0, n, chunk)], axis=0)
+
+
+def band_labels(values: np.ndarray, k: int, iters: int = 12) -> np.ndarray:
+    """Per-row unweighted K-means labels; centers sorted ascending (canonical)."""
+    v = np.asarray(values, dtype=np.float64)
+    if k <= 1:
+        return np.zeros(v.shape, dtype=np.int64)
+
+    def _run(v: np.ndarray) -> np.ndarray:
+        qs = (np.arange(k) + 0.5) / k
+        c = np.quantile(v, qs, axis=1).T
+        a = None
+        for _ in range(iters):
+            a = np.abs(v[:, :, None] - c[:, None, :]).argmin(axis=2)
+            mask = a[:, :, None] == np.arange(k)[None, None, :]
+            num = (v[:, :, None] * mask).sum(axis=1)
+            cnt = mask.sum(axis=1)
+            c = np.sort(num / np.maximum(cnt, 1), axis=1)
+        return np.abs(v[:, :, None] - c[:, None, :]).argmin(axis=2)
+
+    return _rowwise(_run, v).astype(np.int64)
+
+
+def make_order(labels: np.ndarray, values: np.ndarray | None = None) -> np.ndarray:
+    """Per-row slot->original index: band-major layout, position-ascending
+    inside each band (stable argsort).  With this layout the label channel
+    (one label per ORIGINAL position) inverts the permutation exactly, and
+    per-16 ladder error depends only on band width, not intra-band order.
+    ``values`` is accepted but ignored (kept for call-site compatibility)."""
+
+    def _run(labels: np.ndarray) -> np.ndarray:
+        return np.argsort(labels, axis=1, kind="stable").astype(np.int64)
+
+    n = labels.shape[0]
+    if n <= _ROW_CHUNK:
+        return _run(labels)
+    outs = []
+    for s in range(0, n, _ROW_CHUNK):
+        e = min(s + _ROW_CHUNK, n)
+        outs.append(_run(labels[s:e]))
+    return np.concatenate(outs, axis=0)
+
+
+def apply_order(a: np.ndarray, order: np.ndarray) -> np.ndarray:
+    """Slot view: out[b, s] = a[b, order[b, s]]."""
+    return np.take_along_axis(a, order, axis=1)
+
+
+def invert_order(order: np.ndarray) -> np.ndarray:
+    """Inverse permutation: inv[b, orig] = slot."""
+    inv = np.empty_like(order)
+    rows = np.arange(order.shape[0])[:, None]
+    cols = np.arange(order.shape[1])[None, :]
+    inv[rows, order] = cols
+    return inv
+
+
+def recover_order(labels: np.ndarray) -> np.ndarray:
+    """Exact inversion given a label channel (band-major, position-ascending)."""
+    return make_order(labels)
+
+
+def scatter_to_original(xh_slots: np.ndarray, order: np.ndarray) -> np.ndarray:
+    """Inverse of apply_order: out[b, order[b, s]] = xh_slots[b, s]."""
+    out = np.empty(xh_slots.shape, dtype=xh_slots.dtype)
+    np.put_along_axis(out, order, xh_slots, axis=1)
+    return out
+
+
+def segment_stage(
+    X: np.ndarray,
+    W: np.ndarray | None,
+    cfg: HPCConfig,
+    proxy,
+    k: int,
+) -> 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.
+    """
+    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
+
+
+def q2k_proxy_cfg(clamp_val: float, cfg: HPCConfig):
+    """Cheap phase1+2 proxy for Q2_K label iteration (GGML seeds + WLS)."""
+    from .q2k import codes_from_params_batch, dequantize_params_batch
+    from .seeds import build_seeds
+
+    def _proxy(xg: np.ndarray, wg: np.ndarray | None) -> np.ndarray:
+        seeds = build_seeds(xg, wg, cfg, clamp_val)
+        d, m = _wls(xg, wg, seeds, cfg)
+        sd = np.where(d > 1.0e-15, d, 1.0).astype(np.float32)
+        sm = np.where(m > 1.0e-15, m, 1.0).astype(np.float32)
+        ls = np.where(
+            d[:, None] > 1.0e-15,
+            np.clip(nearest_int(seeds["scales"] / sd[:, None]), 0, 15),
+            0,
+        ).astype(np.uint8)
+        lm = np.where(
+            m[:, None] > 1.0e-15,
+            np.clip(nearest_int(seeds["mins"] / sm[:, None]), 0, 15),
+            0,
+        ).astype(np.uint8)
+        codes = codes_from_params_batch(xg, d, m, ls, lm)
+        return dequantize_params_batch(d, m, ls, lm, codes).astype(np.float64)
+
+    return _proxy
+
+
+def _wls(xg, wg, seeds, cfg):
+    """Row-level WLS (same iteration as pipeline._phase2_wls, kept local to
+    avoid a circular import)."""
+    from .fp16 import round_fp16
+
+    n = xg.shape[0]
+    per_blk_clamp = np.zeros((n, 1), dtype=np.float32)
+    clipped = xg.astype(np.float32)
+    base_d = seeds["d"].astype(np.float64).copy()
+    base_m = seeds["dmin"].astype(np.float64).copy()
+    ls = seeds["ls"].copy()
+    lm = seeds["lm"].copy()
+    scales = seeds["scales"].astype(np.float64)
+    mins = seeds["mins"].astype(np.float64)
+    ww = np.ones_like(xg, dtype=np.float64) if wg is None else wg.astype(np.float64)
+
+    from .q2k import codes_from_params_batch
+
+    for _ in range(5):
+        codes = codes_from_params_batch(clipped, base_d, base_m, ls, lm).astype(np.float64)
+        A = np.repeat(ls.astype(np.float64), 16, axis=1) * codes
+        B = np.repeat(lm.astype(np.float64), 16, axis=1)
+        Saa = np.sum(ww * A * A, axis=1)
+        Sab = np.sum(ww * A * B, axis=1)
+        Sbb = np.sum(ww * B * B, axis=1)
+        Sxa = np.sum(ww * clipped.astype(np.float64) * A, axis=1)
+        Sxb = np.sum(ww * clipped.astype(np.float64) * B, axis=1)
+        det = Saa * Sbb - Sab * Sab
+        valid = det > 1.0e-30
+        safe_det = np.where(valid, det, 1.0)
+        d_new = (Sbb * Sxa - Sab * Sxb) / safe_det
+        m_new = (Sab * Sxa - Saa * Sxb) / safe_det
+        upd_d = valid & (d_new > 0.0) & np.isfinite(d_new)
+        upd_m = valid & (m_new > 0.0) & np.isfinite(m_new)
+        base_d = np.where(upd_d, round_fp16(d_new.astype(np.float32)).astype(np.float64), base_d)
+        base_m = np.where(upd_m, round_fp16(m_new.astype(np.float32)).astype(np.float64), base_m)
+        safe_d = np.where(base_d > 1.0e-15, base_d, 1.0)
+        safe_m = np.where(base_m > 1.0e-15, base_m, 1.0)
+        ls = np.where(
+            base_d[:, None] > 1.0e-15,
+            np.clip(nearest_int(scales / safe_d[:, None]), 0, 15),
+            0,
+        ).astype(np.uint8)
+        lm = np.where(
+            base_m[:, None] > 1.0e-15,
+            np.clip(nearest_int(mins / safe_m[:, None]), 0, 15),
+            0,
+        ).astype(np.uint8)
+    return base_d, base_m
+
+
+def q3k_proxy(xg: np.ndarray, wg: np.ndarray | None) -> np.ndarray:
+    """Cheap per-16 sub-block 3-bit ladder proxy for Q3_K label iteration."""
+    from .seeds import make_qkx2_batch
+
+    n, B = xg.shape
+    u = np.asarray(xg, dtype=np.float64).reshape(n * (B // 16), 16)
+    if wg is None:
+        wu = np.ones_like(u)
+    else:
+        wu = np.asarray(wg, dtype=np.float64).reshape(n * (B // 16), 16)
+    scale, xmin, codes = make_qkx2_batch(u, wu, nmax=7)
+    deq = scale[:, None] * codes + xmin[:, None]
+    return deq.reshape(n, B)
+
+
+def q4_0_proxy(xg: np.ndarray, wg: np.ndarray | None) -> np.ndarray:
+    """GGML-reference Q4_0 proxy (single scale + min per 32-block)."""
+    x = np.asarray(xg, dtype=np.float64)
+    xmin = x.min(axis=1, keepdims=True)
+    xmax = x.max(axis=1, keepdims=True)
+    d = np.maximum((xmax - xmin) / 15.0, 1.0e-30)
+    q = np.clip(nearest_int(((x - xmin) / d).astype(np.float32)), 0, 15)
+    return xmin + d * q.astype(np.float64)