Fix eager attention selection during model loading (#1375)

## Summary
- select the configured Transformers attention implementation before
`from_pretrained` constructs attention modules
- thread export compatibility through both build and direct export paths
- cover loader and build propagation while keeping export test doubles
compatible

## Motivation
DistilBERT exports could instantiate `DistilBertSdpaAttention` and only
switch the config to eager immediately before ONNX export. This left the
module implementation and config inconsistent, causing integer attention
masks to reach SDPA and fail with:

```
Expected attn_mask dtype to be bool or float or to match query dtype
```

## Validation
- `uv run pytest tests/unit/loader/test_load_hf_model.py
tests/unit/build/test_hf.py
tests/unit/export/test_htp_exporter_attention_compat.py
tests/unit/commands/test_config_value_priority.py
tests/unit/commands/test_export.py -q --tb=short`
- 175 tests passed
- retried all 6 affected QNN GPU models: the attention-mask dtype error
was eliminated in all 6; 5 completed perf and accuracy successfully,
while one progressed to a separate duplicate `metadata_props` ONNX
validation failure

Author: fangyangci

scripts/e2e_eval/run_eval.py

@@ -2902,11 +2902,11 @@ def parse_args() -> argparse.Namespace:
         "--priority",
         nargs="+",
         choices=["P0", "P1", "P2", "P3"],
-        default=["P0", "P1", "P2"],
+        default=["P0", "P1", "P2", "P3"],
         metavar="{P0,P1,P2,P3}",
         help=(
             "Filter by priority. Pass one or more, e.g. --priority P0 P1. "
-            "Default: P0 P1 P2 (P3 excluded from default runs)."
+            "Default: P0 P1 P2 P3."
         ),
     )
     parser.add_argument("--model-type", help="Filter by model_type")

src/winml/modelkit/build/hf.py

@@ -350,6 +350,9 @@ def _load_model(
             path (PR #719 dedup pattern).
     """
     task = config.loader.task
+    attn_implementation = (
+        config.export.compatibility.transformers_attention if config.export else None
+    )
 
     if random_init:
         from transformers import AutoConfig
@@ -402,7 +405,10 @@ def _load_model(
 
         model_label = model_id or config.loader.model_type
         logger.info("Creating random-weight model: %s (from %s)", model_class.__name__, model_label)
-        return model_class.from_config(hf_config)
+        model_kwargs: dict[str, Any] = {}
+        if attn_implementation is not None:
+            model_kwargs["attn_implementation"] = attn_implementation
+        return model_class.from_config(hf_config, **model_kwargs)
 
     if model_id is not None:
         from ..loader import load_hf_model
@@ -415,6 +421,7 @@ def _load_model(
             trust_remote_code=effective_trust,
             hf_config=hf_config,
             model_type=model_type,
+            attn_implementation=attn_implementation,
         )
         return pytorch_model
 

src/winml/modelkit/commands/export.py

@@ -464,7 +464,11 @@ def _run_component_export(component_task: str | None, out_path: Path) -> None:
             raise click.ClickException(f"Configuration error: {e}") from e
 
         # Load model with task detection (CLI is the orchestration layer)
-        pytorch_model, _, detected_task = load_hf_model(model, task=component_task)
+        pytorch_model, _, detected_task = load_hf_model(
+            model,
+            task=component_task,
+            attn_implementation=cfg.compatibility.transformers_attention,
+        )
         if component_task:
             console.print(f"[dim]Task (override): {detected_task}[/dim]")
         else:

src/winml/modelkit/export/htp/exporter.py

@@ -213,7 +213,10 @@ def export(
                     raise ValueError("Either 'model' or 'model_name_or_path' must be provided.")
                 from ...loader import load_hf_model
 
-                model, _, _ = load_hf_model(model_name_or_path)
+                model, _, _ = load_hf_model(
+                    model_name_or_path,
+                    attn_implementation=export_config.compatibility.transformers_attention,
+                )
 
             # Step 1: Model Preparation
             model.eval()

src/winml/modelkit/loader/hf.py

@@ -147,6 +147,7 @@ def load_hf_model(
     model_type: str | None = None,
     *,
     torch_dtype: Any | None = None,
+    attn_implementation: str | None = None,
 ) -> tuple[nn.Module, PretrainedConfig, str]:
     """Load, detect task, and prepare HuggingFace model.
 
@@ -176,6 +177,8 @@ def load_hf_model(
             pattern as ``resolve_loader_config(hf_config=...)`` from PR #719.
         torch_dtype: Optional dtype policy forwarded to ``from_pretrained``.
             Pass ``"auto"`` to preserve the checkpoint's stored dtype.
+        attn_implementation: Optional Transformers attention implementation
+            forwarded to ``from_pretrained`` before attention modules are created.
 
     Returns:
         Tuple of (model, hf_config, task)
@@ -290,6 +293,8 @@ def load_hf_model(
     }
     if torch_dtype is not None:
         load_kwargs["torch_dtype"] = torch_dtype
+    if attn_implementation is not None:
+        load_kwargs["attn_implementation"] = attn_implementation
     model = loader_cls.from_pretrained(model_name_or_path, **load_kwargs)
 
     # [5] Export Preparation

tests/unit/build/test_hf.py

@@ -379,6 +379,59 @@ def test_pretrained_load_threads_model_class(self, sample_config) -> None:
         m_load.assert_called_once()
         assert m_load.call_args.kwargs["model_class"] == "AutoModelForImageClassification"
 
+    def test_pretrained_load_threads_attention_compatibility(self, sample_config) -> None:
+        """Attention compatibility is selected before model construction."""
+        from winml.modelkit.build.hf import _load_model
+        from winml.modelkit.export.policy import ExportCompatibilityConfig
+
+        sample_config.export.compatibility = ExportCompatibilityConfig(
+            transformers_attention="eager"
+        )
+        with patch("winml.modelkit.loader.load_hf_model") as m_load:
+            m_load.return_value = (MagicMock(), MagicMock(), "image-classification")
+            _load_model(sample_config, "test-model", trust_remote_code=False)
+
+        m_load.assert_called_once()
+        assert m_load.call_args.kwargs["attn_implementation"] == "eager"
+
+    def test_random_init_selects_eager_attention_before_construction(self) -> None:
+        """Random-init builds construct the requested attention implementation."""
+        from transformers import DistilBertConfig
+
+        from winml.modelkit.build.hf import _load_model
+        from winml.modelkit.config import WinMLBuildConfig
+
+        config = WinMLBuildConfig.from_dict(
+            {
+                "loader": {
+                    "task": "text-classification",
+                    "model_class": "AutoModelForSequenceClassification",
+                },
+                "export": {
+                    "compatibility": {"transformers_attention": "eager"},
+                },
+                "optim": {},
+                "quant": None,
+                "compile": None,
+            }
+        )
+        hf_config = DistilBertConfig(n_layers=1, dim=32, hidden_dim=64, n_heads=4)
+
+        model = _load_model(
+            config,
+            model_id=None,
+            trust_remote_code=False,
+            random_init=True,
+            hf_config=hf_config,
+        )
+
+        attention_classes = {
+            type(module).__name__
+            for module in model.modules()
+            if "Attention" in type(module).__name__
+        }
+        assert attention_classes == {"MultiHeadSelfAttention"}
+
     def test_pre_loaded_model_skips_load(
         self, tmp_path: Path, sample_config, mock_pipeline
     ) -> None:

tests/unit/commands/test_export.py

@@ -1090,7 +1090,7 @@ def test_composite_exports_one_onnx_per_component(
         ):
             # Echo the per-component task back as the detected task so it flows
             # through to export_onnx and can be verified per sub-model.
-            mock_load.side_effect = lambda _model, task=None: (MagicMock(), None, task)
+            mock_load.side_effect = lambda _model, task=None, **_kwargs: (MagicMock(), None, task)
             mock_resolve_cfg.return_value = (
                 WinMLExportConfig(),
                 WinMLLoaderConfig(task="text-generation"),
@@ -1247,7 +1247,7 @@ def fake_export_onnx(**kwargs):
             patch("winml.modelkit.export.resolve_export_config") as mock_resolve_cfg,
             patch("winml.modelkit.export.export_pytorch", side_effect=fake_export_onnx),
         ):
-            mock_load.side_effect = lambda _model, task=None: (MagicMock(), None, task)
+            mock_load.side_effect = lambda _model, task=None, **_kwargs: (MagicMock(), None, task)
             mock_resolve_cfg.return_value = (
                 WinMLExportConfig(),
                 WinMLLoaderConfig(task="text-generation"),
@@ -1295,7 +1295,7 @@ def test_submodel_exports_only_requested_component(
             patch("winml.modelkit.loader.load_hf_model") as mock_load,
             patch("winml.modelkit.export.resolve_export_config") as mock_resolve_cfg,
         ):
-            mock_load.side_effect = lambda _model, task=None: (MagicMock(), None, task)
+            mock_load.side_effect = lambda _model, task=None, **_kwargs: (MagicMock(), None, task)
             mock_resolve_cfg.return_value = (
                 WinMLExportConfig(),
                 WinMLLoaderConfig(task="text-generation"),
@@ -1412,7 +1412,7 @@ def test_submodel_allows_input_specs(
             patch("winml.modelkit.loader.load_hf_model") as mock_load,
             patch("winml.modelkit.export.resolve_export_config") as mock_resolve_cfg,
         ):
-            mock_load.side_effect = lambda _model, task=None: (MagicMock(), None, task)
+            mock_load.side_effect = lambda _model, task=None, **_kwargs: (MagicMock(), None, task)
             mock_resolve_cfg.return_value = (
                 WinMLExportConfig(),
                 WinMLLoaderConfig(task="text-generation"),

tests/unit/export/test_htp_exporter_attention_compat.py

@@ -7,6 +7,7 @@
 from __future__ import annotations
 
 from typing import TYPE_CHECKING
+from unittest.mock import patch
 
 import torch
 import torch.nn as nn
@@ -143,6 +144,28 @@ def _export_config(*, eager_attention: bool) -> WinMLExportConfig:
     )
 
 
+def test_htp_exporter_auto_load_threads_attention_compatibility(tmp_path: Path) -> None:
+    export_config = _export_config(eager_attention=True)
+
+    with patch("winml.modelkit.loader.load_hf_model") as mock_load:
+        mock_load.side_effect = RuntimeError("stop after model loading")
+        try:
+            HTPExporter().export(
+                output_path=str(tmp_path / "model.onnx"),
+                export_config=export_config,
+                model_name_or_path="fake/model",
+            )
+        except RuntimeError as exc:
+            assert str(exc) == "stop after model loading"
+        else:
+            raise AssertionError("Expected the loader sentinel to stop export")
+
+    mock_load.assert_called_once_with(
+        "fake/model",
+        attn_implementation="eager",
+    )
+
+
 def test_htp_exporter_uses_eager_attention_when_policy_requests_it(
     monkeypatch: pytest.MonkeyPatch,
     tmp_path: Path,

tests/unit/loader/test_load_hf_model.py

@@ -233,6 +233,44 @@ def test_dtype_is_forwarded_to_task_resolved_class(self, monkeypatch):
             torch_dtype="auto",
         )
 
+    def test_attention_implementation_is_forwarded_before_model_construction(self, monkeypatch):
+        """Export compatibility selects attention while constructing the model."""
+        from types import SimpleNamespace
+        from unittest.mock import MagicMock
+
+        import winml.modelkit.loader.resolution as resolution_module
+
+        task_class = MagicMock()
+        task_class.__name__ = "TaskModel"
+        task_class.config_class = None
+        task_model = MagicMock()
+        task_class.from_pretrained.return_value = task_model
+        config = SimpleNamespace(model_type="unit", architectures=["CheckpointModel"])
+
+        monkeypatch.setattr(
+            resolution_module,
+            "resolve_task",
+            lambda *_a, **_kw: SimpleNamespace(
+                task="fill-mask",
+                model_class=task_class,
+            ),
+        )
+
+        model, _, task = load_hf_model(
+            "fake/model",
+            hf_config=config,
+            attn_implementation="eager",
+        )
+
+        assert model is task_model
+        assert task == "fill-mask"
+        task_class.from_pretrained.assert_called_once_with(
+            "fake/model",
+            trust_remote_code=False,
+            config=config,
+            attn_implementation="eager",
+        )
+
     def test_bert_tiny_uses_model_specific_default_task(self, monkeypatch):
         """bert-tiny should use model-specific default task when task is omitted."""
         from unittest.mock import MagicMock