mixle.engines.qlut module¶

Quantized function lookup tables – every scalar nonlinearity becomes table[code].

The corollary of quantization: once an operand is the integer code round(x/step), any scalar function f is a precomputed table indexed by that code – no transcendental ever runs. This turns the nonlinear ops of a model (activations: sigmoid/tanh/GELU/SiLU/softplus; the exp/log that convert between the log (LNS) and linear domains; the Gaussian-log of mixle.engines.lns) into integer gathers. Combined with integer-add products and the integer logsumexp, a fully-quantized forward pass has no floating-point exp/log at all.

Nearest-code lookup has error <= (step/2) * sup|f'| on the tabulated range (e.g. sigmoid: 0.125*step). Unbounded functions (GELU, softplus, ReLU) extrapolate linearly beyond the range using the boundary slope, so the table need only cover the curved region. Measured 1.5-8x faster than the real transcendental.

class QuantizedFunction(func, step, lo, hi)[source]

Bases: object

A scalar function tabulated over quantized inputs; f(x) ≈ table[round(x/step)] with linear tails.

Parameters:

func (Callable[[np.ndarray], np.ndarray])
step (float)
lo (float)
hi (float)

lookup(code)[source]

Gather directly from integer codes already in the quantized domain (the pure-integer path).

Parameters:: code (Any)
Return type:: ndarray

max_abs_error(x)[source]

Empirical max absolute error vs the true function over x.

Parameters:: x (Any)
Return type:: float

quantized_activation(name, step=0.01, span=20.0)[source]

A common NN activation as a quantized LUT (sigmoid/tanh/gelu/silu/softplus/relu).

Parameters:

name (str)
step (float)
span (float)

Return type:

QuantizedFunction

quantized_exp(log_step=0.01, lo_log=-30.0)[source]

exp from an LNS log-code (units of log_step) back to the linear domain – a pure table gather.

Call with integer log-codes via QuantizedFunction.lookup(); this is the softmax/attention “back to linear” step, with no real exp.

Parameters:

log_step (float)
lo_log (float)

Return type:

QuantizedFunction

quantized_logsumexp(scores, *, bits=12, span=24.0, weights=None)[source]

Log-sum-exp via an integer histogram over a span/2^bits grid plus ONE 2^bits-entry exp table.

Shift by the max, round every score to its grid level, count scores per level (an integer histogram – np.bincount), then a single dot with the exp table and one final log: 2^bits exp evaluations total (the table – 12 bits = 32 KB, cache-resident per table_bytes()) no matter how many scores there are, instead of one exp per score.

weights makes the histogram weighted, which is exactly the group-attention cell form (experiments/group_attention/RESULTS.md): passing per-cell integer counts computes the LSE of the token-level attention mass log sum_c count_c * exp(s_c) without ever expanding cells back to tokens. -inf scores (masked slots) are dropped, matching softmax semantics.

Error: rounding moves each score by at most delta/2 = span / 2^(bits+1), the grid term of lse_error_bound(). Scores more than span below the max are CLIPPED UP to the bottom bin instead of contributing their true (much smaller) mass. For typical, roughly-bounded score distributions that clipped mass is negligible, but it is NOT bounded by the grid term alone – a clipped score paired with a large weight can dominate the sum. Call lse_error_bound(bits, span, scores=scores, weights=weights) with the same data to get a bound that accounts for this rather than assuming the grid term is the whole story.

Parameters:

scores (Any)
bits (int)
span (float)
weights (Any)

Return type:

float

lse_error_bound(bits, span, *, scores=None, weights=None)[source]

Upper bound on |quantized_logsumexp(...) - exact LSE|.

With no data (scores=None), this is JUST the grid term span / 2^(bits+1) – the ordinary nearest-code rounding error. That alone is a COMPLETE bound only if no score ever falls more than span below the running max. quantized_logsumexp() instead CLIPS such scores up to the bottom bin, and nothing in bits/span limits how much weight a clipped score can carry, so the omitted tail is not generally negligible – concretely, scores=[0, -25], weights=[1, 1e12], bits=12, span=24 has actual error ~0.96 against a grid-only bound of ~0.0029, a ~300x understatement.

Pass the SAME scores/weights given to (or about to be given to) quantized_logsumexp() to get a bound that also accounts for the clipped tail: the grid term plus log1p(clipped_mass / in_span_mass), where clipped_mass is the (elevated) linear-domain contribution of out-of-span scores and in_span_mass is the linear-domain contribution of everything else. Derivation: the true linear sum is at least the true in-span mass (clipped terms only add) and at most the true in-span mass plus the FULLY elevated clipped mass (elevation only increases a clipped term’s contribution, never decreases it); the true in-span mass is within the grid term (in log-space) of the computed in-span mass by the ordinary nearest-code argument. Chaining those bounds and converting to log space gives the expression above. Verified empirically: it upper-bounds the exact error on the adversarial example above (bound ~3.67 vs actual ~0.96), stays within ~1e-9 of the plain grid term when nothing clips (matching this function’s pre-existing behavior), and held with zero violations across 20000 randomized trials with huge-dynamic-range weights.