From 65f9e4affbbe2db7f794d7feeec47eb5a552a5cc Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Mon, 27 Jul 2026 08:42:04 -0700 Subject: [PATCH 1/2] Bump vendored XNNPACK to include qd8_bf16_qb4w subgraph FC Bumps backends/xnnpack/third-party/XNNPACK from 1adaa7c to 92a7ad5 (google/XNNPACK master, PR #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. --- backends/xnnpack/cmake/Dependencies.cmake | 4 ++++ backends/xnnpack/third-party/XNNPACK | 2 +- backends/xnnpack/third-party/xnnpack.buck.bzl | 6 +++++- 3 files changed, 10 insertions(+), 2 deletions(-) diff --git a/backends/xnnpack/cmake/Dependencies.cmake b/backends/xnnpack/cmake/Dependencies.cmake index 4fb80d5c18c..d02ac4a7926 100644 --- a/backends/xnnpack/cmake/Dependencies.cmake +++ b/backends/xnnpack/cmake/Dependencies.cmake @@ -66,6 +66,10 @@ if(WIN32) set_overridable_option(XNNPACK_ENABLE_AVX256VNNI OFF) set_overridable_option(XNNPACK_ENABLE_AVX256VNNIGFNI OFF) set_overridable_option(XNNPACK_ENABLE_AVX512BF16 OFF) + # clang-cl (reported by CMake as Clang, not MSVC, so XNNPACK's own MSVC gate + # does not catch it) crashes with an internal codegen error compiling some of + # the AVX512-FP16 micro-kernels, so disable them on Windows. + set_overridable_option(XNNPACK_ENABLE_AVX512FP16 OFF) endif() set(XNNPACK_BUILD_ALL_MICROKERNELS diff --git a/backends/xnnpack/third-party/XNNPACK b/backends/xnnpack/third-party/XNNPACK index 1adaa7c709d..92a7ad501b9 160000 --- a/backends/xnnpack/third-party/XNNPACK +++ b/backends/xnnpack/third-party/XNNPACK @@ -1 +1 @@ -Subproject commit 1adaa7c709d4839d29e1f219cb962b01c9e6a905 +Subproject commit 92a7ad501b9516f9fecae119e0146dd1f58e54c1 diff --git a/backends/xnnpack/third-party/xnnpack.buck.bzl b/backends/xnnpack/third-party/xnnpack.buck.bzl index ac861435af8..3aa82f89cd4 100644 --- a/backends/xnnpack/third-party/xnnpack.buck.bzl +++ b/backends/xnnpack/third-party/xnnpack.buck.bzl @@ -70,7 +70,10 @@ def define_xnnpack(): # @lint-ignore BUCKLINT: native and fb_native are explicitly forbidden in fbcode. native.cxx_library( name = "subgraph", - srcs = SUBGRAPH_SRCS + ["XNNPACK/src/datatype.c"], + srcs = SUBGRAPH_SRCS + [ + "XNNPACK/src/datatype.c", + "XNNPACK/src/subgraph/rewrites/cvt_to_fp32.cc", + ], compiler_flags = [ "-Wno-error=missing-braces", # required since the SGX toolchain does not have this by default ], @@ -1150,6 +1153,7 @@ def define_xnnpack(): "XNNPACK/src/init.c", "XNNPACK/src/params.c", "XNNPACK/src/configs/hardware-config.c", + "XNNPACK/src/xnnpack/init-once.c", "XNNPACK/src/microparams-init.c", "XNNPACK/src/microkernel-utils.c", "XNNPACK/src/reference/binary-elementwise.cc", From f6b35650221e6c4530fdf737eabed0768146a197 Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Tue, 21 Jul 2026 11:03:27 -0700 Subject: [PATCH 2/2] [xnnpack] Support bf16 delegation for fully-connected 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. --- backends/xnnpack/operators/node_visitor.py | 15 +++- backends/xnnpack/operators/op_linear.py | 5 ++ .../partition/config/xnnpack_config.py | 1 + backends/xnnpack/runtime/XNNCompiler.cpp | 78 ++++++++++++++++++- 4 files changed, 97 insertions(+), 2 deletions(-) diff --git a/backends/xnnpack/operators/node_visitor.py b/backends/xnnpack/operators/node_visitor.py index a0f03205ed5..7032b444e6a 100644 --- a/backends/xnnpack/operators/node_visitor.py +++ b/backends/xnnpack/operators/node_visitor.py @@ -271,6 +271,12 @@ def get_per_channel_dtype( if force_fp32 else XNNDatatype.xnn_datatype_fp16 ) + elif node_dtype is not None and node_dtype == torch.bfloat16: + dtype = ( + XNNDatatype.xnn_datatype_fp32 + if force_fp32 + else XNNDatatype.xnn_datatype_bf16 + ) return dtype @@ -591,7 +597,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 +718,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( diff --git a/backends/xnnpack/operators/op_linear.py b/backends/xnnpack/operators/op_linear.py index dda1d3e53ef..31eb0d7950f 100644 --- a/backends/xnnpack/operators/op_linear.py +++ b/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), diff --git a/backends/xnnpack/partition/config/xnnpack_config.py b/backends/xnnpack/partition/config/xnnpack_config.py index 817f9d1cf50..a9ad94ff438 100644 --- a/backends/xnnpack/partition/config/xnnpack_config.py +++ b/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): diff --git a/backends/xnnpack/runtime/XNNCompiler.cpp b/backends/xnnpack/runtime/XNNCompiler.cpp index 3f2fb5a1aaa..60b4ed6dcf4 100644 --- a/backends/xnnpack/runtime/XNNCompiler.cpp +++ b/backends/xnnpack/runtime/XNNCompiler.cpp @@ -791,6 +791,33 @@ 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,6 +837,44 @@ 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 out_dims = + flatbufferDimsToVector(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, @@ -817,7 +882,7 @@ Error defineFullyConnectedNode( 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 +891,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; };