Qualcomm AI Engine Direct - Preserve partition metadata on single-output delegate getitem nodes

[DannyYuyang-quic]

Summary

The single-output path in create_submodule_from_nodes copied only "val" onto the manually-built getitem, dropping other partition meta (e.g. QNN's q_tensor_io tag) that the multi-output path keeps via propagate_meta, leaving backend IO-tagging passes with wrong delegate dtype/spec.

Fix by propagating full meta then dropping nn_module_stack/source_fn_stack, matching the multi-output path.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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:

• Unrelated MacOS MPS test jobs time out (30 min) due to stale cpp_extension lock file poisoning self-hosted runners

✅ You can merge normally! (1 Unrelated Failure)

As of commit 3b99a81 with merge base 2265ab8 ():

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

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

• pull / android / build-android (gh) (trunk failure)
  ImportError: /opt/hostedtoolcache/Python/3.11.15/x64/lib/python3.11/site-packages/executorch/extension/training/pybindings/_training_lib.cpython-311-x86_64-linux-gnu.so: undefined symbol: _ZN3c104impl3cow23materialize_cow_storageERNS_11StorageImplE

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

[DannyYuyang-quic]

@psiddh @larryliu0820 @JacobSzwejbka Hi,
We ran into a node metadata loss issue while sharding Gemma4 into multiple shards, and it traces back to create_submodule_from_nodes in exir/lowered_backend_module.py. Sharding produces multiple delegates, some single-output, some multi-output, and the single-output ones are exactly what triggers this bug:
The getitem node for a single-output delegate drops the metadata carried by its upstream delegate output, whereas multi-output delegates preserve it. I put together a couple of before/after diagrams and wanted to run my understanding (and the fix) by you.

What's happening and Root cause

We use TagQuantIO Pass to tag the partition input/output with q_tensor_io=uint16 in node's metadata before to_backend, and BuildQuantIO Pass uses that tag to assign delegate IO's dtype/spec after to_backend.

single-output and multi-output take two different paths:

• multi-output: the getitem is created with propagate_meta=True, so it inherits the full meta, and then just pops nn_module_stack / source_fn_stack afterwards the q_tensor_io stays. reference

• single-output: insert_subgm doesn't create a getitem node, so create_submodule_from_nodes inserts one separately and only carries val key, dropping everything else → q_tensor_io is lost (see Graph 1).

Fix

Bring the single-output branch in line with multi-output: use replace_all_uses_with(proxy_out, propagate_meta=True) to carry the full meta, then pop nn_module_stack / source_fn_stack. Now both paths behave the same, and we still make sure the source-tracking meta doesn't leak onto the delegate getitem. (see Graph 2).

[Graph 1] Before the fix: getitem is missing q_tensor_io	[Graph 2] After the fix: getitem carries q_tensor_io=torch.uint16

Question

I wanted to check with you, since you have more historical background here, was the single-output branch only copying val on purpose? My change makes it "copy everything + pop the source-tracking keys" to match the multi-output path, but I might be missing some context. Are there any side effects or edge cases I should be aware of? Happy to adjust the approach if there's a better way to handle this.

Thanks so much for taking a look!!

cc: @shewu-quic @haowhsu-quic @winskuo-quic

[DannyYuyang-quic]

@pytorchbot label "release notes: qualcomm"

[DannyYuyang-quic]

@psiddh Hi, please note that this change is part of the Gemma4 support work in #20979.
It would be great if you could also help evaluate whether this change has any side effects or impacts on existing functionality.

Thanks!

[psiddh]

the fix is in generic exir code that runs for every backend, but the restored key has no backend users outside QNN, So iI feel its a safe fix,The only thing missing I feel is a regression test so a future refactor can't silently drop q_tensor_io again. Can you do it in a follow up PR ?

cc @JacobSzwejbka @digantdesai

[DannyYuyang-quic]

the fix is in generic exir code that runs for every backend, but the restored key has no backend users outside QNN, So iI feel its a safe fix,The only thing missing I feel is a regression test so a future refactor can't silently drop q_tensor_io again. Can you do it in a follow up PR ?

Thanks for the feedback! I’ll add a regression test in a follow-up PR.