[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 1 commit 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. ✅ You can merge normally! (1 Unrelated Failure)As of commit 8874028 with merge base aceeb40 ( 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.
Differential Revision: D111009173