QuantizedLinearMLX

Weight-only int4/int8 linear backed by a real Metal low-precision GEMM.

The compute-and-memory quantized linear, and the layer lucid.quantization.convert routes the Linear family to when the MLX backend is active. Where lucid.nn.quantized.Linear dequantizes its int8 weight back to float and runs an ordinary float matmul (a memory-only win), QuantizedLinearMLX stores the weight in MLX's group-wise packed format and runs the genuine low-precision kernel — the packed weight is dequantized inside the GEMM, in one pass, so a full float weight is never materialized. Build it from a float lucid.nn.Linear via from_float, choosing bits (4 or 8).

Group-wise packed representation. The float weight (out_features, in_features) is split, along the input dim, into contiguous groups of group_size elements; every (output-channel, group) block gets its own affine scale $s$ and bias $\beta$ (an offset, distinct from the layer's additive bias term), and the codes are bit-packed into a uint32 tensor (tagged I32). Finer groups track a channel's local dynamic range more tightly than one per-row scale, at the cost of more metadata. This is the packing MLX's quantized_matmul consumes directly.

Encoding happens once (at from_float); the packed decode + matmul run fused on every forward:

where $o$ indexes out_features, $i$ the input dim, $b$ the bit-width (grid 0 .. 2^b - 1), and $g$ maps input index $i$ to its group. A ReLU is optionally fused after + b. In the kernel, $\hat{w}$ is never fully realized — the packed codes and the per-group $(s, \beta)$ are streamed into the matmul directly.