Repository navigation
[XNNPACK] Lower fp32 <-> bf16 casts to XNNPACK converts - #24745
Merged
Merged
Conversation
ToDimOrderCopyConfig already lowers fp32 <-> fp16 casts to XNNPACK convert nodes; allow fp32 <-> bf16 too (when bf16 is enabled in the partitioner). fp16 <-> bf16 stays unsupported. This keeps bf16 models that upcast a step to fp32 (softmax, RMSNorm) in one delegate instead of running the casts in portable kernels between partitions. executor_runner, single thread pinned to the fastest core, before -> after: upcast softmax [1,16,512,512] 5.5-8.0x, Llama RMSNorm [1,128,2048] 2.4-5.2x, elementwise upcast [1,128,4096] 5.4-18x (Apple M4, iPhone 17, Pixel 10, Pixel 6). Test plan: python -m pytest backends/xnnpack/test/ops/test_to_copy_dtype.py backends/xnnpack/test/ops/test_to_copy.py
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24745
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit b87ab61 with merge base a0d6be8 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
This PR needs a
|
Contributor
Author
|
@claude review effort=max |
|
Claude finished @JakeStevens's task in 19m 15s —— View job This PR needs a few small fixes: the two new bf16 runtime tests are missing the PR Review: #24745SummaryThis PR lets
What I couldn't check:
Infrastructure
Testing
Performance
Specific Comments
RecommendationRequest Changes The fbcode skips are what blocks this PR, because the internal import will probably fail without them. The exact-tolerance cast test is cheap and covers the one numerical change this PR makes. The |
Skip the two bf16 runtime tests in fbcode like the other XNNPACK bf16 runtime tests, add an exact-tolerance fp32 -> bf16 cast test with values halfway between two bf16 numbers, and add the bf16 softmax upcast as a single-delegate test. ToDimOrderCopyConfig's docstring now covers dtype casts, its comment says why fp16 <-> bf16 is excluded and that bf16 needs enable_bf16, and a dim order change gets its own why() reason. Authored with AI using Claude Code.
Contributor
Author
|
@claude review effort=max |
|
Claude finished @JakeStevens's task in 22m 32s —— View job Approve: I found no problems in the partitioner change, and what's left is minor: two bf16 tests could check their outputs more tightly, the rounding test is missing some inputs, and the PR body should disclose AI authorship. PR Review: #24745SummaryThis PR lets
What I couldn't check:
Code Quality
Testing
Specific Comments
RecommendationApprove Nothing blocks merging. Both test fixes are cheap, and they would make the new tests catch the numerical differences their names describe. |
shoumikhin
approved these changes
Nov 6, 2026
JakeStevens
added a commit
to JakeStevens/executorch
that referenced
this pull request
Nov 9, 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` as a convert to fp32, the fp32 `XNNStaticMean` from pytorch#24741, and a convert back to bf16, all inside the delegate. **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 Stacked on pytorch#24741, which adds `XNNStaticMean`; only the last commit is new. - The converts are written into the program rather than added by the runtime at load time, as bf16 matmul results are. That keeps the change in Python and matches eager bf16 exactly, but a program exported now keeps the converts if XNNPACK's `xnn_define_static_reduce` later accepts bf16, which is the only thing that rejects a bf16 mean today. 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` goes through the same converts as any other bf16 mean. - bf16 means are lowered at any size, as fp16 and fp32 means are in pytorch#24741. fp32 and fp16 programs are byte-identical to pytorch#24741. - 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 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
to JakeStevens/executorch
that referenced
this pull request
Nov 9, 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` as a convert to fp32, the fp32 `XNNStaticMean` from pytorch#24741, and a convert back to bf16, all inside the delegate. **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 Stacked on pytorch#24741, which adds `XNNStaticMean`; only the last commit is new. - The converts are written into the program rather than added by the runtime at load time, as bf16 matmul results are. That keeps the change in Python and matches eager bf16 exactly, but a program exported now keeps the converts if XNNPACK's `xnn_define_static_reduce` later accepts bf16, which is the only thing that rejects a bf16 mean today. 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` goes through the same converts as any other bf16 mean. - bf16 means are lowered at any size, as fp16 and fp32 means are in pytorch#24741. fp32 and fp16 programs are byte-identical to pytorch#24741. - 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
to JakeStevens/executorch
that referenced
this pull request
Nov 9, 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` as a convert to fp32, the fp32 `XNNStaticMean` from pytorch#24741, and a convert back to bf16, all inside the delegate. **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 Stacked on pytorch#24741, which adds `XNNStaticMean`; only the last commit is new. - The converts are written into the program rather than added by the runtime at load time, as bf16 matmul results are. That keeps the change in Python and matches eager bf16 exactly, but a program exported now keeps the converts if XNNPACK's `xnn_define_static_reduce` later accepts bf16, which is the only thing that rejects a bf16 mean today. 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` goes through the same converts as any other bf16 mean. - bf16 means are lowered at any size, as fp16 and fp32 means are in pytorch#24741. fp32 and fp16 programs are byte-identical to pytorch#24741. - 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
to JakeStevens/executorch
that referenced
this pull request
Nov 9, 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 Stacked on pytorch#24741, which adds `XNNStaticMean`; only the last commit is new. - `defineGenericReduceNode` adds the converts when the serialized input and output are bf16, next to `defineBf16MatMulResult`. The fp32 tensor definition is shared with the matmul path as `defineFp32Like`, and `convertBf16MatMulResult` becomes `convertBf16Result`. 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 pytorch#24741. - 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
to JakeStevens/executorch
that referenced
this pull request
Nov 9, 2026
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. **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. Stacked on pytorch#24741, which adds `XNNStaticMean`; only the last commit is new. - `defineGenericReduceNode` adds the converts when the serialized input and output are bf16, next to `defineBf16MatMulResult`. The fp32 tensor definition is shared with the matmul path as `defineFp32Like`, and `convertBf16MatMulResult` becomes `convertBf16Result`. 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 pytorch#24741. - 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
to JakeStevens/executorch
that referenced
this pull request
Nov 9, 2026
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. **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. Stacked on pytorch#24741, which adds `XNNStaticMean`; only the last commit is new. - `defineGenericReduceNode` adds the converts when the serialized input and output are bf16, next to `defineBf16MatMulResult`. The fp32 tensor definition is shared with the matmul path as `defineFp32Like`, and `convertBf16MatMulResult` becomes `convertBf16Result`. 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 pytorch#24741. - 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
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.
Ulterior Motive
Improving overall BF16 support in light of #24667
Summary
ToDimOrderCopyConfigalready lowers fp32 <-> fp16 casts to XNNPACK convert nodes. This allows fp32 <-> bf16 as well, when bf16 is enabled in the partitioner. fp16 <-> bf16 stays unsupported, because XNNPACK has no direct convert between them. bf16 models that upcast a step to fp32 (HF softmax withdtype=torch.float32, Llama-style RMSNorm) now stay in one delegate, instead of running the casts in portable kernels between partitions. ExecuTorch's current XNNPACK pin already has the bf16 <-> fp32 convert kernels, so no runtime or pin changes are needed.The config matches
_to_dim_order_copy, which export emits when dim order is kept, theto_edge_transform_and_lowerdefault. Programs exported withskip_dim_order=True, asLLMEdgeManagerandexport_llamado by default, keepaten._to_copy, which no XNNPACK partitioner config matches yet, so their casts (fp16 ones included) still run in portable kernels. Partitioning_to_copydtype casts under the same rules would be a follow-up.Benchmarks
Measured with
executor_runner(Release; XNNPACK + portable ops), median of 50 iterations. "Before" exports the same model with the cast lowering disabled. Shapes:softmax(x, -1, dtype=float32).to(bf16)on[1,16,S,S].[1,T,2048].(x.float()*2+1).to(bf16)on[1,128,4096].Single thread (Android pinned to the fastest core with
taskset), ms, before -> after:Default
executor_runnerthreads (M4: 16, iPhone: 6, Pixel 10: 6, Pixel 6: 4; unpinned), ms, before -> after:* Unpinned, this ~10 us graph runs on a core that doesn't ramp its clock up, so the run hits a ~0.1 ms floor even with 1 thread. Pinned, it's 0.012 ms (table above).
RMSNorm still has
meanoutside the delegate (mean.dimover the last dim isn't lowered yet; #24741 lowers it). The casts no longer split it further (3 -> 2 delegates).Test plan
New tests cover:
enable_bf16;The tests that run bf16 in XNNPACK are skipped in fbcode, like the other XNNPACK bf16 runtime tests. The delegation tests fail without the config change.