Skip to content

Add fused_quant, a backend-neutral quantization frontend (#23091) - #23091

Open
mcremon-meta wants to merge 1 commit into
mainfrom
export-D119511656
Open

mcremon-meta wants to merge 1 commit into
mainfrom
export-D119511656

Conversation

@mcremon-meta

@mcremon-meta mcremon-meta commented Sep 23, 2026 •

Copy link
Copy Markdown
Contributor

Summary:

fused_quant represents a dequantize -> aten op -> quantize triple as a single
torch.ops.fused_quant.* operator that carries its quantization parameters
inline. Instead of a quantized graph being a float graph sprinkled with q/dq
pairs that each backend has to pattern-match and fold back together, the
quantized operator is the unit the graph is built from.

Taking a quantized convolution as the example -- before, the quantized op is
implicit. Five nodes, and the one doing the arithmetic is float:

  x_i8        w_i8       bias
    |           |          |
    v           v          |
dequantize  dequantize     |      <- scale, zero_point per edge
    |           |          |
    +-----------+----------+
                |
                v
        aten.convolution           <- float compute
                |
                v
             quantize              <- scale, zero_point
                |
                v
              y_i8

After, it is one node, and the qparams ride along with the operands they
describe:

  x_i8        w_i8       bias
    |           |          |
    +-----------+----------+
                |
                v
    fused_quant.convolution
        inp    : scale, zero_point, dtype, quant_min, quant_max
        weight : scale, zero_point, dtype, quant_min, quant_max
        bias   : scale, zero_point, dtype, quant_min, quant_max
        out    : scale, zero_point, dtype, quant_min, quant_max
                |
                v
              y_i8

Nothing is lost in the round trip: DecomposeFusedQuant turns the second form
back into the first, so a backend that cannot run a fused op still gets the
portable one. test_op_numerics checks this exactly, for every registered op.

This adds that representation under backends/fused_quant/. It is
ahead-of-time only: each op has a reference (eager) kernel that defines its
numerics, but no runtime kernel, so a fused_quant op never reaches
to_executorch. Either a backend claims it and lowers it to its own
instruction, or DecomposeFusedQuant turns it back into the q/dq + ATen form
above and the ordinary portable kernels run it. AssertNoFusedQuantOps fails
the compile if one leaks.

Operator representation and contract (ops.py, ops_utils.py, graph_utils.py)

Every fused_quant op takes its tensor operands first, then one qparam block
(scale, zero_point, dtype, quant_min, quant_max) per tensor operand,
then the out block, then the ATen op's remaining arguments in ATen order. A
block whose scale is None means that edge stays in float, so one schema
covers the fully quantized, partially quantized and float instances of an op
rather than needing a variant for each. For an input block dtype is the
dequantized float dtype; for out it is the quantized storage dtype.

Granularity is carried by the shape of the scale tensor -- per-tensor,
per-channel, per-group and blockwise are the same op with differently shaped
qparams. QuantParamsStruct.is_per_channel(tensor_shape) /
is_per_group(tensor_shape) take the quantized tensor's shape so a group scale
such as [1, K // g] over [1, K] is not mistaken for per-channel.

The README writes down what is stable for backends: the argument layout and
field meanings are frozen, a new capability is a new overload or a new trailing
argument with a behavior-preserving default, and the ops are never serialized.

Fake kernels report the dtype the eager kernel actually returns, including for a
partially quantized op (quantized input, float output), where they previously
reported the int8 storage dtype.

Quantizer (quantizer/)

A torchao PT2E Quantizer that both annotates and fuses:

  • OpQuantizer annotates one ATen op type and folds matched instances into a
    fused_quant op, given a QuantizationConfig and the schema names of its
    quantized activations, weights and passthrough operands.
  • NoopQuantizer claims nodes with an empty annotation so a later quantizer
    skips them, leaving the op in float.
  • MaxPoolQuantizer shares the input observer with output 0 of
    max_pool2d_with_indices.
  • FusedQuantQuantizer composes them. Annotation is first-match-wins, so a
    narrower quantizer is listed before the general one it overrides.

Annotations are marked _annotated, the flag other PT2E quantizers (XNNPACK,
Arm) check before claiming a node, and fuse() only folds nodes the same
quantizer annotated -- never one claimed by another quantizer.

preserved_ops() reports which ATen ops the quantizer matches, so the caller can
hold exactly those back from decomposition instead of maintaining a hardcoded
list. quantizer/qspecs.py builds per-tensor, per-channel and blockwise static
specs (dynamic qspecs are not offered, since fusion needs constant qparams);
quantizer/defaults.py provides a default op set, with get_core_quantizers
for backends that lack per-channel support or the wider op set. Weights default
to symmetric in both make_qconfig and make_fused_quant_quantizer. A
transposed convolution gets its per-channel weight scales on axis 1, where its
output channels are, or per-tensor weights when groups > 1.

Coloring (colorer.py)

ColorerBase is the step between optimization and lowering: a backend declares
which fused ops it can run (supported_ops) and whether a given instance is
within its limits (is_legal), and the base class tags the rest with
delegation_tag. Colorers compose by priority: the first to claim a node keeps
it, and a later colorer skips already-tagged nodes.

The tag is a claim of ownership and the generic passes honour it -- neither
ToChannelsFirst nor DecomposeFusedQuant will touch a colored node. So the
pipeline is colorer -> ToChannelsFirst -> DecomposeFusedQuant: a backend keeps
what it claimed, and everything it did not claim falls back to q/dq plus ATen
automatically. A colored node the partitioner then rejects is caught by running
DecomposeFusedQuant again after to_backend.

A backend that claims nothing is just the degenerate case of this: every fused
op decomposes and the program runs entirely on the ordinary portable ATen
kernels. test_compile.py has that written out end to end.

Passes

  • pre_quantize_passes/ runs on the traced program before annotation: batch-norm
    folding, replacing leftover batch norms with convolutions, and recomposing
    add(mask) + softmax into _masked_softmax so the very-negative mask sentinel
    never enters the quantized domain.
  • optimization_passes/ runs on the fused edge graph: layout canonicalization,
    activation hoisting, folding add/mul into linear, quantization absorption,
    embedding prequantization, linear splitting, requantize formation.
    IterativePassGroup logs a warning when a group hits its step cap.
  • decompose_fused_quant.py is the inverse: it turns any fused op back into
    dequantize + ATen + quantize. It is schema-driven, so it covers every
    registered fused op rather than a fixed list, and handles multi-output ops
    through getitem. It validates every node before rewriting any, so an
    unsupported one (grouped/blockwise qparams) fails the pass without leaving a
    partly decomposed graph. Per-channel qparams are emitted as 1-D constants
    rather than a view_copy of the full-rank one, and the qparam constants left
    dead are removed.

Shared passes moved to backends/transforms

The optimization group no longer depends on any backend. It used 22 passes from
backends/cadence/aot; none were Cadence-specific in what they do:

  • Three (FuseCascadedTransposeOrPermuteOps, FuseCascadedViewOps,
    ReplaceNopTransposeOrPermuteWithViewPass) were already empty subclasses of a
    backends/transforms pass, which is now imported directly.
  • The other 19 move to backends/transforms, one module and target each. The
    ones that listed Cadence's q/dq ops alongside quantized_decomposed's now
    take those op sets as class attributes holding only the standard ops.
    Cadence keeps each pass under its old name as a thin subclass that adds its
    own ops back, so Cadence's behavior is unchanged.
  • FuseQuantDequantToRequantizePass needs a requantize op to fuse a pair whose
    qparams differ, and there is no backend-neutral one. The shared pass takes a
    requantize_op (default None) and defaults to allow_requantize=False;
    Cadence sets its op and keeps True.
  • permute_pass_utils gains get_overload_packet and get_placeholders, and
    quantize_reorder_utils holds the two op sets the reorder passes share.

Tests no longer import Cadence either: test/helpers.trace replaces
cadence.aot.compiler.trace, and the unused ops_registrations imports are gone.
buck2 uquery finds no backends/cadence target in the dependency closure of
any fused_quant target.

Pipeline (frontend.py, passes.py, compile.py)

compile_to_fused_quant runs trace -> prepare -> calibrate -> convert -> fuse ->
export -> to_edge -> passes, and stops. That is the seam a backend plugs into,
and it is the only entry point the package offers: from there a backend runs its
own colorer, lowering and partitioner, and hands whatever it did not claim to
DecomposeFusedQuant.

trace(..., is_qat=True) exports in training mode and keeps batch norms, for
prepare_qat_pt2e; by default it exports in eval mode with batch norms folded.
Either way the caller's model is left in the mode it was passed in.

The pass list is a parameter, defaulting to get_fused_quant_passes(). Callers
that want a different pipeline pass their own, and an empty list skips
optimization entirely.

One wrinkle worth knowing when writing a backend: DecomposeFusedQuant re-emits
the ATen ops that were held back from decomposition, and several of them
(aten.linear, aten.silu, ...) are not core ATen, so the edge verifier rejects
them. compile.py feeds quantizer.preserved_ops() into
EdgeCompileConfig._core_aten_ops_exception_list, which is exact by
construction: the ops held back are precisely the ops re-emitted.

Known gaps

  • The package depends on torchao, which is not yet a core ExecuTorch dependency.
  • DecomposeFusedQuant rejects grouped/blockwise qparams rather than emitting
    dequantize_affine, so an unclaimed blockwise op fails the compile.
  • Bias stays float in the default configs; int32 bias is neither configurable
    nor documented yet.
  • Op registration happens on import, and passes.py looks up
    exir_ops.edge.fused_quant.* at module scope.
  • Follow-ups from review, deliberately not in this diff: splitting the pass
    pipeline into opt-in groups (which takes ToChannelsLast out of the default),
    a colored-node Partitioner adapter, fitting the quantizer and frontend into
    ComposableQuantizer / export() recipes, a QAT fake-quant path, generating
    schemas from the ATen ones, and the core/transforms/quantizer placement split.

Differential Revision: D119511656

@pytorch-bot

pytorch-bot Bot commented Sep 23, 2026 •

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/23091

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 8321a9b with merge base f2575be (image):
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Sep 23, 2026
@meta-codesync

meta-codesync Bot commented Sep 23, 2026

Copy link
Copy Markdown
Contributor

@mcremon-meta has exported this pull request. If you are a Meta employee, you can view the originating Diff in D119511656.

@github-actions

Copy link
Copy Markdown

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

meta-codesync Bot pushed a commit that referenced this pull request Sep 24, 2026
Summary:
Pull Request resolved: #23091

`fused_quant` represents a `dequantize -> aten op -> quantize` triple as a single
`torch.ops.fused_quant.*` operator that carries its quantization parameters
inline. Instead of a quantized graph being a float graph sprinkled with q/dq
pairs that each backend has to pattern-match and fold back together, the
quantized operator is the unit the graph is built from.

Taking a quantized convolution as the example -- before, the quantized op is
implicit. Five nodes, and the one doing the arithmetic is float:

      x_i8        w_i8       bias
        |           |          |
        v           v          |
    dequantize  dequantize     |      <- scale, zero_point per edge
        |           |          |
        +-----------+----------+
                    |
                    v
            aten.convolution           <- float compute
                    |
                    v
                 quantize              <- scale, zero_point
                    |
                    v
                  y_i8

After, it is one node, and the qparams ride along with the operands they
describe:

      x_i8        w_i8       bias
        |           |          |
        +-----------+----------+
                    |
                    v
        fused_quant.convolution
            inp    : scale, zero_point, dtype, quant_min, quant_max
            weight : scale, zero_point, dtype, quant_min, quant_max
            bias   : scale, zero_point, dtype, quant_min, quant_max
            out    : scale, zero_point, dtype, quant_min, quant_max
                    |
                    v
                  y_i8

Nothing is lost in the round trip: `DecomposeFusedQuant` turns the second form
back into the first, so a backend that cannot run a fused op still gets the
portable one.

This adds that representation under `backends/fused_quant/`. It is
ahead-of-time only: a `fused_quant` op has no kernel and is never executed
directly. Either a backend claims it and lowers it to its own instruction, or
`DecomposeFusedQuant` turns it back into the q/dq + ATen form above and the
ordinary portable kernels run it.

**Operator representation** (`ops.py`, `ops_utils.py`, `graph_utils.py`)

Every `fused_quant` op takes its tensor operands first, then one qparam block
(`scale`, `zero_point`, `dtype`, `quant_min`, `quant_max`) per quantizable edge,
including the output. A block whose `scale` is `None` means that edge stays in
float, so one schema covers the fully quantized, partially quantized and float
instances of an op rather than needing a variant for each. Granularity is
carried by the shape of the `scale` tensor -- per-tensor, per-channel, per-group
and blockwise are the same op with differently shaped qparams.

**Quantizer** (`quantizer/`)

A torchao PT2E `Quantizer` that both annotates and fuses:

- `OpQuantizer` annotates one ATen op type and folds matched instances into a
  `fused_quant` op, given a `QuantizationConfig` and the schema names of its
  quantized activations, weights and passthrough operands.
- `NoopQuantizer` claims nodes with an empty annotation so a later quantizer
  skips them, leaving the op in float.
- `MaxPoolQuantizer` shares the input observer with output 0 of
  `max_pool2d_with_indices`.
- `FusedQuantQuantizer` composes them. Annotation is first-match-wins, so a
  narrower quantizer is listed before the general one it overrides.

`preserved_ops()` reports which ATen ops the quantizer matches, so the caller can
hold exactly those back from decomposition instead of maintaining a hardcoded
list. `quantizer/qspecs.py` builds per-tensor, per-channel and blockwise specs;
`quantizer/defaults.py` provides a default op set, with `get_core_quantizers`
for backends that lack per-channel support or the wider op set.

**Coloring** (`colorer.py`)

`ColorerBase` is the step between optimization and lowering: a backend declares
which fused ops it can run (`supported_ops`) and whether a given instance is
within its limits (`is_legal`), and the base class tags the rest with
`delegation_tag`.

The tag is a claim of ownership and the generic passes honour it -- neither
`ToChannelsFirst` nor `DecomposeFusedQuant` will touch a colored node. So the
pipeline is colorer -> ToChannelsFirst -> DecomposeFusedQuant: a backend keeps
what it claimed, and everything it did not claim falls back to q/dq plus ATen
automatically. A backend only has to describe what it supports, not handle what
it doesn't.

A backend that claims nothing is just the degenerate case of this: every fused
op decomposes and the program runs entirely on the ordinary portable ATen
kernels. `test_compile.py` has that written out end to end.

**Passes**

- `pre_quantize_passes/` runs on the traced program before annotation: batch-norm
  folding, replacing leftover batch norms with convolutions, and recomposing
  `add(mask) + softmax` into `_masked_softmax` so the very-negative mask sentinel
  never enters the quantized domain.
- `optimization_passes/` runs on the fused edge graph: layout canonicalization,
  activation hoisting, folding add/mul into linear, quantization absorption,
  embedding prequantization, linear splitting, requantize formation.
- `decompose_fused_quant.py` is the inverse: it turns any fused op back into
  dequantize + ATen + quantize. It is schema-driven, so it covers every
  registered fused op rather than a fixed list, handles multi-output ops through
  `getitem`, and flattens full-rank per-channel qparams to the 1-D form
  `quantized_decomposed` expects.

**Pipeline** (`frontend.py`, `passes.py`, `compile.py`)

`compile_to_fused_quant` runs trace -> prepare -> calibrate -> convert -> fuse ->
export -> to_edge -> passes, and stops. That is the seam a backend plugs into,
and it is the only entry point the package offers: from there a backend runs its
own colorer, lowering and partitioner, and hands whatever it did not claim to
`DecomposeFusedQuant`.

The pass list is a parameter, defaulting to `get_fused_quant_passes()`. Callers
that want a different pipeline pass their own, and an empty list skips
optimization entirely.

One wrinkle worth knowing when writing a backend: `DecomposeFusedQuant` re-emits
the ATen ops that were held back from decomposition, and several of them
(`aten.linear`, `aten.silu`, ...) are not core ATen, so the edge verifier rejects
them. `compile.py` feeds `quantizer.preserved_ops()` into
`EdgeCompileConfig._core_aten_ops_exception_list`, which is exact by
construction: the ops held back are precisely the ops re-emitted.

**Known gaps**, also in the README: the package depends on `torchao`, which is
not yet a core ExecuTorch dependency; and the optimization pass group still
imports about fifteen passes from `backends/cadence/aot` that are not
Cadence-specific and would be better placed in `backends/transforms`.

Differential Revision: D119511656
@meta-codesync meta-codesync Bot changed the title Add fused_quant, a backend-neutral quantization frontend Add fused_quant, a backend-neutral quantization frontend (#23091) Sep 24, 2026
meta-codesync Bot pushed a commit that referenced this pull request Sep 24, 2026
Summary:
Pull Request resolved: #23091

`fused_quant` represents a `dequantize -> aten op -> quantize` triple as a single
`torch.ops.fused_quant.*` operator that carries its quantization parameters
inline. Instead of a quantized graph being a float graph sprinkled with q/dq
pairs that each backend has to pattern-match and fold back together, the
quantized operator is the unit the graph is built from.

Taking a quantized convolution as the example -- before, the quantized op is
implicit. Five nodes, and the one doing the arithmetic is float:

      x_i8        w_i8       bias
        |           |          |
        v           v          |
    dequantize  dequantize     |      <- scale, zero_point per edge
        |           |          |
        +-----------+----------+
                    |
                    v
            aten.convolution           <- float compute
                    |
                    v
                 quantize              <- scale, zero_point
                    |
                    v
                  y_i8

After, it is one node, and the qparams ride along with the operands they
describe:

      x_i8        w_i8       bias
        |           |          |
        +-----------+----------+
                    |
                    v
        fused_quant.convolution
            inp    : scale, zero_point, dtype, quant_min, quant_max
            weight : scale, zero_point, dtype, quant_min, quant_max
            bias   : scale, zero_point, dtype, quant_min, quant_max
            out    : scale, zero_point, dtype, quant_min, quant_max
                    |
                    v
                  y_i8

Nothing is lost in the round trip: `DecomposeFusedQuant` turns the second form
back into the first, so a backend that cannot run a fused op still gets the
portable one.

This adds that representation under `backends/fused_quant/`. It is
ahead-of-time only: a `fused_quant` op has no kernel and is never executed
directly. Either a backend claims it and lowers it to its own instruction, or
`DecomposeFusedQuant` turns it back into the q/dq + ATen form above and the
ordinary portable kernels run it.

**Operator representation** (`ops.py`, `ops_utils.py`, `graph_utils.py`)

Every `fused_quant` op takes its tensor operands first, then one qparam block
(`scale`, `zero_point`, `dtype`, `quant_min`, `quant_max`) per quantizable edge,
including the output. A block whose `scale` is `None` means that edge stays in
float, so one schema covers the fully quantized, partially quantized and float
instances of an op rather than needing a variant for each. Granularity is
carried by the shape of the `scale` tensor -- per-tensor, per-channel, per-group
and blockwise are the same op with differently shaped qparams.

**Quantizer** (`quantizer/`)

A torchao PT2E `Quantizer` that both annotates and fuses:

- `OpQuantizer` annotates one ATen op type and folds matched instances into a
  `fused_quant` op, given a `QuantizationConfig` and the schema names of its
  quantized activations, weights and passthrough operands.
- `NoopQuantizer` claims nodes with an empty annotation so a later quantizer
  skips them, leaving the op in float.
- `MaxPoolQuantizer` shares the input observer with output 0 of
  `max_pool2d_with_indices`.
- `FusedQuantQuantizer` composes them. Annotation is first-match-wins, so a
  narrower quantizer is listed before the general one it overrides.

`preserved_ops()` reports which ATen ops the quantizer matches, so the caller can
hold exactly those back from decomposition instead of maintaining a hardcoded
list. `quantizer/qspecs.py` builds per-tensor, per-channel and blockwise specs;
`quantizer/defaults.py` provides a default op set, with `get_core_quantizers`
for backends that lack per-channel support or the wider op set.

**Coloring** (`colorer.py`)

`ColorerBase` is the step between optimization and lowering: a backend declares
which fused ops it can run (`supported_ops`) and whether a given instance is
within its limits (`is_legal`), and the base class tags the rest with
`delegation_tag`.

The tag is a claim of ownership and the generic passes honour it -- neither
`ToChannelsFirst` nor `DecomposeFusedQuant` will touch a colored node. So the
pipeline is colorer -> ToChannelsFirst -> DecomposeFusedQuant: a backend keeps
what it claimed, and everything it did not claim falls back to q/dq plus ATen
automatically. A backend only has to describe what it supports, not handle what
it doesn't.

A backend that claims nothing is just the degenerate case of this: every fused
op decomposes and the program runs entirely on the ordinary portable ATen
kernels. `test_compile.py` has that written out end to end.

**Passes**

- `pre_quantize_passes/` runs on the traced program before annotation: batch-norm
  folding, replacing leftover batch norms with convolutions, and recomposing
  `add(mask) + softmax` into `_masked_softmax` so the very-negative mask sentinel
  never enters the quantized domain.
- `optimization_passes/` runs on the fused edge graph: layout canonicalization,
  activation hoisting, folding add/mul into linear, quantization absorption,
  embedding prequantization, linear splitting, requantize formation.
- `decompose_fused_quant.py` is the inverse: it turns any fused op back into
  dequantize + ATen + quantize. It is schema-driven, so it covers every
  registered fused op rather than a fixed list, handles multi-output ops through
  `getitem`, and flattens full-rank per-channel qparams to the 1-D form
  `quantized_decomposed` expects.

**Pipeline** (`frontend.py`, `passes.py`, `compile.py`)

`compile_to_fused_quant` runs trace -> prepare -> calibrate -> convert -> fuse ->
export -> to_edge -> passes, and stops. That is the seam a backend plugs into,
and it is the only entry point the package offers: from there a backend runs its
own colorer, lowering and partitioner, and hands whatever it did not claim to
`DecomposeFusedQuant`.

The pass list is a parameter, defaulting to `get_fused_quant_passes()`. Callers
that want a different pipeline pass their own, and an empty list skips
optimization entirely.

One wrinkle worth knowing when writing a backend: `DecomposeFusedQuant` re-emits
the ATen ops that were held back from decomposition, and several of them
(`aten.linear`, `aten.silu`, ...) are not core ATen, so the edge verifier rejects
them. `compile.py` feeds `quantizer.preserved_ops()` into
`EdgeCompileConfig._core_aten_ops_exception_list`, which is exact by
construction: the ops held back are precisely the ops re-emitted.

**Known gaps**, also in the README: the package depends on `torchao`, which is
not yet a core ExecuTorch dependency; and the optimization pass group still
imports about fifteen passes from `backends/cadence/aot` that are not
Cadence-specific and would be better placed in `backends/transforms`.

Differential Revision: D119511656
## Directory structure

```
backends/fused_quant/

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
backends/fused_quant/
backends/transforms/fused_quant/

@@ -0,0 +1,68 @@
# fused_quant

`fused_quant` is a backend-neutral quantization frontend for ExecuTorch. It represents a

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What are your thoughts on making this extensible so backends can customize to their needs? Because that's what they do today to match their hardware capabilities, hence different quantizers etc. There might be overlap which this can help dedup, but I do expect some need for specialization per backend if this were to adopted widely, IMHO.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

of course, that's what all of the "adoption" follow ups will do. This is just landing the infra layers

@@ -0,0 +1,68 @@
# fused_quant

`fused_quant` is a backend-neutral quantization frontend for ExecuTorch. It represents a

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: rather than calling this fused_quant, this should reflect the quant_scheme(s) it reflects or something else like "some_name_standard_int8_symmetric_quant_IR" (obv not that) but something like this?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

well the whole point of this layer is that it supports all of those schemes natively, together. Including fp32 actually, where you would "just" get the optimizations passes and layout changes. cc @JakeStevens since we talked about that briefly

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe a bit of bike shedding but the fused just leads my thoughts to fusing functional ops. Maybe just quantized_ir ?

@rascani
rascani self-requested a review September 24, 2026 17:09
Comment thread backends/fused_quant/graph_utils.py Outdated
return val.fake_mode


def add_constant(

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Probably should reuse constant creation functions in backends/transforms/utils.py

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Probably should reuse constant creation functions in backends/transforms/utils.py

yes, let me see if we have more of those opportunities as well

@@ -0,0 +1,25 @@
# Copyright (c) Meta Platforms, Inc. and affiliates.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why not put this in exir.passes?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we can do that for some of those for sure. Let me see if there are more candidates

Comment thread backends/fused_quant/frontend.py Outdated

The rewrites run here rather than in the caller because they must happen
after tracing (they match on ATen ops) but before annotation (the quantizer
must see the rewritten ops). Which ones run is derived from the quantizer:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

curious why?

Comment thread backends/fused_quant/frontend.py Outdated
fold_quantize_into_mutable_buffers=fold_quantize_into_mutable_buffers,
)
_normalize_affine_overloads(gm)
# convert_pt2e rewrites the graph without refreshing meta["val"], so shape and

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

perhaps we should fix it in the convert?



@torch.no_grad()
def trace(

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

users?

@digantdesai

Copy link
Copy Markdown
Contributor

@claude review this PR. Focus on architectural aspects of the design and how this is setup to be adopted by existing backends as an IR Toolkit for efficient lowering, while allowing them to be opinionated about their hardware. Also other standard ET features like Quantizer composability and same for partitioner - if someone were to write a recipe for a composed setup. Lastly custom op versioning, BC contract, QAT support, for the dialect.

How can optimization passes can be split and opt-in at a finer granularity? Also does exir dir makes a better home for this?

@claude

claude Bot commented Oct 1, 2026 •

Copy link
Copy Markdown

Claude finished @digantdesai's task in 1m 38s —— View job


Summary: The fused-op IR idea is sound, but this PR isn't ready to serve as a shared toolkit for backends yet. It only works through its own compile_to_fused_quant driver, and several pieces don't fit together: the colorer, partitioner, quantizer and recipe paths. The ops have no BC contract and no QAT path. The "backend-neutral" passes enforce NHWC layout and depend on Cadence. One correctness bug: fake kernels report the wrong dtype for partially quantized ops.

Detailed review

PR Review: #23091

Summary

This adds backends/fused_quant/: a fused_quant::* op library that carries qparams inline, a torchao PT2E quantizer that also fuses, a colorer base class, about 38 edge optimization passes, and DecomposeFusedQuant as the fallback for unclaimed ops. The IR is useful as a concept. The problems are that a backend can't adopt it piece by piece, and two of the PR's own promises break at the edges: the fallback (blockwise ops, per-channel ops into XNNPACK) and the backend-neutral claim (NHWC, Cadence).

API Design

Op dialect, versioning and BC

  • Fake-kernel dtype bug when only the output is float. For input qparam blocks, dtype means the dequantized float dtype (ops.py:65-66). Most fake kernels, however, fall back to inp.dtype, which is the int8 storage dtype, when out_qp is None. Examples: get_out_dtype(out_qp, inp.dtype) at ops.py:298 (binary), :319 (unary), :345 (scalar), :1361 (linear). The same pattern appears for bmm, conv, conv_channels_last and the pools. The eager kernel returns float, so meta["val"] is int8 for a quantized-input/float-output node. Downstream passes and partitioners read that meta. native_layer_norm gets it right (ops.py:2283, compute_dtype = inp_qp.dtype if inp_qp is not None else inp.dtype), and the other ops should do the same. Please add a test for this case, since the README advertises partially quantized ops.

  • No BC contract or versioning. The schemas are positional: tensor operands, then 5-field qparam blocks, then extra args. There's no schema version, no overload naming rule and no deprecation path. Hidden assumptions:

    • split_fused_arg_names (graph_utils.py:32-53) locates the qparam blocks by the first name ending in _scale and allows exactly one output block named out.
    • fuse_aten checks extra-arg names but emits them in ATen positional order. If a fused schema's extras are ever ordered differently from ATen, args bind to the wrong parameters with no error.
    • Adding a field to the qparam block, or quantizing a second output, breaks every backend lowering and the decomposer at once.

    Before backends depend on this, please document:

    • what is frozen: the schema layout and the meaning of each field;
    • how a field is added (new overload vs. a new optional argument at the end);
    • that fused_quant ops are AOT-only and must never be serialized.
  • The schemas are hand-written and repetitive. ops.py is 2572 lines, and each op restates 11–27 parameters for both the impl and the fake kernel. Generating the fused schema from aten_op._schema plus a list of quantizable edges (_register_conv_nd at ops.py:1539 already heads that way) would remove most of the file and the ordering risk above. The registration styles are also mixed: Library("fused_quant","DEF") is used only for add.Scalar/mul.Scalar (ops.py:1146,1202), with Scalar other, while sub.Scalar uses custom_op (ops.py:1214) with float other, despite the comment at :1092-1096.

  • Qparam ambiguities that backends will hit:

    • QuantParamsStruct.is_per_channel / channel_axis (ops.py:99-124) count the non-unary scale dims but don't check block_size == 1 on that axis. A 1-D group scale such as [1, K/g] on [1, K] is reported as per-channel.
    • A C=1 per-channel scale is reported as per-tensor.
    • Per-tensor scales come in two shapes: () from fuse_aten.py:55 and [1,…,1] from affine q/dq.
  • The README contradicts the code. It says "no kernels", but every op has a real eager reference kernel (e.g. ops.py:353-374). That kernel is the obvious numerics oracle. The doc should say "reference kernels, not runtime kernels", and a test should compare the reference kernel against the DecomposeFusedQuant output.

Adopting it from a backend (colorer → partitioner)

  • No adapter from colorer tags to a Partitioner. ColorerBase writes delegation_tag outside any Partitioner, and nothing in the PR turns those tags into a PartitionResult. The docstring at colorer.py:15-17 says the partitioner "will act on" the tag. That's only true once each backend writes its own adapter. The colorer also:

    • assigns one tag string to every claimed node (colorer.py:85). A colored → uncolored → colored path then becomes a cyclic partition in fuse_as_graphmodule.
    • doesn't tag the lifted qparam constants or the getitem users of multi-output ops.

    Please ship a ColoredPartitioner that splits the colored nodes into convex groups, assigns per-group tags with a DelegationSpec, and calls tag_constant_data. Otherwise "a backend only describes what it supports" doesn't hold.

  • Backends can't be composed. colorer.py:79-83 raises on any node that's already tagged, whatever the tag, and the test pins this (test_colorer.py:139-144). So an NPU colorer followed by an XNNPACK colorer fails as soon as both claim the same op, which is the normal case. The error text ("non-delegated graphs") is misleading because nothing has been delegated yet. First-claim-wins plus a priority-ordered list of colorers is what a composed recipe needs.

  • No fallback for colored but unlowered nodes, and no leak guard.

    • DecomposeFusedQuant skips colored nodes (decompose_fused_quant.py:318). If a backend partitioner then rejects one, the fused op stays in the graph.
    • Nothing asserts that no fused_quant op remains before to_executorch. The failure surfaces late, as ToOutVarPass reporting "Missing out variants: fused_quant::…".
    • to_backend also strips every delegation_tag (exir/backend/backend_api.py:432-433), which silently erases colors meant for a second backend.

    Add a terminal check pass, and make "decompose after to_backend" the documented fallback.

Quantizer composability

  • The quantizer can't be composed with other quantizers in either direction.

    • FusedQuantQuantizer rejects children that aren't QuantizerBase (quantizer.py:290-292), so XNNPACKQuantizer can't be nested inside it.
    • Wrapping it in torchao's ComposableQuantizer hides fuse()/preserved_ops(). compile.py:54,68 then fail with an AttributeError. QuantizeStage always wraps torchao quantizers this way (export/stages.py:545-547), so the recipe path hits this every time.
  • The annotation-ownership protocol doesn't match other quantizers. fused_quant never sets QuantizationAnnotation._annotated=True (quantizer.py:155-168, :230). XNNPACK and Arm only treat a node as taken when _annotated is True (backends/xnnpack/quantizer/xnnpack_quantizer_utils.py:109-120, backends/arm/quantizer/arm_quantizer_utils.py:47-59). So:

    • when fused_quant runs first, they overwrite its annotations, including NoopQuantizer's suppression;
    • when fused_quant runs second, fuse() (quantizer.py:178-190) fuses any matching node that has any live annotation, including one written by another quantizer. It should only fuse nodes it annotated itself.
  • Silent degradation under plain prepare_pt2e/convert_pt2e.

    • Fusion is a non-standard step that runs after convert.
    • The pre-quantize rewrites live in frontend.prepare (frontend.py:58-74), not in transform_for_annotation.
    • The matchers assume the core-ATen graph from trace() (native_layer_norm, max_pool2d_with_indices, bmm, _masked_softmax).

    As a result, a standard flow with this quantizer annotates less and never fuses, with no warning. The parts that must run for correctness belong in transform_for_annotation.

Recipe / to_edge_transform_and_lower fit

  • compile_to_fused_quant blocks the standard lowering entry points. It calls to_edge internally (compile.py:77), so the caller can't use to_edge_transform_and_lower. Partitioner ops_to_not_decompose is lost: SDPA, upsample and so on are decomposed by trace() and again by to_edge. The edge config is also hardcoded:

    • _skip_dim_order=True (compile.py:80), which diverges from the dim-order default other backends rely on;
    • the verifier exception list (compile.py:81) covers the whole program rather than only the re-emitted ops.

    There is also no dynamic_shapes parameter, and the second export (compile.py:78) is non-strict while trace is strict. A composed recipe needs these steps exposed separately, e.g. transform_for_annotation + post_convert_passes=[fuse] + edge_transform_passes=[...], so the standard export() stages can drive them.

QAT

  • QAT is impossible as written. trace calls model.eval() under @torch.no_grad() (frontend.py:26,38), which also mutates the caller's model. FoldBatchNorm runs unconditionally and folds BN into constants (frontend.py:60). There's no prepare_qat_pt2e path, and qspecs.py builds observers only, with no fake-quant constructors or is_qat. The fused-op IR itself doesn't prevent QAT: fusion runs after convert. Keeping QAT out of the frontend means:
    • making BN folding and eval() the caller's choice, or tied to an is_qat flag;
    • adding fake-quant qspec builders.
  • Dynamic quantization looks supported but isn't fused. is_dynamic is exposed (qspecs.py:29,67), but only .default q/dq overloads are matched (graph_utils.py:56-69), and fuse_aten requires constant scales (fuse_aten.py:90,112-114). Either reject dynamic qspecs or support them.

Pass granularity (how to make passes opt-in)

  • The optimization passes are monolithic. get_optimization_passes (passes.py:137-209) is one 38-pass IterativePassGroup with a fixed 5-step cap. Hitting the cap is silent (pass_base.py, passes.py:146-148), and there are only two knobs. Suggested split into named, separately exported groups:

    1. canonicalize (safe default): ToConvolution, CSEPass, RemoveAlias/Clone/NopExpand/NopSliceOrView, FuseCascaded{Transpose,View}, ReplaceSqueezeUnsqueezeWithView, MergeSplitConcatChain, ReplaceSplitWithSlice, ReplaceSelectWithView, ReplaceNopTransposeOrPermuteWithView, RemoveCatFromSliceCopy, RemoveZeroSizedCatArgs, RemoveBranchedQuantDequant, ConstantFold.
    2. layout (opt-in, backend picks NHWC or NCHW): ToChannelsLast, RemovePermutesAroundElementwiseOps, MovePermuteAfterConcat, MoveSlice*, RemovePermuteBeforeMean.
    3. quant-boundary policy: QuantAbsorption, HoistActivations, AdvanceQuantize*, PostponeDequantize*, SplitDequantizedCat, SinkConstantCat, ReplaceScalarMulWithDequantQuant, ReplaceDequantQuantWithRequantize.
    4. numerics/cost rewrites (opt-in): FuseAddIntoLinear, FuseMulIntoLinear, SplitLinearAtSlices, PrequantizeEmbedding.

    Each group should be its own IterativePassGroup factory, so backends compose canonicalize + their_layout + …. Log a warning when a group hits its step cap. The pre-quantize passes (frontend.py:60-72, plus ReplaceMmWithBmm hidden in BmmQuantizer.transform_for_annotation) should become an explicit list the caller can pass, not something derived from preserved_ops().

  • The NHWC default contradicts "backend-neutral". ToChannelsLast runs unconditionally before coloring (passes.py:155), so its is_colored guards (to_channels_last.py:295,304,313) never fire in this pipeline. A backend that wants NCHW can't claim convolution.default because none survive the pass. It has to run ToChannelsFirst first, which contradicts the documented colorer → ToChannelsFirst → decompose order. The test_compile.py:56 tail has to undo this layout change, and nothing re-runs permute fusion afterwards (to_channels_first.py:215-218).

Infrastructure

  • Dependency direction is inverted. A layer described as neutral imports 22 passes from backends/cadence/aot (passes.py:11-40, targets.bzl:149-152). Several of those passes reference exir_ops.edge.cadence.* in class attributes, for example FuseQuantDequantToRequantizePass at backends/cadence/aot/fuse_ops.py:614-619. That's why their modules import cadence.aot.ops_registrations, so importing fused_quant.passes registers the whole cadence:: library. It also rules out cadence → fused_quant later without a cycle. Some fixes are cheap:

    • FuseCascadedTransposeOrPermuteOps, FuseCascadedViewOps, ReplaceNopTransposeOrPermuteWithViewPass and RemoveCloneOpsTransformImported are thin wrappers over backends/transforms; import those directly.
    • Move the rest to backends/transforms, with the op sets as parameters, as RemovePermutesAroundElementwiseOps already does.

    Tests also use cadence.aot.compiler.trace instead of fused_quant.frontend.trace (e.g. test_fold_batch_norm.py:12).

  • Likely cause of the 4 red unittest jobs (log not read). pytest.ini adds backends/fused_quant/test to the generic unittest run, and that run now pulls in cadence.aot.compiler and its registrations for the first time. Cadence tests normally run in their own workflow. Another candidate: Library("fused_quant","DEF")/custom_op re-registration if the module is imported under both executorch.backends.… and backends.…. Also, test/ has no __init__.py, while xnnpack/test and native/test do. Please check the job logs.

  • Import-time side effects. Op registration happens on import, and every pass module needs import executorch.backends.fused_quant.ops # noqa (passes.py:9, to_channels_last.py:11, …). passes.py:102-117 dereferences exir_ops.edge.fused_quant.* at module scope. Any backend reusing a single pass inherits both registrations, and the order-dependence that comes with them.

  • Placement (exir vs. backends). Neither location fits the whole package as it stands. exir/ already depends on torchao lazily (exir/operator/util.py:58, exir/passes/BUCK:203), but ops.py:13 imports torchao eagerly. Recommended split:

    1. op library + QuantParamsStruct + fuse_aten + DecomposeFusedQuant + graph_utils → a neutral core location (e.g. exir/passes/fused_quant/ or executorch/quantization/), with lazy torchao imports and a leak-check pass;
    2. generic canonicalization passes, including the de-Cadenced ones → backends/transforms/;
    3. quantizer, frontend and pre-quantize passes stay together, with the opinionated pass groups opt-in.

    backends/ implies a delegate, which the README explicitly says this isn't.

  • Target-specific wording in shared code. quantizer.py:239-242 and replace_mm_with_bmm.py:20-22 mention "the DSP" and "Turing TCE delegation". prequantize_embedding.py routes to "the cadence quantized_embedding_byte lowering".

Code Quality

  • The fallback doesn't hold at the edges.
    • DecomposeFusedQuant raises on grouped/blockwise qparams (decompose_fused_quant.py:55-59). An unclaimed 4-bit grouped linear, which the quantizer produces (qspecs.py:90-136), fails the whole compile instead of falling back. Emitting torchao.dequantize_affine would cover it. The pass also mutates in place before failing, leaving a partly decomposed graph.
    • Per-channel scale/zp are emitted as view_copy(scale, [-1]) (decompose_fused_quant.py:70-75). XNNPACK asserts that a per-channel scale is a placeholder (backends/xnnpack/operators/quant_params.py:259-260), so the decomposed graph can't be partitioned by XNNPACK unless a ConstantFold runs afterwards. The documented tail (test_compile.py:56) doesn't run one. The same applies to permute_copy(const) weights from ToChannelsFirst (to_channels_first.py:51-73). Medium confidence; I haven't run this end to end.
  • Default quantizer configs carry hardware assumptions.
    • Bias stays float (qspecs.py:193), while XNNPACK, QNN and Arm expect int32 bias.
    • The symmetric-weight default disagrees: False in make_qconfig (qspecs.py:142), True in make_fused_quant_quantizer (defaults.py:201).
    • Per-channel conv uses ch_axis=0 for aten.convolution without excluding transposed=True (defaults.py:127-130), which gives the wrong axis for transposed conv.

Testing

  • No numerical-equivalence test. Nothing compares a fused op's eager reference kernel, its DecomposeFusedQuant output and the portable-kernel run on the same inputs. That one parametrized test over all registered ops would back the "nothing is lost in the round trip" claim, and it would have caught the fake-dtype bug above.

Recommendation

Request Changes

Blocking:

  • the fake-kernel dtype bug;
  • the red unittest jobs;
  • the inverted Cadence dependency.

Needs discussion before backends adopt this:

  • a written BC contract for the op schemas;
  • a shipped tag → Partitioner adapter, with first-wins coloring;
  • making the quantizer and frontend steps composable with ComposableQuantizer/export() recipes;
  • an opt-in pass grouping;
  • the core/transforms/quantizer split in place of backends/fused_quant/.

@meta-codesync meta-codesync Bot changed the title Add fused_quant, a backend-neutral quantization frontend (#23091) Add fused_quant, a backend-neutral quantization frontend Oct 9, 2026
@meta-codesync
meta-codesync Bot force-pushed the export-D119511656 branch from f16d602 to 1f5540f Compare October 9, 2026 18:16
meta-codesync Bot pushed a commit that referenced this pull request Oct 9, 2026
Summary:
Pull Request resolved: #23091

`fused_quant` represents a `dequantize -> aten op -> quantize` triple as a single
`torch.ops.fused_quant.*` operator that carries its quantization parameters
inline. Instead of a quantized graph being a float graph sprinkled with q/dq
pairs that each backend has to pattern-match and fold back together, the
quantized operator is the unit the graph is built from.

Taking a quantized convolution as the example -- before, the quantized op is
implicit. Five nodes, and the one doing the arithmetic is float:

      x_i8        w_i8       bias
        |           |          |
        v           v          |
    dequantize  dequantize     |      <- scale, zero_point per edge
        |           |          |
        +-----------+----------+
                    |
                    v
            aten.convolution           <- float compute
                    |
                    v
                 quantize              <- scale, zero_point
                    |
                    v
                  y_i8

After, it is one node, and the qparams ride along with the operands they
describe:

      x_i8        w_i8       bias
        |           |          |
        +-----------+----------+
                    |
                    v
        fused_quant.convolution
            inp    : scale, zero_point, dtype, quant_min, quant_max
            weight : scale, zero_point, dtype, quant_min, quant_max
            bias   : scale, zero_point, dtype, quant_min, quant_max
            out    : scale, zero_point, dtype, quant_min, quant_max
                    |
                    v
                  y_i8

Nothing is lost in the round trip: `DecomposeFusedQuant` turns the second form
back into the first, so a backend that cannot run a fused op still gets the
portable one. `test_op_numerics` checks this exactly, for every registered op.

This adds that representation under `backends/fused_quant/`. It is
ahead-of-time only: each op has a reference (eager) kernel that defines its
numerics, but no runtime kernel, so a `fused_quant` op never reaches
`to_executorch`. Either a backend claims it and lowers it to its own
instruction, or `DecomposeFusedQuant` turns it back into the q/dq + ATen form
above and the ordinary portable kernels run it. `AssertNoFusedQuantOps` fails
the compile if one leaks.

**Operator representation and contract** (`ops.py`, `ops_utils.py`, `graph_utils.py`)

Every `fused_quant` op takes its tensor operands first, then one qparam block
(`scale`, `zero_point`, `dtype`, `quant_min`, `quant_max`) per tensor operand,
then the `out` block, then the ATen op's remaining arguments in ATen order. A
block whose `scale` is `None` means that edge stays in float, so one schema
covers the fully quantized, partially quantized and float instances of an op
rather than needing a variant for each. For an input block `dtype` is the
dequantized float dtype; for `out` it is the quantized storage dtype.

Granularity is carried by the shape of the `scale` tensor -- per-tensor,
per-channel, per-group and blockwise are the same op with differently shaped
qparams. `QuantParamsStruct.is_per_channel(tensor_shape)` /
`is_per_group(tensor_shape)` take the quantized tensor's shape so a group scale
such as `[1, K // g]` over `[1, K]` is not mistaken for per-channel.

The README writes down what is stable for backends: the argument layout and
field meanings are frozen, a new capability is a new overload or a new trailing
argument with a behavior-preserving default, and the ops are never serialized.

Fake kernels report the dtype the eager kernel actually returns, including for a
partially quantized op (quantized input, float output), where they previously
reported the int8 storage dtype.

**Quantizer** (`quantizer/`)

A torchao PT2E `Quantizer` that both annotates and fuses:

- `OpQuantizer` annotates one ATen op type and folds matched instances into a
  `fused_quant` op, given a `QuantizationConfig` and the schema names of its
  quantized activations, weights and passthrough operands.
- `NoopQuantizer` claims nodes with an empty annotation so a later quantizer
  skips them, leaving the op in float.
- `MaxPoolQuantizer` shares the input observer with output 0 of
  `max_pool2d_with_indices`.
- `FusedQuantQuantizer` composes them. Annotation is first-match-wins, so a
  narrower quantizer is listed before the general one it overrides.

Annotations are marked `_annotated`, the flag other PT2E quantizers (XNNPACK,
Arm) check before claiming a node, and `fuse()` only folds nodes the same
quantizer annotated -- never one claimed by another quantizer.

`preserved_ops()` reports which ATen ops the quantizer matches, so the caller can
hold exactly those back from decomposition instead of maintaining a hardcoded
list. `quantizer/qspecs.py` builds per-tensor, per-channel and blockwise static
specs (dynamic qspecs are not offered, since fusion needs constant qparams);
`quantizer/defaults.py` provides a default op set, with `get_core_quantizers`
for backends that lack per-channel support or the wider op set. Weights default
to symmetric in both `make_qconfig` and `make_fused_quant_quantizer`. A
transposed convolution gets its per-channel weight scales on axis 1, where its
output channels are, or per-tensor weights when `groups > 1`.

**Coloring** (`colorer.py`)

`ColorerBase` is the step between optimization and lowering: a backend declares
which fused ops it can run (`supported_ops`) and whether a given instance is
within its limits (`is_legal`), and the base class tags the rest with
`delegation_tag`. Colorers compose by priority: the first to claim a node keeps
it, and a later colorer skips already-tagged nodes.

The tag is a claim of ownership and the generic passes honour it -- neither
`ToChannelsFirst` nor `DecomposeFusedQuant` will touch a colored node. So the
pipeline is colorer -> ToChannelsFirst -> DecomposeFusedQuant: a backend keeps
what it claimed, and everything it did not claim falls back to q/dq plus ATen
automatically. A colored node the partitioner then rejects is caught by running
`DecomposeFusedQuant` again after `to_backend`.

A backend that claims nothing is just the degenerate case of this: every fused
op decomposes and the program runs entirely on the ordinary portable ATen
kernels. `test_compile.py` has that written out end to end.

**Passes**

- `pre_quantize_passes/` runs on the traced program before annotation: batch-norm
  folding, replacing leftover batch norms with convolutions, and recomposing
  `add(mask) + softmax` into `_masked_softmax` so the very-negative mask sentinel
  never enters the quantized domain.
- `optimization_passes/` runs on the fused edge graph: layout canonicalization,
  activation hoisting, folding add/mul into linear, quantization absorption,
  embedding prequantization, linear splitting, requantize formation.
  `IterativePassGroup` logs a warning when a group hits its step cap.
- `decompose_fused_quant.py` is the inverse: it turns any fused op back into
  dequantize + ATen + quantize. It is schema-driven, so it covers every
  registered fused op rather than a fixed list, and handles multi-output ops
  through `getitem`. It validates every node before rewriting any, so an
  unsupported one (grouped/blockwise qparams) fails the pass without leaving a
  partly decomposed graph. Per-channel qparams are emitted as 1-D constants
  rather than a `view_copy` of the full-rank one, and the qparam constants left
  dead are removed.

**Shared passes moved to `backends/transforms`**

The optimization group no longer depends on any backend. It used 22 passes from
`backends/cadence/aot`; none were Cadence-specific in what they do:

- Three (`FuseCascadedTransposeOrPermuteOps`, `FuseCascadedViewOps`,
  `ReplaceNopTransposeOrPermuteWithViewPass`) were already empty subclasses of a
  `backends/transforms` pass, which is now imported directly.
- The other 19 move to `backends/transforms`, one module and target each. The
  ones that listed Cadence's q/dq ops alongside `quantized_decomposed`'s now
  take those op sets as class attributes holding only the standard ops.
  Cadence keeps each pass under its old name as a thin subclass that adds its
  own ops back, so Cadence's behavior is unchanged.
- `FuseQuantDequantToRequantizePass` needs a requantize op to fuse a pair whose
  qparams differ, and there is no backend-neutral one. The shared pass takes a
  `requantize_op` (default `None`) and defaults to `allow_requantize=False`;
  Cadence sets its op and keeps `True`.
- `permute_pass_utils` gains `get_overload_packet` and `get_placeholders`, and
  `quantize_reorder_utils` holds the two op sets the reorder passes share.

Tests no longer import Cadence either: `test/helpers.trace` replaces
`cadence.aot.compiler.trace`, and the unused `ops_registrations` imports are gone.
`buck2 uquery` finds no `backends/cadence` target in the dependency closure of
any `fused_quant` target.

**Pipeline** (`frontend.py`, `passes.py`, `compile.py`)

`compile_to_fused_quant` runs trace -> prepare -> calibrate -> convert -> fuse ->
export -> to_edge -> passes, and stops. That is the seam a backend plugs into,
and it is the only entry point the package offers: from there a backend runs its
own colorer, lowering and partitioner, and hands whatever it did not claim to
`DecomposeFusedQuant`.

`trace(..., is_qat=True)` exports in training mode and keeps batch norms, for
`prepare_qat_pt2e`; by default it exports in eval mode with batch norms folded.
Either way the caller's model is left in the mode it was passed in.

The pass list is a parameter, defaulting to `get_fused_quant_passes()`. Callers
that want a different pipeline pass their own, and an empty list skips
optimization entirely.

One wrinkle worth knowing when writing a backend: `DecomposeFusedQuant` re-emits
the ATen ops that were held back from decomposition, and several of them
(`aten.linear`, `aten.silu`, ...) are not core ATen, so the edge verifier rejects
them. `compile.py` feeds `quantizer.preserved_ops()` into
`EdgeCompileConfig._core_aten_ops_exception_list`, which is exact by
construction: the ops held back are precisely the ops re-emitted.

**Known gaps**

- The package depends on `torchao`, which is not yet a core ExecuTorch dependency.
- `DecomposeFusedQuant` rejects grouped/blockwise qparams rather than emitting
  `dequantize_affine`, so an unclaimed blockwise op fails the compile.
- Bias stays float in the default configs; int32 bias is neither configurable
  nor documented yet.
- Op registration happens on import, and `passes.py` looks up
  `exir_ops.edge.fused_quant.*` at module scope.
- Follow-ups from review, deliberately not in this diff: splitting the pass
  pipeline into opt-in groups (which takes `ToChannelsLast` out of the default),
  a colored-node `Partitioner` adapter, fitting the quantizer and frontend into
  `ComposableQuantizer` / `export()` recipes, a QAT fake-quant path, generating
  schemas from the ATen ones, and the core/transforms/quantizer placement split.

Differential Revision: D119511656
@meta-codesync meta-codesync Bot changed the title Add fused_quant, a backend-neutral quantization frontend Add fused_quant, a backend-neutral quantization frontend (#23091) Oct 9, 2026
@meta-codesync
meta-codesync Bot force-pushed the export-D119511656 branch from 1f5540f to 3e4896f Compare October 9, 2026 18:24
meta-codesync Bot pushed a commit that referenced this pull request Oct 10, 2026
Summary:
Pull Request resolved: #23091

`fused_quant` represents a `dequantize -> aten op -> quantize` triple as a single
`torch.ops.fused_quant.*` operator that carries its quantization parameters
inline. Instead of a quantized graph being a float graph sprinkled with q/dq
pairs that each backend has to pattern-match and fold back together, the
quantized operator is the unit the graph is built from.

Taking a quantized convolution as the example -- before, the quantized op is
implicit. Five nodes, and the one doing the arithmetic is float:

      x_i8        w_i8       bias
        |           |          |
        v           v          |
    dequantize  dequantize     |      <- scale, zero_point per edge
        |           |          |
        +-----------+----------+
                    |
                    v
            aten.convolution           <- float compute
                    |
                    v
                 quantize              <- scale, zero_point
                    |
                    v
                  y_i8

After, it is one node, and the qparams ride along with the operands they
describe:

      x_i8        w_i8       bias
        |           |          |
        +-----------+----------+
                    |
                    v
        fused_quant.convolution
            inp    : scale, zero_point, dtype, quant_min, quant_max
            weight : scale, zero_point, dtype, quant_min, quant_max
            bias   : scale, zero_point, dtype, quant_min, quant_max
            out    : scale, zero_point, dtype, quant_min, quant_max
                    |
                    v
                  y_i8

Nothing is lost in the round trip: `DecomposeFusedQuant` turns the second form
back into the first, so a backend that cannot run a fused op still gets the
portable one. `test_op_numerics` checks this exactly, for every registered op.

This adds that representation under `backends/fused_quant/`. It is
ahead-of-time only: each op has a reference (eager) kernel that defines its
numerics, but no runtime kernel, so a `fused_quant` op never reaches
`to_executorch`. Either a backend claims it and lowers it to its own
instruction, or `DecomposeFusedQuant` turns it back into the q/dq + ATen form
above and the ordinary portable kernels run it. `AssertNoFusedQuantOps` fails
the compile if one leaks.

**Operator representation and contract** (`ops.py`, `ops_utils.py`, `graph_utils.py`)

Every `fused_quant` op takes its tensor operands first, then one qparam block
(`scale`, `zero_point`, `dtype`, `quant_min`, `quant_max`) per tensor operand,
then the `out` block, then the ATen op's remaining arguments in ATen order. A
block whose `scale` is `None` means that edge stays in float, so one schema
covers the fully quantized, partially quantized and float instances of an op
rather than needing a variant for each. For an input block `dtype` is the
dequantized float dtype; for `out` it is the quantized storage dtype.

Granularity is carried by the shape of the `scale` tensor -- per-tensor,
per-channel, per-group and blockwise are the same op with differently shaped
qparams. `QuantParamsStruct.is_per_channel(tensor_shape)` /
`is_per_group(tensor_shape)` take the quantized tensor's shape so a group scale
such as `[1, K // g]` over `[1, K]` is not mistaken for per-channel.

The README writes down what is stable for backends: the argument layout and
field meanings are frozen, a new capability is a new overload or a new trailing
argument with a behavior-preserving default, and the ops are never serialized.

Fake kernels report the dtype the eager kernel actually returns, including for a
partially quantized op (quantized input, float output), where they previously
reported the int8 storage dtype.

**Quantizer** (`quantizer/`)

A torchao PT2E `Quantizer` that both annotates and fuses:

- `OpQuantizer` annotates one ATen op type and folds matched instances into a
  `fused_quant` op, given a `QuantizationConfig` and the schema names of its
  quantized activations, weights and passthrough operands.
- `NoopQuantizer` claims nodes with an empty annotation so a later quantizer
  skips them, leaving the op in float.
- `MaxPoolQuantizer` shares the input observer with output 0 of
  `max_pool2d_with_indices`.
- `FusedQuantQuantizer` composes them. Annotation is first-match-wins, so a
  narrower quantizer is listed before the general one it overrides.

Annotations are marked `_annotated`, the flag other PT2E quantizers (XNNPACK,
Arm) check before claiming a node, and `fuse()` only folds nodes the same
quantizer annotated -- never one claimed by another quantizer.

`preserved_ops()` reports which ATen ops the quantizer matches, so the caller can
hold exactly those back from decomposition instead of maintaining a hardcoded
list. `quantizer/qspecs.py` builds per-tensor, per-channel and blockwise static
specs (dynamic qspecs are not offered, since fusion needs constant qparams);
`quantizer/defaults.py` provides a default op set, with `get_core_quantizers`
for backends that lack per-channel support or the wider op set. Weights default
to symmetric in both `make_qconfig` and `make_fused_quant_quantizer`. A
transposed convolution gets its per-channel weight scales on axis 1, where its
output channels are, or per-tensor weights when `groups > 1`.

**Coloring** (`colorer.py`)

`ColorerBase` is the step between optimization and lowering: a backend declares
which fused ops it can run (`supported_ops`) and whether a given instance is
within its limits (`is_legal`), and the base class tags the rest with
`delegation_tag`. Colorers compose by priority: the first to claim a node keeps
it, and a later colorer skips already-tagged nodes.

The tag is a claim of ownership and the generic passes honour it -- neither
`ToChannelsFirst` nor `DecomposeFusedQuant` will touch a colored node. So the
pipeline is colorer -> ToChannelsFirst -> DecomposeFusedQuant: a backend keeps
what it claimed, and everything it did not claim falls back to q/dq plus ATen
automatically. A colored node the partitioner then rejects is caught by running
`DecomposeFusedQuant` again after `to_backend`.

A backend that claims nothing is just the degenerate case of this: every fused
op decomposes and the program runs entirely on the ordinary portable ATen
kernels. `test_compile.py` has that written out end to end.

**Passes**

- `pre_quantize_passes/` runs on the traced program before annotation: batch-norm
  folding, replacing leftover batch norms with convolutions, and recomposing
  `add(mask) + softmax` into `_masked_softmax` so the very-negative mask sentinel
  never enters the quantized domain.
- `optimization_passes/` runs on the fused edge graph: layout canonicalization,
  activation hoisting, folding add/mul into linear, quantization absorption,
  embedding prequantization, linear splitting, requantize formation.
  `IterativePassGroup` logs a warning when a group hits its step cap.
- `decompose_fused_quant.py` is the inverse: it turns any fused op back into
  dequantize + ATen + quantize. It is schema-driven, so it covers every
  registered fused op rather than a fixed list, and handles multi-output ops
  through `getitem`. It validates every node before rewriting any, so an
  unsupported one (grouped/blockwise qparams) fails the pass without leaving a
  partly decomposed graph. Per-channel qparams are emitted as 1-D constants
  rather than a `view_copy` of the full-rank one, and the qparam constants left
  dead are removed.

**Shared passes moved to `backends/transforms`**

The optimization group no longer depends on any backend. It used 22 passes from
`backends/cadence/aot`; none were Cadence-specific in what they do:

- Three (`FuseCascadedTransposeOrPermuteOps`, `FuseCascadedViewOps`,
  `ReplaceNopTransposeOrPermuteWithViewPass`) were already empty subclasses of a
  `backends/transforms` pass, which is now imported directly.
- The other 19 move to `backends/transforms`, one module and target each. The
  ones that listed Cadence's q/dq ops alongside `quantized_decomposed`'s now
  take those op sets as class attributes holding only the standard ops.
  Cadence keeps each pass under its old name as a thin subclass that adds its
  own ops back, so Cadence's behavior is unchanged.
- `FuseQuantDequantToRequantizePass` needs a requantize op to fuse a pair whose
  qparams differ, and there is no backend-neutral one. The shared pass takes a
  `requantize_op` (default `None`) and defaults to `allow_requantize=False`;
  Cadence sets its op and keeps `True`.
- `permute_pass_utils` gains `get_overload_packet` and `get_placeholders`, and
  `quantize_reorder_utils` holds the two op sets the reorder passes share.

Tests no longer import Cadence either: `test/helpers.trace` replaces
`cadence.aot.compiler.trace`, and the unused `ops_registrations` imports are gone.
`buck2 uquery` finds no `backends/cadence` target in the dependency closure of
any `fused_quant` target.

**Pipeline** (`frontend.py`, `passes.py`, `compile.py`)

`compile_to_fused_quant` runs trace -> prepare -> calibrate -> convert -> fuse ->
export -> to_edge -> passes, and stops. That is the seam a backend plugs into,
and it is the only entry point the package offers: from there a backend runs its
own colorer, lowering and partitioner, and hands whatever it did not claim to
`DecomposeFusedQuant`.

`trace(..., is_qat=True)` exports in training mode and keeps batch norms, for
`prepare_qat_pt2e`; by default it exports in eval mode with batch norms folded.
Either way the caller's model is left in the mode it was passed in.

The pass list is a parameter, defaulting to `get_fused_quant_passes()`. Callers
that want a different pipeline pass their own, and an empty list skips
optimization entirely.

One wrinkle worth knowing when writing a backend: `DecomposeFusedQuant` re-emits
the ATen ops that were held back from decomposition, and several of them
(`aten.linear`, `aten.silu`, ...) are not core ATen, so the edge verifier rejects
them. `compile.py` feeds `quantizer.preserved_ops()` into
`EdgeCompileConfig._core_aten_ops_exception_list`, which is exact by
construction: the ops held back are precisely the ops re-emitted.

**Known gaps**

- The package depends on `torchao`, which is not yet a core ExecuTorch dependency.
- `DecomposeFusedQuant` rejects grouped/blockwise qparams rather than emitting
  `dequantize_affine`, so an unclaimed blockwise op fails the compile.
- Bias stays float in the default configs; int32 bias is neither configurable
  nor documented yet.
- Op registration happens on import, and `passes.py` looks up
  `exir_ops.edge.fused_quant.*` at module scope.
- Follow-ups from review, deliberately not in this diff: splitting the pass
  pipeline into opt-in groups (which takes `ToChannelsLast` out of the default),
  a colored-node `Partitioner` adapter, fitting the quantizer and frontend into
  `ComposableQuantizer` / `export()` recipes, a QAT fake-quant path, generating
  schemas from the ATen ones, and the core/transforms/quantizer placement split.

Differential Revision: D119511656
meta-codesync Bot pushed a commit that referenced this pull request Oct 11, 2026
Summary:
Pull Request resolved: #23091

`fused_quant` represents a `dequantize -> aten op -> quantize` triple as a single
`torch.ops.fused_quant.*` operator that carries its quantization parameters
inline. Instead of a quantized graph being a float graph sprinkled with q/dq
pairs that each backend has to pattern-match and fold back together, the
quantized operator is the unit the graph is built from.

Taking a quantized convolution as the example -- before, the quantized op is
implicit. Five nodes, and the one doing the arithmetic is float:

      x_i8        w_i8       bias
        |           |          |
        v           v          |
    dequantize  dequantize     |      <- scale, zero_point per edge
        |           |          |
        +-----------+----------+
                    |
                    v
            aten.convolution           <- float compute
                    |
                    v
                 quantize              <- scale, zero_point
                    |
                    v
                  y_i8

After, it is one node, and the qparams ride along with the operands they
describe:

      x_i8        w_i8       bias
        |           |          |
        +-----------+----------+
                    |
                    v
        fused_quant.convolution
            inp    : scale, zero_point, dtype, quant_min, quant_max
            weight : scale, zero_point, dtype, quant_min, quant_max
            bias   : scale, zero_point, dtype, quant_min, quant_max
            out    : scale, zero_point, dtype, quant_min, quant_max
                    |
                    v
                  y_i8

Nothing is lost in the round trip: `DecomposeFusedQuant` turns the second form
back into the first, so a backend that cannot run a fused op still gets the
portable one. `test_op_numerics` checks this exactly, for every registered op.

This adds that representation under `backends/fused_quant/`. It is
ahead-of-time only: each op has a reference (eager) kernel that defines its
numerics, but no runtime kernel, so a `fused_quant` op never reaches
`to_executorch`. Either a backend claims it and lowers it to its own
instruction, or `DecomposeFusedQuant` turns it back into the q/dq + ATen form
above and the ordinary portable kernels run it. `AssertNoFusedQuantOps` fails
the compile if one leaks.

**Operator representation and contract** (`ops.py`, `ops_utils.py`, `graph_utils.py`)

Every `fused_quant` op takes its tensor operands first, then one qparam block
(`scale`, `zero_point`, `dtype`, `quant_min`, `quant_max`) per tensor operand,
then the `out` block, then the ATen op's remaining arguments in ATen order. A
block whose `scale` is `None` means that edge stays in float, so one schema
covers the fully quantized, partially quantized and float instances of an op
rather than needing a variant for each. For an input block `dtype` is the
dequantized float dtype; for `out` it is the quantized storage dtype.

Granularity is carried by the shape of the `scale` tensor -- per-tensor,
per-channel, per-group and blockwise are the same op with differently shaped
qparams. `QuantParamsStruct.is_per_channel(tensor_shape)` /
`is_per_group(tensor_shape)` take the quantized tensor's shape so a group scale
such as `[1, K // g]` over `[1, K]` is not mistaken for per-channel.

The README writes down what is stable for backends: the argument layout and
field meanings are frozen, a new capability is a new overload or a new trailing
argument with a behavior-preserving default, and the ops are never serialized.

Fake kernels report the dtype the eager kernel actually returns, including for a
partially quantized op (quantized input, float output), where they previously
reported the int8 storage dtype.

**Quantizer** (`quantizer/`)

A torchao PT2E `Quantizer` that both annotates and fuses:

- `OpQuantizer` annotates one ATen op type and folds matched instances into a
  `fused_quant` op, given a `QuantizationConfig` and the schema names of its
  quantized activations, weights and passthrough operands.
- `NoopQuantizer` claims nodes with an empty annotation so a later quantizer
  skips them, leaving the op in float.
- `MaxPoolQuantizer` shares the input observer with output 0 of
  `max_pool2d_with_indices`.
- `FusedQuantQuantizer` composes them. Annotation is first-match-wins, so a
  narrower quantizer is listed before the general one it overrides.

Annotations are marked `_annotated`, the flag other PT2E quantizers (XNNPACK,
Arm) check before claiming a node, and `fuse()` only folds nodes the same
quantizer annotated -- never one claimed by another quantizer.

`preserved_ops()` reports which ATen ops the quantizer matches, so the caller can
hold exactly those back from decomposition instead of maintaining a hardcoded
list. `quantizer/qspecs.py` builds per-tensor, per-channel and blockwise static
specs (dynamic qspecs are not offered, since fusion needs constant qparams);
`quantizer/defaults.py` provides a default op set, with `get_core_quantizers`
for backends that lack per-channel support or the wider op set. Weights default
to symmetric in both `make_qconfig` and `make_fused_quant_quantizer`. A
transposed convolution gets its per-channel weight scales on axis 1, where its
output channels are, or per-tensor weights when `groups > 1`.

**Coloring** (`colorer.py`)

`ColorerBase` is the step between optimization and lowering: a backend declares
which fused ops it can run (`supported_ops`) and whether a given instance is
within its limits (`is_legal`), and the base class tags the rest with
`delegation_tag`. Colorers compose by priority: the first to claim a node keeps
it, and a later colorer skips already-tagged nodes.

The tag is a claim of ownership and the generic passes honour it -- neither
`ToChannelsFirst` nor `DecomposeFusedQuant` will touch a colored node. So the
pipeline is colorer -> ToChannelsFirst -> DecomposeFusedQuant: a backend keeps
what it claimed, and everything it did not claim falls back to q/dq plus ATen
automatically. A colored node the partitioner then rejects is caught by running
`DecomposeFusedQuant` again after `to_backend`.

A backend that claims nothing is just the degenerate case of this: every fused
op decomposes and the program runs entirely on the ordinary portable ATen
kernels. `test_compile.py` has that written out end to end.

**Passes**

- `pre_quantize_passes/` runs on the traced program before annotation: batch-norm
  folding, replacing leftover batch norms with convolutions, and recomposing
  `add(mask) + softmax` into `_masked_softmax` so the very-negative mask sentinel
  never enters the quantized domain.
- `optimization_passes/` runs on the fused edge graph: layout canonicalization,
  activation hoisting, folding add/mul into linear, quantization absorption,
  embedding prequantization, linear splitting, requantize formation.
  `IterativePassGroup` logs a warning when a group hits its step cap.
- `decompose_fused_quant.py` is the inverse: it turns any fused op back into
  dequantize + ATen + quantize. It is schema-driven, so it covers every
  registered fused op rather than a fixed list, and handles multi-output ops
  through `getitem`. It validates every node before rewriting any, so an
  unsupported one (grouped/blockwise qparams) fails the pass without leaving a
  partly decomposed graph. Per-channel qparams are emitted as 1-D constants
  rather than a `view_copy` of the full-rank one, and the qparam constants left
  dead are removed.

**Shared passes moved to `backends/transforms`**

The optimization group no longer depends on any backend. It used 22 passes from
`backends/cadence/aot`; none were Cadence-specific in what they do:

- Three (`FuseCascadedTransposeOrPermuteOps`, `FuseCascadedViewOps`,
  `ReplaceNopTransposeOrPermuteWithViewPass`) were already empty subclasses of a
  `backends/transforms` pass, which is now imported directly.
- The other 19 move to `backends/transforms`, one module and target each. The
  ones that listed Cadence's q/dq ops alongside `quantized_decomposed`'s now
  take those op sets as class attributes holding only the standard ops.
  Cadence keeps each pass under its old name as a thin subclass that adds its
  own ops back, so Cadence's behavior is unchanged.
- `FuseQuantDequantToRequantizePass` needs a requantize op to fuse a pair whose
  qparams differ, and there is no backend-neutral one. The shared pass takes a
  `requantize_op` (default `None`) and defaults to `allow_requantize=False`;
  Cadence sets its op and keeps `True`.
- `permute_pass_utils` gains `get_overload_packet` and `get_placeholders`, and
  `quantize_reorder_utils` holds the two op sets the reorder passes share.

Tests no longer import Cadence either: `test/helpers.trace` replaces
`cadence.aot.compiler.trace`, and the unused `ops_registrations` imports are gone.
`buck2 uquery` finds no `backends/cadence` target in the dependency closure of
any `fused_quant` target.

**Pipeline** (`frontend.py`, `passes.py`, `compile.py`)

`compile_to_fused_quant` runs trace -> prepare -> calibrate -> convert -> fuse ->
export -> to_edge -> passes, and stops. That is the seam a backend plugs into,
and it is the only entry point the package offers: from there a backend runs its
own colorer, lowering and partitioner, and hands whatever it did not claim to
`DecomposeFusedQuant`.

`trace(..., is_qat=True)` exports in training mode and keeps batch norms, for
`prepare_qat_pt2e`; by default it exports in eval mode with batch norms folded.
Either way the caller's model is left in the mode it was passed in.

The pass list is a parameter, defaulting to `get_fused_quant_passes()`. Callers
that want a different pipeline pass their own, and an empty list skips
optimization entirely.

One wrinkle worth knowing when writing a backend: `DecomposeFusedQuant` re-emits
the ATen ops that were held back from decomposition, and several of them
(`aten.linear`, `aten.silu`, ...) are not core ATen, so the edge verifier rejects
them. `compile.py` feeds `quantizer.preserved_ops()` into
`EdgeCompileConfig._core_aten_ops_exception_list`, which is exact by
construction: the ops held back are precisely the ops re-emitted.

**Known gaps**

- The package depends on `torchao`, which is not yet a core ExecuTorch dependency.
- `DecomposeFusedQuant` rejects grouped/blockwise qparams rather than emitting
  `dequantize_affine`, so an unclaimed blockwise op fails the compile.
- Bias stays float in the default configs; int32 bias is neither configurable
  nor documented yet.
- Op registration happens on import, and `passes.py` looks up
  `exir_ops.edge.fused_quant.*` at module scope.
- Follow-ups from review, deliberately not in this diff: splitting the pass
  pipeline into opt-in groups (which takes `ToChannelsLast` out of the default),
  a colored-node `Partitioner` adapter, fitting the quantizer and frontend into
  `ComposableQuantizer` / `export()` recipes, a QAT fake-quant path, generating
  schemas from the ATen ones, and the core/transforms/quantizer placement split.

Differential Revision: D119511656
Summary:
Pull Request resolved: #23091

`fused_quant` represents a `dequantize -> aten op -> quantize` triple as a single
`torch.ops.fused_quant.*` operator that carries its quantization parameters
inline. Instead of a quantized graph being a float graph sprinkled with q/dq
pairs that each backend has to pattern-match and fold back together, the
quantized operator is the unit the graph is built from.

Taking a quantized convolution as the example -- before, the quantized op is
implicit. Five nodes, and the one doing the arithmetic is float:

      x_i8        w_i8       bias
        |           |          |
        v           v          |
    dequantize  dequantize     |      <- scale, zero_point per edge
        |           |          |
        +-----------+----------+
                    |
                    v
            aten.convolution           <- float compute
                    |
                    v
                 quantize              <- scale, zero_point
                    |
                    v
                  y_i8

After, it is one node, and the qparams ride along with the operands they
describe:

      x_i8        w_i8       bias
        |           |          |
        +-----------+----------+
                    |
                    v
        fused_quant.convolution
            inp    : scale, zero_point, dtype, quant_min, quant_max
            weight : scale, zero_point, dtype, quant_min, quant_max
            bias   : scale, zero_point, dtype, quant_min, quant_max
            out    : scale, zero_point, dtype, quant_min, quant_max
                    |
                    v
                  y_i8

Nothing is lost in the round trip: `DecomposeFusedQuant` turns the second form
back into the first, so a backend that cannot run a fused op still gets the
portable one. `test_op_numerics` checks this exactly, for every registered op.

This adds that representation under `backends/fused_quant/`. It is
ahead-of-time only: each op has a reference (eager) kernel that defines its
numerics, but no runtime kernel, so a `fused_quant` op never reaches
`to_executorch`. Either a backend claims it and lowers it to its own
instruction, or `DecomposeFusedQuant` turns it back into the q/dq + ATen form
above and the ordinary portable kernels run it. `AssertNoFusedQuantOps` fails
the compile if one leaks.

**Operator representation and contract** (`ops.py`, `ops_utils.py`, `graph_utils.py`)

Every `fused_quant` op takes its tensor operands first, then one qparam block
(`scale`, `zero_point`, `dtype`, `quant_min`, `quant_max`) per tensor operand,
then the `out` block, then the ATen op's remaining arguments in ATen order. A
block whose `scale` is `None` means that edge stays in float, so one schema
covers the fully quantized, partially quantized and float instances of an op
rather than needing a variant for each. For an input block `dtype` is the
dequantized float dtype; for `out` it is the quantized storage dtype.

Granularity is carried by the shape of the `scale` tensor -- per-tensor,
per-channel, per-group and blockwise are the same op with differently shaped
qparams. `QuantParamsStruct.is_per_channel(tensor_shape)` /
`is_per_group(tensor_shape)` take the quantized tensor's shape so a group scale
such as `[1, K // g]` over `[1, K]` is not mistaken for per-channel.

The README writes down what is stable for backends: the argument layout and
field meanings are frozen, a new capability is a new overload or a new trailing
argument with a behavior-preserving default, and the ops are never serialized.

Fake kernels report the dtype the eager kernel actually returns, including for a
partially quantized op (quantized input, float output), where they previously
reported the int8 storage dtype.

**Quantizer** (`quantizer/`)

A torchao PT2E `Quantizer` that both annotates and fuses:

- `OpQuantizer` annotates one ATen op type and folds matched instances into a
  `fused_quant` op, given a `QuantizationConfig` and the schema names of its
  quantized activations, weights and passthrough operands.
- `NoopQuantizer` claims nodes with an empty annotation so a later quantizer
  skips them, leaving the op in float.
- `MaxPoolQuantizer` shares the input observer with output 0 of
  `max_pool2d_with_indices`.
- `FusedQuantQuantizer` composes them. Annotation is first-match-wins, so a
  narrower quantizer is listed before the general one it overrides.

Annotations are marked `_annotated`, the flag other PT2E quantizers (XNNPACK,
Arm) check before claiming a node, and `fuse()` only folds nodes the same
quantizer annotated -- never one claimed by another quantizer.

`preserved_ops()` reports which ATen ops the quantizer matches, so the caller can
hold exactly those back from decomposition instead of maintaining a hardcoded
list. `quantizer/qspecs.py` builds per-tensor, per-channel and blockwise static
specs (dynamic qspecs are not offered, since fusion needs constant qparams);
`quantizer/defaults.py` provides a default op set, with `get_core_quantizers`
for backends that lack per-channel support or the wider op set. Weights default
to symmetric in both `make_qconfig` and `make_fused_quant_quantizer`. A
transposed convolution gets its per-channel weight scales on axis 1, where its
output channels are, or per-tensor weights when `groups > 1`.

**Coloring** (`colorer.py`)

`ColorerBase` is the step between optimization and lowering: a backend declares
which fused ops it can run (`supported_ops`) and whether a given instance is
within its limits (`is_legal`), and the base class tags the rest with
`delegation_tag`. Colorers compose by priority: the first to claim a node keeps
it, and a later colorer skips already-tagged nodes.

The tag is a claim of ownership and the generic passes honour it -- neither
`ToChannelsFirst` nor `DecomposeFusedQuant` will touch a colored node. So the
pipeline is colorer -> ToChannelsFirst -> DecomposeFusedQuant: a backend keeps
what it claimed, and everything it did not claim falls back to q/dq plus ATen
automatically. A colored node the partitioner then rejects is caught by running
`DecomposeFusedQuant` again after `to_backend`.

A backend that claims nothing is just the degenerate case of this: every fused
op decomposes and the program runs entirely on the ordinary portable ATen
kernels. `test_compile.py` has that written out end to end.

**Passes**

- `pre_quantize_passes/` runs on the traced program before annotation: batch-norm
  folding, replacing leftover batch norms with convolutions, and recomposing
  `add(mask) + softmax` into `_masked_softmax` so the very-negative mask sentinel
  never enters the quantized domain.
- `optimization_passes/` runs on the fused edge graph: layout canonicalization,
  activation hoisting, folding add/mul into linear, quantization absorption,
  embedding prequantization, linear splitting, requantize formation.
  `IterativePassGroup` logs a warning when a group hits its step cap.
- `decompose_fused_quant.py` is the inverse: it turns any fused op back into
  dequantize + ATen + quantize. It is schema-driven, so it covers every
  registered fused op rather than a fixed list, and handles multi-output ops
  through `getitem`. It validates every node before rewriting any, so an
  unsupported one (grouped/blockwise qparams) fails the pass without leaving a
  partly decomposed graph. Per-channel qparams are emitted as 1-D constants
  rather than a `view_copy` of the full-rank one, and the qparam constants left
  dead are removed.

**Shared passes moved to `backends/transforms`**

The optimization group no longer depends on any backend. It used 22 passes from
`backends/cadence/aot`; none were Cadence-specific in what they do:

- Three (`FuseCascadedTransposeOrPermuteOps`, `FuseCascadedViewOps`,
  `ReplaceNopTransposeOrPermuteWithViewPass`) were already empty subclasses of a
  `backends/transforms` pass, which is now imported directly.
- The other 19 move to `backends/transforms`, one module and target each. The
  ones that listed Cadence's q/dq ops alongside `quantized_decomposed`'s now
  take those op sets as class attributes holding only the standard ops.
  Cadence keeps each pass under its old name as a thin subclass that adds its
  own ops back, so Cadence's behavior is unchanged.
- `FuseQuantDequantToRequantizePass` needs a requantize op to fuse a pair whose
  qparams differ, and there is no backend-neutral one. The shared pass takes a
  `requantize_op` (default `None`) and defaults to `allow_requantize=False`;
  Cadence sets its op and keeps `True`.
- `permute_pass_utils` gains `get_overload_packet` and `get_placeholders`, and
  `quantize_reorder_utils` holds the two op sets the reorder passes share.

Tests no longer import Cadence either: `test/helpers.trace` replaces
`cadence.aot.compiler.trace`, and the unused `ops_registrations` imports are gone.
`buck2 uquery` finds no `backends/cadence` target in the dependency closure of
any `fused_quant` target.

**Pipeline** (`frontend.py`, `passes.py`, `compile.py`)

`compile_to_fused_quant` runs trace -> prepare -> calibrate -> convert -> fuse ->
export -> to_edge -> passes, and stops. That is the seam a backend plugs into,
and it is the only entry point the package offers: from there a backend runs its
own colorer, lowering and partitioner, and hands whatever it did not claim to
`DecomposeFusedQuant`.

`trace(..., is_qat=True)` exports in training mode and keeps batch norms, for
`prepare_qat_pt2e`; by default it exports in eval mode with batch norms folded.
Either way the caller's model is left in the mode it was passed in.

The pass list is a parameter, defaulting to `get_fused_quant_passes()`. Callers
that want a different pipeline pass their own, and an empty list skips
optimization entirely.

One wrinkle worth knowing when writing a backend: `DecomposeFusedQuant` re-emits
the ATen ops that were held back from decomposition, and several of them
(`aten.linear`, `aten.silu`, ...) are not core ATen, so the edge verifier rejects
them. `compile.py` feeds `quantizer.preserved_ops()` into
`EdgeCompileConfig._core_aten_ops_exception_list`, which is exact by
construction: the ops held back are precisely the ops re-emitted.

**Known gaps**

- The package depends on `torchao`, which is not yet a core ExecuTorch dependency.
- `DecomposeFusedQuant` rejects grouped/blockwise qparams rather than emitting
  `dequantize_affine`, so an unclaimed blockwise op fails the compile.
- Bias stays float in the default configs; int32 bias is neither configurable
  nor documented yet.
- Op registration happens on import, and `passes.py` looks up
  `exir_ops.edge.fused_quant.*` at module scope.
- Follow-ups from review, deliberately not in this diff: splitting the pass
  pipeline into opt-in groups (which takes `ToChannelsLast` out of the default),
  a colored-node `Partitioner` adapter, fitting the quantizer and frontend into
  `ComposableQuantizer` / `export()` recipes, a QAT fake-quant path, generating
  schemas from the ATen ones, and the core/transforms/quantizer placement split.

Differential Revision: D119511656

This branch was successfully deployed

1 active deployment
cadence — 8321a9be Deployed Oct 11, 2026 by meta-codesync[bot] via hifi-op-test / hifi4 #32876
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. meta-exported

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants