Fold prototype

hexstate_requantize.py

@@ -659,436 +659,98 @@ def dequant_q2k_fast(q2k_bytes, n_blocks):
 # across all fold pairs, so the spectral objective's vesica/DC terms
 # can be driven to zero by the unchanged HPC pipeline.
 #
-# Memory-efficient: permutes f32_data IN-PLACE in small chunks (no
-# full-tensor copy).  The argsort runs on ~1 M elements at a time
-# instead of the whole tensor.  RMSE is computed in permuted order
-# (SE is permutation-invariant, so the result is mathematically
-# identical).  Sidecar perm indices are uint8 (indices 0..255).
-
-_FI_CHUNK_BLOCKS = 4096  # blocks per argsort chunk (~1 M elements)
-
-def apply_fold_interleave(f32_data, block_size=QK_K):
-    """Sort-interleave every block IN-PLACE.  Returns the uint8 perm
-    index array (flat, same length as f32_data) for the sidecar.
-    f32_data is modified in-place — no copy is made."""
-    n = len(f32_data)
-    n_blocks = n // block_size
-    half = block_size // 2
-    all_perms = np.empty(n, dtype=np.uint8)
-
-    for cb in range(0, n_blocks, _FI_CHUNK_BLOCKS):
-        ce = min(cb + _FI_CHUNK_BLOCKS, n_blocks)
-        s = cb * block_size
-        e = ce * block_size
-        chunk = f32_data[s:e].reshape(ce - cb, block_size)
-
-        sorted_idx = np.argsort(chunk, axis=1)
-        perm = np.empty_like(sorted_idx)
-        perm[:, :half] = sorted_idx[:, :half]
-        perm[:, half:] = sorted_idx[:, -1:half-1:-1]
-        del sorted_idx
-
-        f32_data[s:e] = np.take_along_axis(chunk, perm, axis=1).reshape(-1)
-        all_perms[s:e] = perm.astype(np.uint8).reshape(-1)
-        del perm, chunk
-
-    return all_perms
-
-
-def apply_fold_interleave_importance(importance, perm_flat, block_size=QK_K):
-    """Permute an importance vector with the same fold permutation.
-    Returns a new array (importance is typically much smaller than
-    weights for tiled imatrix, so the copy is fine)."""
-    n = len(importance)
-    n_blocks = n // block_size
-    blocks = importance.reshape(n_blocks, block_size)
-    perms = perm_flat[:n].reshape(n_blocks, block_size).astype(np.intp)
-    permuted = np.take_along_axis(blocks, perms, axis=1)
-    return permuted.reshape(-1)
-
-
-def is_attention_tensor(name):
-    """Detect attention Q/K/V/O projection tensors.
-    These are the most sensitive to quantization and get promoted to Q4_0."""
-
-# ─── Fold-Basis Change ────────────────────────────────────────────────────
-# Whole-dimension permutations that propagate losslessly through the graph.
-# P_hidden:  one global permutation of the residual-stream dimension.
-#            Permutes columns of everything that reads from it (Q/K/V,
-#            gate, up, embeddings) and rows of everything that writes
-#            to it (O, down).  Layer norms are 1D — just permute entries.
-# P_inter_N: one permutation per layer of the FFN intermediate dimension.
-#            Permutes columns of W_down (the dimension Q2_K blocks span)
-#            and rows of W_gate/W_up (output dim, no block effect).
-#
-# The permutations are chosen by stride-interleaving the columns sorted
-# by L2 energy: the k-th lowest-energy column is paired with the k-th
-# highest at fold-complementary positions (k, k+128) within each QK_K
-# block.  This gives every block balanced fold pairs for vesica/DC
-# cancellation — the same spectral benefit as the per-block sort, but
-# baked into the weight matrix so the output GGUF is self-contained.
-
-import re as _re
-
-def _fb_tensor_perms(name):
-    """Return list of (perm_key, axis) for this tensor.
-    axis: 0 → permute dims[0] (columns/input, what blocks span)
-          1 → permute dims[1] (rows/output)
-          '1d' → 1D parameter (norms)
+# The permutation is a pre-processing step — the quantization algorithm
+# sees 256 floats and runs identically.  A sidecar .npz stores the
+# per-block permutation indices; the inference engine applies the
+# inverse permutation after dequantization (or equivalently permutes
+# the activation vector).
+
+def compute_fold_permutation(block_weights):
+    """Sort-interleave permutation for one QK_K block.
+
+    Returns perm such that permuted[k] = block[perm[k]], with
+    perm[0..127]  = argsort ascending  (smallest half)
+    perm[128..255] = argsort descending (largest half, reversed)
+
+    Every fold pair (k, k+128) therefore contains the k-th smallest
+    and the k-th largest weight, making their sum approximately
+    constant across all 128 pairs.
     """
-    perms = []
-    # Hidden dim as input (cols = dims[0])
-    if _re.match(r'blk\.\d+\.attn_(q|k|v|qkv)\.weight', name) or \
-       _re.match(r'blk\.\d+\.ffn_(gate|up)\.weight', name) or \
-       name in ('token_embd.weight', 'output.weight'):
-        perms.append(('hidden', 0))
-    # Hidden dim as output (rows = dims[1])
-    if _re.match(r'blk\.\d+\.attn_output\.weight', name) or \
-       _re.match(r'blk\.\d+\.ffn_down\.weight', name):
-        perms.append(('hidden', 1))
-    # Hidden dim as 1D
-    if _re.match(r'blk\.\d+\.(attn|ffn)_norm\.weight', name) or \
-       name == 'output_norm.weight':
-        perms.append(('hidden', '1d'))
-    # Intermediate dim as input (cols = dims[0])  → THIS is the big win
-    m = _re.match(r'blk\.(\d+)\.ffn_down\.weight', name)
-    if m:
-        perms.append((f'inter_{m.group(1)}', 0))
-    # Intermediate dim as output (rows = dims[1])
-    m = _re.match(r'blk\.(\d+)\.ffn_(gate|up)\.weight', name)
-    if m:
-        perms.append((f'inter_{m.group(1)}', 1))
-    return perms
-
-
-def _fb_compute_perm(col_stats, block_size=QK_K):
-    """Per-block-group fold-optimal permutation.
-
-    For each group of 256 consecutive columns, sorts by the column
-    statistic and interleaves: k-th smallest at position k, k-th
-    largest at position k+128.  This IS the fold principle — each
-    block group gets its own sort-interleave — while remaining a
-    valid column permutation that propagates through the graph.
-    Stock llama.cpp compatible: no modified dequant, no sidecar."""
-    dim = len(col_stats)
-    n_blocks = dim // block_size
-    half = block_size // 2
-    remainder = dim % block_size
-
-    if n_blocks == 0:
-        return np.arange(dim, dtype=np.intp)
-
-    perm = np.arange(dim, dtype=np.intp)  # identity for remainder
-
-    for b in range(n_blocks):
-        s = b * block_size
-        e = s + block_size
-        group_sorted = np.argsort(col_stats[s:e])
-        perm[s:s + half] = s + group_sorted[:half]
-        perm[s + half:e] = s + group_sorted[-1:half-1:-1]
-
+    n = len(block_weights)
+    half = n // 2
+    s = np.argsort(block_weights)
+    perm = np.empty(n, dtype=np.intp)
+    perm[:half] = s[:half]
+    perm[half:] = s[-1:half-1:-1]
     return perm
 
 
-def _fb_apply(f32, dims, perm, axis):
-    """Apply a permutation to one axis of a tensor (flat f32 array).
-    axis: 0 → dims[0] columns, 1 → dims[1] rows, '1d' → 1D vector.
-    Returns original f32 unchanged if perm length doesn't match the dim."""
-    if axis == '1d':
-        if len(perm) != len(f32):
-            return f32
-        return f32[perm].copy()
-    d0 = int(dims[0])
-    d1 = int(dims[1]) if len(dims) > 1 else 1
-    # Dimension safety: skip if perm doesn't match the target axis
-    if axis == 0 and len(perm) != d0:
-        return f32
-    if axis == 1 and len(perm) != d1:
-        return f32
-    M = f32.reshape(d1, d0)
-    if axis == 0:
-        M = M[:, perm]
-    else:
-        M = M[perm, :]
-    return M.reshape(-1).copy()
+def apply_fold_interleave(f32_data, block_size=QK_K, chunk_blocks=32768):
+    """Apply fold-optimal permutation to every block in a flat tensor.
 
+    Returns (permuted_data, perm_indices) where perm_indices is a flat
+    int16 array of the same length — block b's permutation lives at
+    [b*block_size : (b+1)*block_size].  int16 suffices since indices
+    are 0..255.
 
-# ─── Fold-Hadamard ────────────────────────────────────────────────────────
-# Per-block Walsh-Hadamard transform before quantization.
-#
-# The WHT decorrelates weights within each QK_K block, making
-# quantization errors statistically independent.  Independent errors
-# have zero expected DC and zero expected vesica — the spectral
-# objective is satisfied structurally, not by chasing it.
-#
-# Key properties:
-#   H is its own inverse:  H·H = n·I  →  H⁻¹ = H/n
-#   H preserves L2 norm:   ‖Hx‖₂ = √n·‖x‖₂  (with unnormalised H)
-#   H is deterministic:    defined entirely by block size
-#   Zero storage overhead: no permutation indices, no sidecar
-#
-# Baked into the GGUF: a metadata key signals "Hadamard Q2_K".
-# The inference dequant kernel applies the same transform after
-# unpacking:  deq_orig = (1/n) · H · deq_packed
-#
-# This is QuIP#'s insight applied per-block instead of per-row,
-# composed with the full HPC sieve pipeline.
-
-def _wht_block(block):
-    """In-place unnormalised Walsh-Hadamard transform of a power-of-2 array."""
-    n = len(block)
-    h = 1
-    while h < n:
-        for i in range(0, n, h * 2):
-            for j in range(i, i + h):
-                x = block[j]
-                y = block[j + h]
-                block[j] = x + y
-                block[j + h] = x - y
-        h *= 2
-
-
-def apply_fold_hadamard(f32_data, block_size=QK_K):
-    """Apply normalised per-block WHT to a flat tensor IN-PLACE.
-    Each block is transformed independently.  O(n log n) per block."""
+    Memory-optimized: chunks the vectorised argsort (int64 8 bytes/idx)
+    to avoid a 400 MB full-tensor allocation on 50M-element tensors
+    (Qwen 27B attn_qkv), and permutes in-place so peak is ~1.7x
+    less.  Squared error is permutation-invariant, so RMSE needs no
+    inverse permutation.
+    """
     n = len(f32_data)
     n_blocks = n // block_size
-    scale = 1.0 / np.sqrt(block_size)
-
-    # Vectorised butterfly: process all blocks in parallel
-    blocks = f32_data.reshape(n_blocks, block_size).astype(np.float64)
-    h = 1
-    while h < block_size:
-        # Reshape to expose butterfly pairs
-        view = blocks.reshape(n_blocks, -1, 2 * h)
-        left = view[:, :, :h].copy()
-        right = view[:, :, h:].copy()
-        view[:, :, :h] = left + right
-        view[:, :, h:] = left - right
-        h *= 2
-
-    blocks *= scale
-    f32_data[:] = blocks.astype(np.float32).reshape(-1)
-
-
-# ─── Sort-Quantize-Unsort-Refit ───────────────────────────────────────────
-# The full fold principle baked into standard Q2_K:
-#
-#   1. Sort each block by value (the fold-interleave)
-#   2. Quantize sorted blocks with full HPC pipeline
-#      (sieve, Viterbi, greedy descent, Floyd-Steinberg — everything)
-#   3. Unsort the 2-bit codes back to original positions
-#   4. Refit sub-block scales (Ls, Lm) via least-squares
-#   5. Repack as standard Q2_K — stock llama.cpp compatible
-#
-# The HPC pipeline sees sorted blocks with narrow sub-block ranges
-# and balanced fold pairs.  The code assignments it produces are
-# globally optimal in sorted space.  After unsorting, the codes
-# retain their relative quality — element k still gets the code
-# that the HPC pipeline chose for it.  Only the sub-block scales
-# need refitting because elements moved across sub-block boundaries.
-#
-# Combined with fold-basis (column permutation), the sort-refit
-# captures both coarse structure (consistent across rows) and fine
-# per-block structure (row-specific), all in standard Q2_K format.
-
-# Sub-block index mapping for Q2_K:
-# Element i maps to sub-block j as follows:
-#   half = i // 128
-#   sub  = (i % 128) // 32
-#   k    = (i % 128) % 32
-#   j    = half * 8 + sub * 2 + (k >= 16)
-# Sub-block j covers 16 elements starting at:
-#   base = (j//8)*128 + ((j%8)//2)*32 + 16*(j%2)
-
-_Q2K_SUB_BASES = np.array([
-    (j // 8) * 128 + ((j % 8) // 2) * 32 + 16 * (j % 2)
-    for j in range(16)
-], dtype=np.int32)
-
-
-def _sort_refit_q2k(f32_orig, q2k_sorted_bytes, sort_perms, n_blocks):
-    """Unsort HPC-quantized codes and refit sub-block scales.
-
-    f32_orig:         [n_blocks, 256] original (unsorted) weight values
-    q2k_sorted_bytes: packed Q2_K data from HPC quantization of sorted blocks
-    sort_perms:       [n_blocks, 256] argsort permutations used to sort
-    n_blocks:         number of Q2_K blocks
-
-    Returns repacked Q2_K bytes with unsorted codes and refit scales.
-    Fully vectorized over blocks — only loops over 16 sub-blocks."""
-
-    BLOCK_BYTES = 84
-
-    q2k_arr = np.frombuffer(q2k_sorted_bytes, dtype=np.uint8).reshape(n_blocks, BLOCK_BYTES).copy()
-
-    # ── Extract d, dmin ──
-    d_arr = np.frombuffer(q2k_arr[:, 0:2].tobytes(),
-                          dtype=np.float16).reshape(n_blocks).astype(np.float32)
-    dmin_arr = np.frombuffer(q2k_arr[:, 2:4].tobytes(),
-                             dtype=np.float16).reshape(n_blocks).astype(np.float32)
-
-    # ── Unpack 2-bit codes: [n_blocks, 256] ──
-    qs_raw = q2k_arr[:, 20:84]  # [n_blocks, 64]
-    q_sorted = np.zeros((n_blocks, 256), dtype=np.int32)
-    for h in range(2):
-        qs_half = qs_raw[:, h * 32:(h + 1) * 32]
-        for s in range(4):
-            q_sorted[:, h * 128 + s * 32: h * 128 + (s + 1) * 32] = (qs_half >> (s * 2)) & 3
-
-    # ── Unsort codes ──
-    inv_perms = np.argsort(sort_perms, axis=1)
-    q_orig = np.take_along_axis(q_sorted, inv_perms, axis=1)  # [n_blocks, 256]
-
-    # ── Refit sub-block scales via OLS ──
-    Ls_new = np.zeros((n_blocks, 16), dtype=np.int32)
-    Lm_new = np.zeros((n_blocks, 16), dtype=np.int32)
-
-    for j in range(16):
-        base = int(_Q2K_SUB_BASES[j])
-        x_sub = f32_orig[:, base:base + 16].astype(np.float64)
-        q_sub = q_orig[:, base:base + 16].astype(np.float64)
-
-        Sq = np.sum(q_sub, axis=1)
-        Sqq = np.sum(q_sub ** 2, axis=1)
-        Sx = np.sum(x_sub, axis=1)
-        Sqx = np.sum(q_sub * x_sub, axis=1)
-
-        n_el = 16.0
-        denom = n_el * Sqq - Sq * Sq
-        safe_denom = np.where(np.abs(denom) > 1e-10, denom, 1.0)
-
-        alpha = (n_el * Sqx - Sq * Sx) / safe_denom
-        beta = (Sx - alpha * Sq) / n_el
-
-        safe_d = np.where(d_arr > 1e-20, d_arr, 1.0)
-        safe_dmin = np.where(dmin_arr > 1e-20, dmin_arr, 1.0)
-
-        Ls_f = alpha / safe_d
-        Lm_f = -beta / safe_dmin
-
-        valid = (np.abs(denom) > 1e-10) & (d_arr > 1e-20) & (dmin_arr > 1e-20)
-        Ls_new[:, j] = np.where(valid, np.clip(np.round(Ls_f), 0, 15), 0).astype(np.int32)
-        Lm_new[:, j] = np.where(valid, np.clip(np.round(Lm_f), 0, 15), 0).astype(np.int32)
-
-    # ── Repack scales ──
-    for j in range(16):
-        q2k_arr[:, 4 + j] = (Ls_new[:, j].astype(np.uint8) & 0x0F) | \
-                             ((Lm_new[:, j].astype(np.uint8) & 0x0F) << 4)
-
-    # ── Repack qs with unsorted codes ──
-    new_qs = np.zeros((n_blocks, 64), dtype=np.uint8)
-    for h in range(2):
-        for s in range(4):
-            codes = q_orig[:, h * 128 + s * 32: h * 128 + (s + 1) * 32].astype(np.uint8)
-            new_qs[:, h * 32:(h + 1) * 32] |= codes << (s * 2)
-    q2k_arr[:, 20:84] = new_qs
-
-    return q2k_arr.tobytes()
-
-
-def _fb_prepass(fin, tensor_infos, data_section_start):
-    """Pre-pass: compute fold-basis permutations from weight statistics.
-
-    Memory-lean: reads only a small sample of tensors.
-    - P_hidden:  from the first 3 layers' W_q only (3 tensor reads)
-    - P_inter_N: from W_down[N] only (1 read per layer)
-
-    Each tensor is loaded, its column means computed, and freed before
-    the next is loaded.  Peak memory: one tensor in F32 at a time.
-
-    Returns dict: perm_key → permutation array."""
-    import gc
-    print("  ┌── Fold-Basis pre-pass: computing dimension permutations ──┐")
-
-    perms = {}
-    hidden_stat = None
-    hidden_dim = None
-    hidden_count = 0
-
-    # ── Pass 1: P_hidden from first 3 layers' W_q ──
-    _HIDDEN_SAMPLE = 3
-    for ti in tensor_infos:
-        if hidden_count >= _HIDDEN_SAMPLE:
-            break
-        if not _re.match(r'blk\.\d+\.attn_q\.weight', ti['name']):
-            continue
-        if ti['n_dims'] < 2:
-            continue
-        d0 = int(ti['dims'][0])
-        if d0 % QK_K != 0:
-            continue
-
-        abs_offset = data_section_start + ti['offset']
-        fin.seek(abs_offset)
-        raw = fin.read(ti['data_size'])
-        if ti['type'] == GGML_TYPE_BF16:
-            f32 = bf16_to_f32(raw, ti['n_elements'])
-        elif ti['type'] == GGML_TYPE_F16:
-            f32 = f16_to_f32(raw, ti['n_elements'])
-        elif ti['type'] == GGML_TYPE_F32:
-            f32 = np.frombuffer(raw, dtype=np.float32)
-        else:
-            continue
-
-        d1 = int(ti['dims'][1]) if ti['n_dims'] > 1 else 1
-        M = f32.reshape(d1, d0)
-        col_mean = np.mean(M.astype(np.float64), axis=0)
-        del f32, M, raw
-        gc.collect()
-
-        if hidden_stat is None:
-            hidden_dim = d0
-            hidden_stat = col_mean
-        else:
-            hidden_stat += col_mean
-        hidden_count += 1
-        print(f"  │  P_hidden sample: {ti['name']}")
-
-    if hidden_stat is not None:
-        perms['hidden'] = _fb_compute_perm(hidden_stat)
-        print(f"  │  P_hidden:  dim={hidden_dim}, from {hidden_count} W_q tensors")
-
-    # ── Pass 2: P_inter from each layer's W_down ──
-    n_inter = 0
-    for ti in tensor_infos:
-        m = _re.match(r'blk\.(\d+)\.ffn_down\.weight', ti['name'])
-        if not m:
-            continue
-        if ti['n_dims'] < 2:
-            continue
-        d0 = int(ti['dims'][0])
-        if d0 % QK_K != 0:
-            continue
-
-        layer = m.group(1)
-        abs_offset = data_section_start + ti['offset']
-        fin.seek(abs_offset)
-        raw = fin.read(ti['data_size'])
-        if ti['type'] == GGML_TYPE_BF16:
-            f32 = bf16_to_f32(raw, ti['n_elements'])
-        elif ti['type'] == GGML_TYPE_F16:
-            f32 = f16_to_f32(raw, ti['n_elements'])
-        elif ti['type'] == GGML_TYPE_F32:
-            f32 = np.frombuffer(raw, dtype=np.float32)
-        else:
-            continue
-
-        d1 = int(ti['dims'][1]) if ti['n_dims'] > 1 else 1
-        M = f32.reshape(d1, d0)
-        col_mean = np.mean(M.astype(np.float64), axis=0)
-        del f32, M, raw
-        gc.collect()
-
-        perms[f'inter_{layer}'] = _fb_compute_perm(col_mean)
-        n_inter += 1
-
-    if n_inter > 0:
-        print(f"  │  P_inter:   {n_inter} layers (from W_down column means)")
+    # Ensure contiguous + writeable for in-place chunk updates
+    if not f32_data.flags.writeable or not f32_data.flags.c_contiguous:
+        f32_data = np.ascontiguousarray(f32_data)
+    perm_flat = np.empty(n, dtype=np.int16)
+    blocks = f32_data.reshape(n_blocks, block_size)
+    perm_blocks = perm_flat.reshape(n_blocks, block_size)
+    half = block_size // 2
+    for start in range(0, n_blocks, chunk_blocks):
+        end = min(start + chunk_blocks, n_blocks)
+        chunk = blocks[start:end]  # view into f32_data
+        sorted_idx = np.argsort(chunk, axis=1)  # chunk-sized int64 only
+        perm_chunk = np.empty((end - start, block_size), dtype=np.intp)
+        perm_chunk[:, :half] = sorted_idx[:, :half]
+        perm_chunk[:, half:] = sorted_idx[:, -1:half-1:-1]
+        tmp = np.take_along_axis(chunk, perm_chunk, axis=1)
+        chunk[:, :] = tmp
+        perm_blocks[start:end] = perm_chunk.astype(np.int16, copy=False)
+    return f32_data, perm_flat
+
+
+def apply_fold_interleave_importance(importance, perm_flat, block_size=QK_K, chunk_blocks=32768):
+    """Permute an importance vector with the same fold permutation.
 
-    print(f"  └── {len(perms)} permutations computed ──┘")
-    print()
-    return perms
+    Chunked + in-place: avoids a full-tensor int64 perm copy
+    (400 MB at 50M elements) and a duplicate float32 buffer.
+    """
+    n = len(importance)
+    n_blocks = n // block_size
+    if not importance.flags.writeable or not importance.flags.c_contiguous:
+        importance = np.ascontiguousarray(importance)
+    blocks = importance.reshape(n_blocks, block_size)
+    perm_blocks = perm_flat.reshape(n_blocks, block_size)
+    for start in range(0, n_blocks, chunk_blocks):
+        end = min(start + chunk_blocks, n_blocks)
+        chunk = blocks[start:end]
+        perm_chunk = perm_blocks[start:end].astype(np.intp, copy=False)
+        tmp = np.take_along_axis(chunk, perm_chunk, axis=1)
+        chunk[:, :] = tmp
+    return importance
+
+
+def invert_fold_interleave(data, perm_flat, block_size=QK_K):
+    """Apply the inverse fold permutation to restore original order."""
+    n = len(data)
+    n_blocks = n // block_size
+    blocks = data.reshape(n_blocks, block_size)
+    perms = perm_flat.reshape(n_blocks, block_size).astype(np.intp)
+    inv_perms = np.argsort(perms, axis=1)
+    result = np.take_along_axis(blocks, inv_perms, axis=1)
+    return result.reshape(-1)
 
 
 def is_attention_tensor(name):
@@ -1232,7 +894,7 @@ def main():
     if len(sys.argv) < 3:
         print("Usage: python3 hexstate_requantize.py <input.gguf> <output.gguf>"
               " [--keep-metadata] [--imatrix FILE] [--keep-embd] [--q2all]"
-              " [--fold-interleave] [--fold-basis] [--fold-hadamard]")
+              " [--fold-interleave]")
         sys.exit(1)
 
     input_path = sys.argv[1]
@@ -1242,9 +904,6 @@ def main():
     q2all = '--q2all' in sys.argv
     keep_embd = '--keep-embd' in sys.argv   # keep tied embedding at source precision instead of Q8_0
     fold_interleave = '--fold-interleave' in sys.argv  # fold-optimal weight interleave for DC/vesica cancellation
-    fold_basis = '--fold-basis' in sys.argv  # whole-dimension fold-optimal basis change (baked into weights)
-    fold_hadamard = '--fold-hadamard' in sys.argv  # per-block Walsh-Hadamard transform (baked, zero storage)
-    fold_sort_refit = '--fold-sort-refit' in sys.argv  # sort→HPC→unsort→refit (baked, stock llama compatible)
 
     # Check for imatrix
     imatrix_data = None
@@ -1275,12 +934,6 @@ def main():
         print("  ║  Engine: Python (numpy vectorized)                           ║")
     if fold_interleave:
         print("  ║  Fold Interleave: ON  (sort-pair DC/vesica cancellation)     ║")
-    if fold_basis:
-        print("  ║  Fold Basis:      ON  (whole-dim permutation, baked in)      ║")
-    if fold_hadamard:
-        print("  ║  Fold Hadamard:   ON  (per-block WHT, zero storage)          ║")
-    if fold_sort_refit:
-        print("  ║  Fold Sort-Refit: ON  (sort→HPC→unsort→refit, stock llama)   ║")
     print("  ╚════════════════════════════════════════════════════════════════╝")
     print()
 
@@ -1412,16 +1065,6 @@ def main():
             print(f"  Tensors to keep as-is:          {total_keep}")
         print()
 
-        # ── Fold-Basis pre-pass ──
-        # Compute whole-dimension permutations from weight statistics.
-        # This reads relevant tensors from the input to gather column
-        # energies, then computes the fold-optimal permutation for each
-        # shared dimension.  The permutations are applied during the
-        # main quantization loop below.
-        fb_perms = {}
-        if fold_basis:
-            fb_perms = _fb_prepass(fin, tensor_infos, data_section_start)
-
         # ── Compute output tensor sizes and offsets ──
         out_tensor_infos = []
         out_data_offset = 0
@@ -1566,24 +1209,22 @@ def main():
                 else:
                     updated_kv.append((key, vtype, raw_value))
 
+        # ── Pre-compute fold-interleave permutations for GGUF KV pairs ──
+        # Instead of per-tensor permutations (expensive to pre-compute),
+        # just store a flag indicating fold-interleave was used.
+        # FoldLlama will apply inverse permutation to all Q2_K tensors.
+        if fold_interleave:
+            updated_kv.append(("fold_interleave", 8, struct.pack('<Q', 4) + b"true"))
+
         # ── Write output GGUF ──
         print("  Writing output GGUF...")
-
-        # ── Inject fold-hadamard metadata ──
-        n_kv_out = n_kv
-        if fold_hadamard:
-            # GGUF KV type 8 = BOOL (stored as int8: 1=true)
-            updated_kv.append(('hexstate.fold_hadamard', 8,
-                               struct.pack('<b', 1)))
-            n_kv_out += 1
-            print("  Added hexstate.fold_hadamard=true to GGUF metadata")
-
+        n_kv = len(updated_kv)
         with open(output_path, 'wb') as fout:
             # Header
             fout.write(struct.pack('<I', GGUF_MAGIC))
             fout.write(struct.pack('<I', GGUF_VERSION))
             fout.write(struct.pack('<Q', n_tensors))
-            fout.write(struct.pack('<Q', n_kv_out))
+            fout.write(struct.pack('<Q', n_kv))
 
             # KV pairs (passthrough)
             for key, vtype, raw_value in updated_kv:
@@ -1614,9 +1255,6 @@ def main():
             q2k_rmse_sum = 0.0
             q2k_tensor_count = 0
 
-            # Fold-interleave sidecar: tensor_name → int16 permutation indices
-            fold_perm_sidecar = {} if fold_interleave else None
-
             for i, ti in enumerate(tensor_infos):
                 # Progress bar
                 pct = (i + 1) / n_tensors * 100
@@ -1651,12 +1289,6 @@ def main():
                     n_el = ti['n_elements']
                     n_blocks_q8 = n_el // 32
 
-                    # ── Fold-Basis ──
-                    if fold_basis and fb_perms:
-                        for perm_key, axis in _fb_tensor_perms(ti['name']):
-                            if perm_key in fb_perms:
-                                f32 = _fb_apply(f32, ti['dims'], fb_perms[perm_key], axis)
-
                     if use_hpc and hasattr(_HEXSTATE_LIB, 'hexstate_quantize_tensor_q8_0_hpc'):
                         output_buf = np.zeros(n_blocks_q8 * 34, dtype=np.uint8)
                         error = ctypes.c_float(0.0)
@@ -1702,13 +1334,6 @@ def main():
                         continue
 
                     n_el = len(f32)
-
-                    # ── Fold-Basis ──
-                    if fold_basis and fb_perms:
-                        for perm_key, axis in _fb_tensor_perms(ti['name']):
-                            if perm_key in fb_perms:
-                                f32 = _fb_apply(f32, ti['dims'], fb_perms[perm_key], axis)
-
                     if n_el % 32:
                         raise ValueError(f"Q4_0 tensor {ti['name']} is not row/block aligned")
                     n_blocks_q4 = n_el // 32
@@ -1772,46 +1397,15 @@ def main():
                             fout.write(b'\x00' * pad)
                         continue
 
-                    # ── Fold-Basis pre-processing ───────────────────────
-                    # Apply whole-dimension permutations computed in the
-                    # pre-pass.  These are baked into the weight matrix —
-                    # the output GGUF is self-contained, no sidecar.
-                    if fold_basis and fb_perms:
-                        for perm_key, axis in _fb_tensor_perms(ti['name']):
-                            if perm_key in fb_perms:
-                                f32 = _fb_apply(f32, ti['dims'], fb_perms[perm_key], axis)
-
-                    # ── Fold-Hadamard pre-processing ────────────────────
-                    # Per-block Walsh-Hadamard transform.  Decorrelates
-                    # weights so quantization errors are independent →
-                    # DC and vesica go to zero structurally.  No storage
-                    # overhead — H is defined by block size alone.
-                    # The dequant kernel applies H⁻¹ = H/n after unpacking.
-                    if fold_hadamard:
-                        apply_fold_hadamard(f32)
-
-                    # ── Fold-Sort-Refit pre-processing ────────────────
-                    # Sort each block → quantize sorted with HPC → unsort
-                    # codes → refit sub-block scales.  Produces standard
-                    # Q2_K blocks compatible with stock llama.cpp.
-                    sort_perms_for_refit = None
-                    f32_presort = None
-                    if fold_sort_refit:
-                        n_blk_sr = len(f32) // QK_K
-                        f32_presort = f32.reshape(n_blk_sr, QK_K).copy()
-                        sort_perms_for_refit = np.argsort(f32_presort, axis=1)
-                        f32 = np.take_along_axis(f32_presort, sort_perms_for_refit, axis=1).reshape(-1).astype(np.float32)
-
-                    # ── Fold-Interleave pre-processing ──────────────────
+                    # ── Fold-Interleave pre-processing ──
                     # Sort-pair weights within each QK_K block so that the
                     # k-th smallest and k-th largest are at fold-complementary
                     # positions (k, k+128).  The HPC pipeline then sees blocks
                     # whose vesica/DC structure is pre-aligned for cancellation.
-                    # f32_orig is kept for RMSE reporting against true weights.
-                    fold_perm_flat = None
-                    f32_orig = f32  # alias — not copied unless fold is active
+                    # No f32_orig copy — squared error is invariant under
+                    # permutation, so RMSE on permuted order is identical.
+                    # apply_fold_interleave permutes f32 in-place chunked.
                     if fold_interleave:
-                        f32_orig = f32.copy()
                         f32, fold_perm_flat = apply_fold_interleave(f32)
 
                     # Quantize to Q2_K — always use HPC with chunked processing
@@ -1823,7 +1417,7 @@ def main():
                     imat_full = _expand_imatrix(ti, imatrix_data)
 
                     # If fold-interleave is active, permute the imatrix the same way
-                    if fold_interleave and imat_full is not None and fold_perm_flat is not None:
+                    if fold_interleave and imat_full is not None:
                         imat_full = apply_fold_interleave_importance(imat_full, fold_perm_flat)
 
                     n_el = ti['n_elements']
@@ -1869,25 +1463,10 @@ def main():
                         n_blocks = n_el // QK_K
                     fout.write(q2k_data)
 
-                    # ── Sort-Refit post-processing ────────────────────
-                    # Unsort the HPC-optimized codes back to original
-                    # weight positions and refit sub-block scales.
-                    if fold_sort_refit and sort_perms_for_refit is not None:
-                        # Seek back — we need to overwrite what we just wrote
-                        fout.seek(fout.tell() - len(q2k_data))
-                        q2k_data = _sort_refit_q2k(
-                            f32_presort, q2k_data,
-                            sort_perms_for_refit, n_blocks)
-                        fout.write(q2k_data)
-
-                    # ── Store fold-interleave permutation in sidecar ──
-                    if fold_interleave and fold_perm_flat is not None:
-                        fold_perm_sidecar[ti['name']] = fold_perm_flat
-
                     # ── Compute and report exact per-tensor RMSE ──
-                    # When fold-interleave is active, dequantized values are
-                    # in permuted order — inverse-permute before comparing
-                    # against the original (un-permuted) weights.
+                    # RMSE is invariant under permutation: with fold-interleave
+                    # both f32 and deq_chunk are in permuted order, so direct
+                    # comparison is identical and avoids the inverse perm + copy.
                     try:
                         CHUNK_BLK = 100_000  # blocks per chunk to bound memory
                         total_se = 0.0
@@ -1896,10 +1475,7 @@ def main():
                             ce = min(ci + CHUNK_BLK, n_blocks)
                             chunk_q = q2k_data[ci*84:ce*84]
                             deq_chunk = dequant_q2k_fast(chunk_q, ce - ci)
-                            if fold_interleave and fold_perm_flat is not None:
-                                perm_slice = fold_perm_flat[ci*QK_K:ce*QK_K]
-                                deq_chunk = invert_fold_interleave(deq_chunk, perm_slice)
-                            orig_chunk = f32_orig[ci*QK_K:ce*QK_K]
+                            orig_chunk = f32[ci*QK_K:ce*QK_K]
                             n_valid = min(len(orig_chunk), len(deq_chunk))
                             diff = orig_chunk[:n_valid] - deq_chunk[:n_valid]
                             total_se += np.sum(diff ** 2)
@@ -1907,7 +1483,7 @@ def main():
                         tensor_rmse = np.sqrt(total_se / max(total_n, 1))
                         q2k_rmse_sum += tensor_rmse
                         q2k_tensor_count += 1
-                        fi_tag = "·SR" if fold_sort_refit else ("·FH" if fold_hadamard else ("·FB" if fold_basis else ("·FI" if fold_interleave else "")))
+                        fi_tag = "·FI" if fold_interleave else ""
                         print(f"\n  [Q2_K{fi_tag}] {ti['name'][:52]}  RMSE={tensor_rmse:.6e}")
                     except Exception as e:
                         print(f"\n  [Q2_K] {ti['name'][:55]}  RMSE=err({e})")
@@ -1915,31 +1491,7 @@ def main():
                     quant_count += 1
                     total_quant_bytes += len(q2k_data)
                 else:
-                    # Keep as-is (passthrough) — but apply fold-basis to
-                    # 1D norm weights so the residual stream is consistent
-                    if fold_basis and fb_perms:
-                        tp = _fb_tensor_perms(ti['name'])
-                        if tp:
-                            if ti['type'] == GGML_TYPE_BF16:
-                                f32 = bf16_to_f32(raw_data, ti['n_elements'])
-                            elif ti['type'] == GGML_TYPE_F16:
-                                f32 = f16_to_f32(raw_data, ti['n_elements'])
-                            elif ti['type'] == GGML_TYPE_F32:
-                                f32 = np.frombuffer(raw_data, dtype=np.float32).copy()
-                            else:
-                                f32 = None
-                            if f32 is not None:
-                                for perm_key, axis in tp:
-                                    if perm_key in fb_perms:
-                                        f32 = _fb_apply(f32, ti['dims'], fb_perms[perm_key], axis)
-                                # Write permuted data in original type
-                                if ti['type'] == GGML_TYPE_BF16:
-                                    raw_data = f32.astype(np.float32).view(np.uint32)
-                                    raw_data = (((raw_data >> 16) & 0xFFFF).astype(np.uint16)).tobytes()
-                                elif ti['type'] == GGML_TYPE_F16:
-                                    raw_data = f32.astype(np.float16).tobytes()
-                                elif ti['type'] == GGML_TYPE_F32:
-                                    raw_data = f32.astype(np.float32).tobytes()
+                    # Keep as-is (passthrough)
                     fout.write(raw_data)
                     total_keep_bytes += len(raw_data)
 
@@ -1978,20 +1530,15 @@ def main():
     print()
     print(f"  Output: {output_path}")
 
-    # ── Save fold-interleave sidecar ──
-    if fold_interleave and fold_perm_sidecar:
-        sidecar_path = output_path + '.fold_perm.npz'
-        np.savez_compressed(sidecar_path, **fold_perm_sidecar)
-        sidecar_size = os.path.getsize(sidecar_path)
-        n_fi_tensors = len(fold_perm_sidecar)
-        print(f"  Fold-interleave sidecar: {sidecar_path}")
-        print(f"    {n_fi_tensors} tensors, {sidecar_size/1024:.0f} KB compressed")
+    if fold_interleave:
+        print(f"  Fold-interleave: ON — permutations applied in-quad, no sidecar")
         print()
         print("  ┌──────────────────────────────────────────────────────────────┐")
-        print("  │  Inference: for each tensor in the sidecar, apply the       │")
-        print("  │  inverse permutation after Q2_K dequant (per QK_K block),   │")
-        print("  │  or equivalently permute the activation vector before the   │")
-        print("  │  matmul.  inv_perm = np.argsort(perm) per block.            │")
+        print("  │  Fold-permuted weights are stored directly in the GGUF    │")
+        print("  │  — no separate .npz sidecar needed. The fold permutation  │")
+        print("  │  must be inverted after dequant (per QK_K block) by the   │")
+        print("  │  inference engine, or equivalently permute the activation  │")
+        print("  │  vector before the matmul.                                  │")
         print("  └──────────────────────────────────────────────────────────────┘")
     print()