[executorch][webgpu] Test linear embedding dynamic reuse

[JCNTH]

Stack from ghstack (oldest at bottom):

• -> [executorch][webgpu] Test linear embedding dynamic reuse #21406
• [executorch][webgpu] Preserve embedding packing on resize #21405
• [executorch][vulkan] Add payload-bounded constant sharding #21404

Dynamic-shape coverage now locks linear-packed Q4 embedding reuse.

The fixture exports a discriminating linear-packed model and reuses one loaded graph across N=16,8,1,16, covering the resize transition that previously reset nibble order.

Key changes:

• test_dynamic_shape_export.py — export deterministic true-packing artifacts and torch goldens for N=16,8,1.
• test_dynamic_shape.cpp — add one-graph linear-packed reuse coverage while preserving the nonlinear control.

The explicit true-vs-false output guard prevents a non-discriminating fixture.

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

Differential Revision: D113627869

Differential Revision: D113627869

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 1 New Failure, 11 Unrelated Failures

As of commit e975971 with merge base 2439920 ():

NEW FAILURE - The following job has failed:

• Build Presets / windows (pybind) / build (gh)
  Process completed with exit code 1.

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

• pull / test-binary-size-linux-gcc / linux-job (gh) (matched linux rule in flaky-rules.json)
  The process '/usr/bin/git' failed with exit code 1
• pull / test-build-wasm-linux / linux-job (gh) (matched linux rule in flaky-rules.json)
  The process '/usr/bin/git' failed with exit code 128
• pull / test-models-linux (emformer_join, portable, linux.4xlarge.memory) / linux-job (gh) (matched linux rule in flaky-rules.json)
  The process '/usr/bin/git' failed with exit code 128
• pull / test-parakeet-xnnpack-linux / linux-job (gh) (matched linux rule in flaky-rules.json)
  The process '/usr/bin/git' failed with exit code 1
• pull / test-qnn-models-linux (mv3) / linux-job (gh) (matched linux rule in flaky-rules.json)
  The process '/usr/bin/git' failed with exit code 128
• pull / test-qnn-passes-linux / linux-job (gh) (matched linux rule in flaky-rules.json)
  The process '/usr/bin/git' failed with exit code 1
• pull / test-qnn-testsuite-linux / test-backend-linux (qnn, operators) / linux-job (gh) (matched linux rule in flaky-rules.json)
  The process '/usr/bin/git' failed with exit code 1
• pull / test-static-llama-qnn-linux (stories_110m) / linux-job (gh) (matched linux rule in flaky-rules.json)
  The process '/usr/bin/git' failed with exit code 128
• pull / test-voxtral-realtime-xnnpack-linux / linux-job (gh) (matched linux rule in flaky-rules.json)
  E: Failed to fetch https://deb.nodesource.com/node_16.x/dists/nodistro/InRelease 403 Forbidden [IP: 104.20.45.190 443]
• pull / unittest-wasm-bindings (--enable-etdump) / linux-job (gh) (matched linux rule in flaky-rules.json)
  The process '/usr/bin/git' failed with exit code 1
• pull / unittest-wasm-bindings / linux-job (gh) (matched linux rule in flaky-rules.json)
  The process '/usr/bin/git' failed with exit code 1

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

[github-actions]

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]

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

[claude]

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

---

Review: redundancy, modularization & high-value changes

• Read PR diff and surrounding context
• Analyze test_dynamic_shape.cpp changes
• Analyze test_dynamic_shape_export.py changes
• Post review feedback

Overall this is a clean, well-scoped test addition. The run_embedding parameterization is a tidy reuse (no copy-pasted reader), and the true-vs-false output guard is a nice piece of defensive fixture design — it makes a non-discriminating fixture fail loudly at export time rather than silently passing at runtime. A few observations below.

---

🟢 High-value: the import-path fix is a real correctness change

test_dynamic_shape_export.py:341 changes the import from ...ops.embedding_q4gsw.test_embedding_q4gsw to ...ops.test_embedding_q4gsw. The module actually lives at backends/webgpu/test/ops/test_embedding_q4gsw.py (there is no embedding_q4gsw/ subdir), so the base was importing from a non-existent path. Worth calling out explicitly in the PR description since it's more than a test-coverage addition — it repairs _export_dynamic_embedding entirely.

🟡 Redundancy: the golden-writing loop is near-duplicated

The new linear block (test_dynamic_shape_export.py:384-404) mirrors the existing emb_dyn block (:364-375) almost exactly — same [EMB_MAXN, 8, 1] sweep, same idx = (torch.arange(n) * 7) % EMB_VOCAB, same idx/golden tofile dance. A small helper collapses both:

def _write_emb_goldens(out_dir, prefix, ns, golden_fn):
    for n in ns:
        idx = (torch.arange(n, dtype=torch.long) * 7) % EMB_VOCAB
        golden_fn(n, idx)  # raises if non-discriminating / does the equal() guard
        idx.detach().numpy().astype("<i8").tofile(
            os.path.join(out_dir, f"{prefix}.S{n}.idx.bin"))
        golden_fn.golden(n, idx).detach().numpy().astype("<f4").tofile(
            os.path.join(out_dir, f"{prefix}.S{n}.golden.bin"))

That may be over-engineering for two call sites given the differing golden/guard logic — reasonable to leave as-is. But at minimum the repeated idx construction and the <i8/<f4 tofile pattern are the duplicated core. Your call on whether the abstraction earns its keep here. Fix this →

🟡 Minor: two C++ diagnostic messages still hardcode "emb"

run_embedding now takes prefix, and most ASSERT/EXPECT messages were updated to use it — but two were missed:

• test_dynamic_shape.cpp:211 — << "emb N=" << n << " forward failed..."
• test_dynamic_shape.cpp:215-216 — << "emb N=" << n << " output numel mismatch"

If the linear-packed test fails on forward() or numel, the message won't say which fixture (emb_dyn vs emb_dyn_linear) — undercutting the diagnostic value of the parameterization for exactly the failure modes a nibble-order regression would surface. Swap "emb N=" → prefix << " N=" in both. Fix this →

🟢 Optional: the discriminating guard could be hoisted

The torch.equal(linear_golden, nonlinear_golden) check runs per-n (:394). Since is_linear_weight flips the nibble interpretation independent of how many rows are gathered, one representative n (e.g. EMB_MAXN) is sufficient to prove the fixture discriminates. Running it per-n is cheap and arguably more thorough, so this is purely a taste call — leaving it is fine.

---

Summary: No bugs. The import fix is the most consequential change and deserves a mention in the PR description. The two suggestions worth acting on are the C++ diagnostic-message consistency (quick, real diagnostic value) and optionally trimming the duplicated golden-writing loop. This is a read-only review — no files changed.
· branch gh/JCNTH/196/head

[SS-JIA]

Review automatically exported from Phabricator review in Meta.