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

[sahas3]

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 #4600 (comment) 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: Future projects RFC #4499
• torchao PT2E: https://github.com/pytorch/ao
• quantized_decomposed defs: torch/ao/quantization/fx/_decomposed.py

[sahas3]

@zjgarvey, @sjarus, @Lallapallooza, @catcor01 -- please take a look and share your thoughts. Thanks!

[sjarus]

Strongly in support of this! I know there are folks who are working in this space at Arm.I've forwarded internally.

[zjgarvey]

This looks like a good plan to me. Do the quantized_decomposed ops work via the ods gen python scripts, or do we need to manually define these in the torch dialect? Not that I'm really excited about either path, particularly given that the ods gen path itself is tragically out of date.

[sahas3]

Do the quantized_decomposed ops work via the ods gen python scripts, or do we need to manually define these in the torch dialect?

The existing generator scraped the JIT operator registry (torch::jit::getAllOperators()) to produce GeneratedTorchOps.td. Since quantized_decomposed ops are registered in the torch.library the existing ods gen scripts won't handle these ops. Previously the only option was hand-authoring them in TorchOps.td.

I'm trying out a new approach that adds a second scraping path alongside the JIT registry one:

1. JitOperator.from_function_schema() — bridges the gap between FunctionSchema (what torch.ops.<ns>.<op>._schema exposes) and the OP_INFO_DICT format that the existing emit_op() machinery expects. The only translation needed is type string format: Optional[int] → int?, List[float] → float[], etc.

2. emit_torch_library_namespace(namespace, import_module, ops=[...]) — imports the module that registers the ops, enumerates overloads via torch.ops.<ns>, constructs a JitOperator per overload via from_function_schema(), and calls the same emit_op() path. The ops allowlist ensures only explicitly approved ops appear in GeneratedTorchOps.td — new ops added to the namespace in a future torch nightly won't silently appear.

So we can now add support for a new torch.library namespace via the existing emit_ops() + a update_torch_ods.sh run, instead of hand-authoring TableGen in TorchOps.td. Seems to be working well, but I am not very familiar with the intricacies of pytorch registry mechanics; so we can discuss more on the PR.

[zjgarvey]

I'm trying out a new approach that adds a second scraping path alongside the JIT registry one:

1. **`JitOperator.from_function_schema()`** — bridges the gap between `FunctionSchema` (what `torch.ops.<ns>.<op>._schema` exposes) and the `OP_INFO_DICT` format that the existing `emit_op()` machinery expects. The only translation needed is type string format: `Optional[int]` → `int?`, `List[float]` → `float[]`, etc.

2. **`emit_torch_library_namespace(namespace, import_module, ops=[...])`** — imports the module that registers the ops, enumerates overloads via `torch.ops.<ns>`, constructs a `JitOperator` per overload via `from_function_schema()`, and calls the same `emit_op()` path. The `ops` allowlist ensures only explicitly approved ops appear in `GeneratedTorchOps.td` — new ops added to the namespace in a future torch nightly won't silently appear.

So we can now add support for a new torch.library namespace via the existing emit_ops() + a update_torch_ods.sh run, instead of hand-authoring TableGen in TorchOps.td. Seems to be working well, but I am not very familiar with the intricacies of pytorch registry mechanics; so we can discuss more on the PR.

Honestly, if you can completely replace the JIT IR Importer path, that is the biggest blocker to number 4 on #4499 (remove pt1 legacy code). The ods gen (and abstract interp lib gen) scripts are based on old pytorch. I would be extremely happy to help push along anything which can help with that.

If you add a parallel path to reproduce the current GeneratedTorchOps.td- specifically, one which doesn't rely on having the legacy torchscript JIT-IR-importer path built- that's a huge win on it's own. I think what you are proposing above still relies on that JIT IR Importer tech via emit_op, so I'm not sure if that will be a good idea to invest time into. Better to have a chatbot handwrite the 6 quantized ops in tablegen.

[sahas3]

Honestly, if you can completely replace the JIT IR Importer path, that is the biggest blocker to number 4 on #4499 (remove pt1 legacy code).

That makes sense. I'll proceed with updating TorchOps.td for now and will try to look into this as a separate PR once the new quant path is stable.

[sahas3]

PR for Milestone 1 is here: #4611

[sahas3]

The following is in relation to the discussion happening in #4499 (comment) that is related to this RFC.

===========================================================

Thanks @ziereis. I think it was worth moving the discussion to the RFC thread so that it's captured with the design — so I am posting my reply here instead.

===========================================================

I do agree with the core of what you're getting at: these ops share a common legalization that really factors into just two patterns — but the split is by granularity, not by op name or operand type:

• Scalar-qparam (per-tensor): one scale/zp for the whole tensor. This covers quantize_per_tensor and its .tensor/.tensor2 overloads (those are numel-1 tensors, so effectively scalar — just extract the element), and torchao.quantize_affine when block_size spans the full input shape.
• Tensor-qparam (sub-tensor granularity): multiple qparams indexed by the granularity — per_channel, per_group, per_token, and torchao.quantize_affine when block_size is smaller than the full shape.

So the same torchao affine op lands in either pattern depending on block_size, and the .tensor overloads land in the scalar one despite carrying tensor operands. The point is the backend really only writes those two conversions, which lines up with your "backends shouldn't reimplement N QDQ forms" argument.

On the shared representation (your point 1): I think a single block_size-parameterized affine op is feasible. The tradeoff I see: the named per-granularity ops buy cleaner per-op verification and shape/dtype inference (each op has a fixed qparam rank, so the invariants are checkable), and they line up with the existing quantized_decomposed naming that PT2E actually emits. The single affine op is more uniform and almost certainly less op-definition code, but it pushes granularity classification (reading block_size against the input shape) into every backend matcher — e.g. a per-channel matmul fusion still has to decode block_size to know it's per-channel. So I'd frame it as "fewer ops, but each backend does more work per op" vs "more ops, but each is trivially matchable." Curious which way you weigh that — I am still leaning towards the latter.

Regarding point 2 about decomposition: It seems that those decompositions don't exist natively — neither PyTorch's default_decompositions() nor torch._decomp has the quantized_decomposed quantize ops, and torchao's affine ops only decompose downward into aten primitives (which we do not want). I hadn't thought of using a custom decomposition implementation as you shared; my interpretation of leveraging decomposition infra was to reuse existing pytorch decompositions so I don't think having all the ODS representations is in direct violation of proposal 1 in #4499. I think it's a matter of choice of implementing new decompositions either in C++ (ONNX path can leverage this) or python. As proposed the pre-import route means torch-mlir owns/maintains those mappings against torch nightly, and it only covers the FX path (ONNX never runs PyTorch decompositions). This is certainly doable but depends on what we decide as per the previous question.

I am not much familiar with the ONNX path yet, but my understanding for com.microsoft.MatMulIntegerToFloat is that it also falls in the same two op granularity of scalar/tensor qparams -- so we can still map that to the whatever torch op representation we choose to support for the actual legalizations in the torch-onnx-to-torch pass?

The RFC currently defers the ONNX route but that seems like a natural point to start the collaboration if you are familiar with that route once we land on the design?

I hadn't originally considered <8bit quantization in the RFC but the channel-wise quantization is slated for future milestones to keep the PR changes minimal. Any help in all these is definitely welcome :)

[zjgarvey]

I think I am going to suggest the following:

We just don't handle any qdq fusion in torch-mlir at all.

Just lower any quantization/dequantization ops directly, and we require subsequent compilers to pattern-match and fuse to a quantized/low-precision implementation if they want. If, for example, linalg consumers like IREE don't find it easy to pattern match the quantize/dequantize generics, they can make a named op and we can lower to that instead.

We don't need to do anything too sophisticated in torch-mlir. The primary purpose of torch-mlir is to faithfully lower pytorch code into a digestible core-mlir dialect. It would be good if we do that in an efficient way, but fusing qdq sequences during all "quantizable" op lowerings into all target dialects is a burden this project probably shouldn't take on.

[sahas3]

Hi @zjgarvey, so my understanding of your proposal is that in torch-mlir we only add ODS for all the variants of q/dq ops and their legalizations in the 3 different backends. We don't perform the fusion of dq->matmul->q to quant_matmul and let downstream users perform that fusion if they want to. For our (MathWorks) use-case, this fusion is a requirement : since we already have a custom pass-pipeline for legalizing torch to other dialects, I think we can add a pass that performs this fusion in our repo before either torch-to-tosa/linalg passes run. However, it's not clear how other consumers like IREE will support this fusion. I think they don't have a custom pass-pipeline, so trying to pattern match the already generic patterns of q/dq and matmul (and other ops) between them will be quite tricky.

I'll update #4611 to only add the ODS + legalization of the q/dq ops, while we discuss the details on the fusion implementation, which can always land as a separate PR if we decide to do so.

[ziereis]

Hello everyone, sorry for getting back to here so late.

On the shared representation (your point 1): I think a single block_size-parameterized affine op is feasible. The tradeoff I see: the named per-granularity ops buy cleaner per-op verification and shape/dtype inference (each op has a fixed qparam rank, so the invariants are checkable), and they line up with the existing quantized_decomposed naming that PT2E actually emits. The single affine op is more uniform and almost certainly less op-definition code, but it pushes granularity classification (reading block_size against the input shape) into every backend matcher — e.g. a per-channel matmul fusion still has to decode block_size to know it's per-channel. So I'd frame it as "fewer ops, but each backend does more work per op" vs "more ops, but each is trivially matchable." Curious which way you weigh that — I am still leaning towards the latter.

Generally don't have a strong opinion here, i just wanted to bring it up so we can consider it. If you think the named ops are easier to deal with then lets keep it this way.

Regarding point 2 about decomposition: It seems that those decompositions don't exist natively — neither PyTorch's default_decompositions() nor torch._decomp has the quantized_decomposed quantize ops, and torchao's affine ops only decompose downward into aten primitives (which we do not want). I hadn't thought of using a custom decomposition implementation as you shared; my interpretation of leveraging decomposition infra was to reuse existing pytorch decompositions so I don't think having all the ODS representations is in direct violation of proposal 1 in #4499. I think it's a matter of choice of implementing new decompositions either in C++ (ONNX path can leverage this) or python. As proposed the pre-import route means torch-mlir owns/maintains those mappings against torch nightly, and it only covers the FX path (ONNX never runs PyTorch decompositions). This is certainly doable but depends on what we decide as per the previous question.

I think even if there are no upstream decomposition for this this and we would have to maintain the custom decomps i generally think its less code and easier to maintain then the c++/mlir conversions. Either way this only really applies for converting from pt2e -> torchao since not all torchao ops can be converted to pt2e so we would need some conditional logic in the torch decomposition which is not possible as far i know. So if we go with the named/pt2e representation i agree i think its easier to do the conversion in c++/mlir.

I am not much familiar with the ONNX path yet, but my understanding for com.microsoft.MatMulIntegerToFloat is that it also falls in the same two op granularity of scalar/tensor qparams

Yes however i did take a quick look at your PR for the initial implementation for the scalar qparams. As far as i understood your current implementation rejects emitting the fused representation if the dimensionality of the LHS and RHS don't match. For the onnx case we would always have a 3D LHS + 2D RHS so this is a important case we need to support. However if we decide to emit it in QDQ form this comment might no longer apply.

We just don't handle any qdq fusion in torch-mlir at all. Just lower any quantization/dequantization ops directly, and we require subsequent compilers to pattern-match and fuse to a quantized/low-precision implementation if they want. If, for example, linalg consumers like IREE don't find it easy to pattern match the quantize/dequantize generics, they can make a named op and we can lower to that instead.

From what i have seen now i feel similarly in that handling the qdq fusions adds to much complexity (i can only speak the the linalg lowering i don't have enough knowledge about other backends). I think the main reason for this is that the linalg.quantized_matmul is just not expressive enough to support the kind of quantized model we see originating from torch and onnx.

Fundamentally all cases that have the following properties prevent us from using the linalg.(batch)_quantized_matmul:

-> arbitrary and mismatching numbers of batch dims of the LHS and RHS
-> all non scalar qparam cases
-> different kinds of quantization granularity for LHS and RHS (for example the LHS could have tensor qparams while the RHS has scalar qparams)
-> qparams with different signdness (currently we work around this with the shiftSign function but i also feel like this is not something that should be decided at the level of torch-mlir, as this is also very hard to reverse.

This might not sound like a lot here but from my experience most of the quantized models i encounter in my daily work fall into the categories from above.

So generally its a + for me in regards to keeping the QDQ form. It might even be possible to get some linalg.quantize_affine/ linalg.dequantize_affine op upstream, what do you think about this? Together with using linalg.contraction for all the higher/mismatching dimensionality case in the lowerings here should make matching the QDQ representation downstream quite feasible.

[sahas3]

It might even be possible to get some linalg.quantize_affine/ linalg.dequantize_affine op upstream, what do you think about this?

Yeah I think if there were some ops in the backend dialects that captures the semantics of the q/dq ops then it'll be easier to pattern match the dq->op->q chain and fuse that as qOp. The main challenge I see now is to faithfully keep the legalization of q/dq op in torch-mlir and in the downstream consumer passes tightly in sync to be able to infer a sequence of ops as performing quantize/dequantize operations.