Module inference_awq 

W1/AWQ — activation-statistics capture for Activation-aware Weight Quantization (no external libs).

AWQ’s premise: a weight’s salience is set by the magnitude of the activation it multiplies, so scaling each input channel by s_j = max|X_j| before quantizing preserves the channels that matter — letting aggressive (ternary) quantization survive. This module is AWQ step 1, the forward hook: during a calibration forward over the eval corpus it records, per FFN layer, the per-input-channel max |activation| at the FFN input (post-ffn_norm, the input to the gate/up projections).

Lock-free + gated: off in production (one relaxed atomic load on the FFN path). The accumulator is fetch_max over |x|.to_bits() — valid because |x| >= 0, so u32 bit-order matches float order. Heap (the stats buffer) is calibration-only, allocated once, never on a production hot path.

Functions

begin_forward

Call once at the start of each token’s forward so the layer cursor tracks layers 0..n_layer-1. No-op (one atomic load) when capture is off.

disable

enable

Begin an AWQ calibration capture for a model with n_layer FFN layers of n_chan input channels.

is_enabled

record_ffn_input

Record one FFN layer’s post-norm input channels. Called from the FFN forward; no-op when off.

reset

Zero the accumulators + reset the per-forward layer cursor.

snapshot

Per-layer per-channel max |activation| ([layer][chan]). Snapshot after the calibration pass.