# Copyright (C) 2026 Embedl AB
"""Lattice-specific activation quantization module."""
from collections.abc import Callable
from typing import Protocol, cast
import torch
from torch import nn
from embedl_deploy._internal.core.quantize.config import (
CalibrationMethod,
Precision,
)
from embedl_deploy._internal.core.quantize.stubs import QuantStub
class _SymbolicValue(Protocol):
"""Protocol for symbolic values used by ONNX export graph builders."""
def setType(self, tensor_type: object) -> None: # noqa: N802
"""Set tensor type metadata on the symbolic value."""
...
def type(self) -> object:
"""Return symbolic tensor type metadata."""
...
class _SymbolicGraph(Protocol):
"""Protocol for ONNX symbolic graph builders used by ``symbolic``."""
def op(
self,
name: str,
*args: object,
**kwargs: object,
) -> _SymbolicValue:
"""Create a symbolic operation node and return its output value."""
...
class _LatticeQuantFunction(torch.autograd.Function):
"""Autograd bridge for runtime fake-quant math and ONNX export.
The ONNX node has 8 tensor inputs ``(x + 7 params)``. We also pass
duplicated Python scalars to ``symbolic`` so the same values can be emitted
as ONNX attributes (``*_f`` and ``*_i``), matching the backend contract.
In ``backward``, autograd expects one returned entry per ``forward`` input.
Only ``x`` is differentiable here, so the implementation returns
``grad_output`` for ``x`` and ``None`` for every parameter/attribute input.
"""
@staticmethod
def forward( # noqa: PLR0913
ctx: object,
x: torch.Tensor,
min_rng: torch.Tensor,
max_rng: torch.Tensor,
num_bits: torch.Tensor,
resolution: torch.Tensor,
q_step_size: torch.Tensor,
is_lsq: torch.Tensor,
is_signed: torch.Tensor,
min_rng_f: float,
max_rng_f: float,
num_bits_i: int,
resolution_f: float,
q_step_size_f: float,
is_signed_b: bool,
) -> torch.Tensor:
"""Run runtime fake-quant math while carrying export-only args.
The scalar duplicates are needed only by ``symbolic`` during ONNX
export. They are intentionally ignored in eager execution.
"""
del (
ctx,
num_bits,
resolution,
is_lsq,
is_signed,
min_rng_f,
max_rng_f,
num_bits_i,
resolution_f,
q_step_size_f,
is_signed_b,
)
q = torch.round(x / q_step_size)
q = torch.clamp(q, min=min_rng, max=max_rng)
return q * q_step_size
@staticmethod
def backward(
ctx: object,
*grad_outputs: object,
) -> object:
del ctx
grad_output = cast(torch.Tensor, grad_outputs[0])
# Straight-through estimator for QAT-style gradient flow.
return (
grad_output,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
None,
)
@staticmethod
def symbolic( # noqa: PLR0913
g: _SymbolicGraph,
x: _SymbolicValue,
min_rng: _SymbolicValue,
max_rng: _SymbolicValue,
num_bits: _SymbolicValue,
resolution: _SymbolicValue,
q_step_size: _SymbolicValue,
is_lsq: _SymbolicValue,
is_signed: _SymbolicValue,
min_rng_f: float,
max_rng_f: float,
num_bits_i: int,
resolution_f: float,
q_step_size_f: float,
is_signed_b: bool,
) -> _SymbolicValue:
"""Emit a single contrib op with 8 tensor inputs and scalar attrs.
Input count and op count are different concepts: this emits one
``ai.onnx.contrib::LatticeQuant`` node that consumes 8 inputs. The
extra scalar arguments become ONNX attributes, not extra inputs.
"""
output = g.op(
"ai.onnx.contrib::LatticeQuant",
x,
min_rng,
max_rng,
num_bits,
resolution,
q_step_size,
is_lsq,
is_signed,
is_lsq_i=1,
is_signed_i=int(is_signed_b),
max_rng_f=max_rng_f,
min_rng_f=min_rng_f,
num_bits_i=num_bits_i,
q_step_size_f=q_step_size_f,
resolution_f=resolution_f,
)
output.setType(x.type())
return output
[docs]
class LatticeQuant(QuantStub):
"""Quantize activations and keep Lattice export identity metadata.
Extends ``QuantStub`` with two Lattice-specific capabilities:
- Tracks per-channel calibration maxima from observed activations.
- Stores stable identity strings used when remapping metadata to ONNX
node names during export.
:param consumers:
Set of modules that consume this quantizer's output.
:param precision:
Quantized format.
:param symmetric:
Whether quantization is symmetric.
:param calibration_method:
Activation calibration algorithm.
:param channel_axis:
Channel axis used when tracking per-channel calibration maxima.
:param fixed_calibration:
Optional fixed ``(scale, zero_point)`` pair.
:param canonical_name:
Optional canonical metadata key for this quantizer.
:param producer_qualified_name:
Optional qualified name of the producer module.
"""
def __init__(
self,
consumers: set[nn.Module],
precision: Precision = Precision.INT8,
symmetric: bool = True,
calibration_method: CalibrationMethod = CalibrationMethod.MINMAX,
*,
channel_axis: int = 1,
fixed_calibration: tuple[float, int] | None = None,
canonical_name: str | None = None,
producer_qualified_name: str | None = None,
) -> None:
super().__init__(
consumers,
precision=precision,
symmetric=symmetric,
calibration_method=calibration_method,
fixed_calibration=fixed_calibration,
)
self.channel_axis = channel_axis
self.canonical_name = canonical_name
self.producer_qualified_name = producer_qualified_name
self._calibration_max: torch.Tensor | None = None
self._calibration_min_scalar: torch.Tensor | None = None
self._calibration_max_scalar: torch.Tensor | None = None
num_bits = self.config.precision.n_bits
self.register_buffer(
"num_bits",
torch.tensor(num_bits, dtype=torch.int64),
)
self.register_buffer(
"resolution",
torch.tensor(2.0 ** float(num_bits), dtype=torch.float32),
)
self.register_buffer(
"q_step_size",
torch.tensor(1.0, dtype=torch.float32),
)
self.register_buffer(
"min_rng",
torch.tensor(float(self.config.quant_min), dtype=torch.float32),
)
self.register_buffer(
"max_rng",
torch.tensor(float(self.config.quant_max), dtype=torch.float32),
)
self.register_buffer(
"is_lsq",
torch.tensor([True], dtype=torch.bool),
)
self.register_buffer(
"is_signed",
torch.tensor(self.config.quant_min < 0, dtype=torch.bool),
)
@property
def calibration_max(self) -> torch.Tensor | None:
"""Return tracked per-channel calibration maxima, if available."""
if self._calibration_max is None:
return None
return self._calibration_max.clone()
[docs]
def reset_calibration_max(self) -> None:
"""Reset tracked per-channel calibration maxima."""
self._calibration_max = None
[docs]
def set_export_identity(
self,
canonical_name: str,
producer_qualified_name: str | None = None,
) -> None:
"""Set canonical identity strings used during metadata export.
:param canonical_name:
Canonical key used for export metadata maps.
:param producer_qualified_name:
Optional fully-qualified producer module name.
"""
self.canonical_name = canonical_name
self.producer_qualified_name = producer_qualified_name
def _observe_calibration_max(self, x: torch.Tensor) -> None:
"""Update per-channel calibration maxima from activation tensor `x`."""
x_abs = x.detach().abs()
if x_abs.ndim <= 1:
current = x_abs.amax().reshape(1)
else:
if self.channel_axis < 0:
axis = x_abs.ndim + self.channel_axis
else:
axis = self.channel_axis
if axis < 0 or axis >= x_abs.ndim:
raise ValueError(
"channel_axis is out of bounds for activation tensor"
)
reduce_dims = tuple(i for i in range(x_abs.ndim) if i != axis)
current = x_abs.amax(dim=reduce_dims)
if self._calibration_max is None:
self._calibration_max = current
else:
self._calibration_max = torch.max(self._calibration_max, current)
def _use_native_calibration(self) -> bool:
"""Return whether calibration should bypass ``torch.ao`` observers."""
return self.config.precision.n_bits > 8
def _observe_calibration_range(self, x: torch.Tensor) -> None:
"""Track scalar min/max needed to compute fake-quant parameters."""
x_detached = x.detach()
current_min = x_detached.amin()
current_max = x_detached.amax()
if self._calibration_min_scalar is None:
self._calibration_min_scalar = current_min
else:
self._calibration_min_scalar = torch.minimum(
self._calibration_min_scalar,
current_min,
)
if self._calibration_max_scalar is None:
self._calibration_max_scalar = current_max
else:
self._calibration_max_scalar = torch.maximum(
self._calibration_max_scalar,
current_max,
)
def _compute_native_parameters(self) -> tuple[torch.Tensor, torch.Tensor]:
"""Derive scale and zero-point without ``torch.ao`` observers."""
if (
self._calibration_min_scalar is None
or self._calibration_max_scalar is None
):
raise RuntimeError(
"QuantStub did not observe valid data during calibration"
)
min_val = self._calibration_min_scalar
max_val = self._calibration_max_scalar
if (
not torch.isfinite(min_val).item()
or not torch.isfinite(max_val).item()
):
raise RuntimeError(
"QuantStub did not observe valid data during calibration"
)
q_min = self.config.quant_min
q_max = self.config.quant_max
if self.config.symmetric:
max_abs = torch.max(min_val.abs(), max_val.abs())
max_abs = torch.clamp(max_abs, min=1e-8)
scale = max_abs / ((q_max - q_min) / 2)
zero_point = torch.zeros_like(scale, dtype=torch.int32)
return scale, zero_point
zeros = torch.zeros_like(min_val)
min_val = torch.minimum(min_val, zeros)
max_val = torch.maximum(max_val, zeros)
max_val = torch.maximum(max_val, min_val + 1e-8)
scale = (max_val - min_val) / (q_max - q_min)
zero_point = (
torch.round(-min_val / scale).clamp(q_min, q_max).to(torch.int32)
)
return scale, zero_point
[docs]
def compute_parameters(self) -> None:
"""Compute quant parameters and synchronize Lattice export buffers.
:raises ValueError:
If the computed zero-point is nonzero. Lattice hardware
does not support asymmetric quantization with a nonzero
zero-point; the eager and ONNX export paths both assume
zero-point is zero.
"""
if self._use_native_calibration():
scale, zero_point = self._compute_native_parameters()
self.scale.copy_(scale.squeeze())
self.zero_point.copy_(zero_point.squeeze())
else:
super().compute_parameters()
if self.zero_point.item() != 0:
raise ValueError(
"LatticeQuant does not support nonzero zero_point. "
"Use symmetric=True (signed range) for activations that "
"can be negative (e.g. post-Conv, post-LeakyReLU), or "
"symmetric=False (unsigned range) only for non-negative "
"activations (e.g. post-ReLU)."
)
# Note: symmetric=False (unsigned) is valid for non-negative
# activations at any bit width, including INT16. Post-ReLU data
# has min ≈ 0, so the asymmetric calibration path naturally
# produces zero_point ≈ 0 and passes the check above. The
# STATE_PREP patterns ensure symmetric is only set to False when
# the fused module's activation guarantees non-negative output.
scale = torch.as_tensor(self.scale, dtype=torch.float32)
with torch.no_grad():
q_step = cast(torch.Tensor, self.q_step_size)
q_step[...] = scale
def _forward_eager_quant(self, x: torch.Tensor) -> torch.Tensor:
"""Apply Lattice fake-quant math in eager mode (no ONNX export)."""
q_step_size = cast(torch.Tensor, self.q_step_size)
min_rng = cast(torch.Tensor, self.min_rng)
max_rng = cast(torch.Tensor, self.max_rng)
q = torch.round(x / q_step_size)
q = torch.clamp(q, min=min_rng, max=max_rng)
return q * q_step_size
def _forward_lattice_quant(self, x: torch.Tensor) -> torch.Tensor:
"""Apply Lattice quantization op using export-compatible parameters.
This path goes through ``_LatticeQuantFunction`` specifically to
trigger ONNX export's ``symbolic()`` method. Do not use it for ordinary
eager execution.
"""
apply_fn = cast(
Callable[..., torch.Tensor], _LatticeQuantFunction.apply
)
result = apply_fn(
x,
cast(torch.Tensor, self.min_rng),
cast(torch.Tensor, self.max_rng),
cast(torch.Tensor, self.num_bits),
cast(torch.Tensor, self.resolution),
cast(torch.Tensor, self.q_step_size),
cast(torch.Tensor, self.is_lsq),
cast(torch.Tensor, self.is_signed),
float(cast(torch.Tensor, self.min_rng)),
float(cast(torch.Tensor, self.max_rng)),
int(cast(torch.Tensor, self.num_bits)),
float(cast(torch.Tensor, self.resolution)),
float(cast(torch.Tensor, self.q_step_size)),
bool(cast(torch.Tensor, self.is_signed)),
)
return result
[docs]
def forward(self, x: torch.Tensor) -> torch.Tensor:
use_native_calibration = self._use_native_calibration()
if self.calibrating:
self._observe_calibration_max(x)
if use_native_calibration:
self._observe_calibration_range(x)
if not self.enabled:
if use_native_calibration:
return x
return super().forward(x)
# Keep current Embedl runtime behavior outside ONNX export while
# emitting ai.onnx.contrib::LatticeQuant when exporting.
if torch.onnx.is_in_onnx_export():
return self._forward_lattice_quant(x)
if use_native_calibration:
return self._forward_eager_quant(x)
return super().forward(x)
