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/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..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, @@ -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, diff --git a/optimum/exporters/executorch/recipes/xnnpack.py b/optimum/exporters/executorch/recipes/xnnpack.py index ef37cc7f..82ef462b 100644 --- a/optimum/exporters/executorch/recipes/xnnpack.py +++ b/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+`, 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,