[ExecuTorch][WebGPU] Add et_vk.prepack (constant-tensor packing) for E2E weight loading#20265
[ExecuTorch][WebGPU] Add et_vk.prepack (constant-tensor packing) for E2E weight loading#20265JulianCloudNTH wants to merge 1 commit into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20265
Note: Links to docs will display an error until the docs builds have been completed. ❗ 1 Active SEVsThere are 1 currently active SEVs. If your PR is affected, please view them below: ❌ 21 New Failures, 2 Unrelated FailuresAs of commit bf5c220 with merge base 5526971 ( NEW FAILURES - The following jobs have failed:
FLAKY - The following job failed but was likely due to flakiness present on trunk:
BROKEN TRUNK - The following job failed but was 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. |
This PR needs a
|
Stack from ghstack (oldest at bottom):
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
srcandoutmatch (dims,elem_size,nbytes, non-null buffers; every check throws fail-loud) and records a one-timesrc->outbuffer-to-buffer copy via the newWebGPUGraph::add_prepack_copy. The recorded copies run once in a newbuild()Phase 4 (after the op-dispatch chain is recorded), mirroring the Vulkan delegate's separateprepack()init phase (distinct from per-inferenceexecute()). Ordering is guaranteed by the WebGPU queue -- the prepack submit precedes the firstexecute()submit on the same queue, so the copied data is visible without an explicit device poll (Dawn has nowgpuDevicePoll, and the backend relies on queue ordering plus the output-map wait elsewhere).src.elem_sizeis theWebGPUTensorfield added by the embedding op lower in this stack, so prepack stacks above it.Differential Revision: D108428754