Skip to content

Vulkan: fix embedding_texture for non width-packed outputs - #22827

Open
msluszniak wants to merge 1 commit into
pytorch:mainfrom
msluszniak:ms/vulkan-embedding-texture-packed-dim
Open

msluszniak wants to merge 1 commit into
pytorch:mainfrom
msluszniak:ms/vulkan-embedding-texture-packed-dim

Conversation

@msluszniak

@msluszniak msluszniak commented Sep 15, 2026

Copy link
Copy Markdown
Contributor

embedding_texture.glsl assumed the output texel always packs along the embedding dim: it looked up one index per texel, then read 4 consecutive weight elements. When the output is height packed, which the memory layout pass picks once the weight has more than 16384 rows, each texel holds 4 different rows instead, so every group of 4 output rows got the first row's data shifted by one element.

Branch on the output's packed dim and load each component from its own weight row when the packed dim is not the embedding dim.

Fixes #22333.

nn.Embedding(16385, 4), ids [[0, 1, 16384], [7, 16000, 12345]], Mali-G76, before:

row 0  -0.0766   0.3599  -0.7820   0.0715   <- correct
row 1   0.3599  -0.7820   0.0715   0.6648   <- row 0 shifted by 1
row 2  -0.7820   0.0715   0.6648  -0.2868   <- row 0 shifted by 2

After: bit-exact against eager. CLIP text encoder at force_fp16 (vocab 49408) goes from cos 0.0625 to 0.9999995. Width-packed outputs are byte-identical before and after.

cc @SS-JIA @manuelcandales @digantdesai @cbilgin

@msluszniak
msluszniak requested a review from SS-JIA as a code owner September 15, 2026 10:21
@pytorch-bot pytorch-bot Bot added the module: vulkan Issues related to the Vulkan delegate and code under backends/vulkan/ label Sep 15, 2026
@pytorch-bot

pytorch-bot Bot commented Sep 15, 2026

Copy link
Copy Markdown

🔗 Helpful Links

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

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

❌ 1 New Failure

As of commit efafdef with merge base 161fbd5 (image):

NEW FAILURE - The following job has failed:

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 15, 2026
@msluszniak

Copy link
Copy Markdown
Contributor Author

@pytorchbot label "module: vulkan"

@msluszniak

Copy link
Copy Markdown
Contributor Author

@pytorchbot label "release notes: vulkan"

@pytorch-bot pytorch-bot Bot added the release notes: vulkan Changes to the Vulkan backend delegate label Sep 15, 2026
msluszniak added a commit to software-mansion-labs/executorch that referenced this pull request Sep 15, 2026
Backport of pytorch/executorch#22827, which fixes #22333. Above 16384 weight
rows the memory layout pass picks a height-packed texture, where each texel
holds 4 different weight rows rather than 4 elements of one row, and the shader
returned the first row shifted. CLIP text at force_fp16 went cos 0.0625 ->
0.9999994.
embedding_texture.glsl assumed the output texel always packs along the
embedding dim: it looked up one index per texel and then read 4 consecutive
weight elements. When the output is height packed, which the memory layout
pass picks once the weight has more than 16384 rows, each texel holds 4
different rows instead, so every group of 4 output rows got the first row's
data shifted by one element.

Branch on the output's packed dim and load each component from its own weight
row when the packed dim is not the embedding dim.

Fixes pytorch#22333.

(cherry picked from commit 2e5b071)
@msluszniak
msluszniak force-pushed the ms/vulkan-embedding-texture-packed-dim branch from 2e5b071 to efafdef Compare September 17, 2026 15:35
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. module: vulkan Issues related to the Vulkan delegate and code under backends/vulkan/ release notes: vulkan Changes to the Vulkan backend delegate

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[ET-VK] Embedding with vocab > 16384 silently returns wrong values on the texture path; force_fp16 breaks all-MiniLM-L6-v2

3 participants