Support bf16 delegation for fully-connected (#21493)

Summary:

A reland of D113977357/#21409


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.



Key addition for the reland: an enable_bf16 flag in the xnnpack_config, which is default false. This fixes the internal failures, where the XNNPACK pin is too old.

Differential Revision: D114220918

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

@@ -50,6 +50,11 @@ def __init__(self, **kwargs):
         self.force_non_static_weights_for_f32_linear = kwargs.get(
             "force_non_static_weights_for_f32_linear", False
         )
+        # Opt-in flag for bf16 delegation (e.g. bf16 fully-connected). XNNPACK
+        # only supports bf16 fully-connected on new enough revisions, so this
+        # stays off by default and bf16 nodes fall back to the portable op
+        # unless the caller explicitly enables it.
+        self.enable_bf16 = kwargs.get("enable_bf16", False)
 
     def get_partition(
         self, node: torch.fx.Node, ep: ExportedProgram
@@ -229,6 +234,8 @@ def _check_node_has_valid_dtype(self, node):
             torch.float32,
             torch.float16,
         }
+        if self.enable_bf16:
+            valid_dtypes.add(torch.bfloat16)
         # Only allow int8 and quant dtypes for quant operations
         if is_quant(node) or is_dequant(node) or is_qparam(node):
             valid_dtypes.update(

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

@@ -46,6 +46,12 @@
     torchao_installed = False
 
 
+def is_fbcode() -> bool:
+    # torch.version.git_version is only set in OSS PyTorch; the internal
+    # Buck-built torch omits it.
+    return not hasattr(torch.version, "git_version")
+
+
 # Pytorch Modules Used for Testing
 class BaseLinear(torch.nn.Module):
     def __init__(
@@ -64,7 +70,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
 
@@ -388,6 +398,7 @@ def _test_groupwise_dq_linear(
         num_linears: int = 1,
         atol: float = 5e-3,  # TODO(T212995726): Investigate right atol for rand[n] inputs
         rtol: float = 5e-3,  # TODO(T212995726): Investigate right rtol for rand[n] inputs
+        enable_bf16: bool = False,
     ):
         """
         Helper function to test groupwise dynamic quantized linear op with different configurations.
@@ -404,6 +415,7 @@ def _test_groupwise_dq_linear(
         DynamicallyQuantizedPartitioner = XnnpackPartitioner(
             config_precisions=ConfigPrecisionType.DYNAMIC_QUANT,
             per_op_mode=True,
+            enable_bf16=enable_bf16,
         )
         tester = (
             Tester(mod, inputs)
@@ -706,11 +718,22 @@ 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
+                        lin_mod,
+                        inputs,
+                        group_size=bl,
+                        use_bias=use_bias,
+                        atol=atol,
+                        enable_bf16=dtype == torch.bfloat16,
                     )
 
     def test_fp16_linear(self):
@@ -839,6 +862,17 @@ 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"
+    )
+    @unittest.skipIf(
+        is_fbcode(),
+        "wait for XNNPACK pin bump to enable",
+    )
+    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"
     )