Merge #458: fmt=5 index codec tools — deployable bytes for the E8/IQ3 container
Tools-only step of the #452 plan: iq3_pack.py index codec + IQ3_XXS grid + quant_ablation hook, with round-trip tests (6/6 locally, CI 8/8). No engine changes. Co-authored-by: ZacharyZcR
This commit is contained in:
Vincenzo Fornaro
@@ -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
|
||||
Reference in New Issue
Block a user