From 6015e741a7077ea7f2b8dab240c9724260fb2183 Mon Sep 17 00:00:00 2001 From: Gregory Comer Date: Wed, 8 Jul 2026 16:58:49 -0700 Subject: [PATCH 1/5] Allow f16/bf16 custom_sdpa --- optimum/executorch/attentions/custom_kv_cache.py | 2 +- optimum/executorch/attentions/custom_sdpa.py | 4 ---- 2 files changed, 1 insertion(+), 5 deletions(-) diff --git a/optimum/executorch/attentions/custom_kv_cache.py b/optimum/executorch/attentions/custom_kv_cache.py index 64b7322d..515d1c7f 100644 --- a/optimum/executorch/attentions/custom_kv_cache.py +++ b/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"): diff --git a/optimum/executorch/attentions/custom_sdpa.py b/optimum/executorch/attentions/custom_sdpa.py index 0f5d0fc7..d59e418f 100644 --- a/optimum/executorch/attentions/custom_sdpa.py +++ b/optimum/executorch/attentions/custom_sdpa.py @@ -81,11 +81,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) From 09000dd5040f71228446185b5737b11719ecef74 Mon Sep 17 00:00:00 2001 From: jrstevens Date: Tue, 21 Jul 2026 11:24:54 -0700 Subject: [PATCH 2/5] Fall back to fp32 custom_sdpa when the op rejects f16/bf16 ExecuTorch's custom_sdpa only gained f16/bf16 support in 1.4; before that it asserts float32, so the previous commit's removal of the unconditional upcast breaks half-dtype exports on older runtimes. Attempt the native call and upcast only if the op actually rejects the dtype, rather than comparing versions -- source builds report versions such as 1.4.0a0+, which PEP 440 sorts after every 1.4.0.devN nightly and would therefore be misread as new enough. The recovery is deliberately narrow: it requires a half input dtype and an assertion mentioning float32 (executorch 1.3.1 raises "Expected query to be float32 but got ..."), so every unrelated assertion still propagates. --- optimum/executorch/attentions/custom_sdpa.py | 38 ++++++++++++++------ 1 file changed, 28 insertions(+), 10 deletions(-) diff --git a/optimum/executorch/attentions/custom_sdpa.py b/optimum/executorch/attentions/custom_sdpa.py index d59e418f..aa80b2b0 100644 --- a/optimum/executorch/attentions/custom_sdpa.py +++ b/optimum/executorch/attentions/custom_sdpa.py @@ -106,16 +106,34 @@ def custom_sdpa_with_start_pos_forward( else: start_pos = 0 - output = torch.ops.llama.custom_sdpa( - query, - key, - value, - start_pos=start_pos, - attn_mask=attn_mask, - drpout_p=0.0, - is_causal=is_causal, - scale=scaling, - ) + # Try the input dtype natively. ExecuTorch's custom_sdpa only gained + # f16/bf16 support in 1.4; older versions assert float32. Rather than + # sniff the version, we attempt the native call and fall back to an fp32 + # upcast if the op rejects a half dtype. + try: + output = torch.ops.llama.custom_sdpa( + query, + key, + value, + start_pos=start_pos, + attn_mask=attn_mask, + drpout_p=0.0, + is_causal=is_causal, + scale=scaling, + ) + except AssertionError as error: + if input_dtype not in (torch.float16, torch.bfloat16) or "float32" not in str(error): + raise + output = torch.ops.llama.custom_sdpa( + query.to(torch.float32), + key.to(torch.float32), + value.to(torch.float32), + start_pos=start_pos, + attn_mask=attn_mask, + drpout_p=0.0, + is_causal=is_causal, + scale=scaling, + ) return output.to(input_dtype), None From 95cf17deaf418d17ffe4e64935ba05adfa893df7 Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Thu, 30 Jul 2026 07:34:15 -0700 Subject: [PATCH 3/5] Enable bf16 XNNPACK delegation for bf16 models bf16 fully-connected delegation is opt-in in the XNNPACK backend (it requires a new enough XNNPACK), so it must be explicitly requested via XnnpackPartitioner(enable_bf16=...). Enable it only when the exported model carries bf16 tensors (i.e. exported with --dtype bfloat16; a bf16 model always has bf16 tensors such as RMSNorm weights, even when linears are quantized), leaving non-bf16 exports unchanged. --- optimum/exporters/executorch/recipes/xnnpack.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/optimum/exporters/executorch/recipes/xnnpack.py b/optimum/exporters/executorch/recipes/xnnpack.py index ef37cc7f..1b32a4a6 100644 --- a/optimum/exporters/executorch/recipes/xnnpack.py +++ b/optimum/exporters/executorch/recipes/xnnpack.py @@ -15,6 +15,7 @@ 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 @@ -83,9 +84,17 @@ 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() + ) + 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, From 36341958e210210be7cf1753be7feeb6a3a916fe Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Mon, 3 Aug 2026 06:55:03 -0700 Subject: [PATCH 4/5] Guard bf16 XNNPACK delegation and bump nightly pins XnnpackPartitioner(enable_bf16=...) only exists in executorch 1.4.0.dev20260801 and later. The pinned nightly stack predated it (dev20260714), so the flag added in the previous commit was accepted and silently ignored: CI exercised the no-op path and every bf16 matmul stayed off the delegate. Bump the nightly stack to dev20260802, which was validated end to end (gemma-3-1b-it, bf16 + 8da4w + 8w, exports and generates). Also fail loudly when the installed executorch cannot delegate bf16, since the downstream symptom is unhelpful. Detect the capability by probing the partitioner rather than comparing versions: source builds report versions such as 1.4.0a0+, and PEP 440 sorts .dev before a0, so a version check would wave them through whether or not they carry the feature. Two severities, because the failure modes differ: * with quantized linears, raise -- torchao's affine quant ops only lower inside an XNNPACK partition, so skipping bf16 linears strands them in the graph and to_executorch dies with "Missing out variants: torchao::..." far from the cause; * without them, warn -- the export still succeeds, it just leaves every bf16 operator on the portable kernels (312 vs 1094 delegated nodes on SmolLM2-135M), so raising would regress a working path. --- install_dev.py | 10 ++-- .../exporters/executorch/recipes/xnnpack.py | 46 +++++++++++++++++++ 2 files changed, 51 insertions(+), 5 deletions(-) diff --git a/install_dev.py b/install_dev.py index 6b1035fd..086183b6 100644 --- a/install_dev.py +++ b/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", ] diff --git a/optimum/exporters/executorch/recipes/xnnpack.py b/optimum/exporters/executorch/recipes/xnnpack.py index 1b32a4a6..82ef462b 100644 --- a/optimum/exporters/executorch/recipes/xnnpack.py +++ b/optimum/exporters/executorch/recipes/xnnpack.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import functools import logging from typing import Dict, Union @@ -22,6 +23,7 @@ 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 ( @@ -42,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+`, 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[ @@ -91,6 +115,28 @@ def _lower_to_executorch( 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, From 7e31fc9c54501a68012c8c852c62a69bbb6e75b4 Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Mon, 3 Aug 2026 13:02:10 -0700 Subject: [PATCH 5/5] Only keep half dtypes in custom_sdpa where the op actually works The try/except fallback around llama.custom_sdpa never fired where it mattered: - Under torch.export, dynamo wraps the meta kernel's AssertionError in a TorchRuntimeError before it reaches the call site, so exporting an f16/bf16 model against ExecuTorch < 1.4 hard-failed instead of falling back (gemma3/qwen3 float16 CI). - In eager, the AOT op library's half kernel doesn't raise at all. It logs 'No temp allocator provided' (1.4) or 'Invalid arguments' (1.3.1) and returns garbage, so bf16 eager generation produced empty output (common CI). Decide up front instead: probe the meta kernel once at import time for half support, and only skip the fp32 upcast when we are building an exported graph, which is the one place half dtypes are known to work end to end. --- optimum/executorch/attentions/custom_sdpa.py | 97 ++++++++++++++------ 1 file changed, 69 insertions(+), 28 deletions(-) diff --git a/optimum/executorch/attentions/custom_sdpa.py b/optimum/executorch/attentions/custom_sdpa.py index aa80b2b0..96e84494 100644 --- a/optimum/executorch/attentions/custom_sdpa.py +++ b/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+` 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, @@ -106,34 +156,25 @@ def custom_sdpa_with_start_pos_forward( else: start_pos = 0 - # Try the input dtype natively. ExecuTorch's custom_sdpa only gained - # f16/bf16 support in 1.4; older versions assert float32. Rather than - # sniff the version, we attempt the native call and fall back to an fp32 - # upcast if the op rejects a half dtype. - try: - output = torch.ops.llama.custom_sdpa( - query, - key, - value, - start_pos=start_pos, - attn_mask=attn_mask, - drpout_p=0.0, - is_causal=is_causal, - scale=scaling, - ) - except AssertionError as error: - if input_dtype not in (torch.float16, torch.bfloat16) or "float32" not in str(error): - raise - output = torch.ops.llama.custom_sdpa( - query.to(torch.float32), - key.to(torch.float32), - value.to(torch.float32), - start_pos=start_pos, - attn_mask=attn_mask, - drpout_p=0.0, - is_causal=is_causal, - scale=scaling, - ) + # 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, + value, + start_pos=start_pos, + attn_mask=attn_mask, + drpout_p=0.0, + is_causal=is_causal, + scale=scaling, + ) return output.to(input_dtype), None