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.

⚠️ 11 Awaiting Approval

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

AWAITING APPROVAL - The following workflows need approval before CI can run:

• Build documentation (gh)
• Build Presets (gh)
• Lint (gh)
• periodic (gh)
• pull (gh)
• Test ARM Backend (gh)
• Test CoreML Backend (gh)
• Test QNN Backend (gh)
• Test Vulkan Backend (gh)
• Test WebGPU Backend (gh)
• Test XNNPACK Backend (gh)

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"