Repository navigation
XNNPACK: delegate bf16 softmax, computing in fp32 - #24805
Merged
Merged
Conversation
XNNPACK has no bf16 softmax kernel, so SoftmaxConfig turned bf16 off and evo1's 24 bf16 softmaxes ran on the single-threaded portable kernels. Remove that gate and, in the runtime, define fp32 values for bf16 softmax inputs and outputs and convert between them around the fp32 softmax, so the delegate boundary stays bf16. Without enable_bf16 the behavior is unchanged, and fp16 and fp32 softmaxes take the same path as before. Add a bf16 softmax test with a no-enable_bf16 negative case, and update the partitioner tests to expect the attention-block softmax inside a delegate. Related to pytorch#24667.
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24805
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 61f0cab with merge base 3f08f3f ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
Contributor
Author
|
@claude review effort:max |
|
Claude finished @JakeStevens's task in 3m 15s —— View job This PR is correct as far as reading the code can tell, and nothing blocks it: the remaining points are two duplicated C++ helpers, a misleading comment, an untested dtype claim and a runtime-version caveat. Detailed reviewPR Review: #24805SummaryThe PR lets the partitioner delegate bf16 Code Quality
Backward Compatibility
RecommendationApprove once CI is green on the two new tests and the attention-block test that weren't run locally. Deduplicating the helpers and fixing the comment are cleanups and don't need to block the merge. |
JakeStevens
added
the
release notes: xnnpack
shoumikhin
approved these changes
Nov 9, 2026
The bf16 softmax helpers duplicated the bf16 matmul ones. Factor out defineFp32ValueLike(), which defines an internal fp32 value shaped like a serialized value, and convertIfDistinct(), so the matmul and softmax paths keep only their own bf16 checks. Also drop an inaccurate comment on SoftmaxConfig. Authored with Claude Code.
JakeStevens
added a commit
to JakeStevens/executorch
that referenced
this pull request
Nov 10, 2026
## Ulterior Motive Keep bf16 RMSNorm in XNNPACK for bf16 models, such as the bf16 LeRobot policies in pytorch#24667, where bf16 ops that XNNPACK does not support fall back to the portable kernels. ## Rationale **What**: With `enable_bf16=True`, lower a bf16 `mean.dim` to the `XNNStaticMean` from pytorch#24741. The serialized graph stays bf16; when the program loads, the runtime reduces in fp32 between a convert from bf16 and a convert back, as it already does for bf16 matmul. **Why**: XNNPACK has no bf16 reduction, so `MeanDimConfig` rejected every bf16 mean, and a bf16 RMSNorm split the delegate around its mean. XNNPACK does convert between bf16 and fp32, and when the mean has lowered neighbours the two converts cost less than leaving the mean outside the delegate. ## Details - `defineGenericReduceNode` adds the converts when a serialized value is bf16, the same way `defineSoftmaxNode` does since pytorch#24805. It reuses that change's helpers: `convertIfDistinct`, and `defineBf16SoftmaxFp32Value`, renamed `defineFp32ValueForBf16` now that reductions use it too. Because the program does not encode the converts, a runtime can switch to a native bf16 reduction if XNNPACK adds one, without re-exporting. The XNNPACK subgraph is the same as with the converts serialized: outputs are bitwise identical and latency is within noise. Each bf16 mean adds an fp32 copy of its input to the XNNPACK workspace (1 MiB for 1x128x2048). - A bf16 mean never takes the Global Average Pooling path, which has no bf16 kernel. A 4D mean over dims [2, 3] with `keepdim=True` is a static mean like any other bf16 mean, so a channels-last input is copied to NCHW for it. - bf16 means are lowered at any size, as fp16 and fp32 means are in pytorch#24741. fp32 and fp16 programs are byte-identical to main. - Latency, branch time over main time (below 1 is faster), each a median over at least 1000 timed runs per program, alternating main and branch programs on one runtime, on an Apple M-series CPU (12 performance + 4 efficiency cores) with the default 16 threads: bf16 mean over the last dim of 1x128x2048 0.42x, bf16 RMSNorm 1x128x2048 0.56x, decode RMSNorm 1x1x2048 0.80x, 4 residual blocks of `x + Linear(RMSNorm(x))` at width 2048 0.87x for prefill (64 tokens) and 0.98x for decode (1 token). bf16 RMSNorm on 13 static shapes from 1x1x128 to 1x32x2048 is 0.47x to 0.96x at 1 thread. At 4 threads the shapes under 8K elements with several outputs are 1.19x to 1.57x (1.6 to 3.0 us more) and the rest 0.57x to 0.96x; at 16 threads those small shapes are 3.8x to 5.4x (22 to 24 us more) and the rest 0.70x to 0.99x, the same threadpool cost as for fp32 in pytorch#24741. - A mean alone in its delegate, with nothing else lowered around it, pays the threadpool wake-up on its own and can be slower than the portable kernel at high thread counts. At 1x128x128 and 16 threads, a lone fp32 mean is 3.2x (as in pytorch#24741) and a lone bf16 mean 5.7x, because of its two converts; at 4 threads the bf16 mean is 1.35x, and at 1 thread it is faster. Any op lowered alone pays this, such as `add` (2.95x at 1x128x128 and 16 threads on main). - Accuracy against fp32 eager is unchanged for single ops. It improves when the means were the only places a model's delegate split: for the 4 blocks above at prefill, max abs error drops from 4.6e-2 to 1.8e-2, because fewer intermediates are rounded to bf16 at delegate boundaries. - An RMSNorm that casts to fp32 and back, as Hugging Face Llama does, already has an fp32 mean. Its remaining cost is the bf16 <-> fp32 casts, which pytorch#24745 lowers for programs exported with dim order, and pytorch#24768 for programs exported with `skip_dim_order`. - `test_enable_bf16_leaves_ops_xnnpack_cannot_define_in_bf16` no longer lists `mean`, which is now delegated in bf16. Test plan: ``` python -m pytest backends/xnnpack/test/ops/test_mean_dim.py python -m pytest backends/xnnpack/test/ops backends/xnnpack/test/passes backends/xnnpack/test/test_xnnpack_partitioner.py backends/xnnpack/test/serialization backends/xnnpack/test/quantizer ``` Authored with AI using Claude Code.
JakeStevens
added a commit
that referenced
this pull request
Nov 10, 2026
## Ulterior Motive Keep bf16 RMSNorm in XNNPACK for bf16 models, such as the bf16 LeRobot policies in #24667, where bf16 ops that XNNPACK does not support fall back to the portable kernels. ## Rationale **What**: With `enable_bf16=True`, lower a bf16 `mean.dim` to the `XNNStaticMean` from #24741. The serialized graph stays bf16; when the program loads, the runtime reduces in fp32 between a convert from bf16 and a convert back, as it already does for bf16 matmul. **Why**: XNNPACK has no bf16 reduction, so `MeanDimConfig` rejected every bf16 mean, and a bf16 RMSNorm split the delegate around its mean. XNNPACK does convert between bf16 and fp32, and when the mean has lowered neighbours the two converts cost less than leaving the mean outside the delegate. ## Details - `defineGenericReduceNode` adds the converts when a serialized value is bf16, the same way `defineSoftmaxNode` does since #24805. It reuses that change's helpers: `convertIfDistinct`, and `defineBf16SoftmaxFp32Value`, renamed `defineFp32ValueForBf16` now that reductions use it too. Because the program does not encode the converts, a runtime can switch to a native bf16 reduction if XNNPACK adds one, without re-exporting. The XNNPACK subgraph is the same as with the converts serialized: outputs are bitwise identical and latency is within noise. Each bf16 mean adds an fp32 copy of its input to the XNNPACK workspace (1 MiB for 1x128x2048). - A bf16 mean never takes the Global Average Pooling path, which has no bf16 kernel. A 4D mean over dims [2, 3] with `keepdim=True` is a static mean like any other bf16 mean, so a channels-last input is copied to NCHW for it. - bf16 means are lowered at any size, as fp16 and fp32 means are in #24741. fp32 and fp16 programs are byte-identical to main. - Latency, branch time over main time (below 1 is faster), each a median over at least 1000 timed runs per program, alternating main and branch programs on one runtime, on an Apple M-series CPU (12 performance + 4 efficiency cores) with the default 16 threads: bf16 mean over the last dim of 1x128x2048 0.42x, bf16 RMSNorm 1x128x2048 0.56x, decode RMSNorm 1x1x2048 0.80x, 4 residual blocks of `x + Linear(RMSNorm(x))` at width 2048 0.87x for prefill (64 tokens) and 0.98x for decode (1 token). bf16 RMSNorm on 13 static shapes from 1x1x128 to 1x32x2048 is 0.47x to 0.96x at 1 thread. At 4 threads the shapes under 8K elements with several outputs are 1.19x to 1.57x (1.6 to 3.0 us more) and the rest 0.57x to 0.96x; at 16 threads those small shapes are 3.8x to 5.4x (22 to 24 us more) and the rest 0.70x to 0.99x, the same threadpool cost as for fp32 in #24741. - A mean alone in its delegate, with nothing else lowered around it, pays the threadpool wake-up on its own and can be slower than the portable kernel at high thread counts. At 1x128x128 and 16 threads, a lone fp32 mean is 3.2x (as in #24741) and a lone bf16 mean 5.7x, because of its two converts; at 4 threads the bf16 mean is 1.35x, and at 1 thread it is faster. Any op lowered alone pays this, such as `add` (2.95x at 1x128x128 and 16 threads on main). - Accuracy against fp32 eager is unchanged for single ops. It improves when the means were the only places a model's delegate split: for the 4 blocks above at prefill, max abs error drops from 4.6e-2 to 1.8e-2, because fewer intermediates are rounded to bf16 at delegate boundaries. - An RMSNorm that casts to fp32 and back, as Hugging Face Llama does, already has an fp32 mean. Its remaining cost is the bf16 <-> fp32 casts, which #24745 lowers for programs exported with dim order, and #24768 for programs exported with `skip_dim_order`. - `test_enable_bf16_leaves_ops_xnnpack_cannot_define_in_bf16` no longer lists `mean`, which is now delegated in bf16. Test plan: ``` python -m pytest backends/xnnpack/test/ops/test_mean_dim.py python -m pytest backends/xnnpack/test/ops backends/xnnpack/test/passes backends/xnnpack/test/test_xnnpack_partitioner.py backends/xnnpack/test/serialization backends/xnnpack/test/quantizer ``` Authored with AI using Claude Code.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Related to #24667. That issue reports that evo1's 24 bf16 softmaxes take 20.0 s of its 35.4 s chunk because their XNNPACK config turns bf16 off, leaving them on the single-threaded portable kernels. This PR addresses only that bf16 softmax item, not the other items tracked in that issue.
XNNPACK defines softmax only on fp32 and fp16 tensors, so flipping the config alone would fail at load. SoftmaxConfig no longer sets
supports_bf16 = False, and the runtime now defines fp32 values for bf16 softmax inputs and outputs and converts between them around the fp32 softmax, sharing the fp32-value and convert helpers with the bf16 matmul path, so the delegate boundary stays bf16. Each side is converted only when it is bf16. Withoutenable_bf16the behavior is unchanged, and fp16 and fp32 softmaxes take the same path as before. The bf16 softmax test and the attention-block partitioner expectation are updated accordingly.Compatibility: a
.pteexported withenable_bf16=Trueafter this change delegates bf16 softmax, and fails to load on a runtime built before it. The serialization format doesn't change, so nothing in the file flags this. The bf16 matmul change had the same constraint.This PR was authored with Muse Code.
Test plan
lintrunneron the four changed files is clean. Partitioner probes confirmed a bf16 softmax delegates withenable_bf16=Trueand stays on the portable kernels without it, and that the attention block still splits into 3 delegates with the softmax inside one.The runtime change was verified with a fresh Release
executor_runnerbuild of this branch running exported.pteprograms: a small bf16 softmax matched eager exactly, a 524288-element bf16 softmax matched to 1.2e-07, a bf16 attention block matched exactly, and fp32 and fp16 softmaxes passed. These cases were rerun after the helper refactor with the same results. On the large bf16 shape, 500 executions took about 0.13 s delegated versus about 1.54 s on the portable kernels on the same binary and inputs.The new
Tester-based bf16 softmax test and the partitioner load-and-compare test were not run locally: this machine's installed Python runtime is an older build that cannot load programs exported from this branch, and an isolated rebuild was blocked by the repository's requirement that the checkout directory be named exactlyexecutorch. Those tests need CI.