[Qualcomm] Don't annotate non-float index_put value for quantization

[psiddh]

[Qualcomm] Don't annotate a non-float index_put value for quantization

Summary

The QNN IndexPut quantizer annotator (HTP and LPAI) put its value argument into the input
qspec map unconditionally. When the value is an integer tensor, prepare_pt2e/convert_pt2e
emit quantize_per_tensor on it and the meta kernel asserts a float input → to_executorch()
fails. Observers only work on float tensors, so integer values must not be annotated.

Symptom

Quantizing Mixtral (MoE) with QnnQuantizer fails:

File ".../torch/ao/quantization/fx/_decomposed.py", quantize_per_tensor_meta
AssertionError: Expecting input to have dtype torch.float32, but got dtype: torch.int64
  While executing quantized_decomposed.quantize_per_tensor.default(%arange_5, ...)

Mixtral routing feeds an int64 arange as the value arg of index_put. XNNPACK is unaffected —
it has no index_put annotator and guards annotated inputs by dtype.

Fix

Guard with _is_float_tensor (the helper QNN already uses elsewhere), in both the HTP and LPAI
annotators:

-        if input_qspec is not None:
+        if input_qspec is not None and _is_float_tensor(value):
             input_qspec_map[value] = input_qspec
             output_qspec = SharedQuantizationSpec((value, node))

(LPAI mirror updated equivalently; output_qspec left None when the value is non-float.)

Testing

test_index_put_int64_value_not_quantized (in test_passes.py): builds an index_put with an
int64 value, runs prepare_pt2e/convert_pt2e with QnnQuantizer, and re-exports. Fails on
the pre-fix code with the dtype torch.int64 assertion; passes after. Verified by file-swap.

Notes

Touches both HTP and LPAI annotators. Surfaced by the HuggingFace transformers ExecuTorch-exporter
QNN work (PR #47747), which currently skips MoE with this exact limitation.

cc @cbilgin

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• [ALERT ONLY] - System Maintenance

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

[github-actions]

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.

[Copilot]

Pull request overview

This PR fixes a PT2E quantization failure in the Qualcomm QNN backend by ensuring the IndexPut annotators (HTP and LPAI) do not annotate index_put’s value argument when it is not a float tensor, preventing quantize_per_tensor from being applied to int tensors during export.

Changes:

• Add a dtype guard (_is_float_tensor) in the HTP IndexPut annotator before adding value to the input qspec map.
• Mirror the same behavior in the LPAI IndexPut annotator, leaving output_qspec=None for non-float value.
• Add a regression test covering the int64 index_put value case that previously triggered a quantize_per_tensor meta-kernel assertion.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.

File	Description
backends/qualcomm/tests/test_passes.py	Adds a regression test to ensure int64 index_put values are not sent through PT2E quantization.
backends/qualcomm/quantizer/annotators/lpai_rules.py	Skips annotating index_put’s value when it’s non-float; avoids generating qspecs that would trigger int quantization.
backends/qualcomm/quantizer/annotators/htp_rules.py	Adds _is_float_tensor(value) guard so non-float index_put values aren’t annotated/quantized.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated no new comments.

Suppressed comments (2)

backends/qualcomm/quantizer/annotators/lpai_rules.py:472

• In the LPAI IndexPut annotator, quantization_config.input_activation is Optional[QuantizationSpec] (e.g. fp16a8w configs set it to None). With the current guard, a float value will still be added to input_qspec_map even when input_activation is None, which can propagate a None qspec into QuantizationAnnotation and break PT2E quantization.

Match the HTP annotator pattern by guarding on input_activation is not None before populating input_qspec_map / output_qspec.

        input_qspec_map = {}
        output_qspec = None
        if _is_float_tensor(value):
            input_qspec_map[value] = quantization_config.input_activation
            output_qspec = SharedQuantizationSpec((value, node))

backends/qualcomm/tests/test_passes.py:478

• Docstring grammar: “to_executorch() fail” should be “to_executorch() to fail” (or similar) to read correctly.

        Regression for MoE (Mixtral) routing, where index_put's value is an int64
        arange: annotating it for per-tensor quant makes quantize_per_tensor assert
        float32 and to_executorch() fail. Only float tensors may be annotated.

[Copilot]

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated no new comments.

Suppressed comments (2)

backends/qualcomm/tests/test_passes.py:514

• Calling self.skipTest() inside the backend loop will skip the entire test as soon as one backend (likely LPAI) is unavailable, so the HTP leg won't be exercised in environments without quantized_aot_lib. Wrap each backend in a subTest and only skip the LPAI subtest when needed; let unexpected failures still surface.

        for backend in (
            QnnExecuTorchBackendType.kHtpBackend,
            QnnExecuTorchBackendType.kLpaiBackend,
        ):
            try:
                quantizer = QnnQuantizer(backend=backend)
            except Exception as e:
                # LPAI needs quantized_aot_lib; skip that leg if it isn't available.
                self.skipTest(f"{backend} quantizer unavailable: {e}")
            quantizer.set_default_quant_config(quant_dtype=QuantDtype.use_8a8w)

backends/qualcomm/tests/test_passes.py:486

• The test relies on torch.arange(4) being int64; making the dtype explicit avoids potential behavior changes and matches the intent being tested.

                idx = torch.arange(4)  # int64 value written by index_put
                buf = buf.index_put((torch.tensor([0, 1, 2, 3]),), idx)

[shewu-quic]

Appreciate you catching that.

[shewu-quic]

If the input is non-float, the output should also be non-float. In this case, I’d expect this node to be unannotated. Could you help add this condition here?

if len(input_qspec_map) > 0 or output_act_qspec is not None:

[qti-horodnic]

Nit: can we make this torch.arange(4, dtype=torch.int64) explicitly?

[qti-horodnic]

LGTM, please address Hutton's comment, the Copilot comment about the test loop and my minor formatting comment before merging.

[Copilot]

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated no new comments.

Suppressed comments (1)

backends/qualcomm/tests/test_passes.py:516

• self.skipTest(...) is called inside the backend loop. If the LPAI quantizer is unavailable, this will skip the entire test (including the already-executed HTP leg), which defeats the intent of “skip that leg”. Instead, skip only the unavailable backend and only mark the whole test skipped if no backend could be exercised.

        for backend in (
            QnnExecuTorchBackendType.kHtpBackend,
            QnnExecuTorchBackendType.kLpaiBackend,
        ):
            try:
                quantizer = QnnQuantizer(backend=backend)
            except Exception as e:
                # LPAI needs quantized_aot_lib; skip that leg if it isn't available.
                self.skipTest(f"{backend} quantizer unavailable: {e}")