Module mlx_int8

Module mlx_int8 

Source
Expand description

2026-09-25: Launches for 8-bit affine weights in the MLX layout: load a (.weight, .scales, .biases) triplet and run the dequant, GEMV, fused GEMV and GEMM kernels of kernels/metal/common/mlx_int8_*.metal.

Layout, as MlxInt8Weight::load checks it:

  • {base}.weight: U32 [out_features, in_features / 4]; each word packs four unsigned bytes, the low byte holding the lowest column.
  • {base}.scales, {base}.biases: BF16 [out_features, in_features / G], where G is group_size.

The kernels dequantise w[r, c] = byte * scales[r, c / G] + biases[r, c / G], where byte is byte c % 4 of packed[r, c / 4].

Owner: gpu-runtime (metal). Invariants: none beyond the types.

Structs§

MlxInt8Weight
2026-09-25: One 8-bit affine linear weight resident on the GPU. There is no Drop: the owner frees the three buffers with Self::release.
MlxQuantConfig
2026-09-25: Quantization metadata from the model’s config.json.

Functions§

gemv_gate_up
2026-09-25: gate_y = gate x and up_y = up x in one launch of mlx_int8_gemv_gate_up, which reads x once for both. Both weights must share (out_features, in_features, group_size); only debug_assert_eq! checks it.