Skip to content

[ET-VK] Fix embedding_q4gsw out-of-bounds access with dynamic shapes#18558

Merged
meta-codesync[bot] merged 1 commit intogh/SS-JIA/513/basefrom
gh/SS-JIA/513/head
Mar 30, 2026
Merged

[ET-VK] Fix embedding_q4gsw out-of-bounds access with dynamic shapes#18558
meta-codesync[bot] merged 1 commit intogh/SS-JIA/513/basefrom
gh/SS-JIA/513/head

Conversation

@SS-JIA
Copy link
Copy Markdown
Contributor

@SS-JIA SS-JIA commented Mar 29, 2026

Stack from ghstack (oldest at bottom):

The embedding_q4gsw shader used push constants for num_indices,
out_height, and embed_dim that were captured at graph build time and
never updated when input tensors were dynamically resized. This caused
out-of-bounds GPU memory reads when the actual input was smaller than
the initial allocation, resulting in VK_ERROR_DEVICE_LOST on Mali GPUs.

The fix derives all shape-dependent values (embed_dim, out_height,
num_indices) from the output tensor's sizes UBO, which is automatically
updated on resize. Only truly constant values (group_size,
is_linear_weight) remain as push constants.

Root cause: With a 7-token input on a graph built for 256 tokens, the
local workgroup rounding created an extra thread (y=7) that passed the
stale bounds check (7 >= 256 == false) and read past the 7-element
indices buffer.

Differential Revision: D98642319

The embedding_q4gsw shader used push constants for num_indices,
out_height, and embed_dim that were captured at graph build time and
never updated when input tensors were dynamically resized. This caused
out-of-bounds GPU memory reads when the actual input was smaller than
the initial allocation, resulting in VK_ERROR_DEVICE_LOST on Mali GPUs.

The fix derives all shape-dependent values (embed_dim, out_height,
num_indices) from the output tensor's sizes UBO, which is automatically
updated on resize. Only truly constant values (group_size,
is_linear_weight) remain as push constants.

Root cause: With a 7-token input on a graph built for 256 tokens, the
local workgroup rounding created an extra thread (y=7) that passed the
stale bounds check (7 >= 256 == false) and read past the 7-element
indices buffer.

Differential Revision: [D98642319](https://our.internmc.facebook.com/intern/diff/D98642319/)

[ghstack-poisoned]
SS-JIA pushed a commit that referenced this pull request Mar 29, 2026
The embedding_q4gsw shader used push constants for num_indices,
out_height, and embed_dim that were captured at graph build time and
never updated when input tensors were dynamically resized. This caused
out-of-bounds GPU memory reads when the actual input was smaller than
the initial allocation, resulting in VK_ERROR_DEVICE_LOST on Mali GPUs.

The fix derives all shape-dependent values (embed_dim, out_height,
num_indices) from the output tensor's sizes UBO, which is automatically
updated on resize. Only truly constant values (group_size,
is_linear_weight) remain as push constants.

Root cause: With a 7-token input on a graph built for 256 tokens, the
local workgroup rounding created an extra thread (y=7) that passed the
stale bounds check (7 >= 256 == false) and read past the 7-element
indices buffer.

Differential Revision: [D98642319](https://our.internmc.facebook.com/intern/diff/D98642319/)

ghstack-source-id: 359350851
Pull Request resolved: #18558
@pytorch-bot
Copy link
Copy Markdown

pytorch-bot Bot commented Mar 29, 2026

🔗 Helpful Links

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

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

❌ 1 New Failure, 1 Cancelled Job, 2 Unrelated Failures

As of commit 33dfb14 with merge base 24751f1 (image):

NEW FAILURE - The following job has failed:

CANCELLED JOB - The following job was cancelled. Please retry:

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

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

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 Mar 29, 2026
@github-actions
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.

@meta-codesync meta-codesync Bot merged commit 2ad79a8 into gh/SS-JIA/513/base Mar 30, 2026
158 of 165 checks passed
@meta-codesync meta-codesync Bot deleted the gh/SS-JIA/513/head branch March 30, 2026 16:42
@meta-codesync meta-codesync Bot temporarily deployed to cherry-pick-bot March 30, 2026 16:42 Inactive
SS-JIA pushed a commit that referenced this pull request Mar 30, 2026
The embedding_q4gsw shader used push constants for num_indices,
out_height, and embed_dim that were captured at graph build time and
never updated when input tensors were dynamically resized. This caused
out-of-bounds GPU memory reads when the actual input was smaller than
the initial allocation, resulting in VK_ERROR_DEVICE_LOST on Mali GPUs.

The fix derives all shape-dependent values (embed_dim, out_height,
num_indices) from the output tensor's sizes UBO, which is automatically
updated on resize. Only truly constant values (group_size,
is_linear_weight) remain as push constants.

Root cause: With a 7-token input on a graph built for 256 tokens, the
local workgroup rounding created an extra thread (y=7) that passed the
stale bounds check (7 >= 256 == false) and read past the 7-element
indices buffer.

Differential Revision: [D98642319](https://our.internmc.facebook.com/intern/diff/D98642319/)

ghstack-source-id: 359350851
Pull Request resolved: #18558
rascani pushed a commit to rascani/executorch that referenced this pull request Apr 1, 2026
The embedding_q4gsw shader used push constants for num_indices,
out_height, and embed_dim that were captured at graph build time and
never updated when input tensors were dynamically resized. This caused
out-of-bounds GPU memory reads when the actual input was smaller than
the initial allocation, resulting in VK_ERROR_DEVICE_LOST on Mali GPUs.

The fix derives all shape-dependent values (embed_dim, out_height,
num_indices) from the output tensor's sizes UBO, which is automatically
updated on resize. Only truly constant values (group_size,
is_linear_weight) remain as push constants.

Root cause: With a 7-token input on a graph built for 256 tokens, the
local workgroup rounding created an extra thread (y=7) that passed the
stale bounds check (7 >= 256 == false) and read past the 7-element
indices buffer.

Differential Revision: [D98642319](https://our.internmc.facebook.com/intern/diff/D98642319/)

ghstack-source-id: 359350851
Pull Request resolved: pytorch#18558
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