[mcremon-meta]

Summary:
X-link: pytorch/executorch#24993

AffineQuantizedObserverBase.convert inserts torchao.choose_qparams_affine, quantize_affine and dequantize_affine as OpOverloadPackets. The per-tensor and per-channel q/dq ops that convert_pt2e inserts are OpOverloads, and torch.export turns these packets into .default when it re-exports a converted model, so the same op shows up in two forms depending on whether a graph was re-exported. Each of these ops has a single overload, so convert now inserts .default.

The bug it fixes. convert_pt2e(..., use_reference_representation=True) never applies the 4-bit groupwise dynamic linear rewrite to the output of torchao's own affine observers. The rewrite's patterns are traced with torch.export, so they carry .default, and the subgraph matcher compares targets exactly. remove_tensor_overload_for_qdq_ops lines the two forms up for the quantized_decomposed ops only, and the rewrite's filter checks for .default explicitly. The new test_dynamic_linear_4bit_groupwise (8-bit dynamic per-token activations, 4-bit per-group weights) fails without the observer change, because reference_dqlinear_int4 is not in the graph, and passes with it.

Matchers that move with it. Code that matches the packet on convert_pt2e output would silently stop matching, so it changes in the same diff:

• torchao: _QUANT_OPS (which quantize nodes convert folds into constants), constant_fold's DEQUANT_OPS, and port_metadata_pass's op lists. test_quantize_pt2e now expects .default.
• Helios post_convert_transform: Inductor's add_dont_constant_fold matches the exact target. With only the packet listed, constant_fold would fold the weight dequantize of AffineQuantizedMinMaxObserver weights into a float constant.
• ExecuTorch QNN (genai_pipeline/quant_utilities.py, and the llama llm_wrappers.py and mix_precision_analyzer.py): the q/dq target sets accept both forms, because OSS ExecuTorch pins torchao (third-party/ao) and keeps getting the packet until that pin moves past this change.
• Bolt per_block_param_observer_test compares against .default.

Unchanged: eager calls such as torch.ops.torchao.dequantize_affine(...) resolve either way; substring matches on str(node.target) (XNNPACK is_affine_qdq, Arm, the native backend tests, Bolt's FileCheck) see both forms; and code that runs on exported or edge graphs already sees .default.

SAS's _normalize_affine_overloads, which rewrites the packets to .default after convert, becomes a no-op. Removing it is a follow-up.

Differential Revision: D130716026

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/5101

• 📄 Preview Python docs built from this PR

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

✅ You can merge normally! (1 Unrelated Failure)

As of commit 164f886 with merge base 4ab9d61 ():

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Run Regression Tests (aarch64) / test-cpu-ops (macos-14) (gh) (trunk failure)

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

[meta-codesync]

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

[vkuzo]

@claude review and also explain

[claude]

Claude finished @vkuzo's task in 1m 43s —— View job

---

Review of #5101: insert .default overloads for the affine q/dq ops in convert

• Gather context (diff, related files)
• Explain the change
• Review for correctness and missed call sites
• Post review

---

Explanation

Background: OpOverloadPacket vs OpOverload. torch.ops.torchao.quantize_affine is an OpOverloadPacket, which is a callable namespace that groups every overload of an op. torch.ops.torchao.quantize_affine.default is an OpOverload, the concrete schema. Both can be used as an FX node's target, and both run the same kernel in eager mode. Graph matching, though, compares node.target by identity or equality, and a packet never equals its .default.

What was happening. AffineQuantizedObserverBase.convert (torchao/quantization/pt2e/observer.py:1736-1789) built nodes with the packet. Everything else in a converted graph uses OpOverloads: the quantized_decomposed.*_per_tensor/per_channel.default nodes, and any graph that goes through torch.export again (re-export turns packets into .default). So the same affine op could appear in two forms depending on how the graph was produced.

The concrete bug. In reference_representation_rewrite (reference_representation_rewrite.py:1166-1189), the 4-bit groupwise dynamic-linear pattern (_qdq_dynamic_quantized_linear_4bit_groupwise) is traced with _get_aten_graph_module_for_pattern, so its nodes are .default. The steps after tracing then fail to line up the two forms:

• remove_tensor_overload_for_qdq_ops (utils.py:1159) only maps quantized_decomposed ops back to packets. The torchao.*_affine ops are untouched.
• The subgraph matcher compares targets exactly, so pattern .default ≠ graph packet, and nothing matches.
• Even if something matched, _filter_fn_for_dynamic_quantized_linear_4bit_groupwise (:442, :450) checks for .default explicitly.

The result was that convert_pt2e(..., use_reference_representation=True) never rewrote torchao's own affine-observer output into reference_dqlinear_int4. Emitting .default at the source fixes this. It is safe because each of these ops has exactly one overload.

Matchers that move with it. Some lists matched on the packet and would silently stop matching:

• _QUANT_OPS (quantize_pt2e.py:252) decides which weight quantize nodes get folded into constants.
• DEQUANT_OPS (constant_fold.py:183) tells the constant folder not to fold a weight dequantize.
• port_metadata_pass (:36, :43, :49) ports metadata onto q/dq/choose_qparams.

All three are updated. I grepped the repo for any remaining packet-form torch.ops.torchao.{quantize,dequantize,choose_qparams}_affine used as a target or matcher entry and found none. The only remaining packet uses are eager calls inside the pattern and reference functions (reference_representation_rewrite.py:226-320), which is fine: they get traced to .default.

---

Review

Overall this is correct and well-scoped. The root cause is fixed at the source, the matchers are updated in the same change, and there is a regression test that fails before the change and passes after. A few comments:

1. Consider accepting both forms in torchao's own matchers, at least for a while (suggestion).
The PR description says ExecuTorch QNN matchers will accept both forms during the transition. The torchao lists, however, now accept only .default. The packet form can still show up in two cases:

• Out-of-tree observers or fake-quants. Any custom convert() (in a downstream repo, or a user's subclass of AffineQuantizedObserverBase that overrides convert) that still inserts torch.ops.torchao.quantize_affine will now silently:
  • lose weight constant-folding (_QUANT_OPS),
  • lose metadata porting,
  • have its weight dequantize folded to a float constant by constant_fold (the same failure mode the PR calls out for Helios).
• Converted GraphModules saved earlier and then run through these passes again.

Since the packet never collides with anything else, keeping it in these lists costs almost nothing and turns a silent behavior change into a no-op for those callers:

torch.ops.torchao.quantize_affine.default,
torch.ops.torchao.quantize_affine,  # BC: graphs produced by pre-#5101 / out-of-tree convert()

The packet entries could be dropped in a later release. This is optional; if the team prefers a single canonical form, the current change is consistent.
Fix this →

2. Test: test_dynamic_linear_4bit_groupwise (test_representation.py:289).

• It relies on the rewrite pattern's baked-in literal (1, 8) (reference_representation_rewrite.py:1073), which is why PerGroup(group_size=8) is used. A short comment saying so would stop someone from "tidying" it to another group size and getting a confusing failure.
• The numeric check uses torch.randn(4, 32) without a seed, with fixed_output_tol=1e-2 (bound = 2*1e-2 + 1e-5). The reference op rounds weight_scale through bf16 (:338), so it deliberately differs from the q/dq path. The margin looks comfortable for these sizes, but seeding the RNG (or the module init) would remove any chance of rare flakes.
• It is good that ref_node_occurrence asserts dequantize_affine.default: 0 as well as reference_dqlinear_int4.default: 1. That proves the whole pattern was consumed, not just that something extra was added.

3. Pre-existing, out of scope. duplicate_dq_pass.py:33-43 and utils.py:583-594 (_is_connected) never listed the affine ops in either form, so DuplicateDQPass does not duplicate dequantize_affine nodes. This PR doesn't change that. It may be worth a follow-up now that there is a single canonical target to add.

4. Minor. The test_quantize_pt2e.py expectation changes are mechanical and correct. They also confirm that checkGraphModuleNodes compares exact targets, which is why they had to change.

On CI: the only unclassified failure is the PR Label Check (it needs a module: label, e.g. module: pt2e_quant). The rest are failures already present on trunk and unrelated to this PR.

Verdict: LGTM. Item 1 is the only thing I'd consider before merging.

[mcremon-meta]

Claude finished @vkuzo's task in 1m 43s —— View job

Review of #5101: insert .default overloads for the affine q/dq ops in convert

• Gather context (diff, related files)
• Explain the change
• Review for correctness and missed call sites
• Post review

Explanation

Background: OpOverloadPacket vs OpOverload. torch.ops.torchao.quantize_affine is an OpOverloadPacket, which is a callable namespace that groups every overload of an op. torch.ops.torchao.quantize_affine.default is an OpOverload, the concrete schema. Both can be used as an FX node's target, and both run the same kernel in eager mode. Graph matching, though, compares node.target by identity or equality, and a packet never equals its .default.

What was happening. AffineQuantizedObserverBase.convert (torchao/quantization/pt2e/observer.py:1736-1789) built nodes with the packet. Everything else in a converted graph uses OpOverloads: the quantized_decomposed.*_per_tensor/per_channel.default nodes, and any graph that goes through torch.export again (re-export turns packets into .default). So the same affine op could appear in two forms depending on how the graph was produced.

The concrete bug. In reference_representation_rewrite (reference_representation_rewrite.py:1166-1189), the 4-bit groupwise dynamic-linear pattern (_qdq_dynamic_quantized_linear_4bit_groupwise) is traced with _get_aten_graph_module_for_pattern, so its nodes are .default. The steps after tracing then fail to line up the two forms:

• remove_tensor_overload_for_qdq_ops (utils.py:1159) only maps quantized_decomposed ops back to packets. The torchao.*_affine ops are untouched.
• The subgraph matcher compares targets exactly, so pattern .default ≠ graph packet, and nothing matches.
• Even if something matched, _filter_fn_for_dynamic_quantized_linear_4bit_groupwise (:442, :450) checks for .default explicitly.

The result was that convert_pt2e(..., use_reference_representation=True) never rewrote torchao's own affine-observer output into reference_dqlinear_int4. Emitting .default at the source fixes this. It is safe because each of these ops has exactly one overload.

Matchers that move with it. Some lists matched on the packet and would silently stop matching:

• _QUANT_OPS (quantize_pt2e.py:252) decides which weight quantize nodes get folded into constants.
• DEQUANT_OPS (constant_fold.py:183) tells the constant folder not to fold a weight dequantize.
• port_metadata_pass (:36, :43, :49) ports metadata onto q/dq/choose_qparams.

All three are updated. I grepped the repo for any remaining packet-form torch.ops.torchao.{quantize,dequantize,choose_qparams}_affine used as a target or matcher entry and found none. The only remaining packet uses are eager calls inside the pattern and reference functions (reference_representation_rewrite.py:226-320), which is fine: they get traced to .default.

Review

Overall this is correct and well-scoped. The root cause is fixed at the source, the matchers are updated in the same change, and there is a regression test that fails before the change and passes after. A few comments:

1. Consider accepting both forms in torchao's own matchers, at least for a while (suggestion). The PR description says ExecuTorch QNN matchers will accept both forms during the transition. The torchao lists, however, now accept only .default. The packet form can still show up in two cases:

• Out-of-tree observers or fake-quants. Any custom convert() (in a downstream repo, or a user's subclass of AffineQuantizedObserverBase that overrides convert) that still inserts torch.ops.torchao.quantize_affine will now silently:

  • lose weight constant-folding (_QUANT_OPS),
  • lose metadata porting,
  • have its weight dequantize folded to a float constant by constant_fold (the same failure mode the PR calls out for Helios).

• Converted GraphModules saved earlier and then run through these passes again.

Since the packet never collides with anything else, keeping it in these lists costs almost nothing and turns a silent behavior change into a no-op for those callers:

torch.ops.torchao.quantize_affine.default,
torch.ops.torchao.quantize_affine,  # BC: graphs produced by pre-#5101 / out-of-tree convert()

The packet entries could be dropped in a later release. This is optional; if the team prefers a single canonical form, the current change is consistent. Fix this →

2. Test: test_dynamic_linear_4bit_groupwise (test_representation.py:289).

• It relies on the rewrite pattern's baked-in literal (1, 8) (reference_representation_rewrite.py:1073), which is why PerGroup(group_size=8) is used. A short comment saying so would stop someone from "tidying" it to another group size and getting a confusing failure.
• The numeric check uses torch.randn(4, 32) without a seed, with fixed_output_tol=1e-2 (bound = 2*1e-2 + 1e-5). The reference op rounds weight_scale through bf16 (:338), so it deliberately differs from the q/dq path. The margin looks comfortable for these sizes, but seeding the RNG (or the module init) would remove any chance of rare flakes.
• It is good that ref_node_occurrence asserts dequantize_affine.default: 0 as well as reference_dqlinear_int4.default: 1. That proves the whole pattern was consumed, not just that something extra was added.

3. Pre-existing, out of scope. duplicate_dq_pass.py:33-43 and utils.py:583-594 (_is_connected) never listed the affine ops in either form, so DuplicateDQPass does not duplicate dequantize_affine nodes. This PR doesn't change that. It may be worth a follow-up now that there is a single canonical target to add.

4. Minor. The test_quantize_pt2e.py expectation changes are mechanical and correct. They also confirm that checkGraphModuleNodes compares exact targets, which is why they had to change.

On CI: the only unclassified failure is the PR Label Check (it needs a module: label, e.g. module: pt2e_quant). The rest are failures already present on trunk and unrelated to this PR.

Verdict: LGTM. Item 1 is the only thing I'd consider before merging.

Updated:

1. The packets are back next to .default in _QUANT_OPS, constant_fold's DEQUANT_OPS and the port_metadata_pass lists. An observer that overrides convert() and still inserts the packet gets the same folding and metadata porting as before. convert itself only inserts .default.

2. The rewrite isn't tied to group size 8. The pattern's (1, 8) is replaced with the group-size placeholder, and the match runs with ignore_literals=True. So instead of a comment, the test now runs group sizes 8, 16 and 32. That turned up a bug in the reference op. When a group spans the whole row (32 here), the observer gives the weight a 1-D scale, and the reference reshaped it to [1, out_features], so every output channel used channel 0's scale. It now reshapes it to [out_features, 1]. Group size 32 fails without that fix.

The RNG is already seeded: QuantizationTestCase.setUp goes through TestCase.setUp, which calls set_rng_seed().