Skip to content

[executorch][webgpu] Preserve embedding packing on resize - #21405

Open
JCNTH wants to merge 2 commits into
gh/JCNTH/195/basefrom
gh/JCNTH/195/head
Open

[executorch][webgpu] Preserve embedding packing on resize#21405
JCNTH wants to merge 2 commits into
gh/JCNTH/195/basefrom
gh/JCNTH/195/head

Conversation

@JCNTH

@JCNTH JCNTH commented Jul 27, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Stack from ghstack (oldest at bottom):

Dynamic Q4 embedding reuse now preserves the complete packing configuration.

A resize used to rebuild the uniform field by field, which dropped is_linear_weight. The resize hook now captures the fully initialized EmbeddingParams and rewrites only num_indices and total_blocks.

Key changes:

  • EmbeddingQ4gsw.cpp — pass a 32-byte base parameter struct by value through the resize callback.
  • Resize recomputation updates only dynamic counts, so current and future invariant fields are preserved by construction.
  • Mirrors Vulkan EmbeddingQ4gsw.cpp, which retains is_linear_weight in push-constant state across resize.

Static and nonlinear embedding behavior is unchanged.

Co-authored-with: Claude Code.
@exported-using-ghexport

Differential Revision: D113627870

Differential Revision: D113627870

@JCNTH
Update
993edd7 
[ghstack-poisoned]
@pytorch-bot

pytorch-bot Bot commented Jul 27, 2026 •
edited
Loading

Copy link
Copy Markdown

🔗 Helpful Links

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

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

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

As of commit 4a9e98f with merge base 2439920 (image):

NEW FAILURE - The following job has failed:

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

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

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

This was referenced Jul 27, 2026
Open
Open
@github-actions

github-actions Bot commented Jul 27, 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.

@JCNTH

JCNTH commented Jul 27, 2026

Copy link
Copy Markdown
Contributor Author

@claude Thoroughly review for redundancy or areas of modularization or any high value changes

@claude

claude Bot commented Jul 27, 2026 •
edited
Loading

Copy link
Copy Markdown

Claude finished @JCNTH's task in 1m 2s —— View job

Review: Preserve embedding packing on resize

The fix is correct and minimal — is_linear is now threaded through the resize hook's capture list and written into p.is_linear_weight, so a dynamic-shape resize no longer silently drops back to the non-linear (zero) packing. It matches the stated Vulkan behavior. ✅

That said, since you asked specifically about redundancy / modularization, this diff is a good illustration of a latent structural issue rather than a problem with the change itself.

High-value observation: the resize path re-derives params instead of reusing them

EmbeddingParams is populated in two places with the same graph-invariant fields:

  • Initial construction — EmbeddingQ4gsw.cpp:181-189
  • Resize rebuild — EmbeddingQ4gsw.cpp:68-76

On resize, only two fields are actually dynamic: num_indices and total_blocks (plus the dispatch count and output dims). Every other field — embed_dim, blocks_per_row, group_size, groups_per_row, bytes_per_row, and now is_linear_weight — is graph-invariant and gets recomputed identically.

This duplication is precisely the bug this PR fixes. is_linear_weight was forgotten because the resize path re-assembles the struct from a parallel argument list rather than reusing the one already built. The +4 here patches the symptom; the same class of bug will recur the next time a field is added to EmbeddingParams.

A more robust shape would be to capture the fully-populated base EmbeddingParams by value and mutate only the dynamic fields on resize:

void resize_embedding_q4gsw(
    WebGPUGraph& g,
    int indices_id,
    int out_id,
    EmbeddingParams base,        // graph-invariant fields already set
    uint32_t wg_size,
    size_t dispatch_idx,
    WGPUBuffer params_buf) {
  const auto& id = g.cur_dims(indices_id);
  const uint64_t ni = utils::numel_of(id);
  if (ni == 0) {
    throw std::runtime_error("WebGPU embedding_q4gsw: zero indices");
  }
  const uint64_t total_blocks = ni * base.blocks_per_row;
  if (total_blocks > UINT32_MAX) {
    throw std::runtime_error("WebGPU embedding_q4gsw: total_blocks exceeds uint32");
  }
  std::vector<int64_t> od = id;
  od.push_back(static_cast<int64_t>(base.embed_dim));
  g.set_cur_dims(out_id, od);

  base.num_indices = static_cast<uint32_t>(ni);
  base.total_blocks = static_cast<uint32_t>(total_blocks);
  wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &base, sizeof(base));
  g.dispatch_at(dispatch_idx).workgroup_count_x =
      utils::compute_1d_workgroup_count(
          g.device(), static_cast<uint32_t>(total_blocks), wg_size,
          "embedding_q4gsw(resize)");
}

Benefits:

  • The 12-parameter signature collapses to ~6, and the lambda capture list at EmbeddingQ4gsw.cpp:239-249 shrinks from 10 captures to ~5.
  • is_linear_weight (and any future field) is preserved by construction — no per-field threading, no recurrence of this bug.
  • The total_blocks > UINT32_MAX guard (currently duplicated at :61-64 and :170-173) and the two struct-population blocks stop drifting.

Fix this →

Minor notes

  • num_indices recompute divergence. The initial path computes num_indices from out_numel / embed_dim (:142), whereas the resize path uses numel_of(indices_dims) (:56). These should agree, but they're two different derivations — capturing the base struct wouldn't fix this (it's genuinely dynamic), just worth confirming the two definitions are equivalent for all supported index shapes.
  • The comment at :184 (std140 layout only; shader derives it) says num_indices is padding the shader ignores. If that's truly the case, the resize path writing it is harmless; if the shader does read it in some path, the two derivations above matter.
  • No test change accompanies this fix. If there's a dynamic-shape embedding test fixture, a case that resizes a is_linear=true embedding and asserts output equality would lock in the regression guard cheaply.

Net: the diff is correct — approve on merit. The refactor above is the higher-value follow-up and would retire this bug class rather than one instance of it.

--- · branch gh/JCNTH/195/head

@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 Jul 27, 2026
@JCNTH
Update
4a9e98f 
[ghstack-poisoned]
SS-JIA
SS-JIA approved these changes Jul 27, 2026
View reviewed changes

@SS-JIA SS-JIA left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review automatically exported from Phabricator review in Meta.

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

Reviewers

@SS-JIA SS-JIA SS-JIA approved these changes

Assignees

No one assigned

Labels

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

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@JCNTH @SS-JIA