[xnnpack] Support bf16 delegation for fully-connected (pytorch#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 pytorch#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 pytorch#21400, so its diff currently
shows the bump commit as well; once pytorch#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).

Author: JakeStevens

backends/xnnpack/operators/node_visitor.py

@@ -265,12 +265,15 @@ def get_per_channel_dtype(
                     )
         else:
             node_dtype = get_node_dtype(node)
-            if node_dtype is not None and node_dtype == torch.float16:
-                dtype = (
-                    XNNDatatype.xnn_datatype_fp32
-                    if force_fp32
-                    else XNNDatatype.xnn_datatype_fp16
-                )
+            # fp16/bf16 tensors keep their datatype unless we've been asked to
+            # force fp32 (e.g. biases for dynamic-quant or bf16 fully-connected),
+            # in which case they fall back to the default fp32.
+            float_dtype_map = {
+                torch.float16: XNNDatatype.xnn_datatype_fp16,
+                torch.bfloat16: XNNDatatype.xnn_datatype_bf16,
+            }
+            if not force_fp32:
+                dtype = float_dtype_map.get(node_dtype, dtype)
 
         return dtype
 
@@ -591,7 +594,7 @@ def get_serialized_buffer_index(
         # Quantize buffer if static data is indeed quantized
         if quant_params is not None and not quant_params.is_dynamic:
             const_val = quant_params.quantize_tensor(const_val).contiguous()
-        elif const_val.dtype != torch.float16 or force_fp32:
+        elif const_val.dtype not in (torch.float16, torch.bfloat16) or force_fp32:
             # ensure that the const is fp32
             const_val = const_val.to(dtype=torch.float32).contiguous()
 
@@ -712,12 +715,19 @@ def define_nodes_tensor_inputs_outputs(
                 bias_quant_params = QuantParams.from_bias(
                     bias_node, weight_quant_params, input_quant_params
                 )
+                # XNNPACK's bf16 fully-connected (bf16_bf16_f32) takes a bf16
+                # activation/weight but an fp32 bias, so force the bias to fp32.
+                weight_val = weight_node.meta.get("val", None)
+                bias_force_fp32 = (
+                    weight_val is not None and weight_val.dtype == torch.bfloat16
+                )
                 self.define_tensor(
                     bias_node,
                     xnn_graph,
                     vals_to_ids,
                     quant_params=bias_quant_params,
                     convert_to_nhwc=False,  # Bias is generally 1d and can not be in NHWC
+                    force_fp32=bias_force_fp32,
                 )
 
     def define_node(

backends/xnnpack/operators/op_linear.py

@@ -73,6 +73,11 @@ def define_node(
             force_fp32 = False
             if input_quant_params is not None and input_quant_params.is_dynamic:
                 force_fp32 = True
+            # XNNPACK's bf16 fully-connected (bf16_bf16_f32) takes a bf16
+            # activation/weight but an fp32 bias, so force the bias to fp32.
+            weight_val = weight_node.meta.get("val", None)
+            if weight_val is not None and weight_val.dtype == torch.bfloat16:
+                force_fp32 = True
 
             self.define_tensor(
                 get_input_node(node, 2),

backends/xnnpack/partition/config/xnnpack_config.py

@@ -228,6 +228,7 @@ def _check_node_has_valid_dtype(self, node):
         valid_dtypes = {
             torch.float32,
             torch.float16,
+            torch.bfloat16,
         }
         # Only allow int8 and quant dtypes for quant operations
         if is_quant(node) or is_dequant(node) or is_qparam(node):

backends/xnnpack/runtime/XNNCompiler.cpp

@@ -791,6 +791,32 @@ Error defineConvertNode(
 
   return Error::Ok;
 };
+/*
+Look up a serialized tensor value (plain or quantized wrapper) by its
+output id. Returns nullptr if not found.
+*/
+const fb_xnnpack::XNNTensorValue* getSerializedTensorValue(
+    const fb_xnnpack::XNNGraph* graph,
+    uint32_t id) noexcept {
+  if (graph == nullptr || graph->xvalues() == nullptr) {
+    return nullptr;
+  }
+  for (auto value : *graph->xvalues()) {
+    const fb_xnnpack::XNNTensorValue* tv = nullptr;
+    if (value->xvalue_union_type() == fb_xnnpack::XValueUnion::XNNTensorValue) {
+      tv = value->xvalue_union_as_XNNTensorValue();
+    } else if (
+        value->xvalue_union_type() ==
+        fb_xnnpack::XValueUnion::XNNQuantizedTensorValue) {
+      tv = value->xvalue_union_as_XNNQuantizedTensorValue()->tensor_value();
+    }
+    if (tv != nullptr && tv->id_out() == id) {
+      return tv;
+    }
+  }
+  return nullptr;
+}
+
 /*
 Define serialized linear(fully-connected) node into the subgraph using
 the remapped ids to map the serialized ids, to the new ids generated
@@ -810,14 +836,52 @@ Error defineFullyConnectedNode(
   REMAP_ID(remapped_ids, graph_node->bias_id(), fc_bias);
   REMAP_ID(remapped_ids, graph_node->output_id(), fc_output);
 
+  // XNNPACK only provides a bf16 fully-connected of type bf16_bf16_f32:
+  // bf16 activation x bf16 weight -> fp32 output. When the serialized graph
+  // asks for a bf16 output (e.g. a fully bf16 model), define the FC with an
+  // fp32 intermediate output and append a convert (fp32 -> bf16) so the
+  // delegate boundary stays bf16.
+  const auto* in_tv = getSerializedTensorValue(graph, graph_node->input1_id());
+  const auto* filt_tv =
+      getSerializedTensorValue(graph, graph_node->filter_id());
+  const auto* out_tv = getSerializedTensorValue(graph, graph_node->output_id());
+  const bool needs_bf16_output_convert = in_tv != nullptr &&
+      filt_tv != nullptr && out_tv != nullptr &&
+      in_tv->datatype() == DataType::xnn_datatype_bf16 &&
+      filt_tv->datatype() == DataType::xnn_datatype_bf16 &&
+      out_tv->datatype() == DataType::xnn_datatype_bf16;
+
+  uint32_t fc_compute_output = fc_output;
+  if (needs_bf16_output_convert) {
+    std::vector<size_t> out_dims =
+        flatbufferDimsToVector<size_t>(out_tv->dims());
+    uint32_t intermediate_id = XNN_INVALID_VALUE_ID;
+    xnn_status ts = xnn_define_tensor_value(
+        subgraph_ptr,
+        xnn_datatype_fp32,
+        out_dims.size(),
+        out_dims.data(),
+        /*data=*/nullptr,
+        /*external_id=*/XNN_INVALID_VALUE_ID,
+        /*flags=*/0,
+        &intermediate_id);
+    ET_CHECK_OR_RETURN_ERROR(
+        ts == xnn_status_success,
+        Internal,
+        "Failed to define fp32 intermediate for bf16 linear node %i: %s",
+        node->debug_handle(),
+        xnn_status_to_string(ts));
+    fc_compute_output = intermediate_id;
+  }
+
   xnn_status status = xnn_define_fully_connected(
       subgraph_ptr,
       min_max.first,
       min_max.second,
       fc_input1,
       fc_filter,
       fc_bias,
-      fc_output,
+      fc_compute_output,
       graph_node->flags());
   ET_CHECK_OR_RETURN_ERROR(
       status == xnn_status_success,
@@ -826,6 +890,17 @@ Error defineFullyConnectedNode(
       node->debug_handle(),
       xnn_status_to_string(status));
 
+  if (needs_bf16_output_convert) {
+    xnn_status cs = xnn_define_convert(
+        subgraph_ptr, fc_compute_output, fc_output, /*flags=*/0);
+    ET_CHECK_OR_RETURN_ERROR(
+        cs == xnn_status_success,
+        Internal,
+        "Failed to define bf16 output convert for linear node %i: %s",
+        node->debug_handle(),
+        xnn_status_to_string(cs));
+  }
+
   return Error::Ok;
 };
 

backends/xnnpack/test/ops/test_linear.py

@@ -64,7 +64,11 @@ def __init__(
         self.ic = input_channels
         self.oc = output_channels
 
-        assert dtype in [torch.float, torch.half], "Unsupported op dtype"
+        assert dtype in [
+            torch.float,
+            torch.half,
+            torch.bfloat16,
+        ], "Unsupported op dtype"
         self.op_dtype = dtype
         self.in_size = in_size
 
@@ -706,9 +710,15 @@ def _test_qd8_per_token_weight_per_channel_group_int4(
                     # Mean: 0.2373046875, 0.237060546875
                     # Max: 1.0078125, 1.0078125
                     # Min: -0.08465576171875, -0.08441162109375
-                    atol = (
-                        1e-2 if dtype == torch.half else 5e-3
-                    )  # TODO(T212995726): Investigate right atol for rand[n] inputs
+                    # bf16 has ~8x coarser mantissa than fp16, so it needs a
+                    # looser atol.
+                    # TODO(T212995726): Investigate right atol for rand[n] inputs
+                    if dtype == torch.bfloat16:
+                        atol = 8e-2
+                    elif dtype == torch.half:
+                        atol = 1e-2
+                    else:
+                        atol = 5e-3
                     self._test_groupwise_dq_linear(
                         lin_mod, inputs, group_size=bl, use_bias=use_bias, atol=atol
                     )
@@ -839,6 +849,13 @@ def test_linear_qd8_f16_per_token_weight_per_channel_group_int4(self):
     def test_linear_qd8_f32_per_token_weight_per_channel_group_int4(self):
         self._test_qd8_per_token_weight_per_channel_group_int4(dtype=torch.float)
 
+    # Tests for q[dp]8-bf16-qb4w
+    @unittest.skipIf(
+        not torchao_installed, "Per Channel Group Quantization Required TorchAO"
+    )
+    def test_linear_qd8_bf16_per_token_weight_per_channel_group_int4(self):
+        self._test_qd8_per_token_weight_per_channel_group_int4(dtype=torch.bfloat16)
+
     @unittest.skipIf(
         not torchao_installed, "Per Channel Group Quantization Required TorchAO"
     )