tools: fmt=5 index codec — the deployable bytes for the E8/IQ3 container (#452 step 2)

c/tests/test_iq3_pack.py

@@ -0,0 +1,96 @@
+"""fmt=5 codec oracle (#452 ladder step 2).
+
+Checks the properties the container, the converter and the decode kernels all
+depend on: exact byte budget, deterministic encode, decode agreeing with a
+straight-from-the-spec reader, sign parity closure, and reconstruction quality
+matching the ablation that chose this scheme (#453).
+"""
+import os
+import sys
+import unittest
+
+# The runtime path is dependency-free by design and CI keeps it that way, so the
+# offline-tooling tests skip rather than fail where numpy is absent.
+np = None
+P = None
+try:
+    import numpy as np
+
+    sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "tools"))
+    import iq3_pack as P
+except ImportError:  # pragma: no cover - exercised only on dependency-free CI
+    pass
+
+
+def ref_decode(packed, K):
+    """Independent reader written straight from the layout comment — deliberately
+    naive and loop-based, so a shared bug in the vectorized path shows up."""
+    g = P.grid()
+    nsb = K // P.QK
+    rows = packed.reshape(-1, nsb * P.BLOCK_BYTES)
+    out = np.zeros((len(rows), K), dtype=np.float32)
+    for r in range(len(rows)):
+        for sb in range(nsb):
+            base = sb * P.BLOCK_BYTES
+            d = float(rows[r, base + 96:base + 98].copy().view(np.float16)[0])
+            for ib in range(P.QK // P.SUB):
+                word = int(np.ascontiguousarray(
+                    rows[r, base + 64 + ib * 4:base + 64 + ib * 4 + 4]).view(np.uint32)[0])
+                db = d * (0.5 + ((word >> 28) & 0xF)) * 0.5
+                for l in range(4):
+                    seven = (word >> (7 * l)) & 0x7F
+                    bits = [(seven >> j) & 1 for j in range(7)]
+                    bits.append(sum(bits) & 1)          # odd parity closes the 8th
+                    idx = int(rows[r, base + ib * 8 + l * 2 + 0])
+                    idx2 = int(rows[r, base + ib * 8 + l * 2 + 1])
+                    mags = list(g[idx]) + list(g[idx2])
+                    for j in range(8):
+                        pos = sb * P.QK + ib * P.SUB + l * 8 + j
+                        out[r, pos] = mags[j] * db * (-1.0 if bits[j] else 1.0)
+    return out
+
+
+@unittest.skipIf(P is None, "numpy not available (offline-tooling test)")
+class TestIq3Pack(unittest.TestCase):
+    def setUp(self):
+        np.random.seed(4242)
+        self.x = (np.random.randn(6, 1024) * 0.05).astype(np.float32)
+
+    def test_byte_budget(self):
+        self.assertEqual(P.BLOCK_BYTES, 98)
+        self.assertAlmostEqual(P.bpw(), 3.0625, places=6)
+        packed = P.encode(self.x)
+        self.assertEqual(packed.shape, (6, 1024 // P.QK * 98))
+        self.assertEqual(packed.dtype, np.uint8)
+
+    def test_encode_is_deterministic(self):
+        self.assertTrue(np.array_equal(P.encode(self.x), P.encode(self.x)))
+
+    def test_decode_matches_spec_reader(self):
+        packed = P.encode(self.x)
+        fast = P.decode(packed, 1024)
+        slow = ref_decode(packed, 1024)
+        self.assertTrue(np.allclose(fast, slow, rtol=1e-6, atol=1e-8),
+                        f"max |Δ| = {np.abs(fast - slow).max()}")
+
+    def test_sign_parity_closes(self):
+        """Every 8-weight block must have an even number of negatives — that is
+        what lets the 8th sign be derived instead of stored."""
+        y = P.decode(P.encode(self.x), 1024)
+        neg = (y < 0).reshape(-1, 8).sum(-1)
+        self.assertTrue(np.all(neg % 2 == 0), "a block stored odd negatives")
+
+    def test_reconstruction_quality(self):
+        y = P.decode(P.encode(self.x), 1024)
+        rel = np.sqrt(((y - self.x) ** 2).mean()) / np.sqrt((self.x ** 2).mean())
+        # the torch model that won the #453 A/B measures ~0.195 on this input class
+        self.assertLess(rel, 0.25, f"rel-RMSE {rel:.4f} — worse than the chosen scheme")
+        self.assertGreater(rel, 0.05, f"rel-RMSE {rel:.4f} — implausibly good, check the test")
+
+    def test_shape_guard(self):
+        with self.assertRaises(ValueError):
+            P.encode(np.zeros((2, 300), dtype=np.float32))
+
+
+if __name__ == "__main__":
+    unittest.main()

c/tools/iq3_pack.py

@@ -0,0 +1,170 @@
+#!/usr/bin/env python3
+"""fmt=5 (E8/IQ3 grouped container) index codec — #452 ladder step 2.
+
+The ablation (#453) proved the SCHEME: an IQ3_XXS-style codebook plus rotation
+matches our simulated E8 ball (51.5% vs 51.5% on OLMoE). That code quantizes to
+lattice points and keeps floats. This module produces the DEPLOYABLE bytes and
+reads them back, so the container, the converter and the decode kernels all
+agree on one layout.
+
+Layout — one 256-weight super-block, 98 bytes, 3.0625 bpw:
+
+    [0  .. 63]  uint8  grid index per 4-dim magnitude block   (64 blocks)
+    [64 .. 95]  uint32 x8, one per 32-weight sub-block:
+                  bits  0..20  three 7-bit sign words (8 weights each,
+                               bit i set => weight i negative; the 8th sign
+                               is implied by odd parity)
+                  bits 21..27  the fourth 7-bit sign word
+                  bits 28..31  4-bit sub-scale code
+    [96 .. 97]  fp16   super-scale d
+
+    value(w) = d * (0.5 + code) * 0.5 * grid[idx][j] * 0.5 * sign
+
+The last 0.5 is the half-unit convention of the published grid (magnitudes are
+stored doubled: 4,12,...,62 mean 2.0,6.0,...,31.0).
+
+Odd-parity signs: llama.cpp stores 7 of every 8 signs and derives the 8th so the
+product of the eight is +1. The encoder therefore flips the smallest-magnitude
+weight of any block whose true signs violate that — the same cost the ablation
+priced in, now applied for real.
+"""
+import json
+import os
+import numpy as np
+
+QK = 256                      # weights per super-block
+SUB = 32                      # weights per sub-block (one uint32 of signs+scale)
+BLOCK_BYTES = QK // 4 + (QK // SUB) * 4 + 2      # 64 + 32 + 2 = 98
+
+_GRID = None
+
+
+def grid():
+    """[256,4] float32 magnitudes in weight units (published table is doubled)."""
+    global _GRID
+    if _GRID is None:
+        path = os.path.join(os.path.dirname(__file__), "iq3xxs_grid.json")
+        _GRID = np.asarray(json.load(open(path)), dtype=np.float32) * 0.5
+    return _GRID
+
+
+def _nearest(mag4):
+    """[N,4] magnitudes -> [N] grid indices, argmin ||m-g||^2 without cdist."""
+    g = grid()
+    g2 = (g * g).sum(1)
+    out = np.empty(len(mag4), dtype=np.uint8)
+    for i in range(0, len(mag4), 1 << 16):          # bounded working set
+        c = mag4[i:i + (1 << 16)]
+        out[i:i + len(c)] = np.argmin(g2 - 2.0 * (c @ g.T), axis=1).astype(np.uint8)
+    return out
+
+
+def encode(x):
+    """float32 [..., K] (K % 256 == 0) -> packed uint8 [..., K//256 * 98]."""
+    x = np.ascontiguousarray(x, dtype=np.float32)
+    K = x.shape[-1]
+    if K % QK:
+        raise ValueError(f"fmt=5 needs K % {QK} == 0, got {K}")
+    rows = x.reshape(-1, K)
+    nsb = K // QK
+    out = np.empty((len(rows), nsb * BLOCK_BYTES), dtype=np.uint8)
+
+    for sb in range(nsb):
+        blk = rows[:, sb * QK:(sb + 1) * QK]                     # [R,256]
+        sign = np.where(blk < 0, -1.0, 1.0).astype(np.float32)
+        mag = np.abs(blk)
+
+        # parity fix: flip the smallest magnitude of every 8 whose product is -1
+        s8 = sign.reshape(len(rows), QK // 8, 8)
+        m8 = mag.reshape(len(rows), QK // 8, 8)
+        viol = s8.prod(-1) < 0                                   # [R,32]
+        amin = m8.argmin(-1)
+        r, b = np.nonzero(viol)
+        s8[r, b, amin[r, b]] *= -1.0
+        sign = s8.reshape(len(rows), QK)
+
+        base = sb * BLOCK_BYTES
+        # super-scale: RMS anchor, same statistic the ablation searches around
+        d = np.sqrt((mag * mag).mean(-1, keepdims=True)) / 20.0 + 1e-12
+        out[:, base + 96:base + 98] = d.astype(np.float16).view(np.uint8)
+        d = d.astype(np.float16).astype(np.float32)              # encode what we store
+
+        g = grid()
+        for ib in range(QK // SUB):
+            m = mag[:, ib * SUB:(ib + 1) * SUB]                  # [R,32]
+            best_err = None
+            best = None
+            for code in range(16):
+                db = d * (0.5 + code) * 0.5
+                q = (m / np.maximum(db, 1e-20)).reshape(-1, 4)
+                idx = _nearest(q)
+                rec = g[idx].reshape(len(rows), SUB) * db
+                err = ((rec - m) ** 2).sum(-1, keepdims=True)
+                if best_err is None:
+                    best_err, best = err, (idx.reshape(len(rows), SUB // 4), code)
+                else:
+                    take = (err < best_err)[:, 0]
+                    if take.any():
+                        keep_idx, keep_code = best
+                        ni = idx.reshape(len(rows), SUB // 4)
+                        keep_idx = np.where(take[:, None], ni, keep_idx)
+                        # per-row code: store alongside, resolved below
+                        keep_code = np.where(take, code, keep_code) if isinstance(
+                            keep_code, np.ndarray) else np.where(
+                            take, code, np.full(len(rows), keep_code))
+                        best = (keep_idx, keep_code)
+                        best_err = np.where(take[:, None], err, best_err)
+            bidx, bcode = best
+            if not isinstance(bcode, np.ndarray):
+                bcode = np.full(len(rows), bcode)
+            out[:, base + ib * 8:base + (ib + 1) * 8] = bidx.astype(np.uint8)
+
+            # signs: four 7-bit words for this sub-block + the 4-bit code
+            s = sign[:, ib * SUB:(ib + 1) * SUB].reshape(len(rows), 4, 8)
+            neg = (s < 0).astype(np.uint32)
+            word = np.zeros(len(rows), dtype=np.uint32)
+            for l in range(4):
+                seven = np.zeros(len(rows), dtype=np.uint32)
+                for j in range(7):
+                    seven |= neg[:, l, j] << j
+                word |= seven << (7 * l)
+            word |= (bcode.astype(np.uint32) & 0xF) << 28
+            off = base + QK // 4 + ib * 4
+            out[:, off:off + 4] = word.view(np.uint8).reshape(len(rows), 4) if False else \
+                np.ascontiguousarray(word).view(np.uint8).reshape(len(rows), 4)
+    return out.reshape(*x.shape[:-1], nsb * BLOCK_BYTES)
+
+
+def decode(packed, K):
+    """packed uint8 [..., K//256*98] -> float32 [..., K]. The kernels' reference."""
+    packed = np.ascontiguousarray(packed, dtype=np.uint8)
+    nsb = K // QK
+    rows = packed.reshape(-1, nsb * BLOCK_BYTES)
+    out = np.empty((len(rows), K), dtype=np.float32)
+    g = grid()
+    for sb in range(nsb):
+        base = sb * BLOCK_BYTES
+        d = rows[:, base + 96:base + 98].copy().view(np.float16).astype(np.float32)
+        for ib in range(QK // SUB):
+            idx = rows[:, base + ib * 8:base + (ib + 1) * 8]                  # [R,8]
+            off = base + QK // 4 + ib * 4
+            word = np.ascontiguousarray(rows[:, off:off + 4]).view(np.uint32).reshape(-1)
+            code = (word >> 28) & 0xF
+            db = d[:, 0] * (0.5 + code) * 0.5                                 # [R]
+            mag = g[idx].reshape(len(rows), SUB)                              # [R,32]
+            sgn = np.ones((len(rows), 4, 8), dtype=np.float32)
+            for l in range(4):
+                seven = (word >> (7 * l)) & 0x7F
+                par = 0
+                for j in range(7):
+                    bit = (seven >> j) & 1
+                    sgn[:, l, j] = np.where(bit == 1, -1.0, 1.0)
+                    par ^= bit
+                sgn[:, l, 7] = np.where(par == 1, -1.0, 1.0)   # odd parity closes the block
+            out[:, sb * QK + ib * SUB:sb * QK + (ib + 1) * SUB] = \
+                mag * sgn.reshape(len(rows), SUB) * db[:, None]
+    return out.reshape(*packed.shape[:-1], K)
+
+
+def bpw():
+    return BLOCK_BYTES * 8 / QK

c/tools/iq3xxs_grid.json

@@ -0,0 +1 @@
+[[4, 4, 4, 4], [20, 4, 4, 4], [36, 4, 4, 4], [12, 12, 4, 4], [28, 12, 4, 4], [62, 12, 4, 4], [4, 20, 4, 4], [20, 20, 4, 4], [12, 28, 4, 4], [20, 36, 4, 4], [28, 62, 4, 4], [44, 62, 4, 4], [12, 4, 12, 4], [28, 4, 12, 4], [4, 12, 12, 4], [20, 12, 12, 4], [12, 20, 12, 4], [44, 20, 12, 4], [4, 28, 12, 4], [20, 28, 12, 4], [12, 36, 12, 4], [36, 44, 12, 4], [4, 62, 12, 4], [4, 4, 20, 4], [20, 4, 20, 4], [36, 4, 20, 4], [12, 12, 20, 4], [4, 20, 20, 4], [20, 20, 20, 4], [12, 28, 20, 4], [28, 28, 20, 4], [62, 28, 20, 4], [12, 44, 20, 4], [62, 44, 20, 4], [44, 62, 20, 4], [12, 4, 28, 4], [62, 4, 28, 4], [4, 12, 28, 4], [20, 12, 28, 4], [44, 20, 28, 4], [4, 62, 28, 4], [28, 12, 36, 4], [62, 28, 36, 4], [36, 36, 36, 4], [62, 44, 36, 4], [28, 62, 36, 4], [44, 62, 36, 4], [12, 4, 44, 4], [62, 4, 44, 4], [20, 28, 44, 4], [20, 44, 44, 4], [44, 28, 52, 4], [36, 52, 52, 4], [4, 12, 62, 4], [36, 12, 62, 4], [52, 12, 62, 4], [28, 36, 62, 4], [12, 52, 62, 4], [12, 4, 4, 12], [28, 4, 4, 12], [4, 12, 4, 12], [20, 12, 4, 12], [12, 20, 4, 12], [28, 20, 4, 12], [4, 28, 4, 12], [20, 28, 4, 12], [36, 28, 4, 12], [62, 36, 4, 12], [4, 44, 4, 12], [4, 4, 12, 12], [20, 4, 12, 12], [12, 12, 12, 12], [4, 20, 12, 12], [20, 20, 12, 12], [12, 4, 20, 12], [28, 4, 20, 12], [4, 12, 20, 12], [20, 12, 20, 12], [12, 20, 20, 12], [4, 28, 20, 12], [20, 62, 20, 12], [4, 4, 28, 12], [20, 4, 28, 12], [4, 20, 28, 12], [12, 28, 28, 12], [52, 36, 28, 12], [52, 52, 28, 12], [12, 4, 36, 12], [44, 4, 36, 12], [4, 44, 36, 12], [4, 20, 44, 12], [36, 20, 44, 12], [52, 36, 44, 12], [12, 62, 44, 12], [44, 4, 52, 12], [20, 20, 62, 12], [4, 36, 62, 12], [4, 4, 4, 20], [20, 4, 4, 20], [12, 12, 4, 20], [28, 12, 4, 20], [4, 20, 4, 20], [20, 20, 4, 20], [52, 20, 4, 20], [12, 28, 4, 20], [20, 36, 4, 20], [12, 4, 12, 20], [28, 4, 12, 20], [44, 4, 12, 20], [4, 12, 12, 20], [20, 12, 12, 20], [12, 20, 12, 20], [4, 28, 12, 20], [28, 52, 12, 20], [62, 52, 12, 20], [4, 62, 12, 20], [4, 4, 20, 20], [20, 4, 20, 20], [12, 12, 20, 20], [62, 12, 20, 20], [4, 20, 20, 20], [20, 20, 20, 20], [62, 28, 20, 20], [4, 36, 20, 20], [44, 44, 20, 20], [12, 4, 28, 20], [4, 12, 28, 20], [36, 12, 28, 20], [4, 62, 28, 20], [36, 62, 28, 20], [44, 28, 36, 20], [28, 44, 36, 20], [28, 4, 44, 20], [62, 20, 44, 20], [12, 36, 44, 20], [36, 62, 44, 20], [12, 4, 62, 20], [28, 4, 62, 20], [52, 12, 62, 20], [44, 36, 62, 20], [12, 4, 4, 28], [4, 12, 4, 28], [20, 12, 4, 28], [12, 20, 4, 28], [28, 20, 4, 28], [4, 44, 4, 28], [44, 52, 4, 28], [20, 62, 4, 28], [4, 4, 12, 28], [20, 4, 12, 28], [4, 20, 12, 28], [12, 28, 12, 28], [36, 36, 12, 28], [52, 36, 12, 28], [12, 4, 20, 28], [28, 4, 20, 28], [4, 12, 20, 28], [44, 20, 20, 28], [20, 44, 20, 28], [20, 62, 20, 28], [12, 12, 28, 28], [28, 28, 28, 28], [4, 28, 36, 28], [62, 36, 36, 28], [20, 62, 36, 28], [4, 4, 44, 28], [52, 4, 44, 28], [20, 20, 44, 28], [44, 44, 44, 28], [36, 12, 52, 28], [52, 28, 52, 28], [28, 52, 52, 28], [28, 28, 62, 28], [4, 52, 62, 28], [36, 4, 4, 36], [62, 12, 4, 36], [44, 28, 4, 36], [62, 28, 4, 36], [28, 44, 4, 36], [62, 44, 4, 36], [36, 62, 12, 36], [4, 20, 20, 36], [62, 28, 20, 36], [4, 36, 20, 36], [4, 52, 20, 36], [52, 52, 20, 36], [62, 4, 28, 36], [44, 36, 28, 36], [36, 4, 36, 36], [12, 44, 36, 36], [36, 52, 36, 36], [44, 20, 44, 36], [28, 36, 44, 36], [4, 62, 44, 36], [44, 4, 62, 36], [4, 12, 62, 36], [20, 12, 62, 36], [4, 28, 62, 36], [20, 12, 4, 44], [12, 36, 4, 44], [4, 62, 4, 44], [4, 4, 12, 44], [52, 4, 12, 44], [52, 20, 12, 44], [44, 44, 12, 44], [36, 12, 20, 44], [20, 28, 20, 44], [20, 62, 20, 44], [20, 4, 28, 44], [28, 44, 28, 44], [4, 12, 36, 44], [28, 20, 36, 44], [62, 20, 36, 44], [20, 62, 36, 44], [20, 4, 44, 44], [12, 28, 44, 44], [4, 44, 52, 44], [36, 20, 62, 44], [20, 36, 62, 44], [36, 20, 4, 52], [36, 36, 4, 52], [52, 36, 4, 52], [36, 52, 4, 52], [12, 20, 12, 52], [12, 52, 12, 52], [62, 12, 20, 52], [36, 52, 20, 52], [4, 28, 28, 52], [52, 28, 28, 52], [36, 36, 36, 52], [44, 4, 44, 52], [20, 44, 44, 52], [28, 28, 52, 52], [28, 4, 62, 52], [12, 20, 62, 52], [28, 4, 4, 62], [44, 4, 4, 62], [62, 4, 4, 62], [4, 12, 4, 62], [20, 28, 4, 62], [20, 44, 4, 62], [52, 20, 12, 62], [4, 36, 12, 62], [20, 12, 20, 62], [44, 36, 20, 62], [20, 44, 20, 62], [4, 4, 28, 62], [44, 12, 28, 62], [28, 28, 28, 62], [4, 52, 28, 62], [12, 20, 36, 62], [12, 36, 36, 62], [4, 4, 44, 62], [20, 4, 44, 62], [36, 20, 44, 62], [4, 28, 52, 62]]