Fix QNN E2E evaluation failures

[DingmaomaoBJTU]

Summary

• Add JSON-backed export compatibility rules with wildcard EP/device matching and a global eager Transformers attention policy for export portability.
• Apply export compatibility through centralized config/export policy resolution while keeping empty compatibility out of serialized configs and cache identity.
• Make eager-attention export overrides restore nested Transformers configs losslessly.
• Validate JSON rules against supported EP/device targets so typos cannot create silent dead rules.
• Update docs and targeted tests for policy resolution, config serialization, HTP export attention behavior, and CI type/model-stub fixes.

Validation

• uv run pytest tests\unit\config\test_build.py tests\unit\export\test_export_policy.py tests\unit\export\test_config_validation.py tests\unit\export\test_htp_exporter_attention_compat.py tests\unit\models\auto\test_auto_onnx.py tests\unit\models\auto\test_composite_model_type_routing.py tests\unit\models\winml\test_composite_from_pretrained.py tests\unit\commands\test_build.py tests\unit\commands\test_export.py tests\unit\commands\test_perf_cli.py tests\unit\commands\test_perf_module.py tests\unit\test_uv_lock.py -q — 548 passed
• uv run --with ruff ruff check ... — passed
• uv run --with ruff ruff format --check ... — passed for Python changes
• uv run --with mypy --with types-colorama mypy -p winml.modelkit — passed
• uv lock --check — passed
• git diff --check — passed
• QNN E2E revalidation with scripts\e2e_eval\run_eval.py --eval-type perf — passed:
  • openai/clip-vit-base-patch32, feature-extraction, QNN GPU
  • openai/clip-vit-base-patch32, zero-shot-image-classification, QNN GPU
  • google-bert/bert-base-multilingual-cased, feature-extraction, QNN GPU
  • google-bert/bert-base-multilingual-cased, fill-mask, QNN GPU
  • google-bert/bert-base-multilingual-cased, masked-lm, QNN GPU
  • google-bert/bert-base-multilingual-cased, fill-mask, QNN NPU
  • google/flan-t5-base, QNN GPU

[KayMKM]

Reviewed the branch locally (24 files, +1144/-46), ran the new tests (234 passed), and empirically verified two of the findings below.

What's good

Single choke point (HTPExporter._convert_model_to_onnx is the only torch.onnx.export caller), data-driven JSON rules instead of hardcoded model logic, compatibility persisted into the config so it participates in cache keys, and solid serialization/validation coverage.

Findings

1. use_eager_attention_for_export restore is lossy (real bug)

src/winml/modelkit/transformers_compat.py:49

Transformers 5's _attn_implementation setter recurses into sub_configs (configuration_utils.py:401-417). Setting the parent config overwrites its children before the loop reaches them, so the children's original values are never recorded in restored. Verified against transformers 5.14.1 with a CLIP config:

before: sdpa flash_attention_2
inside: eager eager
after : sdpa sdpa      # text_config silently downgraded, not restored

Suggested fix: snapshot every reachable config's value in a first pass, mutate in a second pass, then restore all of them. Alternatively use the public model.set_attn_implementation() (present in 5.14.1, modeling_utils.py:2176) instead of poking the private attribute.

2. Auto path no longer forwards the resolved target into build-config generation

src/winml/modelkit/models/auto.py:375-380, src/winml/modelkit/commands/build.py:1027-1038

When neither device nor ep is explicit, generate_hf_build_config now receives device="auto", ep=None. That makes _resolve_policy_target (config/build.py:324) re-resolve the target via auto_detect_device() + EP_SUPPORTED_DEVICES, while the runtime target is resolved via resolve_device() (session/ep_device.py:686), which iterates EP_DEVICE_SPECS in catalog priority order and never consults auto_detect_device().

Two different algorithms now decide the same thing, so build-time quant/compile policy can select a different device than the session actually runs on. The comment deleted in this PR ("Forwarding both concrete axes keeps analyzer/build output aligned with the target selected by the catalog-backed resolver") was guarding exactly this case.

Could you confirm the two resolvers cannot diverge, or thread the resolved runtime target through separately from the export-policy target? The export policy needs the request; quant/compile policy still needs the resolved target.

3. PR description is stale

Commit b782d681 ("chore: prune qnn export policy changes") reverted the E2E scripts, _native_ep_registration.py, analyze/analyzer.py, loader/hf.py, and optracing/qnn/profiler.py. Summary bullets 3 and 4 describe code that is no longer in the diff — worth trimming so reviewers and the changelog match reality.

4. The "portable default" is effectively "always eager"

QNN is in EP_DEVICE_SPECS, so resolve_export_compatibility(None) always resolves to eager. Net effect:

• winml build -m X → eager attention
• winml build -m X --device gpu --ep dml → SDPA

Two different ONNX graphs (and different numerics) for what resolves to the same runtime target on a GPU machine. The same flip happens on eval's CPU fallback (eval/evaluate.py:343), which sets device="cpu"/ep="cpu", turning the request explicit and forcing a full re-export on the retry.

If eager is the portable answer for unspecified targets, is the explicit-target branch earning its complexity? Worth stating the intended contract explicitly.

5. Cache-key churn from {} vs. absent

Verified locally:

unresolved default  : 6c13c2bc15df31cf
resolved-empty ({}) : cdb71a905c9bb3c7

Both produce byte-identical exports, but the is_resolved bookkeeping leaks into artifact identity because to_dict() emits "compatibility": {} for a resolved-but-empty config and omits the key otherwise. Consider tracking "policy already applied" out of band so it does not participate in the cache key.

6. Docs not updated

docs/reference/index.md:47-60 documents the export config fields; compatibility is missing from the table.

Nits

• TransformersAttentionPolicy = Literal["eager"] (export/policy.py:24) is defined but never used; the field is typed str | None.
• EXPORT_COMPATIBILITY_RULES = load_export_compatibility_rules() runs at module scope, so importing export/config.py (widely imported) does JSON file I/O, and a malformed resource becomes an import-time crash. A lazy functools.cache accessor would be safer.
• is_resolved participates in __eq__, and the __bool__ overload doubling as "policy already resolved" is subtle enough that an explicit is_resolved check at the three call sites would read better.
• _coerce_target / Sequence[object] is dead generality now that only ExportPolicyTarget is ever passed (after 4c5e8a06).
• Vestigial assertions that are always true and test nothing: assert "export_policy_targets" not in mock_gen.call_args.kwargs (tests/unit/commands/test_build.py, tests/unit/commands/test_perf_module.py) and assert "export_target_was_explicit" not in received (tests/unit/commands/test_perf_cli.py) — those kwargs no longer exist anywhere in the codebase.

[KayMKM]

1. Finding 2 is still live on the winml build CLI path

auto.py and perf.py now split the runtime target from the export-policy request, but commands/build.py does not. The resolve_device() call that used to run before config generation was moved after it (runtime_device / runtime_ep_value), so generate_build_config() still receives the raw device="auto", ep=None. generate_build_config() also has no export_policy_target parameter to forward, so _resolve_policy_target() re-resolves the target through auto_detect_device() + EP_SUPPORTED_DEVICES, while the session is resolved through resolve_device() + EP_DEVICE_SPECS.

Two different algorithms still decide the same thing on this path, so winml build -m X (no --device/--ep) can generate an NPU-targeted quant/compile policy and cache entry for a build that later runs build_hf_model(device=runtime_device, ...) on GPU.

Suggested fix: mirror the auto.py split here — pass runtime_device / runtime_ep_value to generate_build_config() for quant/compile policy, and thread the raw request through a new export_policy_target passthrough on generate_build_config().

2. use_eager_attention_for_export() still has a lossy edge case

The two-pass snapshot fixes the parent-overwrites-child case. But configs whose current value is None or "eager" are skipped and never recorded:

if current in (None, "eager"):
    continue

If a parent config is sdpa and a child sub-config is explicitly eager, the child is not snapshotted. On restore, setting the parent back to sdpa recurses into sub_configs and silently rewrites the child to sdpa, with no entry in restored to put it back.

Suggested fix: snapshot every reachable config unconditionally in the first pass, and only use the None/"eager" check to decide whether a write is needed — not whether to record the original value.

3. Nit: _coerce_target / Sequence[object] is still dead generality

resolve_export_compatibility() is typed Sequence[object] | None and routes every element through _coerce_target(), but the only producer (export_policy_targets_for_request()) already returns tuple[ExportPolicyTarget, ...]. Narrowing the signature to Sequence[ExportPolicyTarget] | None and dropping _coerce_target() would remove the duck-typing branch and its error path.

[KayMKM]

59ed9b44 addresses all three of my previous points — thanks. The build-path target split, the lossless nested-config restore (including sub-configs not attached to any module, and the parent-before-child restore order), and the _coerce_target removal all look right, and the three new attention tests cover exactly the cascade cases I was worried about.

One spot looks like it was missed by the same split.

Composite sub-model branch still generates configs from the raw request

src/winml/modelkit/commands/build.py:1339

component_config = gen_cfg(
    model,
    task=component_task,
    trust_remote_code=trust_remote_code,
    device=device,
    precision=precision,
    ep=ep_value,
    ...
)

Every other generate_build_config() / generate_hf_build_config() call site in build() now passes runtime_device / runtime_ep_value plus export_policy_target=(request_device, request_ep_value), but this one still passes the raw device / ep_value.

Before this PR that was harmless, because device / ep_value were reassigned in place by the resolve_device() block near the top of build(). Now that the resolved values live in separate variables, this branch receives device="auto", ep=None whenever the user does not pin --device / --ep, so each component's quant/compile policy is re-resolved through _resolve_policy_target() → auto_detect_device() while the actual build/session target comes from resolve_device(). That is the same runtime/config divergence the rest of the commit just fixed, only scoped to composite models.

Suggested fix:

component_config = gen_cfg(
    model,
    task=component_task,
    trust_remote_code=trust_remote_code,
    device=runtime_device,
    precision=precision,
    ep=runtime_ep_value,
    export_policy_target=(request_device, request_ep_value),
    ...
)

A composite regression test asserting the component configs see the resolved runtime target would keep this from drifting again.