Enable bf16 export for XNNPACK backend

install_dev.py

@@ -11,12 +11,12 @@
 ]
 
 NIGHTLY_TORCH_DEPS = [
-    "executorch==1.4.0.dev20260714+cpu",
+    "executorch==1.4.0.dev20260802+cpu",
     # Keep torch aligned with the published torchvision nightly dependency.
-    "torch==2.14.0.dev20260713+cpu",
-    "torchvision==0.29.0.dev20260714+cpu",
-    "torchaudio==2.11.0.dev20260714+cpu",
-    "torchao==0.18.0.dev20260714+cpu",
+    "torch==2.14.0.dev20260801+cpu",
+    "torchvision==0.29.0.dev20260802+cpu",
+    "torchaudio==2.11.0.dev20260802+cpu",
+    "torchao==0.19.0.dev20260802+cpu",
 ]
 
 

optimum/executorch/attentions/custom_kv_cache.py

@@ -174,7 +174,7 @@ def from_legacy_cache(
             dtype = legacy_cache.k_cache.dtype
 
         # assert device is None or device == "cpu"
-        assert dtype is None or dtype == torch.float32
+        assert dtype is None or dtype in (torch.float32, torch.bfloat16, torch.float16)
 
         # Use the legacy cache's max_seq_len if max_cache_len is not specified
         if max_cache_len is None and hasattr(legacy_cache, "max_seq_len"):

optimum/executorch/attentions/custom_sdpa.py

@@ -12,12 +12,62 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
 from typing import Callable, Optional, Tuple, Union
 
 import torch
 from executorch.extension.llm.custom_ops.custom_ops import custom_sdpa  # noqa
 
 
+_HALF_DTYPES = (torch.float16, torch.bfloat16)
+
+
+def _custom_sdpa_traces_half() -> bool:
+    """
+    Whether `llama.custom_sdpa` can be traced with an f16/bf16 query.
+
+    ExecuTorch only taught the op about half dtypes in 1.4; earlier versions assert float32 in the
+    meta kernel, which surfaces as an opaque dynamo failure rather than something we can catch
+    around the call site. Probe the meta kernel once instead of comparing versions, since source
+    builds report versions such as `1.4.0a0+<sha>` that no version predicate can classify.
+    """
+    from torch._subclasses.fake_tensor import FakeTensorMode
+
+    # A failing probe is an expected outcome, so don't let the meta kernel log its traceback.
+    fake_tensor_logger = logging.getLogger("torch._subclasses.fake_tensor")
+    previously_disabled = fake_tensor_logger.disabled
+    fake_tensor_logger.disabled = True
+    try:
+        with FakeTensorMode():
+            qkv = torch.empty(1, 1, 1, 8, dtype=torch.bfloat16)
+            torch.ops.llama.custom_sdpa(
+                qkv,
+                qkv,
+                qkv,
+                start_pos=0,
+                attn_mask=None,
+                drpout_p=0.0,
+                is_causal=False,
+                scale=None,
+            )
+    except Exception:
+        return False
+    finally:
+        fake_tensor_logger.disabled = previously_disabled
+    return True
+
+
+# Evaluated at import time: the probe traces the op, so it must not run inside an export trace.
+_CUSTOM_SDPA_TRACES_HALF = _custom_sdpa_traces_half()
+
+
+def _is_tracing(tensor: torch.Tensor) -> bool:
+    """Whether we are building an exported graph rather than actually computing."""
+    # `is_compiling` is constant-folded by dynamo (strict export); non-strict export never enters
+    # dynamo but does hand the forward fake tensors.
+    return torch.compiler.is_compiling() or isinstance(tensor, torch._subclasses.FakeTensor)
+
+
 def sdpa_mask_passthrough(
     batch_size: int,
     cache_position: torch.Tensor,
@@ -81,11 +131,7 @@ def custom_sdpa_with_start_pos_forward(
     key = key.transpose(1, 2)
     value = value.transpose(1, 2)
 
-    # Convert the hell out of the inputs to fp32 and back
     input_dtype = query.dtype
-    query = query.to(torch.float32)
-    key = key.to(torch.float32)
-    value = value.to(torch.float32)
 
     # Ignore the causal flag from kwargs but use the one in module
     kwargs.pop("is_causal", None)
@@ -110,6 +156,15 @@ def custom_sdpa_with_start_pos_forward(
         else:
             start_pos = 0
 
+    # Keep half dtypes only where they are known to work: inside an exported graph, running against
+    # an ExecuTorch that supports them. Eager calls go through the AOT op library, whose f16/bf16
+    # kernel needs a temp allocator it is never given, so it silently returns garbage. Everywhere
+    # else, upcast to fp32 and cast back.
+    if input_dtype in _HALF_DTYPES and not (_CUSTOM_SDPA_TRACES_HALF and _is_tracing(query)):
+        query = query.to(torch.float32)
+        key = key.to(torch.float32)
+        value = value.to(torch.float32)
+
     output = torch.ops.llama.custom_sdpa(
         query,
         key,

optimum/exporters/executorch/recipes/xnnpack.py

@@ -12,15 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import functools
 import logging
 from typing import Dict, Union
 
+import torch
 from packaging.version import parse
 from tabulate import tabulate
 from torch import __version__ as torch_version
 from torch.export import ExportedProgram
 from torchao.utils import unwrap_tensor_subclass
 
+from executorch import version as executorch_version
 from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner
 from executorch.devtools.backend_debug import get_delegation_info
 from executorch.exir import (
@@ -41,6 +44,28 @@
 from ..recipe_registry import register_recipe
 
 
+# First ExecuTorch nightly whose XNNPACK partitioner understands `enable_bf16`.
+_MIN_ET_FOR_BF16_DELEGATION = "1.4.0.dev20260801"
+
+
+@functools.lru_cache(maxsize=1)
+def _xnnpack_honors_enable_bf16() -> bool:
+    """
+    Whether the installed XNNPACK partitioner actually acts on `enable_bf16`.
+
+    The flag is forwarded through `**kwargs` into every partitioner config, so an ExecuTorch that
+    predates it accepts `enable_bf16=True` and silently ignores it. Probe the behavior rather than
+    comparing versions: source builds report versions such as `1.4.0a0+<sha>`, which PEP 440 sorts
+    *after* every `1.4.0.devN` nightly and would pass a version check whether or not they carry the
+    feature.
+    """
+    try:
+        configs = XnnpackPartitioner(enable_bf16=True).target_partitioner_configs.values()
+    except Exception:  # Partitioner internals differ across versions; assume unsupported.
+        return False
+    return any(getattr(config, "enable_bf16", False) for config in configs)
+
+
 @register_recipe("xnnpack")
 def export_to_executorch_with_xnnpack(
     model: Union[
@@ -83,9 +108,39 @@ def _lower_to_executorch(
         if len(exported_programs) == 1:
             exported_programs = {"forward": next(iter(exported_programs.values()))}
 
+        # bf16 delegation is opt-in in the backend. A bf16 model always carries bf16 tensors
+        # (e.g. RMSNorm weights), even when its linears are quantized.
+        enable_bf16 = any(
+            getattr(tensor, "dtype", None) == torch.bfloat16
+            for exported_program in exported_programs.values()
+            for tensor in exported_program.state_dict.values()
+        )
+        if enable_bf16 and not _xnnpack_honors_enable_bf16():
+            reason = (
+                f"the installed ExecuTorch ({executorch_version.__version__}) cannot delegate bf16 to XNNPACK "
+                f"(that needs ExecuTorch >= {_MIN_ET_FOR_BF16_DELEGATION})"
+            )
+            # ExecuTorch has no portable kernels for these, so leaving them undelegated makes
+            # `to_executorch` die later with a far less obvious "Missing out variants: torchao::...".
+            has_affine_quant_ops = any(
+                node.op == "call_function" and "torchao" in str(node.target) and "affine" in str(node.target)
+                for exported_program in exported_programs.values()
+                for node in exported_program.graph.nodes
+            )
+            if has_affine_quant_ops:
+                raise RuntimeError(
+                    "Quantized linears (--qlinear) lower only through XNNPACK delegation, since torchao's "
+                    f"affine quant ops have no portable ExecuTorch kernels, but {reason}. "
+                    "Upgrade ExecuTorch, or re-export with --dtype float32 or --dtype float16."
+                )
+            logging.warning(
+                f"Exporting a bf16 model but {reason}, so every bf16 operator will fall back to the portable "
+                "kernels and inference will be slow. Upgrade ExecuTorch to delegate bf16 to XNNPACK."
+            )
+
         et_prog = to_edge_transform_and_lower(
             exported_programs,
-            partitioner=[XnnpackPartitioner()],
+            partitioner=[XnnpackPartitioner(enable_bf16=enable_bf16)],
             compile_config=EdgeCompileConfig(
                 _check_ir_validity=False,
                 _skip_dim_order=True,