From 8b7d5346fdf138d606f520011beddb1207e54637 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Sun, 1 Nov 2026 21:28:41 -0800
Subject: [PATCH 1/7] Qualcomm AI Engine Direct - Decompose kept-whole ops that
 HTP leaves on CPU

On HTP, a partition that reads nothing at runtime stays on the CPU,
because HTP refuses a graph without runtime inputs. That is decided after
decomposition. So a constant-only op that QNN keeps whole, such as
linear, could land in such a partition undecomposed, and lowering failed
with "was not decomposed or delegated".

The HTP filter of ops_to_not_decompose now runs the same drop rule before
the kept-whole ops are decomposed, and decomposes only the ones it would
leave on the CPU. The real partitions are not known yet, so it partitions the
graph as if QNN took every op except the ones whose data types or shapes
rule it out (the checks filter_fn already makes, now shared as
must_stay_on_cpu). Ops in partitions that stay delegated keep their
whole form, and with it their quantization encodings.

The caller's skip_node_op_set and skip_node_id_set are left out of that
estimate on purpose. Decomposing adds nodes, and the new nodes take names
that the later nodes had before. If a skip could change what is
decomposed, adding one skip could rename the node that a
skip_node_id_set entry means, and the wrong node would be skipped. Ops
that QNN only rejects while partitioning are not known yet either. So a
kept-whole op that only a skip or such an op cuts off still fails as
before. And when such an op reads only constants and feeds a kept-whole
op, the estimate decomposes the kept-whole op although it ends up
delegated, so it is delegated in pieces.

exir calls ops_to_not_decompose several times on the ATen program and
uses only one of the filters, so the estimate runs on the filter's first
use.

On the lowered program the filter stays plain filter_fn, so exir still
reports kept-whole ops left on the CPU.

A new test stubs the QNN SDK setup, the QNN manager and the compile. It
checks that constant-only linears cut off by bfloat16 casts or by a
masked select are decomposed, that a constant-only linear in a delegated
partition stays whole, that the filter is plain filter_fn on the lowered
program and on GPU, which drops no partitions, and that a node ID skip
still names the same add when an op skip and an ID-skipped feeder are
present. The existing constant-only partition test still runs against
the real SDK.
---
 .../qualcomm/partition/qnn_partitioner.py     |  48 +++++++-
 backends/qualcomm/partition/utils.py          |  14 ++-
 backends/qualcomm/tests/test_passes.py        | 114 +++++++++++++++++-
 3 files changed, 169 insertions(+), 7 deletions(-)

diff --git a/backends/qualcomm/partition/qnn_partitioner.py b/backends/qualcomm/partition/qnn_partitioner.py
index 640d6fbf478..bf86ceb0fdc 100644
--- a/backends/qualcomm/partition/qnn_partitioner.py
+++ b/backends/qualcomm/partition/qnn_partitioner.py
@@ -4,9 +4,10 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 import copy
+import functools
 import logging
 from collections import defaultdict
-from typing import Any, Callable, Dict, List, Optional, Tuple
+from typing import Any, Callable, Dict, List, Optional, Set, Tuple
 
 import torch
 from executorch.backends.qualcomm.builders import node_visitor_manager
@@ -48,7 +49,7 @@
 from executorch.exir.dialects._ops import ops as exir_ops
 from torch.export.exported_program import ExportedProgram
 from torch.fx.passes.infra.partitioner import Partition
-from torch.fx.passes.operator_support import OperatorSupportBase
+from torch.fx.passes.operator_support import create_op_support, OperatorSupportBase
 
 from .common_defs import (
     allow_list_operator,
@@ -62,6 +63,7 @@
     get_skip_decomp_table,
     has_data_dependent_shape,
     has_unsupported_dtype,
+    must_stay_on_cpu,
     warn_if_dilated_conv_may_not_fit_vtcm,
 )
 
@@ -397,6 +399,46 @@ def ops_to_not_decompose(
         Determines which op should not be decomposed during partitioning.
         The list of operators is obtained from `get_skip_decomp_table()`.
         The filter function (`filter_fn`) can be used to further refine which nodes are not decomposed. (advanced use case)
+        On HTP, kept-whole ops that `_drop_partitions_without_runtime_inputs`
+        would leave on CPU, as far as can be told before they are decomposed,
+        are decomposed too.
         """
         do_not_decompose = get_skip_decomp_table()
-        return (do_not_decompose, filter_fn)
+        # On the lowered program exir reports every kept-whole op left there
+        # that the filter accepts, so it stays plain filter_fn.
+        if self.backend != QnnExecuTorchBackendType.kHtpBackend or ep.dialect == "EDGE":
+            return (do_not_decompose, filter_fn)
+
+        @functools.cache
+        def left_on_cpu() -> Set[torch.fx.Node]:
+            # The real partitions are only known once the kept-whole ops are
+            # decomposed, so the drop rule runs on the largest ones possible:
+            # only ops whose dtypes or shapes rule QNN out count as CPU ops.
+            # The caller's skips are left out, since decomposing adds nodes,
+            # and a skip must not rename the nodes that skip_node_id_set names.
+            # Ops that QNN rejects while partitioning are not known yet either.
+            # So when a skipped or rejected op is all that cuts a kept-whole op
+            # off, the op is not decomposed and still fails as before, and when
+            # such an op reads only constants and feeds a kept-whole op, the
+            # kept-whole op is decomposed although it ends up delegated.
+            partitions = generate_partitions_from_list_of_nodes(
+                ep.graph_module,
+                op_support=create_op_support(
+                    lambda _, node: node.op == "call_function"
+                    and not must_stay_on_cpu(node)
+                ),
+            )
+            delegated = self._drop_partitions_without_runtime_inputs(partitions, ep)
+            return {
+                node
+                for partition in partitions
+                if partition not in delegated
+                for node in partition.nodes
+            }
+
+        # exir calls this method several times on the ATen program but uses
+        # only one of the filters, so the partitions are found on first use.
+        return (
+            do_not_decompose,
+            lambda node: node not in left_on_cpu() and filter_fn(node),
+        )
diff --git a/backends/qualcomm/partition/utils.py b/backends/qualcomm/partition/utils.py
index 19fa9b1ad55..93adaaf2e6e 100644
--- a/backends/qualcomm/partition/utils.py
+++ b/backends/qualcomm/partition/utils.py
@@ -199,12 +199,20 @@ def has_data_dependent_shape(node: torch.fx.Node) -> bool:
     )
 
 
+def must_stay_on_cpu(node: torch.fx.Node) -> bool:
+    """
+    Whether the node's data types or shapes keep it off QNN. Both are already
+    known before decomposition.
+    """
+    return (
+        has_unsupported_dtype(node) and not _unsupported_dtypes_are_quantized(node)
+    ) or has_data_dependent_shape(node)
+
+
 # Logic to determine whether to skip decompose and has higher priority than get_skip_decomp_table()
 def filter_fn(node: torch.fx.Node) -> bool:
     # The node will stay on CPU, so it must not be kept from decomposition.
-    if has_unsupported_dtype(node) and not _unsupported_dtypes_are_quantized(node):
-        return False
-    if has_data_dependent_shape(node):
+    if must_stay_on_cpu(node):
         return False
     # QNN does not support int32/int64 IO for the following OPs.
     potential_i32_i64_io_ops = [
diff --git a/backends/qualcomm/tests/test_passes.py b/backends/qualcomm/tests/test_passes.py
index 2c1818f8a17..3042a894cc4 100644
--- a/backends/qualcomm/tests/test_passes.py
+++ b/backends/qualcomm/tests/test_passes.py
@@ -31,7 +31,11 @@
 )
 from executorch.backends.qualcomm.builders.qnn_constants import OpContextLoader
 from executorch.backends.qualcomm.builders.utils import is_parameter
-from executorch.backends.qualcomm.partition.qnn_partitioner import QnnOperatorSupport
+from executorch.backends.qualcomm.partition.qnn_partitioner import (
+    QnnOperatorSupport,
+    QnnPartitioner,
+)
+from executorch.backends.qualcomm.partition.utils import filter_fn
 from executorch.backends.qualcomm.qnn_preprocess import QnnBackend
 from executorch.backends.qualcomm.quantizer.quantizer import QnnQuantizer, QuantDtype
 from executorch.backends.qualcomm.serialization.qc_schema import (
@@ -52,14 +56,18 @@
     QCOM_SCALE,
     QCOM_ZERO_POINTS,
 )
+from executorch.backends.qualcomm.utils.qnn_manager_lifecycle import QnnManagerRegistry
 from executorch.backends.qualcomm.utils.utils import (
+    generate_gpu_compiler_spec,
     generate_htp_compiler_spec,
     generate_qnn_executorch_compiler_spec,
     to_edge_transform_and_lower_to_qnn,
 )
 from executorch.exir import EdgeCompileConfig, to_edge
+from executorch.exir.backend.backend_details import PreprocessResult
 from executorch.exir.debug_handle_utils import DEBUG_HANDLE_KEY
 from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.lowered_backend_module import get_lowered_backend_modules
 from torch.export.exported_program import OutputKind
 from torch.library import Library
 from torchao.quantization import IntxWeightOnlyConfig, PerAxis, quantize_
@@ -1044,6 +1052,110 @@ def forward(self, x):
             raise
         edge.to_executorch()
 
+    def test_htp_decomposes_kept_whole_ops_left_on_cpu(self):
+        class ConstantBranches(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("table", torch.randn(4, 16))
+                self.register_buffer("weight", torch.randn(16, 16))
+                self.register_buffer("mask", torch.rand(4, 16) > 0.5)
+
+            def forward(self, x):
+                linear = torch.nn.functional.linear
+                # QNN has no bfloat16 type, and a masked select has no static
+                # shape, so these branches read nothing at runtime once those
+                # ops run on CPU. Their linears are decomposed.
+                cast = linear(self.table, self.weight).to(torch.bfloat16).float()
+                masked = linear(self.table, self.weight)[self.mask].sum()
+                # This linear reads only constants too, but it shares a
+                # partition with the add that reads x, so it stays whole.
+                delegated = linear(self.table, self.weight)
+                return x + delegated + cast + masked
+
+        class SkippedFeeder(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("table", torch.randn(2, 4, 16))
+                self.linear = torch.nn.Linear(16, 16)
+
+            def forward(self, x):
+                return x + self.linear(self.table.abs()).unsqueeze(0)
+
+        sample_input = (torch.randn(4, 16),)
+        compiler_specs = generate_qnn_executorch_compiler_spec(
+            soc_model=QcomChipset.SM8650,
+            backend_options=generate_htp_compiler_spec(use_fp16=True),
+        )
+        gpu_compiler_specs = generate_qnn_executorch_compiler_spec(
+            soc_model=QcomChipset.SM8650,
+            backend_options=generate_gpu_compiler_spec(),
+        )
+
+        # Runs without the QNN SDK: the backend accepts every op and compiles
+        # nothing.
+        with patch(
+            "executorch.backends.qualcomm.utils.utils.setup_qnn_sdk"
+        ), patch.object(
+            QnnManagerRegistry, "get_or_create_qnn_manager", return_value=MagicMock()
+        ), patch.object(
+            QnnBackend,
+            "preprocess_multimethod",
+            side_effect=lambda programs, _: {
+                method: [PreprocessResult(processed_bytes=b"qnn")] * len(partitions)
+                for method, partitions in programs.items()
+            },
+        ):
+            edge = to_edge_transform_and_lower_to_qnn(
+                ConstantBranches().eval(), sample_input, compiler_specs
+            )
+            # Skips must not change what is decomposed. A decomposed linear
+            # would add a bias add that takes the name of the add skipped here.
+            skipped = to_edge_transform_and_lower_to_qnn(
+                SkippedFeeder().eval(),
+                sample_input,
+                compiler_specs,
+                skip_node_id_set={"aten_abs_default", "aten_add_tensor"},
+                skip_node_op_set={"aten.unsqueeze_copy.default"},
+            )
+        program = edge.exported_program()
+        top_level_ops = {
+            node.target for node in program.graph.nodes if node.op == "call_function"
+        }
+        self.assertNotIn(exir_ops.edge.aten.linear.default, top_level_ops)
+        (lowered,) = get_lowered_backend_modules(program.graph_module)
+        delegated_ops = {
+            node.target
+            for node in lowered.original_module.graph.nodes
+            if node.op == "call_function"
+        }
+        self.assertIn(exir_ops.edge.aten.linear.default, delegated_ops)
+        self.assertIn(exir_ops.edge.aten.add.Tensor, delegated_ops)
+        # On the lowered program the filter must not let a kept-whole op left
+        # on CPU through, so exir still reports it.
+        _, keep_whole = QnnPartitioner(compiler_specs).ops_to_not_decompose(program)
+        self.assertIs(keep_whole, filter_fn)
+        edge.to_executorch()
+        # GPU drops no partitions, so it keeps the same ops whole as before.
+        _, keep_whole = QnnPartitioner(gpu_compiler_specs).ops_to_not_decompose(
+            torch.export.export(ConstantBranches().eval(), sample_input)
+        )
+        self.assertIs(keep_whole, filter_fn)
+
+        program = skipped.exported_program()
+        (unsqueeze,) = (
+            node
+            for node in program.graph.nodes
+            if node.target == exir_ops.edge.aten.unsqueeze_copy.default
+        )
+        # The skipped add is the one that reads the unsqueeze, so it ran on CPU.
+        (add,) = unsqueeze.users
+        self.assertEqual(add.target, exir_ops.edge.aten.add.Tensor)
+        (lowered,) = get_lowered_backend_modules(program.graph_module)
+        self.assertIn(
+            exir_ops.edge.aten.linear.default,
+            [node.target for node in lowered.original_module.graph.nodes],
+        )
+
     def test_partitioner_falls_back_on_bfloat16(self):
         class Bf16Region(torch.nn.Module):
             def __init__(self):

From d1841774b7be418faa31de2bf04d83a61ae45913 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Sun, 1 Nov 2026 21:28:42 -0800
Subject: [PATCH 2/7] Document both calls to Partitioner.ops_to_not_decompose

to_edge_transform_and_lower calls ops_to_not_decompose during decomposition,
on the ATen program, to choose which nodes to keep whole, and after
lowering, on the EDGE program, to report kept-whole nodes that were neither
decomposed nor delegated. The docstring only described the first use.
QnnPartitioner tells the 2 calls apart by ep.dialect, and its HTP filter only
works on the nodes of the program it was given, so the docstring now names
both stages, their dialects, and that rule. It also says that the list
checked after lowering includes the lists of earlier partitioners for the
same method, and that some quantization ops and ops that exir does not
preserve are not reported. Docstring only.
---
 exir/backend/partitioner.py | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/exir/backend/partitioner.py b/exir/backend/partitioner.py
index c0edf0c71d8..a3bf97b7f9e 100644
--- a/exir/backend/partitioner.py
+++ b/exir/backend/partitioner.py
@@ -105,6 +105,18 @@ def ops_to_not_decompose(
         guaranteed that the program that the backend receives will not have any of these ops
         decomposed.
 
+        to_edge_transform_and_lower calls this method at two stages, possibly more than once
+        in each, and `ep.dialect` tells them apart. Each returned filter is only applied to
+        the nodes of the `ep` passed to that same call.
+
+        * During decomposition, on the ATen program (`ep.dialect == "ATEN"`). The list and
+          the filter decide which nodes are kept whole.
+        * After lowering, on the EDGE program (`ep.dialect == "EDGE"`). A node still in the
+          graph whose op is in this list, or in the list of an earlier partitioner for the
+          same method, and that the filter accepts was neither decomposed nor delegated,
+          and raises an error, except for some quantization ops and ops that exir
+          does not preserve.
+
         Returns:
             List[torch._ops.OpOverload]: a list of operator names that should not be decomposed.
             Optional[Callable[[torch.fx.Node], bool]]]: an optional callable, acting as a filter, that users can provide

From e2ba0765a54ea337c5d0ffdd67125b2cf7084807 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Mon, 2 Nov 2026 13:27:43 -0800
Subject: [PATCH 3/7] Qualcomm AI Engine Direct - Find HTP CPU ops without
 partitioning

To decide which kept-whole ops to decompose, the HTP filter of
ops_to_not_decompose estimated which ops HTP leaves on the CPU by running
the capability-based partitioner. That partitioner's time grows with the
square of the graph size: it was about 4 times slower for each doubling,
from 2,000 to 32,000 nodes. On a large vision-language-action model, with
about 43,000 nodes, the export sat in its cycle search for more than 2
hours.

Only a partition whose ops read no runtime value is dropped, so the
estimate now looks only at those ops. Ops whose data types or shapes rule
QNN out, and ops without a tensor output, such as the runtime asserts
export adds, count as CPU ops. An op whose schema returns nothing counts
as one too, since retracing can drop the "val" of an assert. Every
getitem counts as a QNN op, since QNN's op check accepts it. One pass
over the graph in order and one in reverse find, for every op QNN could
take:
- whether it reads a runtime value, directly or through other QNN ops;
- whether it reaches a QNN op through a CPU op, without passing an op
  that reads a runtime value. The partitioner may split such an op off
  to avoid a cycle.
The partitioner also merges ops that share no data, so an op that reads
no runtime value joins a partition that reads one, unless it has such a
CPU detour. Then it is left on the CPU. The estimate needs more in 2
cases. When no op reads a runtime value, every partition is dropped. And
the getitems of a CPU op join partitions while the partitioner merges
them, but are moved out afterwards, so they can keep an op out of every
partition that reads a runtime value. In these cases, a walk from the
ops that read a runtime value also finds the ops linked to them by QNN
ops alone, and an op that nothing links is left on the CPU too.

Each pass visits every node and edge once. The names of the constant
inputs are read from the graph signature once, since rebuilding them for
each input would grow with the square of their number. The drop rule
reads them the same way. The estimate, and its check of which ops may run
on QNN, are static methods next to the drop rule.

Skips, ops that QNN rejects while partitioning, and ops that its passes
remove before then, such as a clone, are still not known here. So when
such ops are all that cut a kept-whole op off, or all that read a
runtime value, the kept-whole op is not decomposed and fails as before.
Apart from those, the estimate errs toward decomposing: a constant-only
op with a CPU detour, or one that nothing links in the cases above, can
still end up delegated. Then it is delegated in pieces.

The constant-only linear that stays whole in the HTP test now also feeds
an add that reaches another QNN op through CPU casts, so the test checks
that such a path does not decompose it. New cases check that these
linears stay whole too:
- one that a masked select on CPU also reads. Nothing after the select
  reaches QNN, since its sum reads a data-dependent size and the runtime
  asserts for that size have no tensor output;
- one that shares no data with the ops that read the input, among them
  the getitems of a split that QNN takes, as before this estimate;
- one that feeds an add that reads the input, while the getitems of a
  CPU split read the input too;
- one that reads the input through an add and reaches a QNN op through
  CPU casts.
The test builds its compiler specs under the SDK patch, so it needs no
QNN SDK.

New cases in the HTP test, where a constant-only linear must be
decomposed: one feeds a delegated add both directly and through CPU
casts, one feeds a delegated add and a CPU split whose getitems QNN takes,
one sits in a model whose input reaches only a CPU op, and one sits in a
model where only the getitems of a CPU split read the input.

A new test checks, on chains of 2 and 4 links whose constants each feed an
op that reads only constants, that the estimate never runs the
partitioner and reads the input specs of the graph signature as often on
both. Partitioning, or going over those specs for each input, would make
its time grow with the square of the graph size. It counts instead of
timing, so it does not depend on the machine, and it needs no QNN SDK.
---
 .../qualcomm/partition/qnn_partitioner.py     | 164 ++++++++++++----
 backends/qualcomm/tests/test_passes.py        | 178 ++++++++++++++++--
 2 files changed, 284 insertions(+), 58 deletions(-)

diff --git a/backends/qualcomm/partition/qnn_partitioner.py b/backends/qualcomm/partition/qnn_partitioner.py
index bf86ceb0fdc..f78e84bf14a 100644
--- a/backends/qualcomm/partition/qnn_partitioner.py
+++ b/backends/qualcomm/partition/qnn_partitioner.py
@@ -6,6 +6,7 @@
 import copy
 import functools
 import logging
+import operator
 from collections import defaultdict
 from typing import Any, Callable, Dict, List, Optional, Set, Tuple
 
@@ -15,10 +16,7 @@
     PER_CHANNEL_GROUP_ENCODING,
 )
 from executorch.backends.qualcomm.builders.qnn_constants import OpContextLoader
-from executorch.backends.qualcomm.builders.utils import (
-    is_mutable_buffer_input,
-    is_parameter,
-)
+from executorch.backends.qualcomm.builders.utils import is_parameter
 from executorch.backends.qualcomm.qnn_preprocess import QnnBackend
 from executorch.backends.qualcomm.serialization.qc_schema import (
     HtpArch,
@@ -48,8 +46,10 @@
 from executorch.exir.backend.utils import tag_constant_data, tag_mutated_buffer
 from executorch.exir.dialects._ops import ops as exir_ops
 from torch.export.exported_program import ExportedProgram
+from torch.export.graph_signature import InputKind
 from torch.fx.passes.infra.partitioner import Partition
-from torch.fx.passes.operator_support import create_op_support, OperatorSupportBase
+from torch.fx.passes.operator_support import OperatorSupportBase
+from torch.utils import _pytree as pytree
 
 from .common_defs import (
     allow_list_operator,
@@ -312,6 +312,26 @@ def _drop_orphaned_parameter_dequants(
                 )
         return [partition for partition in partitions if partition.nodes]
 
+    @staticmethod
+    def _constant_inputs(edge_program: torch.export.ExportedProgram) -> Set[str]:
+        """
+        Names of the inputs that hold the same value on every run: parameters,
+        buffers that are never mutated, and lifted constants.
+        """
+        signature = edge_program.graph_signature
+        mutated = set(signature.buffers_to_mutate.values())
+        return {
+            spec.arg.name
+            for spec in signature.input_specs
+            if spec.kind
+            in (InputKind.PARAMETER, InputKind.BUFFER, InputKind.CONSTANT_TENSOR)
+            and spec.target not in mutated
+        }
+
+    @staticmethod
+    def _is_runtime_input(node: torch.fx.Node, constant_inputs: Set[str]) -> bool:
+        return node.op != "get_attr" and node.name not in constant_inputs
+
     @staticmethod
     def _drop_partitions_without_runtime_inputs(
         partitions: List[Partition], edge_program: torch.export.ExportedProgram
@@ -321,23 +341,114 @@ def _drop_partitions_without_runtime_inputs(
         only by constants once a CPU op splits it from the rest. Such a
         partition stays on CPU.
         """
-
-        def is_runtime_input(node: torch.fx.Node) -> bool:
-            return not is_parameter(node, edge_program) or is_mutable_buffer_input(
-                node, edge_program
-            )
-
+        constant_inputs = QnnPartitioner._constant_inputs(edge_program)
         return [
             partition
             for partition in partitions
             if any(
-                is_runtime_input(arg)
+                QnnPartitioner._is_runtime_input(arg, constant_inputs)
                 for node in partition.nodes
                 for arg in node.all_input_nodes
                 if arg not in partition.nodes
             )
         ]
 
+    @staticmethod
+    def _may_run_on_qnn(node: torch.fx.Node) -> bool:
+        """
+        Whether QNN's op check may accept the node, as far as can be told
+        before the kept-whole ops are decomposed. Ops whose dtypes or shapes
+        rule QNN out, and ops without a tensor output, such as the runtime
+        asserts export adds, stay on CPU. Every getitem is accepted, as in
+        QNN's op check.
+        """
+        if node.op != "call_function":
+            return False
+        if node.target in allow_list_operator:
+            return True
+        # Retracing can drop the "val" of an assert, but not its schema.
+        if (
+            isinstance(node.target, torch._ops.OpOverload)
+            and not node.target._schema.returns
+        ):
+            return False
+        if "val" in node.meta and not any(
+            isinstance(val, torch.Tensor)
+            for val in pytree.tree_leaves(node.meta["val"])
+        ):
+            return False
+        return not must_stay_on_cpu(node)
+
+    @staticmethod
+    def _left_on_cpu(ep: ExportedProgram) -> Set[torch.fx.Node]:
+        """
+        Estimates, on the ATen program, which ops HTP leaves on CPU, so that
+        the kept-whole ones among them can be decomposed.
+        """
+
+        # The real partitions are only known once the kept-whole ops are
+        # decomposed, and partitioning would be too slow here, since its time
+        # grows with the square of the graph size. So only the ops that
+        # _may_run_on_qnn rules out count as CPU ops here.
+        # Only a partition of ops that read no runtime value is dropped. The
+        # partitioner also merges ops that no path links, so an op that reads
+        # no runtime value joins a partition that reads one, unless it reaches
+        # a QNN op through a CPU op without passing an op that reads one,
+        # since the partitioner may then split it off to avoid a cycle.
+        # In 2 cases this is not enough. When no op reads a runtime value,
+        # every partition is dropped. And the getitems of a CPU op join
+        # partitions while the partitioner merges them, but are moved out
+        # afterwards, so they can keep an op out of every partition that reads
+        # a runtime value. In these cases an op is also left on CPU when no
+        # chain of QNN ops links it to an op that reads one. Such an op can
+        # still end up delegated, in pieces.
+        # The caller's skips are left out, since decomposing adds nodes, and a
+        # skip must not rename the nodes that skip_node_id_set names. Ops that
+        # QNN rejects while partitioning, or that its passes remove before
+        # then, are not known yet either. So when such an op is all that cuts
+        # a kept-whole op off, or all that reads a runtime value, the op is
+        # not decomposed and still fails as before. And when such an op reads
+        # only constants and feeds a kept-whole op, the kept-whole op can be
+        # decomposed although it ends up delegated.
+        nodes = list(ep.graph.nodes)
+        qnn = set(filter(QnnPartitioner._may_run_on_qnn, nodes))
+        constant_inputs = QnnPartitioner._constant_inputs(ep)
+        runtime = set()
+        for node in nodes:
+            if node in qnn and any(
+                (
+                    arg in runtime
+                    if arg in qnn
+                    else QnnPartitioner._is_runtime_input(arg, constant_inputs)
+                )
+                for arg in node.all_input_nodes
+            ):
+                runtime.add(node)
+        to_qnn, before_cpu = set(), set()
+        for node in reversed(nodes):
+            if node in qnn or not to_qnn.isdisjoint(node.users):
+                to_qnn.add(node)
+            if any(
+                (user in before_cpu and user not in runtime)
+                or (user in to_qnn and user not in qnn)
+                for user in node.users
+            ):
+                before_cpu.add(node)
+        left_on_cpu = (qnn - runtime) & before_cpu
+        if runtime and not any(
+            node.target is operator.getitem and node.args[0] not in qnn
+            for node in runtime
+        ):
+            return left_on_cpu
+        linked, stack = set(runtime), list(runtime)
+        while stack:
+            node = stack.pop()
+            for neighbor in (*node.all_input_nodes, *node.users):
+                if neighbor in qnn and neighbor not in linked:
+                    linked.add(neighbor)
+                    stack.append(neighbor)
+        return left_on_cpu | (qnn - runtime - linked)
+
     def tag_nodes(
         self, partitions: List[Partition], edge_program: torch.export.ExportedProgram
     ) -> None:
@@ -409,35 +520,10 @@ def ops_to_not_decompose(
         if self.backend != QnnExecuTorchBackendType.kHtpBackend or ep.dialect == "EDGE":
             return (do_not_decompose, filter_fn)
 
-        @functools.cache
-        def left_on_cpu() -> Set[torch.fx.Node]:
-            # The real partitions are only known once the kept-whole ops are
-            # decomposed, so the drop rule runs on the largest ones possible:
-            # only ops whose dtypes or shapes rule QNN out count as CPU ops.
-            # The caller's skips are left out, since decomposing adds nodes,
-            # and a skip must not rename the nodes that skip_node_id_set names.
-            # Ops that QNN rejects while partitioning are not known yet either.
-            # So when a skipped or rejected op is all that cuts a kept-whole op
-            # off, the op is not decomposed and still fails as before, and when
-            # such an op reads only constants and feeds a kept-whole op, the
-            # kept-whole op is decomposed although it ends up delegated.
-            partitions = generate_partitions_from_list_of_nodes(
-                ep.graph_module,
-                op_support=create_op_support(
-                    lambda _, node: node.op == "call_function"
-                    and not must_stay_on_cpu(node)
-                ),
-            )
-            delegated = self._drop_partitions_without_runtime_inputs(partitions, ep)
-            return {
-                node
-                for partition in partitions
-                if partition not in delegated
-                for node in partition.nodes
-            }
+        left_on_cpu = functools.cache(lambda: self._left_on_cpu(ep))
 
         # exir calls this method several times on the ATen program but uses
-        # only one of the filters, so the partitions are found on first use.
+        # only one of the filters, so the estimate runs on first use.
         return (
             do_not_decompose,
             lambda node: node not in left_on_cpu() and filter_fn(node),
diff --git a/backends/qualcomm/tests/test_passes.py b/backends/qualcomm/tests/test_passes.py
index 3042a894cc4..ad96e9d39d8 100644
--- a/backends/qualcomm/tests/test_passes.py
+++ b/backends/qualcomm/tests/test_passes.py
@@ -1,5 +1,5 @@
 import unittest
-from unittest.mock import MagicMock, patch
+from unittest.mock import MagicMock, patch, PropertyMock
 
 import torch
 from executorch.backends.qualcomm._passes import (
@@ -69,6 +69,8 @@
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.lowered_backend_module import get_lowered_backend_modules
 from torch.export.exported_program import OutputKind
+from torch.export.graph_signature import ExportGraphSignature
+from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner
 from torch.library import Library
 from torchao.quantization import IntxWeightOnlyConfig, PerAxis, quantize_
 from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e
@@ -1052,7 +1054,7 @@ def forward(self, x):
             raise
         edge.to_executorch()
 
-    def test_htp_decomposes_kept_whole_ops_left_on_cpu(self):
+    def test_htp_decomposes_kept_whole_ops_left_on_cpu(self):  # noqa: C901
         class ConstantBranches(torch.nn.Module):
             def __init__(self):
                 super().__init__()
@@ -1067,10 +1069,77 @@ def forward(self, x):
                 # ops run on CPU. Their linears are decomposed.
                 cast = linear(self.table, self.weight).to(torch.bfloat16).float()
                 masked = linear(self.table, self.weight)[self.mask].sum()
+                # This linear feeds a delegated add both directly and through
+                # CPU casts, so the partitioner splits it off to avoid a cycle.
+                # It is decomposed.
+                detour = linear(self.table, self.weight)
+                detour = detour + detour.to(torch.bfloat16).float()
                 # This linear reads only constants too, but it shares a
-                # partition with the add that reads x, so it stays whole.
+                # partition with the add that reads x, so it stays whole, even
+                # though that add reaches another QNN op through CPU casts.
                 delegated = linear(self.table, self.weight)
-                return x + delegated + cast + masked
+                y = x + delegated + cast + masked + detour
+                return y.to(torch.bfloat16).float() + 1
+
+        class MaskedReader(ConstantBranches):
+            def forward(self, x):
+                # Nothing after the masked select reaches QNN: the sum reads a
+                # data-dependent size, and the runtime asserts that export adds
+                # for that size have no tensor output. So the linear shares the
+                # add's partition and stays whole.
+                linear = torch.nn.functional.linear(self.table, self.weight)
+                return x + linear, linear[self.mask].sum()
+
+        class NoRuntimeInput(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("table", torch.randn(4, 16))
+                self.register_buffer("weight", torch.randn(16, 16))
+
+            def forward(self, x):
+                # x only reaches a CPU cast, so no partition reads a runtime
+                # value, and the linear is decomposed.
+                linear = torch.nn.functional.linear(self.table, self.weight)
+                return x.to(torch.bfloat16), linear * 2
+
+        class CastSplit(NoRuntimeInput):
+            def forward(self, x):
+                # QNN takes every getitem, so the getitems of this CPU split
+                # can join the add's partition. The linear would then close a
+                # cycle through the split, so it is left out and decomposed.
+                linear = torch.nn.functional.linear(self.table, self.weight)
+                low, high = linear.to(torch.bfloat16).split(8, dim=-1)
+                return x + linear, low, high
+
+        class ConstantOutput(NoRuntimeInput):
+            def forward(self, x):
+                # No path links this linear to the ops that read x, but the
+                # partitioner also merges such ops, so the linear stays whole.
+                # The getitems of the split stay with it, since QNN takes it.
+                linear = torch.nn.functional.linear(self.table, self.weight)
+                return x + 1, *x.split(8, dim=-1), linear
+
+        class CastSplitInput(NoRuntimeInput):
+            def forward(self, x):
+                # Only the getitems of a CPU split read x, and the partitioner
+                # moves them out of the linear's partition, so it reads no
+                # runtime value. The linear is decomposed.
+                low, high = x.to(torch.bfloat16).split(8, dim=-1)
+                return low, high, torch.nn.functional.linear(self.table, self.weight)
+
+        class CastSplitLinked(NoRuntimeInput):
+            def forward(self, x):
+                # Here the linear feeds an add that reads x, so it stays whole.
+                low, high = x.to(torch.bfloat16).split(8, dim=-1)
+                linear = torch.nn.functional.linear(self.table, self.weight)
+                return low, high, x + linear
+
+        class LinearOfInput(NoRuntimeInput):
+            def forward(self, x):
+                # This linear reads x through the add, so it stays whole even
+                # though it reaches a QNN op through CPU casts.
+                linear = torch.nn.functional.linear(x + 1, self.weight)
+                return linear + linear.to(torch.bfloat16).float()
 
         class SkippedFeeder(torch.nn.Module):
             def __init__(self):
@@ -1082,15 +1151,6 @@ def forward(self, x):
                 return x + self.linear(self.table.abs()).unsqueeze(0)
 
         sample_input = (torch.randn(4, 16),)
-        compiler_specs = generate_qnn_executorch_compiler_spec(
-            soc_model=QcomChipset.SM8650,
-            backend_options=generate_htp_compiler_spec(use_fp16=True),
-        )
-        gpu_compiler_specs = generate_qnn_executorch_compiler_spec(
-            soc_model=QcomChipset.SM8650,
-            backend_options=generate_gpu_compiler_spec(),
-        )
-
         # Runs without the QNN SDK: the backend accepts every op and compiles
         # nothing.
         with patch(
@@ -1105,9 +1165,32 @@ def forward(self, x):
                 for method, partitions in programs.items()
             },
         ):
+            compiler_specs = generate_qnn_executorch_compiler_spec(
+                soc_model=QcomChipset.SM8650,
+                backend_options=generate_htp_compiler_spec(use_fp16=True),
+            )
+            gpu_compiler_specs = generate_qnn_executorch_compiler_spec(
+                soc_model=QcomChipset.SM8650,
+                backend_options=generate_gpu_compiler_spec(),
+            )
             edge = to_edge_transform_and_lower_to_qnn(
                 ConstantBranches().eval(), sample_input, compiler_specs
             )
+            for model in (NoRuntimeInput(), CastSplit(), CastSplitInput()):
+                to_edge_transform_and_lower_to_qnn(
+                    model.eval(), sample_input, compiler_specs
+                )
+            kept_whole = [
+                to_edge_transform_and_lower_to_qnn(
+                    model.eval(), sample_input, compiler_specs
+                )
+                for model in (
+                    MaskedReader(),
+                    ConstantOutput(),
+                    CastSplitLinked(),
+                    LinearOfInput(),
+                )
+            ]
             # Skips must not change what is decomposed. A decomposed linear
             # would add a bias add that takes the name of the add skipped here.
             skipped = to_edge_transform_and_lower_to_qnn(
@@ -1122,9 +1205,9 @@ def forward(self, x):
             node.target for node in program.graph.nodes if node.op == "call_function"
         }
         self.assertNotIn(exir_ops.edge.aten.linear.default, top_level_ops)
-        (lowered,) = get_lowered_backend_modules(program.graph_module)
         delegated_ops = {
             node.target
+            for lowered in get_lowered_backend_modules(program.graph_module)
             for node in lowered.original_module.graph.nodes
             if node.op == "call_function"
         }
@@ -1150,11 +1233,68 @@ def forward(self, x):
         # The skipped add is the one that reads the unsqueeze, so it ran on CPU.
         (add,) = unsqueeze.users
         self.assertEqual(add.target, exir_ops.edge.aten.add.Tensor)
-        (lowered,) = get_lowered_backend_modules(program.graph_module)
-        self.assertIn(
-            exir_ops.edge.aten.linear.default,
-            [node.target for node in lowered.original_module.graph.nodes],
-        )
+        self._assert_delegated(
+            exir_ops.edge.aten.linear.default, [*kept_whole, skipped]
+        )
+
+    def _assert_delegated(self, target, edges):
+        for edge in edges:
+            graph_module = edge.exported_program().graph_module
+            self.assertIn(
+                target,
+                [
+                    node.target
+                    for lowered in get_lowered_backend_modules(graph_module)
+                    for node in lowered.original_module.graph.nodes
+                ],
+            )
+
+    def test_htp_kept_whole_estimate_avoids_quadratic_steps(self):
+        class Chain(torch.nn.Module):
+            def __init__(self, length):
+                super().__init__()
+                self.scales = torch.nn.ParameterList(
+                    torch.rand(16) for _ in range(length)
+                )
+
+            def forward(self, x):
+                # Each scale * 2 reads only constants, like the dequantize of a
+                # quantized weight.
+                for scale in self.scales:
+                    x = torch.relu(x * (scale * 2) + x)
+                return x
+
+        with patch("executorch.backends.qualcomm.utils.utils.setup_qnn_sdk"):
+            compiler_specs = generate_qnn_executorch_compiler_spec(
+                soc_model=QcomChipset.SM8650,
+                backend_options=generate_htp_compiler_spec(use_fp16=True),
+            )
+
+        # Large models have tens of thousands of nodes and constants. The time
+        # to partition grows with the square of the graph size, and so does
+        # going over the input specs of the graph signature for each input.
+        def input_spec_reads(length):
+            ep = torch.export.export(Chain(length), (torch.randn(4, 16),))
+            input_specs = ep.graph_signature.input_specs
+            _, keep_whole = QnnPartitioner(compiler_specs).ops_to_not_decompose(ep)
+            node = next(node for node in ep.graph.nodes if node.op == "call_function")
+            with patch.object(
+                CapabilityBasedPartitioner,
+                "propose_partitions",
+                autospec=True,
+                side_effect=CapabilityBasedPartitioner.propose_partitions,
+            ) as partition, patch.object(
+                ExportGraphSignature,
+                "input_specs",
+                new_callable=PropertyMock,
+                return_value=input_specs,
+                create=True,
+            ) as read_input_specs:
+                keep_whole(node)
+            partition.assert_not_called()
+            return read_input_specs.call_count
+
+        self.assertEqual(input_spec_reads(2), input_spec_reads(4))
 
     def test_partitioner_falls_back_on_bfloat16(self):
         class Bf16Region(torch.nn.Module):

From c389cba45dfe5d7b8a2ebdc338d11b4f51e940c6 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Mon, 2 Nov 2026 14:02:00 -0800
Subject: [PATCH 4/7] Test reserve_node_names on its own

reserve_node_names was only tested through the Qualcomm pass manager,
and only in a test that needs the Qualcomm backend. The Samsung pass
manager calls it too. A torch-only test now exports a model whose input
takes the name "mul", so export renames the multiply to "mul_1". After
reserve_node_names, a new node asked to be named "mul_1" gets another
name, and the graph lints. Once export records the names it renames,
the test skips, like the QNN name collision test, since it then no
longer tests the helper.
---
 backends/transforms/targets.bzl               | 11 ++++++
 .../test/test_reserve_node_names.py           | 39 +++++++++++++++++++
 2 files changed, 50 insertions(+)
 create mode 100644 backends/transforms/test/test_reserve_node_names.py

diff --git a/backends/transforms/targets.bzl b/backends/transforms/targets.bzl
index 4dd160aa44d..3722955f1b2 100644
--- a/backends/transforms/targets.bzl
+++ b/backends/transforms/targets.bzl
@@ -607,6 +607,17 @@ def define_common_targets():
         ],
     )
 
+    runtime.python_test(
+        name = "test_reserve_node_names",
+        srcs = [
+            "test/test_reserve_node_names.py",
+        ],
+        deps = [
+            "//caffe2:torch",
+            ":utils",
+        ],
+    )
+
     runtime.python_test(
         name = "test_convert_dtype_pass",
         srcs = [
diff --git a/backends/transforms/test/test_reserve_node_names.py b/backends/transforms/test/test_reserve_node_names.py
new file mode 100644
index 00000000000..20c7b7a779a
--- /dev/null
+++ b/backends/transforms/test/test_reserve_node_names.py
@@ -0,0 +1,39 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import unittest
+
+import torch
+from executorch.backends.transforms.utils import reserve_node_names
+
+
+class TestReserveNodeNames(unittest.TestCase):
+    def test_new_node_does_not_take_an_existing_name(self):
+        class Model(torch.nn.Module):
+            def forward(self, x, mul):
+                return x * 2 + mul
+
+        # The input takes the name "mul", so export renames the multiply to
+        # "mul_1" without recording that name in the graph's namespace.
+        ep = torch.export.export(Model(), (torch.randn(2), torch.randn(2)))
+        graph = ep.graph_module.graph
+        self.assertIn("mul_1", [node.name for node in graph.nodes])
+        # A skip, not a failure: a torch update that records the name is fine.
+        if "mul_1" in graph._graph_namespace._used_names:
+            self.skipTest(
+                "export now records the names it renames, so this no longer "
+                "tests reserve_node_names; see "
+                "https://github.com/pytorch/pytorch/issues/202927"
+            )
+
+        reserve_node_names(ep.graph_module)
+        (add,) = graph.find_nodes(op="call_function", target=torch.ops.aten.add.Tensor)
+        with graph.inserting_after(add):
+            node = graph.create_node(
+                "call_function", torch.ops.aten.neg.default, (add,), name="mul_1"
+            )
+        self.assertNotEqual(node.name, "mul_1")
+        graph.lint()

From 0b23612602dfe1dc304699390ee9efdbdc13c26b Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Mon, 2 Nov 2026 14:02:00 -0800
Subject: [PATCH 5/7] Keep the original error in the QNN name collision test

When the pass manager failed for a reason other than the name
collision, the test reported an assertion about the message, and the
real error showed up only as the context of that assertion. The test
now re-raises any error that is not the collision, with its own
traceback.
---
 backends/qualcomm/tests/test_passes.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/backends/qualcomm/tests/test_passes.py b/backends/qualcomm/tests/test_passes.py
index ad96e9d39d8..44865722b9c 100644
--- a/backends/qualcomm/tests/test_passes.py
+++ b/backends/qualcomm/tests/test_passes.py
@@ -261,7 +261,8 @@ def forward(self, input, k, v, mul):
                         ).graph.lint()
                     except Exception as e:
                         # The pass manager wraps the lint error in its own.
-                        self.assertIn("Node redefined name", str(e.__cause__ or e))
+                        if "Node redefined name" not in str(e.__cause__ or e):
+                            raise
                     else:
                         no_collision.append(type(model).__name__)
                     reserve.assert_called_once()

From 02b56a53970b255d4a0dfc8443804e13b43ac94d Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Mon, 2 Nov 2026 14:02:10 -0800
Subject: [PATCH 6/7] Skip the QNN name collision test only when no case
 collides

The test skips when a torch update removes the name collision, since
the cases then no longer test reserve_node_names. It skipped as soon as
any case stopped colliding, which also hid the cases that still did.
It now skips only when none of them collide.
---
 backends/qualcomm/tests/test_passes.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/backends/qualcomm/tests/test_passes.py b/backends/qualcomm/tests/test_passes.py
index 44865722b9c..b16f2a0e064 100644
--- a/backends/qualcomm/tests/test_passes.py
+++ b/backends/qualcomm/tests/test_passes.py
@@ -266,8 +266,9 @@ def forward(self, input, k, v, mul):
                     else:
                         no_collision.append(type(model).__name__)
                     reserve.assert_called_once()
-        # A skip, not a failure: a torch update that removes the collision is fine.
-        if no_collision:
+        # A skip, not a failure: a torch update that removes the collision is
+        # fine. While any case still collides, the test still means something.
+        if len(no_collision) == len(cases):
             self.skipTest(
                 f"{no_collision} no longer collide without reserve_node_names, so "
                 "those cases no longer test it; see "

From 2a452b300e8cc38dc03ecdaa7b8cf5fd127ce9b3 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Mon, 2 Nov 2026 17:46:18 -0800
Subject: [PATCH 7/7] Qualcomm AI Engine Direct - Cast uint8 to float32 through
 int32

On the HTP simulator (QNN SDK 2.51, SM8650, float16 graph), a cast from
uint8 to float32 gives 1 for every input above 0, as if the input were
bool. So a camera image passed as uint8 reached the model as a mask of
0s and 1s. The cast op is built with the right data types (UINT_8 in,
FLOAT_32 out). HTP casts uint8 to int32 correctly, and int32 to float32
correctly, but not uint8 to float32 in one step. Chains such as uint8 to
float16 to float32 gave the same result, because consecutive casts are
fused into one uint8 to float32 cast before the graph is built.

The cast builder already goes through int32 for casts to or from int64
that it cannot do in one step. A uint8 to float32 cast with a plain
float output now takes the same 2 steps and gives 0 to 255. Casts with
a quantized output are built as before: on the simulator they were
already right.

test_qnn_backend_cast in the floating point and quantized operator tests
now also scales a uint8 tensor holding 0 to 255 to float32 in [0, 1],
the way a camera input is prepared. The quantized test checks its
outputs with a relative tolerance of 0.1 instead of 1, since with 1 a
cast that gives only 0 or 1 would also pass.
---
 backends/qualcomm/builders/op_to.py          | 22 +++++++++++++++-----
 backends/qualcomm/tests/models.py            |  8 +++++++
 backends/qualcomm/tests/test_qnn_delegate.py | 19 ++++++++++++-----
 3 files changed, 39 insertions(+), 10 deletions(-)

diff --git a/backends/qualcomm/builders/op_to.py b/backends/qualcomm/builders/op_to.py
index 9b024307934..ce677fde335 100644
--- a/backends/qualcomm/builders/op_to.py
+++ b/backends/qualcomm/builders/op_to.py
@@ -97,13 +97,25 @@ def define_node(
         # since int32 is the only source that can be cast into int64
         # this is mainly for validation purpose, redundant cast ops will be fused
         # in preprocess stage.
+        # HTP casts uint8 to float32 as if the input were bool, while uint8 to
+        # int32 and int32 to float32 are right, so a plain uint8 to float32
+        # cast goes through int32 too. Here the step is needed for the values,
+        # and nothing fuses it away.
         ops = []
         if (
-            node.meta["val"].dtype == torch.int64
-            and input_node.meta["val"].dtype != torch.int32
-        ) or (
-            input_node.meta["val"].dtype == torch.int64
-            and node.meta["val"].dtype != torch.int32
+            (
+                node.meta["val"].dtype == torch.int64
+                and input_node.meta["val"].dtype != torch.int32
+            )
+            or (
+                input_node.meta["val"].dtype == torch.int64
+                and node.meta["val"].dtype != torch.int32
+            )
+            or (
+                input_node.meta["val"].dtype == torch.uint8
+                and node.meta["val"].dtype == torch.float32
+                and not node.meta.get(QCOM_QUANT_ATTRS)
+            )
         ):
             input_quant_encoding, input_quant_configs = self.get_quant_encoding_conf(
                 input_node, node
diff --git a/backends/qualcomm/tests/models.py b/backends/qualcomm/tests/models.py
index 94d875617c7..2167c1fa3ed 100644
--- a/backends/qualcomm/tests/models.py
+++ b/backends/qualcomm/tests/models.py
@@ -402,6 +402,14 @@ def forward(self, x, y):
         return res + index.to(torch.int32)
 
 
+class CastToFloat(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return x.to(torch.float32) / 255
+
+
 class Cat2(torch.nn.Module):
     def __init__(self):
         super().__init__()
diff --git a/backends/qualcomm/tests/test_qnn_delegate.py b/backends/qualcomm/tests/test_qnn_delegate.py
index eca38e27451..384eead521a 100644
--- a/backends/qualcomm/tests/test_qnn_delegate.py
+++ b/backends/qualcomm/tests/test_qnn_delegate.py
@@ -528,10 +528,11 @@ def test_qnn_backend_bmm(self):
         self.lower_module_and_test_output(module, sample_input)
 
     def test_qnn_backend_cast(self):
-        modules = [Cast(), CastMultiUsers()]  # noqa: F405
+        modules = [Cast(), CastMultiUsers(), CastToFloat()]  # noqa: F405
         sample_inputs = [
             (10 * torch.rand((9, 4, 5, 3)),),
             (torch.randint(0, 3, size=(3, 3)), torch.randn(3, 3)),
+            (torch.arange(256, dtype=torch.uint8).reshape(16, 16),),
         ]
         for i, (module, sample_input) in enumerate(zip(modules, sample_inputs)):
             with self.subTest(i=i):
@@ -3715,10 +3716,18 @@ def test_qnn_backend_bmm(self):
                 self.lower_module_and_test_output(qdq_module, sample_input)
 
     def test_qnn_backend_cast(self):
-        module = CastMultiUsers()  # noqa: F405
-        sample_input = (torch.randint(0, 3, size=(3, 3)), torch.randn(3, 3))
-        module = self.get_qdq_module(module, sample_input)
-        self.lower_module_and_test_output(module, sample_input)
+        # With rtol 1, any output between 0 and twice the reference passes, so
+        # a uint8 cast that gives only 0 or 1 would pass too.
+        self.rtol = 1e-1
+        modules = [CastMultiUsers(), CastToFloat()]  # noqa: F405
+        sample_inputs = [
+            (torch.randint(0, 3, size=(3, 3)), torch.randn(3, 3)),
+            (torch.arange(256, dtype=torch.uint8).reshape(16, 16),),
+        ]
+        for i, (module, sample_input) in enumerate(zip(modules, sample_inputs)):
+            with self.subTest(i=i):
+                module = self.get_qdq_module(module, sample_input)
+                self.lower_module_and_test_output(module, sample_input)
 
     def test_qnn_backend_cat(self):
         modules = [Cat2(), Cat3(), Cat4(), Cat5()]  # noqa: F405