| """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 = [] |
| |
| 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) |
|
|
