[ExecuTorch][WebGPU] Add et_vk.prepack (constant-tensor packing) for E2E weight loading

[JulianCloudNTH]

Stack from ghstack (oldest at bottom):

• [ExecuTorch][WebGPU] et_vk.prepack test suite (export + native golden) #20292
• -> [ExecuTorch][WebGPU] Add et_vk.prepack (constant-tensor packing) for E2E weight loading #20265
• [ExecuTorch][WebGPU] et_vk.apply_rotary_emb test suite (export + native golden) #20290
• [ExecuTorch][WebGPU] Add et_vk.apply_rotary_emb (interleaved RoPE) + ValueList multi-output #20264
• [ExecuTorch][WebGPU] et_vk.embedding_q4gsw test suite (export + native golden) #20289
• [ExecuTorch][WebGPU] Add et_vk.embedding_q4gsw (4-bit groupwise-symmetric quantized embedding) #20263

Adds the WebGPU backend handler for et_vk.prepack.default, the node the VulkanPartitioner wraps around every constant feeding a delegated op so the constant is materialized into its dedicated GPU buffer before inference.

For the WebGPU backend's buffer-flat/fp32 model, prepack is an identity layout (same dims, dtype, and bytes), so the handler runs no compute shader: it validates that src and out match (dims, elem_size, nbytes, non-null buffers; every check throws fail-loud) and records a one-time src->out buffer-to-buffer copy via the new WebGPUGraph::add_prepack_copy. The recorded copies run once in a new build() Phase 4 (after the op-dispatch chain is recorded), mirroring the Vulkan delegate's separate prepack() init phase (distinct from per-inference execute()). Ordering is guaranteed by the WebGPU queue -- the prepack submit precedes the first execute() submit on the same queue, so the copied data is visible without an explicit device poll (Dawn has no wgpuDevicePoll, and the backend relies on queue ordering plus the output-map wait elsewhere).

src.elem_size is the WebGPUTensor field added by the embedding op lower in this stack, so prepack stacks above it.
@exported-using-ghexport

Differential Revision: D108428754

Differential Revision: D108428754

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• [ROCm] MI350 CI jobs will have longer queue times due to CI migration

❌ 3 New Failures, 3 Unrelated Failures

As of commit 69ed368 with merge base 0e65ba6 ():

NEW FAILURES - The following jobs have failed:

• pull / test-arm-backend-no-driver (test_pytest_ops_tosa) / linux-job (gh)
  RuntimeError: Command docker exec -t 427808bdbfaf14627461a0f959db9f020accbb8b64a65fb516dd0830d0d89c5d /exec failed with exit code 1
• pull / unittest / linux / linux-job (gh)
  RuntimeError: Command docker exec -t 5626fadaaf897f0cba04019fba72b8336260f7dd566f7281be66d0600a8a8f7e /exec failed with exit code 1
• pull / unittest-editable / linux / linux-job (gh)
  RuntimeError: Command docker exec -t ca3719045e70afa57e5c31f516cec3f908848357d34899ea7c8e963d50aeb0df /exec failed with exit code 1

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

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

• pull / test-llama-runner-qnn-linux (fp32, qnn_16a16w, qnn) / linux-job (gh) (trunk failure)
• pull / unittest / macos / macos-job (gh) (trunk failure)
  RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1
• pull / unittest-editable / macos / macos-job (gh) (trunk failure)
  RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script 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.

[JulianCloudNTH]

@claude review

[claude]

Claude finished @JulianCloudNTH's task in 2m 23s —— View job

---

Code Review: WebGPU et_vk.prepack handler

• Read PR diff against base
• Analyze WebGPUGraph build/prescan/materialize logic
• Analyze Prepack.cpp handler
• Check correctness, buffer usage, ordering, and memory accounting

---

Overall this is a clean implementation and the deferred-constant design is nicer than what the PR text describes. A few things worth addressing before merge.

1. PR description doesn't match the implementation

The description says prepack records a src→out buffer-to-buffer copy via WebGPUGraph::add_prepack_copy, run in a new build() Phase 4. None of that exists in the diff:

• There is no add_prepack_copy method.
• There is no Phase 4 — the copies happen inline during the Phase 3 op-dispatch loop (WebGPUGraph.cpp:531).
• It's not a buffer→buffer copy. Deferred constants get no eager src buffer; prepack_impl materializes the .pte bytes directly into the out buffer via wgpuQueueWriteBuffer (materialize_constant, WebGPUGraph.cpp:543).

The actual approach (defer + single CPU→GPU write into the consumer buffer) is better — it avoids the redundant src buffer and the extra copy. Please update the PR/commit description so reviewers and git log reflect what landed.

2. Duplicated constant-loading logic

The inline-offset / named-data-map load is now written twice — the eager path (WebGPUGraph.cpp:330-395) and materialize_constant (WebGPUGraph.cpp:543-575) implement the same offset() != UINT64_MAX vs named_key + undersized-check + Free() logic. This is ~20 lines of non-trivial duplication that can drift. Consider extracting a single helper write_constant_to(WGPUBuffer dst, const ConstantSource&) (or have the eager path build a ConstantSource and call materialize_constant), which would also unify the error messages. Fix this →

3. Confirm prepack out buffers are never memory-plan aliased

materialize_constant writes into out.buffer once at build time. If out ever has a mem_obj_id >= 0 (shared buffer), a later execute() write to a transient aliasing the same mem_obj_id would clobber the prepacked weights. This is presumably safe because the partitioner gives prepacked weights dedicated, non-planned buffers — but the invariant isn't asserted anywhere. Worth a one-line comment (or assert) in prepack_impl documenting that out must be a dedicated buffer, since the correctness of write-once depends on it.

Smaller notes

• out.buffer usage flags: the direct wgpuQueueWriteBuffer into out requires CopyDst. Both the dedicated (WebGPUGraph.cpp:319) and shared (:469) buffer paths set CopyDst, so this is fine — just flagging it as a real dependency of the deferred design.
• Prescan allocations: oc->name()->str() == "et_vk.prepack.default" (:258) builds a std::string per op per chain pass. Minor, but oc->name()->string_view()/c_str() comparison avoids the allocation in the hot prescan loop.
• Deferred-constant invariant: correctness hinges on a deferred constant (null buffer) being reachable only through prepack. The prescan enforces this via direct_use_ids + ValueList expansion (:267-277), which looks correct. Just note it doesn't account for a constant appearing directly in graph->output_ids() — almost certainly impossible for a prepack source, but it's the one path not covered by the prescan.
• memory_stats fix (:891) correctly excludes deferred (null-buffer) constants. 👍

Correctness

The defer logic, source recording, and ordering (queue-ordered writeBuffer before the first execute submit) are sound. No blocking bugs found — items 1–3 above are the main things to resolve, and #1 (description) is purely cosmetic.
· branch gh/JulianCloudNTH/27/head