Module quantization 

Quantization-Aware Training (QAT) primitives

This module is a CPU reference implementation of the numerical core used for quantization-aware training. It contains no real GPU calls: every routine operates on host Array1 / Array2 data and exactly simulates the precision loss that the corresponding low-precision GPU kernels would introduce. The intent is that the bit-for-bit rounding behaviour modelled here matches what a fused int8 / fp8 kernel produces, so a network trained with these “fake-quant” operators behaves like the eventually-deployed quantized network.

What “fake quantization” means

A fake-quant operator maps a floating-point value through the quantize/dequantize round trip while staying in floating point:

  fake_quant(x) = dequant(quant(x))

The result is the value the network would see if x were stored in the low-precision format, but it remains an f64 so that the rest of the forward/backward pass runs in full precision. The gradient of this (piecewise-constant, hence a.e. zero-derivative) operator is supplied by the straight-through estimator (see fake_quant_backward).

Integer quantization (int8 / int4)

For an affine integer grid with step scale, integer zero-point zero_point and clamp range [qmin, qmax]:

  quant(x)   = clamp(round(x / scale) + zero_point, qmin, qmax)
  dequant(q) = (q - zero_point) * scale

Two schemes are supported (QuantScheme):

• Symmetric – zero_point = 0. The grid is symmetric about zero and the scale is derived from the absolute maximum. To keep the negative and positive arms equal in length we use the restricted signed range, i.e. int8 uses qmin = -127, qmax = 127 (the -128 code is dropped) and int4 uses qmin = -7, qmax = 7. scale = absmax / qmax.
• Affine (asymmetric) – the scale comes from the real [min, max] interval (nudged to include the real value 0) and zero_point is the integer code onto which the real value 0 maps. Affine uses the full signed range, int8 = [-128, 127], int4 = [-8, 7].

Note that because dequant(quant(x)) = round(x / scale) * scale the reconstructed grid is always a set of integer multiples of scale (the zero_point cancels in the round trip and only affects the asymmetric clamp window). Hence the per-element error is bounded by scale / 2 for round-to-nearest, which the test-suite checks.

fp8 quantization (E4M3 / E5M2)

Two 8-bit floating formats are modelled, decomposing each value into sign / exponent / mantissa and rounding the mantissa to the available bits, correctly handling normals, subnormals and saturation:

format	sign	exp	mantissa	bias	max-normal	min-normal	min-subnormal
E4M3	1	4	3	7	448	2^-6	2^-9
E5M2	1	5	2	15	57344	2^-14	2^-16

• E4M3 follows the OCP / deep-learning “E4M3” variant: there are no infinities, the only NaN encoding is S.1111.111, and the largest finite value is S.1111.110 = 1.75 * 2^8 = E4M3_MAX_NORMAL = 448. Its dynamic range runs from the smallest subnormal 2^-9 up to 448.
• E5M2 is IEEE-like (it has Inf/NaN at exponent field 11111); the largest finite value is S.11110.11 = 1.75 * 2^15 = E5M2_MAX_NORMAL = 57344, with dynamic range from 2^-16 up to 57344.

The cast implemented here is saturating: magnitudes above the format max (and any infinities) clamp to the format max rather than overflowing to Inf. NaN inputs propagate to NaN.

The rounding uses a single unified rule that is continuous across the normal/subnormal boundary. For a magnitude a with binade exponent e = floor(log2(a)), the unit-in-the-last-place is 2^(max(e, emin) - mbits) where emin = 1 - bias is the smallest normal exponent and mbits is the mantissa width; a is rounded to the nearest multiple of that ULP. For e >= emin this reproduces the normal-number grid; for e < emin it freezes at the subnormal granularity 2^(emin - mbits).

Rounding modes

RoundingMode selects between:

• Nearest – round-to-nearest-even (ties to even), the deterministic default.
• Stochastic – round up with probability equal to the fractional part and down otherwise, drawing a uniform [0, 1) variate from the supplied Rng. Stochastic rounding is unbiased: E[round(r)] = r, so the expected reconstruction equals the true value; the test-suite verifies this by Monte-Carlo averaging.

Straight-through estimator (STE)

Because quant is piecewise constant its true derivative is zero almost everywhere, which would block training. The STE replaces it with the identity on the representable interval: the incoming gradient passes through unchanged where the (pre-quant) value lies inside [qmin_real, qmax_real] and is zeroed outside (the clamp saturates, so no gradient flows). See fake_quant_backward.

QAT master-weight scheme

QatOptimizer implements the standard QAT bookkeeping: a full-precision master copy of the weights is what the optimizer (here an Adam/AdamW update) actually integrates, while the fake-quantized view of those master weights is what the forward pass “uses”. On every QatOptimizer::step the FP32 master is updated and the quantized view is re-derived (re-calibrated and re-rounded) from the new master. Keeping the master in full precision is essential: the tiny gradient steps would otherwise vanish under the quantization rounding and the network would never learn.

Structs

QatConfig

Configuration for a QatOptimizer: the quantization target plus the (Adam/AdamW) master-update hyper-parameters.

QatOptimizer

QAT optimizer wrapper maintaining full-precision master weights.

QuantParams

Affine/symmetric integer quantization parameters.

Enums

Fp8Format

8-bit floating-point formats modelled by fake_quant_fp8.

IntDtype

Supported signed-integer quantization widths.

QuantScheme

Quantization grid geometry: symmetric (zero-point pinned to 0) or affine (asymmetric, zero-point derived from the real minimum).

QuantTarget

What a QatOptimizer quantizes its weights to.

RoundingMode

Rounding rule applied when collapsing a real value onto the quantization grid.

Constants

E4M3_MAX_NORMAL

Largest finite (max-normal) magnitude of the E4M3 format, 1.75 * 2^8.

E5M2_MAX_NORMAL

Largest finite (max-normal) magnitude of the E5M2 format, 1.75 * 2^15.

Functions

fake_quant_backward

Straight-through estimator backward for a fake-quant op.

fake_quant_fp8

Fake-quantize a 1-D tensor onto an fp8 (E4M3 / E5M2) grid.

fake_quant_int

Fake-quantize a 1-D tensor through the integer grid (per-tensor).

fake_quant_int_per_channel

Fake-quantize a 2-D tensor with one QuantParams per channel.

per_channel_params

Calibrate one QuantParams per channel along axis of a 2-D tensor.