diff --git a/c/tests/test_iq3_pack.py b/c/tests/test_iq3_pack.py new file mode 100644 index 0000000..fd0a10a --- /dev/null +++ b/c/tests/test_iq3_pack.py @@ -0,0 +1,88 @@ +"""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 + +import numpy as np + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "tools")) +import iq3_pack as P # noqa: E402 + + +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 + + +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() diff --git a/c/tools/iq3_pack.py b/c/tools/iq3_pack.py new file mode 100644 index 0000000..03d139b --- /dev/null +++ b/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