@@ -408,7 +408,6 @@ def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool:
 
 class SoftmaxConfig(GenericNodePartitionerConfig):
     target_name = "_softmax.default"
-    supports_bf16 = False
 
     def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool:
         """
 
@@ -819,35 +819,19 @@ const fb_xnnpack::XNNTensorValue* getSerializedTensorValue(
 }
 
 /*
-XNNPACK multiplies bf16 matrices only into an fp32 result. When the serialized
-graph multiplies bf16 matrices into a bf16 output, define an fp32 value for the
-op to write instead, which convertBf16MatMulResult() then converts into the
-output, so the delegate boundary stays bf16. Otherwise the op writes the output
-itself.
+Defines an internal fp32 value with the shape of a serialized value, for an op
+that XNNPACK computes only in fp32 on bf16 data. The caller converts between
+the two values with convertIfDistinct(), so the delegate boundary stays bf16.
 */
-Error defineBf16MatMulResult(
+Error defineFp32ValueLike(
     xnn_subgraph_t subgraph_ptr,
-    const fb_xnnpack::XNNGraph* graph,
-    uint32_t input1_id,
-    uint32_t input2_id,
-    uint32_t output_id,
-    uint32_t remapped_output_id,
+    const fb_xnnpack::XNNTensorValue* value,
     const NodePtr node,
-    uint32_t* result_id) noexcept {
-  *result_id = remapped_output_id;
-  const auto* input1 = getSerializedTensorValue(graph, input1_id);
-  const auto* input2 = getSerializedTensorValue(graph, input2_id);
-  const auto* output = getSerializedTensorValue(graph, output_id);
-  if (input1 == nullptr || input2 == nullptr || output == nullptr ||
-      input1->datatype() != DataType::xnn_datatype_bf16 ||
-      input2->datatype() != DataType::xnn_datatype_bf16 ||
-      output->datatype() != DataType::xnn_datatype_bf16) {
-    return Error::Ok;
-  }
+    uint32_t* fp32_id) noexcept {
   // Like defineTensor(), treat a null dims vector as rank 0.
   std::vector<size_t> dims;
-  if (output->dims() != nullptr) {
-    dims = flatbufferDimsToVector<size_t>(output->dims());
+  if (value->dims() != nullptr) {
+    dims = flatbufferDimsToVector<size_t>(value->dims());
   }
   xnn_status status = xnn_define_tensor_value(
       subgraph_ptr,
@@ -857,39 +841,63 @@ Error defineBf16MatMulResult(
       /*data=*/nullptr,
       /*external_id=*/XNN_INVALID_VALUE_ID,
       /*flags=*/0,
-      result_id);
+      fp32_id);
   ET_CHECK_OR_RETURN_ERROR(
       status == xnn_status_success,
       Internal,
-      "Failed to define the fp32 result of bf16 node %i: %s",
+      "Failed to define an fp32 value for bf16 node %i: %s",
       node->debug_handle(),
       xnn_status_to_string(status));
   return Error::Ok;
 }
 
-/*
-Converts the fp32 result that defineBf16MatMulResult() defined into the bf16
-output. Does nothing when the op wrote the output itself.
-*/
-Error convertBf16MatMulResult(
+Error convertIfDistinct(
     xnn_subgraph_t subgraph_ptr,
-    uint32_t result_id,
-    uint32_t remapped_output_id,
+    uint32_t from_id,
+    uint32_t to_id,
     const NodePtr node) noexcept {
-  if (result_id == remapped_output_id) {
+  if (from_id == to_id) {
     return Error::Ok;
   }
-  xnn_status status = xnn_define_convert(
-      subgraph_ptr, result_id, remapped_output_id, /*flags=*/0);
+  xnn_status status =
+      xnn_define_convert(subgraph_ptr, from_id, to_id, /*flags=*/0);
   ET_CHECK_OR_RETURN_ERROR(
       status == xnn_status_success,
       Internal,
-      "Failed to convert the result of bf16 node %i to bf16: %s",
+      "Failed to convert a value of bf16 node %i: %s",
       node->debug_handle(),
       xnn_status_to_string(status));
   return Error::Ok;
 }
 
+/*
+XNNPACK multiplies bf16 matrices only into an fp32 result. When the serialized
+graph multiplies bf16 matrices into a bf16 output, define an fp32 value for the
+op to write instead, which the caller converts into the output. Otherwise the
+op writes the output itself.
+*/
+Error defineBf16MatMulResult(
+    xnn_subgraph_t subgraph_ptr,
+    const fb_xnnpack::XNNGraph* graph,
+    uint32_t input1_id,
+    uint32_t input2_id,
+    uint32_t output_id,
+    uint32_t remapped_output_id,
+    const NodePtr node,
+    uint32_t* result_id) noexcept {
+  *result_id = remapped_output_id;
+  const auto* input1 = getSerializedTensorValue(graph, input1_id);
+  const auto* input2 = getSerializedTensorValue(graph, input2_id);
+  const auto* output = getSerializedTensorValue(graph, output_id);
+  if (input1 == nullptr || input2 == nullptr || output == nullptr ||
+      input1->datatype() != DataType::xnn_datatype_bf16 ||
+      input2->datatype() != DataType::xnn_datatype_bf16 ||
+      output->datatype() != DataType::xnn_datatype_bf16) {
+    return Error::Ok;
+  }
+  return defineFp32ValueLike(subgraph_ptr, output, node, result_id);
+}
+
 /*
 Define serialized linear(fully-connected) node into the subgraph using
 the remapped ids to map the serialized ids, to the new ids generated
@@ -934,9 +942,29 @@ Error defineFullyConnectedNode(
       node->debug_handle(),
       xnn_status_to_string(status));
 
-  return convertBf16MatMulResult(subgraph_ptr, fc_result, fc_output, node);
+  return convertIfDistinct(subgraph_ptr, fc_result, fc_output, node);
 };
 
+/*
+XNNPACK defines softmax only on fp32 and fp16 tensors. When a serialized
+softmax value is bf16, define an fp32 value for the softmax to use instead.
+Otherwise the softmax uses the value itself.
+*/
+Error defineBf16SoftmaxFp32Value(
+    xnn_subgraph_t subgraph_ptr,
+    const fb_xnnpack::XNNGraph* graph,
+    uint32_t serialized_id,
+    uint32_t remapped_id,
+    const NodePtr node,
+    uint32_t* fp32_id) noexcept {
+  *fp32_id = remapped_id;
+  const auto* value = getSerializedTensorValue(graph, serialized_id);
+  if (value == nullptr || value->datatype() != DataType::xnn_datatype_bf16) {
+    return Error::Ok;
+  }
+  return defineFp32ValueLike(subgraph_ptr, value, node, fp32_id);
+}
+
 /*
 Define serialized softmax node into the subgraph, using the remapped ids
 to map the serialized ids, to the new ids generated when defining
@@ -947,22 +975,40 @@ Error defineSoftmaxNode(
     const std::unordered_map<uint32_t, uint32_t>& remapped_ids,
     const NodePtr node,
     const fb_xnnpack::XNNGraph* graph) noexcept {
-  MAYBE_UNUSED(graph);
-
   auto graph_node = node->xnode_union_as_XNNSoftmax();
   REMAP_ID(remapped_ids, graph_node->input_id(), sm_input);
   REMAP_ID(remapped_ids, graph_node->output_id(), sm_output);
 
+  uint32_t softmax_input = XNN_INVALID_VALUE_ID;
+  ET_CHECK_OK_OR_RETURN_ERROR(defineBf16SoftmaxFp32Value(
+      subgraph_ptr,
+      graph,
+      graph_node->input_id(),
+      sm_input,
+      node,
+      &softmax_input));
+  uint32_t softmax_output = XNN_INVALID_VALUE_ID;
+  ET_CHECK_OK_OR_RETURN_ERROR(defineBf16SoftmaxFp32Value(
+      subgraph_ptr,
+      graph,
+      graph_node->output_id(),
+      sm_output,
+      node,
+      &softmax_output));
+
+  ET_CHECK_OK_OR_RETURN_ERROR(
+      convertIfDistinct(subgraph_ptr, sm_input, softmax_input, node));
+
   xnn_status status = xnn_define_softmax(
-      subgraph_ptr, sm_input, sm_output, graph_node->flags());
+      subgraph_ptr, softmax_input, softmax_output, graph_node->flags());
   ET_CHECK_OR_RETURN_ERROR(
       status == xnn_status_success,
       Internal,
       "Failed to create softmax node %i with code: %s",
       node->debug_handle(),
       xnn_status_to_string(status));
 
-  return Error::Ok;
+  return convertIfDistinct(subgraph_ptr, softmax_output, sm_output, node);
 }
 
 Error defineGlobalAvgPooling2dNode(
@@ -1793,7 +1839,7 @@ Error defineBatchMatrixMultiplyNode(
       node->debug_handle(),
       xnn_status_to_string(status));
 
-  return convertBf16MatMulResult(subgraph_ptr, bmm_result, bmm_out, node);
+  return convertIfDistinct(subgraph_ptr, bmm_result, bmm_out, node);
 }
 
 /*
 
@@ -7,7 +7,9 @@
 import unittest
 
 import torch
+from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner
 from executorch.backends.xnnpack.test.tester import Tester
+from executorch.backends.xnnpack.test.tester.tester import ToEdgeTransformAndLower
 
 
 class TestSoftmax(unittest.TestCase):
@@ -22,7 +24,7 @@ def __init__(self, dim):
         def forward(self, x):
             return torch.nn.Softmax(dim=self.dim)(x)
 
-    def _test_softmax(self, inputs):
+    def _test_softmax(self, inputs, enable_bf16=False):
         # Dim can be either the last dimension index or -1 (last dimension),
         # as xnnpack only supports softmax on the last dimension.
         valid_dims = [len(inputs[0]) - 1, -1]
@@ -32,19 +34,48 @@ def _test_softmax(self, inputs):
             dynamic_shape[i] = torch.export.Dim(f"dynamic_dim{i}", min=1, max=100)
         dynamic_shape = (dynamic_shape,)
 
+        # XNNPACK computes a bf16 softmax in fp32 and rounds back to bf16,
+        # so it needs the looser bf16 tolerance.
+        atol = rtol = 2e-2 if enable_bf16 else 1e-3
+
         for dim in valid_dims:
-            (
+            tester = (
                 Tester(self.Softmax(dim), inputs, dynamic_shapes=dynamic_shape)
                 .export()
                 .check_count({"torch.ops.aten.softmax": 1})
-                .to_edge_transform_and_lower()
-                .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
+            )
+            if enable_bf16:
+                tester = tester.to_edge_transform_and_lower(
+                    ToEdgeTransformAndLower(
+                        partitioners=[XnnpackPartitioner(enable_bf16=True)]
+                    )
+                )
+            else:
+                tester = tester.to_edge_transform_and_lower()
+            (
+                tester.check_count(
+                    {"torch.ops.higher_order.executorch_call_delegate": 1}
+                )
                 .check_not(["executorch_exir_dialects_edge__ops_aten__softmax_default"])
                 .to_executorch()
                 .serialize()
-                .run_method_and_compare_outputs(num_runs=5)
+                .run_method_and_compare_outputs(num_runs=5, atol=atol, rtol=rtol)
             )
 
+    def test_bf16_softmax(self):
+        inputs = (torch.rand((3, 5, 7)).to(torch.bfloat16),)
+        self._test_softmax(inputs, enable_bf16=True)
+
+    def test_bf16_softmax_not_partitioned_without_enable_bf16(self):
+        inputs = (torch.rand((3, 5, 7)).to(torch.bfloat16),)
+        (
+            Tester(self.Softmax(-1), inputs)
+            .export()
+            .check_count({"torch.ops.aten.softmax": 1})
+            .to_edge_transform_and_lower()
+            .check(["executorch_exir_dialects_edge__ops_aten__softmax_default"])
+        )
+
     def test_fp16_softmax(self):
         inputs = (torch.rand((3, 5, 7)).to(torch.float16),)
         self._test_softmax(inputs)
 
@@ -182,10 +182,6 @@ class Mean(torch.nn.Module):
             def forward(self, x):
                 return x.mean((-1, -2), keepdim=True)
 
-        class Softmax(torch.nn.Module):
-            def forward(self, x):
-                return torch.softmax(x, -1)
-
         def lower(module, x):
             return to_edge_transform_and_lower(
                 export(module, (x,)),
@@ -200,7 +196,6 @@ def lower(module, x):
             ),
             "max_pool2d": torch.nn.MaxPool2d(2),
             "mean": Mean(),
-            "softmax": Softmax(),
             "upsample_bilinear2d": torch.nn.Upsample(scale_factor=2, mode="bilinear"),
         }
         x = torch.randn(1, 4, 8, 8)
@@ -270,8 +265,9 @@ def forward(self, x):
     )
     def test_enable_bf16_runs_an_attention_block(self):
         """
-        The linear layers and both batch matrix multiplications of a bf16 attention
-        block run in XNNPACK, and the program loads and matches eager bf16.
+        The linear layers, both batch matrix multiplications and the softmax
+        of a bf16 attention block run in XNNPACK, and the program loads and
+        matches eager bf16.
         """
 
         class Attention(torch.nn.Module):
@@ -294,7 +290,7 @@ def forward(self, x):
         ops = _ops_outside_delegates(program)
         # The expand copies that matmul adds split the block into 3 delegates.
         self.assertEqual(ops.count("executorch_call_delegate"), 3)
-        self.assertIn("aten._softmax.default", ops)
+        self.assertNotIn("aten._softmax.default", ops)
         matmuls = {
             "aten.addmm.default",
             "aten.bmm.default",