[JakeStevens]

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 XNNPACK: delegate bf16 softmax, computing in fp32 #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 Lower mean.dim beyond 4D pooling as an XNNPACK static mean #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 Lower mean.dim beyond 4D pooling as an XNNPACK static mean #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 Lower mean.dim beyond 4D pooling as an XNNPACK static mean #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 [XNNPACK] Lower fp32 <-> bf16 casts to XNNPACK converts #24745 lowers for programs exported with dim order, and [XNNPACK] Lower aten._to_copy dtype casts for skip_dim_order exports #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.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24746

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit fedc378 with merge base b3e2457 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[JakeStevens]

@claude review effort=max

[claude]

Claude finished @JakeStevens's task in 24m 26s —— View job

---

Nothing blocks this PR, but two things should be settled before merge: whether bf16 means carry their converts in the serialized program or get them from the runtime (as bf16 matmul does), and a bf16 mean case in test_xnnpack_partitioner.py that now passes only because of the size threshold.

PR Review: #24746

Summary

With enable_bf16=True, this PR lowers a bf16 mean.dim as Convert → fp32 XNNStaticMean → Convert in the serialized XNNPACK graph. Nothing blocks it. The main open question is a design choice: the converts are written into the program instead of being added by the runtime. Smaller issues: one stale test, gaps in bf16 test coverage, one inaccurate documented constraint and one duplicated helper.

Coverage: three independent reviewers went through all 13 files. Tests could not be run here. The XNNPACK submodule is not checked out, so XNNPACK-side behaviour (bf16 convert and transpose support) is unverified.

API Design

• The bf16 converts are fixed into every exported program, unlike bf16 matmul.
  • op_mean_dim.py:90-97 serializes Convert → fp32 XNNStaticMean → Convert.
  • For matmul, the runtime already handles the same gap ("XNNPACK has no bf16 kernel") when the program loads: defineBf16MatMulResult / convertBf16MatMulResult in XNNCompiler.cpp:821-891. There the serialized graph keeps its bf16 dtypes.
  • The PR description accepts the cost: programs keep the converts even after xnn_define_static_reduce accepts bf16. But the benefits it gives don't depend on this choice:
    • Adding the converts at load time builds the same Convert → fp32 mean → Convert subgraph, so the numerics are identical.
    • XNNStaticMean is new in Lower mean.dim beyond 4D pooling as an XNNPACK static mean #24741 and no shipped runtime can parse it. Moving bf16 handling into defineGenericReduceNode therefore costs no compatibility today. _DEFINE_REDUCE_NODE already passes graph, so the same getSerializedTensorValue dtype check works there.
  • Once this merges, every bf16 .pte keeps this encoding permanently, and the backend has two mechanisms for the same gap.
  • Suggestion: serialize a bf16 XNNStaticMean and have defineGenericReduceNode add the input and output converts when the serialized input is bf16, next to the matmul helpers. If you keep the converts in the program, say why in the visitor.
• The Global Average Pooling exception is documented as fp16/fp32 only.
  • op-support.csv:32, op_mean_dim.py:41-46 and generic_node_configs.py:710-711 now say only "fp16 or fp32" 4D [2, 3] keepdim means skip the threshold.
  • Static int8 means skip it too. Their input is a dq node whose val is fp32, so is_global_avg_pool_mean returns True. test_qs8_mean_dim lowers a 1x5x4x4 one, which is below the threshold.
  • "a mean that is not bf16" is the accurate wording. Fix this →

Testing

• A partitioner test now passes for the wrong reason.

  • test_xnnpack_partitioner.py:174-221 (test_enable_bf16_leaves_ops_xnnpack_cannot_define_in_bf16) still lists mean. Its docstring says XNNPACK refuses these ops on bf16 tensors at load.
  • With supports_bf16 = False removed from MeanDimConfig (generic_node_configs.py:704-706), the 1x4x8x8 bf16 mean stays off the delegate only because of the size threshold: it has 4 outputs and 256 < 8192 input elements.
  • If the shared x grows to 1x4x64x64, the mean is delegated and this case fails.
  • Fix: drop the mean case. The small_global_avg_pool_shape subtest of test_bf16_mean_dim_not_partitioned already covers the threshold. Fix this →

• Some new bf16 paths have no tests.

  • Channels-last input: utils.py:103 switches the input layout ChannelsLastTaggedReshapePass requires for a bf16 4D [2, 3] keepdim mean from NHWC to NCHW. A channels_last bf16 input now gets a bf16 NHWC→NCHW copy inside the delegate. No test covers this: test_bf16_mean_dim_global_avg_pool_shape passes a contiguous tensor through x + x, and test_channels_last_tagged_reshape.py has no bf16 case.
  • All dims: there is no bf16 version of test_fp32_mean_dim_all_dims. That is the rank-0 output case, where define_fp32_like copies the [1] dims that define_tensor writes (node_visitor.py:434).
  • Dynamic reduced dim: there is no bf16 version of test_fp32_mean_dim_dynamic_reduced_dim. test_bf16_mean_dim_last_dim_3d_dynamic only makes a kept dim dynamic.

  Fix this →

Performance

• The size threshold is the fp32 value, but bf16 adds two converts.
  • min_multi_row_input_numel = 8192 (generic_node_configs.py:706) was tuned in Lower mean.dim beyond 4D pooling as an XNNPACK static mean #24741 for a single fp32 reduce. The bf16 path adds two full-size converts.
  • The PR description's own numbers: a bf16 mean alone in its delegate at 1x128x128 (16384 elements, so it passes the threshold) takes 5.7x the portable kernel's time at 16 threads, against 3.2x for fp32. At 4 threads it is still 1.35x.
  • The RMSNorm sweep over model shapes doesn't regress, so this may be acceptable. If it is, the comment at lines 735-737 should say why one threshold works for both dtypes. Otherwise use a higher value for bf16.

Code Quality

• define_fp32_like duplicates an existing helper.
  • MeanDim.define_fp32_like (op_mean_dim.py:104-122) repeats UpsampleNearest2dVisitor.define_intermediate (op_upsample_nearest2d.py:127-154).
  • Both append an internal XNNTensorValue with id_out=len(xnn_graph.xvalues), XNN_INVALID_VALUE_ID and constant_buffer_idx=0. They are now the only two visitors that do this.
  • Fix: move one helper into NodeVisitor and use it from both. Fix this →
• Possible KeyError: 'val' in the channels-last pass (from Lower mean.dim beyond 4D pooling as an XNNPACK static mean #24741, low confidence).
  • requires_nhwc_input (channels_last_tagged_reshape_pass.py:318-321) used to be a set lookup. It now reads get_input_node(node, 0).meta["val"] through is_global_avg_pool_mean.
  • Nodes the pass creates have no val until the retrace at line 743. insert_copy_q_dq (line 405) moves every user of a dq onto a new dq whose meta is copied from the copy node.
  • So a static-quant graph where one dq feeds an NHWC op before a mean.dim would raise KeyError: 'val'. This is unlikely, because convert_pt2e normally gives each consumer its own dq.
  • Fix: decide from the mean node's own meta["val"]. With keepdim=True its rank and dtype match the input's.

Specific Comments

• op_mean_dim.py:81-97: the code builds XNNStaticMean, then overwrites mean.input_id / mean.output_id and reassigns xnode_unions. Choosing the ids first and building the node once is clearer.
• utils.py:95-96: the docstring describes a rule of the caller ("subjects it to MeanDimConfig's size threshold"). It should only say what the predicate returns.
• generic_node_configs.py:712-715: "Any mean other than Global Average Pooling with a static shape, more than one output and …" can be read as "other than (Global Average Pooling with a static shape)". Suggested wording: "A static-shape mean that does not lower to Global Average Pooling, with more than one output and fewer than min_multi_row_input_numel input elements, stays on the portable kernel."
• generic_node_configs.py:706 (from Lower mean.dim beyond 4D pooling as an XNNPACK static mean #24741): the check is on the number of outputs, not rows. min_multi_output_input_numel fits better than min_multi_row_input_numel.
• test_mean_dim.py:242-261: _export_and_lower / _test_bf16_mean_dim duplicate _test_mean_dim / _test_mean_dim_not_partitioned. An enable_bf16 parameter on the existing helpers would remove the second pair. _export_and_lower also has a generic name but defaults to enable_bf16=True.
• .ci/scripts/test_lora.sh:49-50 and test_lora_multimethod.sh (from Lower mean.dim beyond 4D pooling as an XNNPACK static mean #24741): "can flip between builds" is vaguer than the cause given in Lower mean.dim beyond 4D pooling as an XNNPACK static mean #24741's description: 8da4w rounding tips the first token once RMSNorm and QK-norm run in XNNPACK. Naming that cause in the comment would help the next person who edits the regex.

Recommendation

Needs Discussion

• Decide where the converts go (serialized program or runtime) before merge. Once bf16 .pte files are exported, their encoding can't change. While XNNStaticMean is unreleased, the runtime option costs no compatibility.
• Fix the stale test_xnnpack_partitioner.py mean case in this PR.
• Everything else is minor.

[JakeStevens]

Thanks for the review. Replies to the review above, as of b3d74b7:

• Converts in the program vs the runtime: moved to the runtime, as for bf16 matmul. A bf16 mean now serializes as a bf16 XNNStaticMean. When the serialized input and output are bf16, defineGenericReduceNode adds an fp32 input with a convert, reduces in fp32, and converts the result back. The fp32 tensor definition is shared with the matmul path as defineFp32Like, and convertBf16MatMulResult is now convertBf16Result. Outputs are bitwise identical to the serialized-convert version, and latency is within noise.
• Global Average Pooling wording: the docstrings now say "not bf16", which includes static int8.
• Stale test_xnnpack_partitioner.py case: mean is removed from test_enable_bf16_leaves_ops_xnnpack_cannot_define_in_bf16.
• Missing bf16 tests: added a channels_last input, all dims (keepdim true and false), and a dynamic reduced dim. The bf16 tests also check that no converts are serialized.
• Threshold for bf16: moot. Lower mean.dim beyond 4D pooling as an XNNPACK static mean #24741 dropped the cut-off, so bf16 means are lowered at any size. The description gives the bf16 RMSNorm grid without it, and the lone-mean cost.
• define_fp32_like duplication and the overwritten ids: both gone with the converts.
• Possible KeyError: 'val': fixed in Lower mean.dim beyond 4D pooling as an XNNPACK static mean #24741. is_global_avg_pool_mean reads the mean's own meta["val"].
• Specific comments:
  • The utils.py docstring only says what the predicate returns.
  • The threshold wording and the min_multi_row_input_numel name went away with the threshold.
  • The bf16 helpers are folded into _test_mean_dim, which takes bf16 from the input dtype.
  • The LoRA comment is fixed in Lower mean.dim beyond 4D pooling as an XNNPACK static mean #24741.

Authored with AI using Claude Code.