[RFC] First-class PT2E quantization support in torch-mlir

Problem

Quantized inference is a primary deployment use case, and the modern PyTorch quantization path (PT2E, now shipped via torchao) is poorly served in torch-mlir today. Follows up on Discussion #4499.

What PT2E produces:

PT2E exports graphs where quantization is expressed as explicit quantized_decomposed.* quantize/dequantize ops surrounding normal float compute. A quantized linear looks like:

%x_q  = quantized_decomposed.quantize_per_tensor(%x, scale, zp, -128, 127, si8)   // float → si8
%x_dq = quantized_decomposed.dequantize_per_tensor(%x_q, scale, zp, -128, 127, si8) // si8 → float (models quant error)
%w_dq = quantized_decomposed.dequantize_per_tensor(%weight, scale, zp, -127, 127, si8) // symmetric weight
%lin  = aten.linear(%x_dq, %w_dq, %bias)            // float compute
%out_q = quantized_decomposed.quantize_per_tensor(%lin, scale, zp, -128, 127, si8) // re-quantize for next layer

Three properties drive the whole design:

• Quantization lives in the graph, not the type. scale/zp are explicit op operands and the tensors are plain si8/float — there is no !torch.qint8. This is precisely what lets qparams be runtime SSA values (dynamic quant).
• quant_min/quant_max are explicit (-128,127 for the activation, -127,127 for the symmetric weight) — the range and symmetry are visible, not inferred from dtype.
• There are no quantized compute ops. aten.linear/aten.relu run on float; the compiler is expected to fuse the surrounding dq → float_op → q into integer arithmetic (or fall back to scalar q/dq math).

The representation is fully self-contained. The quantized_decomposed namespace defines 19 ops spanning per-tensor / per-channel / per-token / per-channel-group granularities and runtime-qparam choose_qparams variants (dynamic and LLM quant). See the Op coverage section below for the full list.

What torch-mlir does today:

PT2E ops aren't registered in the Torch dialect, so the FX importer brings them in as generic torch.operator strings. The MatchQuantizedCustomOps shim then rewrites ~3 of the ~19 ops backwards into a legacy representation built on !torch.qint8 types and aten._make_per_tensor_quantized_tensor (MPTQT). Running the shim on the same linear graph above expands each self-contained PT2E op into a 2–3 op legacy chain:

// %x_q  (quantize_per_tensor) → qint8-typed quantize + int_repr + clamp
%q0 = torch.aten.quantize_per_tensor %x, %scale, %zp, %dtype : ...f32... -> ...!torch.qint8...
%r0 = torch.aten.int_repr %q0 : ...!torch.qint8... -> ...si8...
%c0 = torch.aten.clamp %r0, %amin, %amax : ...si8...

// %x_dq  (dequantize_per_tensor) → clamp + re-type to qint8 + dequantize
%c1 = torch.aten.clamp %c0, %amin, %amax : ...si8...
%m1 = torch.aten._make_per_tensor_quantized_tensor %c1, %scale, %zp : ...si8... -> ...!torch.qint8...
%d1 = torch.aten.dequantize.tensor %m1 : ...!torch.qint8... -> ...f32...

// %w_dq  (dequantize_per_tensor) → same clamp + _make_per_tensor_quantized_tensor + dequantize chain
...
%lin  = torch.aten.linear %d1, %w_dq, %bias    // float compute (unchanged)
// %out_q (quantize_per_tensor) → quantize + int_repr + clamp again

The self-contained representation is gone: qparams are re-baked into the !torch.qint8 type, the explicit range collapses into clamps, and redundant ops appear (the q → dq activation pair becomes a quantize/int_repr/clamp and a clamp/make/dequantize). FuseQuantizedOps then matches the resulting MPTQT → dq → op chains. The linalg/TOSA static-quant lowerings do work — but only when fed through this round-trip. The result:

• (a) No first-class representation — PT2E ops exist only as opaque torch.operator strings that this Python-injected shim (fx.py's _module_lowering) immediately rewrites away. Nothing in the Torch dialect understands them, so they get no verification, canonicalization, or shape/dtype inference.
• (b) Coverage gated on a 3-op shim — any op outside {quantize/dequantize_per_tensor, dequantize_per_channel} has no path in.
• (c) No dynamic quant — MPTQT (and any "fold qparams into the type/attribute" scheme) cannot express runtime qparams from choose_qparams, so dynamic/LLM quant is unrepresentable.

Forcing the modern, self-contained representation through a legacy bottleneck loses information and caps coverage. torch-mlir should consume what PyTorch actually emits.

Proposed design

Make quantized_decomposed.* first-class ops in the Torch dialect. Each TorchToX pass matches the dq → op → q structure shown above and lowers it to its backend's quantized target. Quantization parameters stay SSA values throughout — never folded into types or attributes. Because the q/dq ops already carry every parameter inline, the backend needs no side metadata; and because qparams remain SSA, the same path serves static and dynamic quant while preserving each backend's native fast path.

PT2E graph (torchao)
        │  quantized_decomposed.q / dq around float ops
        ▼
FX / ONNX importer  ──►  Torch dialect (quantized_decomposed.* first-class)
        │
        ▼
Simplify  (constant-fold dq(q(const)); dead q/dq elim; scale propagation
        │   through reshape/transpose/slice — NO fusion, NO qparam→attr)
        ▼
TorchTo{Linalg,Tosa,Stablehlo}   each matches  dq → op → q
        │                          static qparams → native quantized target
        │                          dynamic qparams → in-graph rescale math
        ▼
Backend IR

The Torch dialect does propagation/simplification only, not fusion; fusion decisions belong to each backend. q → dq cancellation is forbidden (it would erase intended quantization error); dq → q with matching qparams and dq(q(const)) folding are safe.

Scope

Milestone 1: prove the architecture end-to-end on a minimal slice — static per-tensor quantized matmul on TorchToLinalg.

• First-class ops (2): quantized_decomposed.quantize_per_tensor.default and quantized_decomposed.dequantize_per_tensor.default (scalar/constant qparams). The .tensor overloads and choose_qparams.* are deferred.
  - Lowering: static per-tensor dq → {aten.mm, aten.matmul, aten.bmm} → q fused to linalg.quantized_matmul / linalg.quantized_batch_matmul with an integer rescale epilogue (i32 accumulate → scale multiply → roundeven → clamp → truncate to i8).
• No simplification pass needed for the matmul slice — the lowering matches its own surrounding dq/q. General simplification (constant-fold dq(q(const)), dq motion through shape ops) is deferred.
• Success criterion: static per-tensor quantized matmul models import through FX and lower to linalg, with compiled output matching eager execution via the existing e2e harness.

** Note on the fusion behavior :** Based on https://github.com/llvm/torch-mlir/issues/4600#issuecomment-4939317495 the fusion implementation is now left off for downstream consumers. torch-mlir will only support legalizing a quantized model in the QDQ form but won't implement fusion to ensure the compute happens in float domain as represented in the QDQ form and not in integer domain.

Milestones

The existing e2e test suite already exercises the per-tensor static quantization path through the legacy shim. M1 takes that set of tests as the regression baseline. The new first-class quantized_decomposed ops must keep those models working end-to-end on TorchToLinalg before anything else is added. Coverage then expands outward:

• M1: per-tensor static q/dq + TorchToLinalg lowering (matches current repo coverage, no regressions).
• M2: per-tensor on TOSA and StableHLO (set up base for these backends).
• M3: per-channel (dequantize_per_channel for weights, mixed with per-tensor activation q/dq — the standard conv pattern).
• M4: per-channel-group (GPTQ/AWQ weight quant). No existing tests.
• M5: dynamic quant ops (per-token activations + choose_qparams family). No existing tests.
• M6: ONNX importer retargeted — QuantizeLinear / DequantizeLinear handlers emit quantized_decomposed ops instead of the legacy MPTQT representation. Once done, FuseQuantizedOps, UnpackQuantTensor, and the !torch.qint8 types can be removed.

From M3 onward, backend support will likely diverge based on the needs of each backend's users. Not every granularity needs to land on all three backends within the same milestone — a backend can skip or defer a granularity if there is no user demand for it. Each landed combination (granularity × backend) must have full test coverage (lit tests for ops, passes, and lowering, plus e2e tests) before the next one starts.

Pass removal is a two-track problem:

• MatchQuantizedCustomOps is only invoked in the FX importer pipeline (fx.py). Once all quantized_decomposed ops that appear in PT2E FX graphs have first-class ODS definitions, the shim finds nothing to rewrite and can be removed — no ONNX dependency. Removal can happen after M5 (or earlier, once the covered op set is sufficient for all tested FX models).
• FuseQuantizedOps is in the Linalg and TOSA backend pipelines. It fuses the legacy !torch.qint8 / MPTQT representation that the ONNX importer emits directly (not through torch.operator strings). It cannot be removed until M6 lands.

Note on the shim during transition: once the ODS definitions land, quantized_decomposed ops arrive as first-class typed ops instead of opaque torch.operator strings, so MatchQuantizedCustomOps will no longer find them to rewrite. PT2E tests for ops without a backend lowering yet will be temporarily broken until all backends are supported.

Op coverage

All 19 ops in the quantized_decomposed namespace have to be implemented for full quantization coverage. They are grouped below by granularity and milestone target. Overloads (.tensor, .tensor2) share the same ODS op and are distinguished by operand types.

Per-tensor (M1 / M2)

Op	Schema	Notes
quantize_per_tensor	(Tensor, float scale, int zp, int qmin, int qmax, ScalarType) → Tensor	scalar qparams; M1
quantize_per_tensor.tensor	(Tensor, Tensor scale, Tensor zp, int qmin, int qmax, ScalarType) → Tensor	runtime scalar qparams (dynamic quant prep)
quantize_per_tensor.tensor2	(Tensor, Tensor scale, Tensor zp, Tensor qmin, Tensor qmax, ScalarType) → Tensor	fully dynamic qparams
dequantize_per_tensor	(Tensor, float scale, int zp, int qmin, int qmax, ScalarType, *, ScalarType? out_dtype) → Tensor	scalar qparams; M1
dequantize_per_tensor.tensor	(Tensor, Tensor scale, Tensor zp, int qmin, int qmax, ScalarType, *, ScalarType? out_dtype) → Tensor	runtime scalar qparams
dequantize_per_tensor.tensor2	(Tensor, Tensor scale, Tensor zp, Tensor qmin, Tensor qmax, ScalarType, *, ScalarType? out_dtype) → Tensor	fully dynamic qparams
choose_qparams.tensor	(Tensor, int qmin, int qmax, float eps, ScalarType) → (Tensor scale, Tensor zp)	calibration; dynamic quant
choose_qparams_symmetric.tensor	(Tensor, int qmin, int qmax, float eps, ScalarType) → (Tensor scale, Tensor zp)	symmetric calibration

Per-channel (M3)

Op	Schema	Notes
quantize_per_channel	(Tensor, Tensor scales, Tensor zps, int axis, int qmin, int qmax, ScalarType) → Tensor	weight quant
dequantize_per_channel	(Tensor, Tensor scales, Tensor? zps, int axis, int qmin, int qmax, ScalarType, *, ScalarType? out_dtype) → Tensor	partial support exists in legacy shim

Per-channel-group (M4)

Op	Schema	Notes
quantize_per_channel_group	(Tensor, Tensor scales, Tensor zps, int qmin, int qmax, ScalarType, int group_size) → Tensor	GPTQ/AWQ weight quant
dequantize_per_channel_group	(Tensor, Tensor scales, Tensor? zps, int qmin, int qmax, ScalarType, int group_size, ScalarType output_dtype) → Tensor	GPTQ/AWQ weight dequant

Per-token (M5)

Op	Schema	Notes
quantize_per_token	(Tensor, Tensor scales, Tensor zps, int qmin, int qmax, ScalarType) → Tensor	dynamic activation quant (LLMs)
dequantize_per_token	(Tensor, Tensor scales, Tensor zps, int qmin, int qmax, ScalarType, ScalarType output_dtype) → Tensor	dynamic activation quant
choose_qparams_per_token_asymmetric	(Tensor, ScalarType) → (Tensor scale, Tensor zp)	per-row calibration
choose_qparams_per_token	(Tensor, ScalarType) → (Tensor scale, Tensor zp)	symmetric per-row calibration
_choose_qparams_per_token_asymmetric_impl	(Tensor, ScalarType) → (Tensor, Tensor)	internal impl; not expected in exported graphs

Utility / training (deferred)

Op	Schema	Notes
fake_quant_per_channel	(Tensor, Tensor scales, Tensor zps, int axis, int qmin, int qmax) → Tensor	training only; not needed for inference lowering
convert_element_type.no_fuse	(Tensor, ScalarType) → Tensor	dtype cast helper; defer until needed

Granularity support

The four granularities differ only in the shape of the scale and zero-point arguments:

• Per-tensor: single scalar scale and zp for the whole tensor. Static and dynamic variants (.tensor overloads with runtime Tensor qparams).
• Per-channel: one scale/zp per output channel. The canonical case for weight quantization in conv and linear.
• Per-token: one scale/zp per row (token) of the activation tensor. Critical for dynamic quantization of LLM activations — qparams are computed at runtime from choose_qparams_per_token.
• Per-channel-group: one scale/zp per group of group_size input-channel elements. The LLM-critical weight quant case (GPTQ, AWQ); scales are a 2-D tensor [out_channels, in_channels // group_size].

Open questions

1. Backend priority. Instinct is TorchToLinalg first (most complete quantized lowering infrastructure, exercises the full integer-compute-with-rescale story), then TOSA (cleanest native mapping to quantized integer ops with scale/zp attributes), then StableHLO. If your backend is higher priority for your use case, contributions are very welcome — we'd love to have more people involved.

2. ONNX importer scope. The ONNX importer emits the legacy MPTQT representation directly — it does not go through MatchQuantizedCustomOps. Retargeting QuantizeLinear / DequantizeLinear to emit quantized_decomposed ops is captured as M6 and is a prerequisite for removing FuseQuantizedOps. This is intentionally deferred until the PT2E path is proven out on the FX importer side.

3. FuseQuantizedOps and UnpackQuantTensor downstream users. Both passes cannot be removed until M6 lands. IREE is also relying on these passes. So, we will propose a deprecation window before removal once the PT2E path is fully implemented.

Migration / test impact

Registering an op in ODS is an atomic, global flip of the FX importer: it emits a first-class op instead of torch.operator exactly when is_registered_operation is true (fx_importer.py's _emit_operation). Adding the two per-tensor ops therefore diverts every existing FX-imported model using those op names away from the legacy MatchQuantizedCustomOps → MPTQT → FuseQuantizedOps path.

This means each PR must:

• Scope the registered op set to exactly what it can lower.
• xfail models that are flipped by the ODS registration but not yet covered by a backend lowering — these are part of the change, not a follow-up.
• Leave lit tests that hand-write torch.operator "torch.quantized_decomposed.*" unaffected — the shim still matches those opaque strings.

References

• Future-projects discussion: https://github.com/llvm/torch-mlir/discussions/4499
• torchao PT2E: https://github.com/pytorch/ao
• quantized_decomposed defs: torch/ao/quantization/fx/_decomposed.py

Start with fx.py and the MatchQuantizedCustomOps shim, then inspect the Torch dialect definitions and the existing quantization end-to-end tests. The first milestone is to represent the two static per-tensor quantized_decomposed q/dq ops as first-class operations and keep static quantized matmul models working through TorchToLinalg. Done means FX import and compiled output match eager execution without requiring the legacy shim for this slice.