[ExecuTorch][WebGPU] et_vk.apply_rotary_emb_hf: HuggingFace rotate-half RoPE runtime op (unblocks Qwen3)#20875
[ExecuTorch][WebGPU] et_vk.apply_rotary_emb_hf: HuggingFace rotate-half RoPE runtime op (unblocks Qwen3)#20875JCNTH wants to merge 3 commits into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20875
Note: Links to docs will display an error until the docs builds have been completed. ❌ 129 New Failures, 3 Unrelated Failures, 7 Unclassified FailuresAs of commit 3a8fa4b with merge base c2b273e ( NEW FAILURES - The following jobs have failed:
UNCLASSIFIED FAILURES - DrCI could not classify the following jobs because the workflow did not run on the merge base. The failures may be pre-existing on trunk or introduced by this PR:
FLAKY - The following jobs failed but were likely due to flakiness present on trunk:
BROKEN TRUNK - The following job failed but were 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):
Problem: the WebGPU runtime registers only
et_vk.apply_rotary_emb(the interleaved/Meta RoPE convention). HuggingFace-derived models (Qwen3, etc.) export the rotate-half convention, which fuses under VulkanPartitioner intoet_vk.apply_rotary_emb_hf— an op the runtime graph builder has no handler for, soWebGPUGraph::build()throws and the delegate is rejected at load withDelegateInvalidCompatibility(et_load error 48). The whole model then fails to load on WebGPU.Solution: add the
et_vk.apply_rotary_emb_hfruntime kernel + handler as a rotate-half sibling of the interleaved op.Before: only
apply_rotary_emb(interleaved) is registered; HF-RoPE models throw at load.After: both conventions are handled; HF-RoPE models (Qwen3) load and run.
Implementation:
rotary_embedding_hf.wgsl: one thread per (i, i+half_dim) pair (rotate-half pairing vs the interleaved even/odd), reading a full[max_seq, rotary_dim]freqs table indexed at rowstart_pos + s. Scalar,wg_size64 — structural + optimization parity with the interleaved kernel (RoPE is ~1% of runtime; vec4 is neutral for this elementwise-class op on Apple's scalar ALU).RotaryEmbedding.cpp:apply_rotary_emb_hf_implmirrors the interleaved handler; it parses the extrastart_posarg as a build-time Int (baked) or a runtime SymInt (dynamic KV-cache decode) exactly asSdpa.cpphandlesinput_pos, and registers a seq resize hook (xq/xk) plus a start_pos resize hook (dynamic decode). Full rotary only (rotary_dim == head_dim); partial-rotary passthrough throws (documented follow-up; Qwen3 uses full RoPE). Mirrors Vulkanet_vk.apply_rotary_emb_hf(backends/vulkan/runtime/graph/ops/impl/RotaryEmbedding.cpp).et_vk.apply_rotary_emb_hf.default.Constraints: full rotary only for now; scalar one-thread-per-pair, kept at parity with the interleaved sibling rather than vec4 (neutral for RoPE per the closed vec4 sweep).
Co-authored-with: Claude Code.
@exported-using-ghexport
Differential Revision: D111009173
Differential Revision: D111009173