File size: 2,572 Bytes
3b57fd1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 | """Lossless PQ2_0/PTQ1_0 to MLX affine 2-bit block transcoding.
Preserves folded weights. Runtime Hadamard metadata remains mandatory; this
module alone does not produce a loadable or validated model.
"""
import numpy as np
def transcode(raw: bytes, shape: tuple[int, int], source: str):
"""Return packed uint32 weights and FP16 affine scales/biases, group size 128."""
rows, width = shape
if rows <= 0 or width <= 0 or width % 128:
raise ValueError(
"Expected positive [output, input] shape with input divisible by 128"
)
sizes = {"PQ2_0": 34, "PTQ1_0": 28}
if source not in sizes:
raise ValueError(f"Unsupported format: {source}")
blocks = rows * width // 128
if len(raw) != blocks * sizes[source]:
raise ValueError("Raw byte length does not match shape and format")
data = np.frombuffer(raw, dtype=np.uint8).reshape(blocks, sizes[source])
scale_bytes = data[:, :2] if source == "PQ2_0" else data[:, 26:28]
scales = scale_bytes.copy().view("<f2").reshape(rows, width // 128)
if not np.isfinite(scales).all():
raise ValueError("Non-finite quantization scale")
if source == "PQ2_0":
words = data[:, 2:].copy().view("<u4").reshape(rows, width // 16)
return words, np.ascontiguousarray(scales), np.ascontiguousarray(-scales)
else:
pieces = []
# GGML PTQ stages consume 16 bytes, then 8, then the two-byte tail.
for lo, hi, count in [(0, 16, 5), (16, 24, 5), (24, 26, 4)]:
packed = data[:, lo:hi].astype(np.uint16)
for trit in range(count):
remainder = (packed * (3**trit)) & 255
pieces.append(((remainder * 3) >> 8).astype(np.uint8))
codes = np.concatenate(pieces, axis=1)
words = np.bitwise_or.reduce(
codes.astype(np.uint32).reshape(rows, width // 16, 16)
<< (2 * np.arange(16, dtype=np.uint32)),
axis=-1,
).astype("<u4")
return (
np.ascontiguousarray(words),
np.ascontiguousarray(scales),
np.ascontiguousarray(-scales),
)
def unpack(weight, scales, biases):
"""Independent affine reconstruction in float32 for packing validation."""
rows, words = weight.shape
values = np.empty((rows, words * 16), dtype=np.float32)
for lane in range(16):
values[:, lane::16] = (weight >> (2 * lane)) & 3
groups = values.reshape(rows, -1, 128)
return (
groups * scales.astype(np.float32)[..., None]
+ biases.astype(np.float32)[..., None]
).reshape(rows, -1)
|