Delete segmentation.py

segmentation.py

@@ -1,480 +0,0 @@
-"""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
-
-
-# ---- 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; full per-block or small tensors)
-#   enc=1: stores raw uint8 labels (k > 16 only)
-#   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
-
-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, 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 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, table_size: int = SEG_TABLE_SIZE) -> bytes:
-    k = int(k)
-    table_size = int(table_size)
-    # Fast path when segment_stage attached pre-packed LUT data
-    lut_cache = getattr(labels, "_hpcq_lut_bytes", None)
-    if lut_cache is not None and lut_cache[0] == (k, table_size):
-        return lut_cache[1]
-
-    flat = np.ascontiguousarray(labels, dtype=np.uint8).ravel()
-    n = int(flat.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 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)
-            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((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(idx_dt)
-            else:
-                top_u = np.argsort(-counts)[:table_size]
-                table[:] = uniq_rows[top_u]
-                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 = _struct.pack("<H", table_size) + table.tobytes() + block_idx.tobytes()
-        else:
-            enc = 0
-            data = _pack_nibbles_1d(flat)
-    else:
-        enc = 1
-        if flat.size and (int(labels.max()) > 255 or int(labels.min()) < 0):
-            raise ValueError("labels must be in [0, 255] for 8-bit packing")
-        data = flat.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:
-        out = _unpack_nibbles_1d(data, n)
-    elif enc in (2, 3):
-        n_blocks = n // 256
-        tsize = _struct.unpack("<H", data[:2])[0]
-        tbytes = tsize * 128
-        table = _unpack_nibbles_1d(data[2 : 2 + tbytes], tsize * 256).reshape(
-            tsize, 256
-        )
-        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()
-    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
-
-
-class _LUTLabels(np.ndarray):
-    _hpcq_lut_bytes: tuple[tuple[int, int], bytes] | None = None
-
-
-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.
-
-    For tensors with ``n_blocks <= M`` (or ``k > 16``), refines a per-block
-    label vector directly.  For larger tensors (``n_blocks > M, k <= 16``),
-    constructs an ``M``-entry generalized permutation lookup table (Entry 0 =
-    natural order, Entries 1..2 = column RMS/abs profiles, Entries 3..M-1 =
-    prototype blocks refined by ``_segment_stage_dense``) and assigns each
-    superblock its optimal table index in chunked float32 / CUDA batches.
-    """
-    k = max(1, int(k))
-    n_blocks, B = X.shape
-    M = int(getattr(cfg, "segment_table_size", SEG_TABLE_SIZE))
-    if M <= 0 or n_blocks <= M or B != 256 or k > 16:
-        Xo = np.asarray(X, dtype=np.float64)
-        Wo = None if W is None else np.asarray(W, dtype=np.float64)
-        return _segment_stage_dense(Xo, Wo, cfg, proxy, k)
-
-    Xf = np.ascontiguousarray(X, dtype=np.float32)
-    table_labels = np.empty((M, 256), dtype=np.uint8)
-    # 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).astype(np.uint8)
-    # Entry 1 & 2: column-wide RMS and abs profiles (sampled across up to 4096 blocks)
-    s_step = max(1, n_blocks // 4096)
-    X_sample = Xf[::s_step].astype(np.float64)
-    col_rms = np.sqrt(np.mean(X_sample * X_sample, axis=0))
-    table_labels[1] = _order_to_canonical_labels(np.argsort(col_rms, kind="stable"), k).astype(np.uint8)
-    col_abs = np.mean(np.abs(X_sample), axis=0)
-    table_labels[2] = _order_to_canonical_labels(np.argsort(col_abs, kind="stable"), k).astype(np.uint8)
-
-    # Entries 3..M-1: up to 253 prototype blocks refined with the proxy loop,
-    # plus fast direct band_labels prototypes when M > 256
-    n_need = M - 3
-    n_dense = min(n_need, 253)
-    proto_idx = np.linspace(0, n_blocks - 1, n_dense, dtype=np.int64)
-    Xp = Xf[proto_idx].astype(np.float64)
-    Wp = None if W is None else np.asarray(W[proto_idx], dtype=np.float64)
-    ord_p, _, rounds = _segment_stage_dense(Xp, Wp, cfg, proxy, k)
-    for i in range(n_dense):
-        table_labels[3 + i] = _order_to_canonical_labels(ord_p[i], k).astype(np.uint8)
-    if n_need > n_dense:
-        extra_idx = np.linspace(1, n_blocks - 2, n_need - n_dense, dtype=np.int64)
-        ord_e = make_order(band_labels(Xf[extra_idx].astype(np.float64), k))
-        for i in range(n_need - n_dense):
-            table_labels[3 + n_dense + i] = _order_to_canonical_labels(ord_e[i], k).astype(np.uint8)
-
-    table_orders = make_order(table_labels)  # (M, 256) int64
-
-    # Build sub-block projection matrix P: (256, M * 16)
-    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
-
-    topK = min(4, M)
-    best_m = np.zeros(n_blocks, dtype=np.int64)
-    nmax_sub = 3.0 if str(getattr(cfg, "quant_format", "q2_k")) != "q3_k" else 7.0
-
-    use_torch_cuda = False
-    if getattr(cfg, "device", "auto") != "cpu":
-        try:
-            import torch
-            use_torch_cuda = torch.cuda.is_available()
-        except Exception:
-            use_torch_cuda = False
-
-    if use_torch_cuda:
-        import torch
-        dev_str = "cuda" if cfg.device == "auto" else cfg.device
-        dev = torch.device(dev_str)
-        P_t = torch.from_numpy(P).to(dev)
-        orders_t = torch.from_numpy(table_orders).to(dev)
-        chunk = 16384
-        with torch.no_grad():
-            for s in range(0, n_blocks, chunk):
-                e = min(s + chunk, n_blocks)
-                xb = torch.from_numpy(Xf[s:e]).to(dev)
-                sc = (xb @ P_t).square().view(-1, M, 16).sum(dim=2) + 0.5 * (
-                    (xb.abs() @ P_t).square().view(-1, M, 16).sum(dim=2)
-                )
-                top_idx = torch.topk(sc, k=topK, dim=1).indices
-                top_idx[:, 0] = 0
-                ords = orders_t[top_idx]
-                u = torch.gather(xb.unsqueeze(1).expand(-1, topK, -1), 2, ords).view(
-                    -1, topK, 16, 16
-                )
-                umin = torch.clamp_max(u.amin(dim=3, keepdim=True), 0.0)
-                umax = u.amax(dim=3, keepdim=True)
-                step = torch.clamp_min((umax - umin) * (1.0 / nmax_sub), 1.0e-30)
-                q = torch.clamp(torch.round((u - umin) / step), 0.0, nmax_sub)
-                err = u - (umin + step * q)
-                if W is not None:
-                    wb = torch.from_numpy(np.ascontiguousarray(W[s:e], dtype=np.float32)).to(dev)
-                    wg = torch.gather(wb.unsqueeze(1).expand(-1, topK, -1), 2, ords).view(
-                        -1, topK, 16, 16
-                    )
-                    sse = (wg * err.square()).sum(dim=(2, 3))
-                else:
-                    sse = err.square().sum(dim=(2, 3))
-                best_local = sse.argmin(dim=1)
-                best_m[s:e] = (
-                    top_idx[torch.arange(e - s, device=dev), best_local].cpu().numpy()
-                )
-    else:
-        chunk = 8192
-        for s in range(0, n_blocks, chunk):
-            e = min(s + chunk, n_blocks)
-            xb = Xf[s:e]
-            sc = ((xb @ P) ** 2).reshape(e - s, M, 16).sum(axis=2) + 0.5 * (
-                (np.abs(xb) @ P) ** 2
-            ).reshape(e - s, M, 16).sum(axis=2)
-            top_idx = np.argpartition(-sc, topK - 1, axis=1)[:, :topK]
-            top_idx[:, 0] = 0
-            ords = table_orders[top_idx]
-            u = np.take_along_axis(xb[:, None, :], ords, axis=2).reshape(
-                e - s, topK, 16, 16
-            )
-            umin = np.minimum(u.min(axis=3, keepdims=True), 0.0)
-            umax = u.max(axis=3, keepdims=True)
-            step = np.maximum((umax - umin) * (1.0 / nmax_sub), 1.0e-30)
-            q = np.clip(np.rint((u - umin) / step), 0.0, nmax_sub)
-            err = u - (umin + step * q)
-            if W is not None:
-                wb = np.ascontiguousarray(W[s:e], dtype=np.float32)
-                wg = np.take_along_axis(wb[:, None, :], ords, axis=2).reshape(
-                    e - s, topK, 16, 16
-                )
-                sse = np.sum(wg * err * err, axis=(2, 3))
-            else:
-                sse = np.sum(err * err, axis=(2, 3))
-            best_local = np.argmin(sse, axis=1)
-            best_m[s:e] = top_idx[np.arange(e - s), best_local]
-
-    out_labels = table_labels[best_m].view(_LUTLabels)
-    enc = 2 if M <= 256 else 3
-    idx_dt = np.uint8 if enc == 2 else np.dtype("<u2")
-    lo = table_labels[:, 0::2]
-    hi = table_labels[:, 1::2]
-    table_nib = (lo | (hi << 4)).astype(np.uint8)
-    raw_lut = (
-        SEG_MAGIC
-        + bytes([k, enc])
-        + _struct.pack("<I", n_blocks * 256)
-        + _struct.pack("<H", M)
-        + table_nib.tobytes()
-        + best_m.astype(idx_dt).tobytes()
-    )
-    out_labels._hpcq_lut_bytes = ((k, M), raw_lut)
-    return table_orders[best_m], out_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
-    from .pipeline import _phase2_wls
-
-    def _proxy(xg: np.ndarray, wg: np.ndarray | None) -> np.ndarray:
-        seeds = build_seeds(xg, wg, cfg, clamp_val)
-        d, m = _phase2_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 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)
-    # make_qkx2_batch returns (scale, -min, codes) with the GGML positive-dmin
-    # convention; reconstruction is scale * code - (-min).
-    scale, neg_min, codes = make_qkx2_batch(u, wu, nmax=7)
-    deq = scale[:, None] * codes - neg_min[:, 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)