[xnnpack] Support bf16 delegation for fully-connected - #21409
Open
JakeStevens wants to merge 2 commits into
Open
[xnnpack] Support bf16 delegation for fully-connected#21409JakeStevens wants to merge 2 commits into
JakeStevens wants to merge 2 commits into
Conversation
Bumps backends/xnnpack/third-party/XNNPACK from 1adaa7c to 92a7ad5 (google/XNNPACK master, PR pytorch#10818), adding the qd8_bf16_qb4w fully-connected path to the subgraph layer. Required for bf16 dynamic quant (8da4w) to lower and run via XNNPACK. This XNNPACK revision splits two new source files out of existing translation units, so the ExecuTorch Buck build (which enumerates sources explicitly rather than via CMake) must list them or linking fails with undefined symbols: - src/subgraph/rewrites/cvt_to_fp32.cc (fp16/bf16 -> fp32 fallback rewrite; defines xnn_subgraph_fallback_from_{fp16,bf16}_to_fp32 and xnn_subgraph_alias_fp32_fallback_data) -> added to the 'subgraph' target. - src/xnnpack/init-once.c (defines xnn_init_once_impl / xnn_reset_all_init_guards; in CMake it is part of the xnnpack-hardware-config object lib) -> added to the 'XNNPACK' target next to configs/hardware-config.c. It also adds new AVX512-FP16 micro-kernels (e.g. f16-vapproxgelu) that crash clang-cl 18's codegen on Windows. clang-cl is reported to CMake as Clang (not MSVC), so XNNPACK's own MSVC gate for AVX512FP16 does not apply; disable XNNPACK_ENABLE_AVX512FP16 on Windows in the existing WIN32 block alongside the other options that don't build there.
Enable the XNNPACK delegate to partition and run bf16 fully-connected (aten.linear / aten.addmm) instead of forcing everything through fp32. XNNPACK's only bf16 fully-connected type is bf16_bf16_f32: bf16 activation x bf16 weight -> fp32 output, with an fp32 bias. Wire that up: - partition/config/xnnpack_config.py: allow torch.bfloat16 in the partitioner's valid dtypes so bf16 nodes get delegated. - operators/node_visitor.py: serialize bf16 tensors as xnn_datatype_bf16 (previously defaulted to fp32 while holding bf16 bytes -> NaN); keep bf16 constant weights in bf16; force the FC bias to fp32 in the shared input/weight/bias serializer (covers addmm). - operators/op_linear.py: force the bias to fp32 for bf16 linears. - runtime/XNNCompiler.cpp: when a fully-connected has bf16 in/weight/out, define an fp32 intermediate output and append a convert (fp32 -> bf16) so the delegate boundary stays bf16. Requires a compiler new enough for XNNPACK's AVX512BF16 microkernels (gcc >= 13 / clang >= 15). Verified: bf16 linear/addmm delegate to XNNPACK with runtime parity vs eager (max abs diff 0.008); gemma-3-1b-it bf16 export goes from 262 all-fp32 delegates to 446 delegates carrying bf16 I/O and runs end to end.
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21409
Note: Links to docs will display an error until the docs builds have been completed. ✅ You can merge normally! (1 Unrelated Failure)As of commit f6b3565 with merge base 1a770b8 ( FLAKY - The following job failed but was likely due to flakiness present on trunk:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
This PR needs a
|
meta-cla
Bot
added
the
CLA Signed
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
1 participant
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
Adds XNNPACK delegation support for bf16 fully-connected, including bf16 dynamic-quant (
8da4w). This lets bf16 models (e.g.google/gemma-3-1b-itexported with--dtype bfloat16) lower their linear layers to XNNPACK instead of falling back to portable.Changes (delegation only):
partition/config/xnnpack_config.py— allowtorch.bfloat16as a valid partitioned dtype.operators/node_visitor.py— serialize bf16 tensors asxnn_datatype_bf16.operators/op_linear.py— force the bias to fp32 for bf16 FC (XNNPACK's bf16 FC isbf16_bf16_f32: bf16 activation/weight, fp32 bias).runtime/XNNCompiler.cpp— XNNPACK only provides abf16_bf16_f32fully-connected (bf16 in → fp32 out). When the serialized graph asks for a bf16 output (fully-bf16 model), define the FC with an fp32 intermediate output and append anxnn_define_convert(fp32 → bf16) so the delegate boundary stays bf16.Dependency
The bf16 dynamic-quant path relies on the
qd8_bf16_qb4wsubgraph fully-connected added in google/XNNPACK#10818, which only exists in the bumped revision. This PR is stacked on #21400, so its diff currently shows the bump commit as well; once #21400 lands I'll rebase and this PR will reduce to just the delegation changes above.Test plan
Exported
google/gemma-3-1b-itvia optimum-executorch (XNNPACK recipe, custom SDPA + KV cache):qd8_bf16_qb4w..pteproduces finitebfloat16logits, with argmax matching the fp32 and bf16-no-quant baselines (argmax == 107 on the smoke input).