[ExecuTorch][WebGPU] Add softmax ops (_softmax, _log_softmax) to the WebGPU backend#20919
[ExecuTorch][WebGPU] Add softmax ops (_softmax, _log_softmax) to the WebGPU backend#20919JCNTH wants to merge 3 commits into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20919
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 9e9f7da with merge base 7013c8d ( 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.
Stack from ghstack (oldest at bottom):
Add
aten._softmax+aten._log_softmaxhandlers (max-stable online softmax) for the training tail.Key changes:
runtime/ops/{softmax,log_softmax}/— max-stable softmax + log_softmax WGSL kernels + handlersCMakeLists.txtWEBGPU_SRCS— wire the sourcesReuses the shared Vulkan partitioner (
aten._softmax/_log_softmaxalready registered); WebGPU kernels only.Co-authored-with: Claude Code.
@exported-using-ghexport
Differential Revision: D111755118
Differential Revision: D111755118