Fix QNN E2E evaluation failures

docs/reference/index.md

@@ -55,6 +55,7 @@ stages based on the target device and precision.
 | `do_constant_folding` | `bool` | `true` | Fold constants during export. |
 | `verbose` | `bool` | `false` | Verbose export logging. |
 | `dynamo` | `bool` | `false` | Use PyTorch's TorchDynamo ONNX exporter; the default `false` selects the legacy TorchScript exporter. |
+| `compatibility` | `dict \| null` | omitted | Resolved export compatibility knobs. Currently supports `transformers_attention: "eager"` to request eager Transformers attention during ONNX export. |
 | `enable_hierarchy_tags` | `bool` | `true` | Add module hierarchy tags to ONNX nodes. |
 | `clean_onnx` | `bool` | `false` | Strip hierarchy tags after export. |
 | `hierarchy_tag_format` | `"full" \| "module_only"` | `"full"` | Tag detail level. |

pyproject.toml

@@ -152,6 +152,19 @@ override-dependencies = [
   "onnxruntime ; sys_platform == 'unobtainium'",
 ]
 
+[[tool.uv.index]]
+name = "pytorch-cpu"
+url = "https://download.pytorch.org/whl/cpu"
+explicit = true
+
+[tool.uv.sources]
+torch = [
+  { index = "pytorch-cpu", marker = "sys_platform == 'linux' and platform_machine == 'x86_64'" },
+]
+torchvision = [
+  { index = "pytorch-cpu", marker = "sys_platform == 'linux' and platform_machine == 'x86_64'" },
+]
+
 [dependency-groups]
 dev = [
   "jupyter>=1.1.1",
@@ -186,6 +199,7 @@ include = [ "winml", "winml.*" ]
   "runtime_checker/need_rerun_errors.json",
 ]
 "winml.modelkit.export" = [
+  "compatibility_rules.json",
   "htp/htp_metadata_schema.json",
 ]
 "winml.modelkit.pattern" = [

src/winml/modelkit/commands/build.py

@@ -901,7 +901,6 @@ def build(
         "EPNameOrAlias | None",
         _reject_ep_source(ep, "winml build"),
     )
-
     # Validate mutual exclusion
     if output_dir and use_cache:
         raise click.UsageError("--output-dir and --use-cache are mutually exclusive.")
@@ -933,20 +932,6 @@ def build(
         dynamic_axes=dynamic_axes,
     )
 
-    # Resolve an omitted EP and the requested device as one target. Forwarding
-    # both concrete axes keeps analyzer/build output aligned with the target
-    # selected by the catalog-backed resolver.
-    if ep_value is None:
-        from ..session import EPDeviceTarget, resolve_device
-
-        try:
-            resolved_target = resolve_device(EPDeviceTarget(ep="auto", device=device))
-        except ValueError as e:
-            raise click.UsageError(str(e)) from e
-        device = resolved_target.device
-        ep_value = cast("EPNameOrAlias", resolved_target.ep)
-        logger.info("Auto-resolved device=%s, EP=%s", device, ep_value)
-
     try:
         # Hub-hosted ONNX (e.g. ``onnx-community/sam3-tracker-ONNX/onnx/...``)
         # is downloaded once and treated as a local .onnx file thereafter.
@@ -989,6 +974,9 @@ def build(
                     ]
                 else:
                     config_or_configs = merge_export_overrides(config_or_configs, export_overrides)
+            from ..config.build import apply_export_compatibility_policy
+
+            apply_export_compatibility_policy(config_or_configs, device=device, ep=ep_value)
         else:
             if not model:
                 raise click.UsageError("-m/--model is required when -c is not provided.")
@@ -1036,6 +1024,19 @@ def build(
             if no_compile:
                 config_or_configs.compile = None
 
+        runtime_device = device
+        runtime_ep_value = ep_value
+        if runtime_ep_value is None:
+            from ..session import EPDeviceTarget, resolve_device
+
+            try:
+                resolved_target = resolve_device(EPDeviceTarget(ep="auto", device=runtime_device))
+            except ValueError as e:
+                raise click.UsageError(str(e)) from e
+            runtime_device = resolved_target.device
+            runtime_ep_value = cast("EPNameOrAlias", resolved_target.ep)
+            logger.info("Auto-resolved device=%s, EP=%s", runtime_device, runtime_ep_value)
+
         # If --device or --precision was explicitly provided, patch quant/compile
         # to honor the requested policy. fp16/fp32 clear quant; npu/int8 etc set it.
         if cli_utils.is_cli_provided(ctx, "device") or cli_utils.is_cli_provided(ctx, "precision"):
@@ -1045,7 +1046,7 @@ def _patch_device(cfg: WinMLBuildConfig) -> None:
                 from ..config import resolve_quant_compile_config
 
                 resolved_quant, _ = resolve_quant_compile_config(
-                    device=device, precision=precision, ep=ep_value
+                    device=runtime_device, precision=precision, ep=runtime_ep_value
                 )
                 if cfg.skip_optimize or not quant or resolved_quant is None:
                     cfg.quant = None
@@ -1073,7 +1074,7 @@ def _patch_device(cfg: WinMLBuildConfig) -> None:
                     cfg.precision = precision.lower()  # type: ignore[attr-defined]
                 if cfg.compile is not None and cfg.compile.ep_config is not None:
                     provider = cfg.compile.ep_config.provider
-                    patched = WinMLCompileConfig.for_provider(provider, device=device)
+                    patched = WinMLCompileConfig.for_provider(provider, device=runtime_device)
                     if patched is not None:
                         cfg.compile = patched
 
@@ -1142,8 +1143,8 @@ def _patch_device(cfg: WinMLBuildConfig) -> None:
             preloaded_hf_config=preloaded_hf_config,
             output_dir=output_dir,
             use_cache=use_cache,
-            device=device,
-            ep=ep_value,
+            device=runtime_device,
+            ep=runtime_ep_value,
             precision=precision,
             rebuild=rebuild,
             submodel=submodel,
@@ -1190,8 +1191,8 @@ def _patch_device(cfg: WinMLBuildConfig) -> None:
                 configs=configs,
                 output_dir=resolved_dir,
                 rebuild=rebuild,
-                ep=ep_value,
-                device=device,
+                ep=runtime_ep_value,
+                device=runtime_device,
                 allow_unsupported_nodes=allow_unsupported_nodes,
             )
 
@@ -1399,8 +1400,8 @@ def _patch_device(cfg: WinMLBuildConfig) -> None:
                             resolved_dir=resolved_dir,
                             rebuild=rebuild,
                             cache_key=name,
-                            ep=ep_value,
-                            device=device,
+                            ep=runtime_ep_value,
+                            device=runtime_device,
                             extra_kwargs=dict(extra_kwargs),
                             preloaded_hf_config=preloaded_hf_config,
                         )
@@ -1417,8 +1418,8 @@ def _patch_device(cfg: WinMLBuildConfig) -> None:
                     resolved_dir=resolved_dir,
                     rebuild=rebuild,
                     cache_key=cache_key,
-                    ep=ep_value,
-                    device=device,
+                    ep=runtime_ep_value,
+                    device=runtime_device,
                     extra_kwargs=extra_kwargs,
                     preloaded_hf_config=preloaded_hf_config,
                 )

src/winml/modelkit/commands/export.py

@@ -270,7 +270,12 @@ def export(
         if not cli_utils.is_cli_provided(ctx, "dynamo") and "dynamo" in ec:
             dynamo = ec["dynamo"]
 
-    from ..export import InputTensorSpec, OutputTensorSpec, WinMLExportConfig
+    from ..export import (
+        InputTensorSpec,
+        OutputTensorSpec,
+        WinMLExportConfig,
+        resolve_export_compatibility,
+    )
     from ..export import export_pytorch as export_onnx
     from ..loader import load_hf_model
 
@@ -430,6 +435,8 @@ def _run_component_export(component_task: str | None, out_path: Path) -> None:
 
         try:
             cfg = WinMLExportConfig.from_dict(config_kwargs)
+            if not cfg.compatibility:
+                cfg.compatibility = resolve_export_compatibility()
         except Exception as e:
             console.print(f"[bold red]Configuration error:[/bold red] {e}")
             logger.exception("Failed to create export config")

src/winml/modelkit/commands/perf.py

@@ -987,6 +987,8 @@ def _load_model(self) -> None:
             "task": resolved_task,
             "config": override,
             "ep_device": self._ep_device,
+            "device": self.config.device,
+            "ep": self.config.ep,
             "precision": self.config.precision,
             "provider_options": self.config.ep_options,
             "use_cache": use_cache,
@@ -1346,8 +1348,10 @@ def _perf_modules(
     from ..session import EPDeviceTarget, WinMLEPRegistry, resolve_device
     from .build import _instantiate_parent_model
 
+    request_device = (device or "auto").lower()
+    request_ep = ep
     resolved_target = resolve_device(
-        EPDeviceTarget(ep=ep or "auto", device=device or "auto", source=ep_source)
+        EPDeviceTarget(ep=request_ep or "auto", device=request_device, source=ep_source)
     )
     resolved_ep_device = WinMLEPRegistry.instance().auto_device(resolved_target)
     resolved_device = resolved_target.device
@@ -1360,9 +1364,9 @@ def _perf_modules(
             model_id=hf_model,
             task=task,
             module=module_class,
-            device=resolved_device,
+            device=request_device,
             precision=precision,
-            ep=ep,
+            ep=request_ep,
         )
     except SubmoduleClassNotFoundError as e:
         # User-error: --module pattern didn't match. List what's available so
@@ -2669,7 +2673,6 @@ def perf(
     except Exception as e:
         raise click.ClickException(f"Failed to resolve Hub-hosted ONNX path {model!r}: {e}") from e
     model = hf_model
-
     # AC 11 (mockup spec): --top-k requires --op-tracing. Outside the
     # op-tracing section the flag is meaningless, so reject it explicitly
     # rather than silently ignoring a user's intent.

src/winml/modelkit/config/build.py

@@ -58,6 +58,7 @@
     WinMLExportConfig,
     _resolve_export_config_from_specs,
 )
+from ..export.policy import export_policy_targets_for_request, resolve_export_compatibility
 from ..loader.config import WinMLLoaderConfig, resolve_loader_config
 from ..optim.config import WinMLOptimizationConfig
 from ..quant.config import WinMLQuantizationConfig
@@ -71,7 +72,7 @@
 
 
 if TYPE_CHECKING:
-    from collections.abc import Mapping
+    from collections.abc import Mapping, Sequence
 
     import torch
     from torch import nn
@@ -457,6 +458,34 @@ def _apply_target_policy(
         )
 
 
+def _is_explicit_export_policy_target(*, device: str | None, ep: str | None) -> bool:
+    """Return whether the request named a specific EP/device export target."""
+    return (ep is not None and ep.lower() != "auto") or (
+        device is not None and device.lower() != "auto"
+    )
+
+
+def apply_export_compatibility_policy(
+    config: WinMLBuildConfig | Sequence[WinMLBuildConfig],
+    *,
+    device: str | None = "auto",
+    ep: str | None = None,
+) -> None:
+    """Populate export compatibility when the config has an export stage."""
+    export_policy_targets = export_policy_targets_for_request(
+        ep=ep,
+        device=device,
+        target_was_explicit=_is_explicit_export_policy_target(device=device, ep=ep),
+    )
+    configs = (config,) if isinstance(config, WinMLBuildConfig) else config
+    for cfg in configs:
+        if cfg.export is None:
+            continue
+        if cfg.export.compatibility:
+            continue
+        cfg.export.compatibility = resolve_export_compatibility(export_policy_targets)
+
+
 def resolve_quant_compile_config(
     *,
     device: str = "auto",
@@ -1078,6 +1107,10 @@ class name. Uses torchinfo to discover submodules and infer
     if no_compile:
         parent_config.compile = None
 
+    # Apply export compatibility policy so parent_config.export.compatibility is populated
+    # (used for serialization/cache-key participation and inheritance by submodules).
+    apply_export_compatibility_policy(parent_config, device=device, ep=ep)
+
     # =========================================================================
     # STEP 5: Specialize for submodules if requested
     # =========================================================================
@@ -1305,6 +1338,11 @@ def _input_name(i: int) -> str:
                 if parent_config.export is not None
                 else WinMLExportConfig().dynamo
             ),
+            compatibility=(
+                copy.deepcopy(parent_config.export.compatibility)
+                if parent_config.export is not None
+                else WinMLExportConfig().compatibility
+            ),
             # opset_version and batch_size use dataclass defaults from WinMLExportConfig
         ),
         optim=copy.deepcopy(parent_config.optim),
@@ -1390,6 +1428,11 @@ def _merge_export_config(
             if override.hierarchy_tag_format != defaults.hierarchy_tag_format
             else base.hierarchy_tag_format
         ),
+        compatibility=(
+            copy.deepcopy(override.compatibility)
+            if override.compatibility
+            else copy.deepcopy(base.compatibility)
+        ),
     )
 
 

src/winml/modelkit/eval/evaluate.py

@@ -318,6 +318,8 @@ def _load_model(config: WinMLEvaluationConfig) -> WinMLPreTrainedModel | WinMLCo
             config.model_id,
             ep_device,
             task=config.task,
+            device=config.device,
+            ep=config.ep,
             precision=config.precision,
             allow_unsupported_nodes=config.allow_unsupported_nodes,
             config=build_override,

src/winml/modelkit/export/__init__.py

@@ -19,6 +19,13 @@
     WinMLExportConfig,
     resolve_export_config,
 )
+from .policy import (
+    ExportCompatibilityConfig,
+    ExportCompatibilityRule,
+    ExportPolicyTarget,
+    export_policy_targets_for_request,
+    resolve_export_compatibility,
+)
 
 
 # Static type re-exports for the names exposed by ``__getattr__`` below.
@@ -40,15 +47,20 @@
 __version__ = "2.1.0"
 
 __all__ = [
+    "ExportCompatibilityConfig",
+    "ExportCompatibilityRule",
+    "ExportPolicyTarget",
     "InputTensorSpec",
     "MaxLengthTextInputGenerator",
     "ONNXConfigNotFoundError",
     "OutputTensorSpec",
     "WinMLExportConfig",
     "export_onnx",
+    "export_policy_targets_for_request",
     "export_pytorch",
     "generate_dummy_inputs",
     "register_onnx_overwrite",
+    "resolve_export_compatibility",
     "resolve_export_config",
     "resolve_io_specs",
 ]

src/winml/modelkit/export/compatibility_rules.json

@@ -0,0 +1,15 @@
+{
+  "schema_version": 1,
+  "rules": [
+    {
+      "match": {
+        "ep": null,
+        "device": null
+      },
+      "export": {
+        "transformers_attention": "eager"
+      },
+      "reason": "Transformers SDPA-exported attention guard paths are not broadly portable."
+    }
+  ]
+}