| """Lossless BF16 weight storage with a 12-bit common path and sparse escapes. |
| |
| BF16's low byte is stored verbatim. A nibble encodes the sign and one of eight |
| adjacent high-byte exponents. Rare values outside that window use an exact |
| uint16 escape, located by a sorted-index binary search only on escape lanes. |
| The GPU reconstructs the original BF16 bits before the ordinary Metal GEMM. |
| This reduces resident weights without quantization or a new matrix algorithm. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| from functools import lru_cache |
|
|
| import mlx.core as mx |
| import numpy as np |
|
|
|
|
| @lru_cache(maxsize=1) |
| def _decoder(): |
| return mx.fast.metal_kernel( |
| name="h3_unpack_bf16_palette12_v1", |
| input_names=["low", "nibbles", "escape_indices", "escape_values", "base"], |
| output_names=["out"], |
| source=r""" |
| uint i = thread_position_in_grid.x; |
| if (i >= low_shape[0]) return; |
| uint code = (nibbles[i / 2] >> ((i % 2) * 4)) & 15; |
| uint count = escape_indices_shape[0]; |
| if (code == 0 && count != 0) { |
| uint left = 0, right = count; |
| while (left < right) { |
| uint mid = (left + right) / 2; |
| if (escape_indices[mid] < i) left = mid + 1; |
| else right = mid; |
| } |
| out[i] = escape_values[left]; |
| } else { |
| out[i] = ushort(low[i]) | ushort(((code & 7) + base[0]) << 8) |
| | ushort((code & 8) << 12); |
| } |
| """, |
| ) |
|
|
|
|
| @dataclass |
| class PackedBF16: |
| shape: tuple[int, ...] |
| low: mx.array |
| nibbles: mx.array |
| escape_indices: mx.array |
| escape_values: mx.array |
| base: mx.array |
|
|
| @classmethod |
| def pack(cls, value): |
| if value.dtype != mx.bfloat16: |
| raise ValueError("Lossless palette12 stores BF16 tensors only") |
| raw = np.asarray(mx.contiguous(value).view(mx.uint16)).reshape(-1) |
| high = ((raw >> 8) & 127).astype(np.uint8) |
| base = max(0, int(high.max(initial=0)) - 7) |
| sign = ((raw >> 15) << 3).astype(np.uint8) |
| codes = np.clip(high.astype(np.int16) - base, 0, 7).astype(np.uint8) | sign |
| |
| |
| mask = high < base |
| if base: |
| mask |= (high == base) & (sign == 0) |
| indices = np.flatnonzero(mask).astype(np.uint32) |
| codes[indices] = 0 |
| if codes.size % 2: |
| codes = np.pad(codes, (0, 1)) |
| values = [ |
| mx.array(raw.astype(np.uint8)), |
| mx.array(codes[::2] | (codes[1::2] << 4)), |
| mx.array(indices), |
| mx.array(raw[indices]), |
| mx.array([base], mx.uint32), |
| ] |
| mx.eval(values) |
| return cls(tuple(value.shape), *values) |
|
|
| @property |
| def nbytes(self): |
| return sum( |
| x.nbytes |
| for x in ( |
| self.low, |
| self.nibbles, |
| self.escape_indices, |
| self.escape_values, |
| self.base, |
| ) |
| ) |
|
|
| def state(self): |
| return ( |
| self.low, |
| self.nibbles, |
| self.escape_indices, |
| self.escape_values, |
| self.base, |
| ) |
|
|
| def unpack(self, state=None): |
| raw = _decoder()( |
| inputs=self.state() if state is None else state, |
| grid=(self.low.size, 1, 1), |
| threadgroup=(256, 1, 1), |
| output_shapes=[(self.low.size,)], |
| output_dtypes=[mx.uint16], |
| )[0] |
| return raw.view(mx.bfloat16).reshape(self.shape) |
|
|
| def tensors(self, prefix): |
| return { |
| f"{prefix}.{name}": getattr(self, name) |
| for name in ("low", "nibbles", "escape_indices", "escape_values", "base") |
| } |
|
|
| @classmethod |
| def from_tensors(cls, tensors, prefix, shape): |
| result = cls( |
| tuple(shape), |
| *( |
| tensors[f"{prefix}.{name}"] |
| for name in ( |
| "low", |
| "nibbles", |
| "escape_indices", |
| "escape_values", |
| "base", |
| ) |
| ), |
| ) |
| size = int(np.prod(shape)) |
| if result.low.shape != (size,) or result.nibbles.shape != ((size + 1) // 2,): |
| raise ValueError("Invalid palette12 tensor dimensions") |
| if result.escape_indices.shape != result.escape_values.shape: |
| raise ValueError("Invalid palette12 escape arrays") |
| return result |
|
|
|
|
| class PackedLinear: |
| def __init__(self, weight, bias=None): |
| self.weight, self.bias = weight, bias |
|
|
| @classmethod |
| def from_linear(cls, linear): |
| if linear.scales is not None: |
| raise ValueError("Lossless packing cannot be combined with quantization") |
| return cls(PackedBF16.pack(linear.weight), linear.bias) |
|
|
| @property |
| def nbytes(self): |
| return self.weight.nbytes + (0 if self.bias is None else self.bias.nbytes) |
|
|
| @property |
| def dtype(self): |
| return mx.bfloat16 |
|
|
| def __call__(self, x, *, state=None): |
| x = x.astype(mx.bfloat16) |
| w = self.weight.unpack(state) |
| return x @ w.T if self.bias is None else mx.addmm(self.bias, x, w.T) |
|
|
