Upload segmentation.py

segmentation.py

@@ -269,65 +269,128 @@ def segment_stage(
 ) -> tuple[np.ndarray, np.ndarray, int]:
     """Run the label fixed-point loop.
 
-    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).
+    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.
     """
-    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))
-    n_blocks, B = Xo.shape
+    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)
 
-    table_labels = np.empty((M, 256), dtype=np.int64)
+    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)
-    # 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)
+    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: 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)
+    n_proto = M - 3
+    proto_idx = np.linspace(0, n_blocks - 1, n_proto, 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_proto):
+        table_labels[3 + i] = _order_to_canonical_labels(ord_p[i], k).astype(np.uint8)
 
-    table_orders = make_order(table_labels)  # (M, 256)
+    table_orders = make_order(table_labels)  # (M, 256) int64
 
-    # Prescreen top-8 table entries per block via between-subblock variance GEMM
+    # 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
-    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)
+
+    topK = min(4, M)
     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)
+    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]
 
     return table_orders[best_m], table_labels[best_m], rounds