Skip to main content

Module quant_insert

Module quant_insert 

Source
Expand description

Quantize/dequantize insertion pass.

The IR-rewrite half of post-training quantization. Given a per-tensor or per-channel calibration record (produced by running forward on a sample batch, see rlx_cpu::calibrate), this pass walks the graph and inserts Op::Quantize → Op::Dequantize pairs immediately downstream of each tagged node. Consumers of the original tap node are rewired to read the dequantized result, so everything past the tap sees an INT8 round-tripped activation / weight while the rest of the graph stays in fp32.

Why a Q/DQ pair instead of switching the whole subgraph to INT8? For PTQ this is the standard “fake-quant” pattern — the IR stays coherent in fp32, but each tap loses one quant step of precision to simulate the on-device int8 path. Real INT8-arithmetic kernels (Op::DequantMatMul, etc.) replace specific Q/DQ-bracketed regions later in the pipeline; this pass just produces the canonical form.

Scope intentionally narrow: insert-only, no measurement. The caller is responsible for filling CalibrationRecord from whatever execution path it has access to.

Structs§

CalibrationEntry
One calibrated quant entry per tap. axis = None is per-tensor; axis = Some(d) is per-channel along axis d, in which case scales and zero_points must each have length tap.shape.dim(d).

Functions§

insert_q_dq
Insert Quantize → Dequantize pairs at every tap in record. Returns a graph where each tagged node is followed by a Quantize → Dequantize pair, and every consumer of the original tap reads from the dequantized output instead.

Type Aliases§

CalibrationRecord
Map of tap NodeId → calibrated quant params.