Skip to content

[ET-VK] Wire the missing resize functions for embedding and index_select - #22410

Open
msluszniak wants to merge 1 commit into
pytorch:mainfrom
msluszniak:ms/vulkan-embedding-index-select-resize
Open

[ET-VK] Wire the missing resize functions for embedding and index_select#22410
msluszniak wants to merge 1 commit into
pytorch:mainfrom
msluszniak:ms/vulkan-embedding-index-select-resize

Conversation

@msluszniak

Copy link
Copy Markdown
Contributor

Summary

add_embedding_legacy_node() and both index_select node builders pass nullptr as their resizing logic:

      // Resize Args
      {},
      // Resizing Logic
      nullptr));

even though resize_embedding_node, resize_index_select_channel_node and the local resize_fn are already defined immediately above them and do the right thing. Under dynamic shapes the outputs keep the extents they were built with (the upper bound) instead of tracking the real input sizes, so consumers read them at the wrong size.

register_index_select() also does not set supports_resize, so this patch sets it now that the op honours resize.

Fix

Pass the resize functions that already exist. Three one-line changes plus the registry flag.

How it surfaced

A TTS model on Adreno 840 produced wrong outputs after an unrelated change moved a partition boundary. The embedding output had been keeping its upper-bound extents all along; previously a CPU round-trip happened to re-establish the correct shape downstream, so the defect was masked. Once the gather stayed on the GPU the stale extents reached the consumer.

Verified on a Galaxy S26 Ultra: with these wired up, the affected sub-models return correct shapes and match their CPU references (cosine >= 0.999 at sequence lengths well below the dynamic bound) where before they were wrong at every length except the bound itself.

Note this class of bug is hard to see with executor_runner, which can only run at the dynamic upper bound -- exactly the one shape where a missing resize is harmless.

add_embedding_legacy_node() and both index_select node builders pass
nullptr as their resizing logic, even though the matching resize
functions are already defined right next to them. Under dynamic shapes
their outputs therefore keep the extents they were built with instead of
tracking the real input sizes, and downstream ops read them at the wrong
size.

Pass the resize functions that already exist, and mark index_select as
supporting resize now that it actually does.

Found while chasing a wrong-output bug in a TTS model on Adreno 840: the
embedding output kept its upper-bound extents whenever a graph change
moved the partition boundary so that a CPU round-trip no longer happened
to re-establish the correct shape.
@msluszniak
msluszniak requested a review from SS-JIA as a code owner September 1, 2026 13:45
@pytorch-bot

pytorch-bot Bot commented Sep 1, 2026

Copy link
Copy Markdown

🔗 Helpful Links

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

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

⚠️ 15 Awaiting Approval

As of commit 922bd83 with merge base 39d623e (image):

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

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

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Sep 1, 2026
@github-actions

github-actions Bot commented Sep 1, 2026

Copy link
Copy Markdown

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants