Module Nx_quant

The type for quantised weights on the set 'd. Match on a weight to read its parts; only the constructors, quantize_q8_0 and walk build one.

MXFP4 (OCP Microscaling Formats, 2023) decodes a value as its e2m1 code, a sign bit over a magnitude among 0, 0.5, 1, 1.5, 2, 3, 4 and 6, times 2 ^ (s - 127) for its group's e8m0 scale byte s of 32 values; the scale byte 255 makes its whole group NaN. A GGUF format stores each block of k's values as one run of bytes, in the layout of ggml's block_mxfp4, block_q8_0, block_q4_K and block_q6_K (ggml-common.h), its float16 fields little-endian. In every format a block of zero bytes decodes to zeros, -0. in Q6_K.

mxfp4 ~scales codes is the MXFP4 weight of codes, two per byte, and scales: Mxfp4 { codes; scales }.

Raises Invalid_argument naming the part if codes does not have shape [| …; n; k / 2 |] with k a multiple of 32, or scales does not have shape [| …; n; k / 32 |].

mxfp4_blocks blocks is the MXFP4 weight stored as GGUF's blocks.

Its values are the format's, t's: the scale byte 255 makes its group NaN where ggml reads it as 2 ^ 128, and the code of sign 1 and magnitude 0 is -0. where ggml gives 0..

Raises Invalid_argument if blocks does not have shape [| …; n; k / 32 * 17 |].

q8_0 blocks is the Q8_0 weight stored as blocks.

Raises Invalid_argument if blocks does not have shape [| …; n; k / 32 * 34 |].

q4_k blocks is the Q4_K weight stored as blocks.

Raises Invalid_argument if blocks does not have shape [| …; n; k / 256 * 144 |].

q6_k blocks is the Q6_K weight stored as blocks.

Raises Invalid_argument if blocks does not have shape [| …; n; k / 256 * 210 |].

quantize_q8_0 x is the float weight x, [| …; n; k |], stored as Q8_0, x read at float32. A block of 32 values v has d = amax / 127 for its greatest magnitude amax, stored at float16, rounded to nearest, and the quants round (v * (1 / d)), halves away from zero, the products at float32, 0 where d is 0.

Outside a block holding a NaN, its bytes are those of ggml's quantize_row_q8_0_ref (ggml-quants.c) wherever ggml's conversion of a product to int8 is defined. A product that is not finite has the quant 0, as ggml's conversion gives on x86-64:

• a block whose 1 / d is infinite, amax below about 127 * 2^-128, is zero bytes, which decode to +0.;
• a block holding an infinity has the scale +inf and the quants 0;
• a block holding a NaN has a NaN scale and the quants 0: the NaN reaches its block's every value, where ggml leaves it out of amax unless it is the block's last value.

A block holding an infinity or a NaN, or whose d rounds past float16's largest, decodes to no value of its own: NaN, or an infinity times its quant.

Raises Invalid_argument if x does not have shape [| …; n; k |] with k a multiple of 32.

shape w is w's logical shape, [| …; n; k |].

dequant dt w is the values of w at dt, of shape shape w, each rounded once to dt.

A value of w is a float32 value. An MXFP4 value is exact barring overflow: codes of magnitude 4 or more at scale byte 253, and of 2 or more at 254, are infinite at every dtype. Float32, bfloat16 and float64 hold every other value exactly. A GGUF value is computed as ggml's float32 dequantisation computes it, its products left to right, and equals it bit for bit: a Q8_0 value, a Q6_K value and a Q4_K product are exact, and a Q4_K value rounds once at its subtraction. A block whose float16 scale is infinite or NaN gives infinite or NaN values.

apply w x is w applied to x as a linear map from k inputs to n outputs, summed at acc, float64 for a float64 x and float32 otherwise, which holds each value of w exactly, and rounded once to x's dtype:

Nx.cast (Nx.dtype x)
  (Nx.matmul (Nx.cast acc x) (Nx.matrix_transpose (dequant acc w)))

w is [| …; n; k |], x is [| …; m; k |] or [| k |], batch axes broadcast as Nx.matmul's, and the result is [| …; m; n |], or [| …; n |] for a vector.

A Q8_0 weight, and an Mxfp4 one summed at float32, are read where they lie: the product decodes each block as it sums it and stores no decoded weight. A Q6_K weight's product first unpacks its quants as int8 and its sub-blocks' scales d * sc at float32, and sums them scaled, storing none of its values. Any other product first decodes w at float32 as dequant does: an Mxfp4 value past float32's range is infinite there, and its code times its scale would be finite at float64. Read where it lies or unpacked, the product needs kernels that compute a block-scaled contraction, as nx.cpu's do. Each element times its block's scale is then the value dequant gives, exact in acc, so in kernels that scale each element, as nx.cpu's do, the scales add no rounding: the product holds Nx.contract's bound over w's values. A compiled Q8_0 product that a device's tensor cores take scales each block's partial sum instead, and holds the same bound, barring the partial sum's overflow or underflow.

Raises Invalid_argument if w is read where it lies or unpacked on a set whose kernels decline that contraction, if x is 0-d, if x's last axis is not k, or if the batch axes do not broadcast.

apply_experts ~ids w x is each row of x through the expert its id names: at each position p of ids, the product apply defines of w's matrix ids.{p} with p's row of x, or zeros where ids.{p} is not in [0, e), whatever that row holds. w is [| e; n; k |], ids [| s… |] and x [| s'…; k |], s'… broadcasting to s…; the result is [| s…; n |].

The positions in range are summed in one contraction, each row with the matrix its id names, at the dtype apply sums in, within Nx.contract's bound; a position out of range has no products. The weight is read as apply reads it: where it lies, each block decoded as it is summed, but for a weight apply decodes at float32 first, of which it decodes, for p positions with 0 < p < e, ids whole on each device and experts not cut, p matrices, each position's id's, a position out of range decoding the first. A row's bits are a function of the shapes, the placements, its row and its expert: they do not depend on the other positions' experts, and may differ from apply's. Cut along k, each device sums its window of k for every row and the devices' sums close by partials, as Nx.contract's do, where the result would lie with k whole. ids and x's rows may be cut along the positions together, into windows of one extent, beside a weight whole on each device: each device computes its own positions. x's leading axes are broadcast to ids' and flattened where that keeps where they lie: at most one has more than one index, or only the first such is cut, into windows of one extent. Cut otherwise, along a later such axis or into windows of unequal extents, x is gathered along each such axis in turn, from the longest, which holds each position's row times the extents of the others. The sum needs kernels that compute a segmented contraction, block-scaled where apply's is, as nx.cpu's do.

Raises Invalid_argument if the sum lies on a set whose kernels decline it, if w is not 3-d, if x does not have one axis more than ids, if x's last axis is not k, or if x's leading axes do not broadcast to ids'.

take ~axis p w is the rows or matrices of w at the positions p along axis, gathered from its packed parts with no value decoded, p read as Nx.take reads it: at each position in range, dequant dt (take ~axis p w) is Nx.take ~axis p (dequant dt w). A position out of range gathers zero bytes, which decode to zeros. Where w is cut along axis, each device gathers from its window, as Nx.take does.

Raises Invalid_argument if axis is w's last axis, negative or not, since a gather along it would cut blocks, and if axis is not an axis of w.

walk c w walks w at c's path: its case, "mxfp4", "mxfp4_blocks", "q8_0", "q4_k" or "q6_k", then each part as a fixed array (Nx.Ptree.Walk.tensor) at its field's name, codes then scales, or blocks. It rebuilds the weight as the case's constructor does. A structure holding a weight walks it with field c "w" Nx_quant.walk w.

Raises Invalid_argument as the constructors do, naming Nx_quant.walk, where the walk returns parts of other shapes.

place p w is w with each part at p, as Nx.place.