Uh oh!
There was an error while loading. Please reload this page.
[ExecuTorch][WebGPU] Fuse the SwiGLU elementwise chain - #21129
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21129
Note: Links to docs will display an error until the docs builds have been completed. ❌ 37 New Failures, 22 PendingAs of commit 38d5fa0 with merge base 28a7fac ( NEW FAILURES - The following jobs have failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
This PR needs a |
SS-JIA
left a comment
There was a problem hiding this comment.
Review automatically exported from Phabricator review in Meta.
fbsource master [ghstack-poisoned]
Uh oh!
There was an error while loading. Please reload this page.
Pull Request resolved: #21129 The SwiGLU feed-forward activation lowered to separate gate and up projections followed by a sigmoid and two elementwise multiplies, launching several dispatches over the same tensors. This fuses the sigmoid and the two multiplies into a single dynamically-resized 2D dispatch that computes (gate * sigmoid(gate)) * up in one pass. Fusion applies only when the pattern is exact and its scratch is planner-owned; graph outputs, extra consumers, and separate inputs fall back to the unfused chain, and the landed broad-QKV behavior is left untouched until its dedicated corrective diff. No Vulkan analogue: the Vulkan backend has no fused SiLU/SwiGLU shader and composes it from a sigmoid unary op plus an elementwise multiply at graph level. Key changes: - runtime/ops/mul/silu_mul_fused.wgsl (+ generated header): one-thread-per- element SiLU-gated multiply of gate and up into a single output. - WebGPUGraph.cpp: detect the exact gate/up/sigmoid/mul pattern, own the gate scratch, size the 2D dispatch, and fall back conservatively. ghstack-source-id: 411961441 @exported-using-ghexport Differential Revision: [D113171742](https://our.internmc.facebook.com/intern/diff/D113171742/)
Pull Request resolved: #21129 The SwiGLU feed-forward activation lowered to separate gate and up projections followed by a sigmoid and two elementwise multiplies, launching several dispatches over the same tensors. This fuses the sigmoid and the two multiplies into a single dynamically-resized 2D dispatch that computes (gate * sigmoid(gate)) * up in one pass. Fusion applies only when the pattern is exact and its scratch is planner-owned; graph outputs, extra consumers, and separate inputs fall back to the unfused chain, and the landed broad-QKV behavior is left untouched until its dedicated corrective diff. No Vulkan analogue: the Vulkan backend has no fused SiLU/SwiGLU shader and composes it from a sigmoid unary op plus an elementwise multiply at graph level. Key changes: - runtime/ops/mul/silu_mul_fused.wgsl (+ generated header): one-thread-per- element SiLU-gated multiply of gate and up into a single output. - WebGPUGraph.cpp: detect the exact gate/up/sigmoid/mul pattern, own the gate scratch, size the 2D dispatch, and fall back conservatively. ghstack-source-id: 411961441 @exported-using-ghexport Differential Revision: [D113171742](https://our.internmc.facebook.com/intern/diff/D113171742/)
Stack from ghstack (oldest at bottom):
The SwiGLU feed-forward activation lowered to separate gate and up projections
followed by a sigmoid and two elementwise multiplies, launching several
dispatches over the same tensors. This fuses the sigmoid and the two multiplies
into a single dynamically-resized 2D dispatch that computes
(gate * sigmoid(gate)) * up in one pass. Fusion applies only when the pattern
is exact and its scratch is planner-owned; graph outputs, extra consumers, and
separate inputs fall back to the unfused chain, and the landed broad-QKV
behavior is left untouched until its dedicated corrective diff. No Vulkan
analogue: the Vulkan backend has no fused SiLU/SwiGLU shader and composes it
from a sigmoid unary op plus an elementwise multiply at graph level.
Key changes:
element SiLU-gated multiply of gate and up into a single output.
scratch, size the 2D dispatch, and fall back conservatively.
@exported-using-ghexport
Differential Revision: D113171742
Differential Revision: D113171742