| """Quantized weights for the port (MLX affine quantization), with a per-tensor recipe chosen by measurement. |
| |
| A quantized weight is stored as `<name>.q` (packed uint32), `<name>.scales`, `<name>.biases`; `quant.json` in the |
| model folder records bits and group size per tensor. The runtime's linear layers call `mm`, which uses |
| `mx.quantized_matmul` for a quantized weight and a plain matmul otherwise; embedding tables use `rows`. |
| """ |
| import json |
| import os |
|
|
| import mlx.core as mx |
|
|
|
|
| class QW: |
| """A quantized weight [out, in]: x @ W.T through mx.quantized_matmul.""" |
| __slots__ = ("q", "scales", "biases", "bits", "group_size", "shape") |
|
|
| def __init__(self, q, scales, biases, bits, group_size, shape): |
| self.q, self.scales, self.biases = q, scales, biases |
| self.bits, self.group_size, self.shape = int(bits), int(group_size), tuple(shape) |
|
|
| @classmethod |
| def from_dense(cls, w, bits, group_size): |
| q, s, b = mx.quantize(w, group_size=group_size, bits=bits) |
| return cls(q, s, b, bits, group_size, w.shape) |
|
|
| def dense(self, dtype=None): |
| w = mx.dequantize(self.q, self.scales, self.biases, group_size=self.group_size, bits=self.bits) |
| return w if dtype is None else w.astype(dtype) |
|
|
| def arrays(self): |
| return [self.q, self.scales, self.biases] |
|
|
|
|
| def mm(x, w): |
| """x @ w.T for a dense [out, in] weight or a QW.""" |
| if isinstance(w, QW): |
| return mx.quantized_matmul(x, w.q, w.scales, w.biases, transpose=True, group_size=w.group_size, bits=w.bits) |
| return x @ w.T |
|
|
|
|
| def rows(w, idx): |
| """Embedding lookup w[idx] for a dense table or a QW (dequantizes only the gathered rows).""" |
| if isinstance(w, QW): |
| return mx.dequantize(w.q[idx], w.scales[idx], w.biases[idx], group_size=w.group_size, bits=w.bits) |
| return w[idx] |
|
|
|
|
| def cast(w, dtype): |
| """A dense weight is cast; a QW keeps its packed values and casts its scales and biases.""" |
| if isinstance(w, QW): |
| return QW(w.q, w.scales.astype(dtype), w.biases.astype(dtype), w.bits, w.group_size, w.shape) |
| return w.astype(dtype) |
|
|
|
|
| def arrays(tree): |
| """Every mx.array in a nested dict/list/QW structure (for mx.eval).""" |
| if isinstance(tree, QW): |
| return tree.arrays() |
| if isinstance(tree, dict): |
| return [a for v in tree.values() for a in arrays(v)] |
| if isinstance(tree, (list, tuple)): |
| return [a for v in tree for a in arrays(v)] |
| return [tree] if isinstance(tree, mx.array) else [] |
|
|
|
|
| def load_recipe(path): |
| p = os.path.join(path, "quant.json") |
| return json.load(open(p)) if os.path.exists(p) else None |
|
|
|
|
| def take(w, name, recipe): |
| """Pop `name` from a loaded file dict as a dense array or a QW, as quant.json says.""" |
| if recipe and name in recipe["tensors"]: |
| r = recipe["tensors"][name] |
| return QW(w.pop(name + ".q"), w.pop(name + ".scales"), w.pop(name + ".biases"), r["bits"], r["group_size"], |
| r["shape"]) |
| return w.pop(name) |
|
|
