Affine group-wise int2/int4/int8 quantization for Apple Silicon inference,
plus microscaled floating-point modes ("mxfp4"/"mxfp8"/"nvfp4", via
:mode on quantize/2 — see EMLX.quantize/2).
Quantized weights are represented as annotated Nx.Tensor values — the
tensor carries the original logical shape and type (e.g. {:s, 4} for
4-bit), while the EMLX.Backend struct stores the packed uint32 data and
a EMLX.Quantization.Config with scales, biases, group_size, bits, and
mode. Microscaled modes have no biases (Config.biases is nil — MLX's
fp_quantize returns only (wq, scales) for them).
Nx.dot automatically dispatches to mx::quantized_matmul when it detects
a quantized operand on EMLX.Backend — no explicit call site changes needed.
Basic usage
# Quantize a dense weight
weight = Nx.iota({512, 4096}, type: :f32)
weight = Nx.backend_transfer(weight, {EMLX.Backend, device: :gpu})
qw = EMLX.Quantization.from_dense(weight)
# Standard Nx.dot dispatches to mx::quantized_matmul automatically
input = Nx.iota({1, 8, 4096}, type: :f32)
result = Nx.dot(input, [2], qw, [1])Inside defn
quantize/2, dequantize/1, and quantized_matmul/2 are all
deftransform functions backed by Nx.runtime_call, so they are safe
to call inside Nx.Defn.jit-traced forward passes:
defn my_layer(x, qw) do
dense = EMLX.Quantization.dequantize(qw)
Nx.dot(x, [2], dense, [1])
endNx.dot with a quantized tensor also dispatches transparently via
EMLX.Backend.dot/7 at execution time, so explicit dequantize is
only needed when you want the dense weight for a non-dot operation.
Limitation: freshly-quantized values can't feed a traced Nx.dot
Nx.dot's mx::quantized_matmul dispatch (see
EMLX.Native.Expr.quantizable_param_positions/1) only recognizes a
quantized tensor that started life as a bound top-level parameter of the
Nx.Defn.jit-compiled function — it's resolved via a call-time
quant_signature keyed by parameter position, baked into a specialized
compiled program. A value produced by quantize/2 inside the traced
program (e.g. dequantize/1 → some op → quantize/2 again) has no such
position and isn't recognized as a quantized :dot operand; it may only
be used as a direct output of the traced function (see quantize/2's doc).
See also
EMLX.Quantization.Config— internal metadata struct
Summary
Functions
Dequantize a quantized Nx.Tensor via Nx.runtime_call.
At execution time the callback receives the real tensor, unpacks
quantization_config, and calls mx::dequantize. Safe to call inside
Nx.Defn.jit-traced forward passes.
The output has the same shape as the input (the quantized tensor's logical
shape equals the dense shape). Supports every mode EMLX.quantize/2
accepts.
Quantize a dense 2-D tensor via Nx.runtime_call.
At execution time the callback receives the real tensor and runs
mx::quantize. Returns an annotated quantized Nx.Tensor with the same
logical shape as the input and type {:s, N}. Safe to call inside
Nx.Defn.jit-traced forward passes — including on a value that is itself
freshly computed in-graph (e.g. dequantize/1 followed by some op followed
by quantize/2 again), not just a bare top-level parameter.
Note: the result of a traced quantize/2 call may only be used as a
direct output of the Nx.Defn.jit-compiled function (returned as-is, or
passed through unchanged). Feeding it into another op in the same traced
program (e.g. Nx.dot) is not yet supported and raises during lowering —
only a quantized tensor that started life as a bound top-level parameter
dispatches to mx::quantized_matmul today.
Options
:type— Nx storage type:{:s, 2},{:s, 4}(default), or{:s, 8}.:group_size— 32, 64, or 128 (default 64). Must evenly divide the last dimension oftensor.:mode—"affine"(default), or a microscaled mode ("mxfp4","mxfp8","nvfp4") — seeEMLX.quantize/2.
Returns true if the tensor has quantization metadata on its backend.
Run a quantized matmul via Nx.runtime_call.
At execution time the callback unpacks quantization_config from qw and
calls mx::quantized_matmul. Safe to call inside Nx.Defn.jit-traced
forward passes.
Output shape is {batch_dims..., out_features} where out_features is the
first dimension of qw. Output type matches the scales dtype for
"affine" (historically the activation's own float type); for microscaled
modes scales are :u8 (a packed exponent byte), so the output type
follows the activation's dtype instead.
Construct an annotated quantized Nx.Tensor from pre-computed device refs.
Use this when you already have packed weights from a checkpoint. For
quantizing a dense tensor from scratch, prefer from_dense/2.
Options
:type— Nx storage type:{:s, 2},{:s, 4}(default), or{:s, 8}.:group_size— quantization group size (default 64).