[xnnpack] Support bf16 delegation for fully-connected (#21409)
## Summary Adds XNNPACK delegation support for **bf16** fully-connected, including bf16 dynamic-quant (`8da4w`). This lets bf16 models (e.g. `google/gemma-3-1b-it` exported with `--dtype bfloat16`) lower their linear layers to XNNPACK instead of falling back to portable. Changes (delegation only): - **`partition/config/xnnpack_config.py`** — allow `torch.bfloat16` as a valid partitioned dtype. - **`operators/node_visitor.py`** — serialize bf16 tensors as `xnn_datatype_bf16`. - **`operators/op_linear.py`** — force the bias to fp32 for bf16 FC (XNNPACK's bf16 FC is `bf16_bf16_f32`: bf16 activation/weight, fp32 bias). - **`runtime/XNNCompiler.cpp`** — XNNPACK only provides a `bf16_bf16_f32` fully-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 an `xnn_define_convert` (fp32 → bf16) so the delegate boundary stays bf16. ## Dependency > **Depends on #21400** (vendored XNNPACK bump to `92a7ad5`) — **merge that first.** The bf16 dynamic-quant path relies on the `qd8_bf16_qb4w` subgraph 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-it` via optimum-executorch (XNNPACK recipe, custom SDPA + KV cache): ``` optimum-cli export executorch --model google/gemma-3-1b-it --task text-generation \ --recipe xnnpack --use_custom_sdpa --use_custom_kv_cache \ --qlinear 8da4w --qembedding 8w --dtype bfloat16 --output_dir <out> ``` - Lowering succeeds; the FC nodes delegate to XNNPACK as `qd8_bf16_qb4w`. - Runtime forward on the resulting `.pte` produces finite `bfloat16` logits, with argmax matching the fp32 and bf16-no-quant baselines (argmax == 107 on the smoke input).
J
Jacob Stevens committed
4019e4c0addc6eda8e1af6cf4c4fac5d02684720
Parent: 8a5153b
Committed by GitHub <noreply@github.com>
on 7/28/2026, 11:12:04 PM