From 317fbbf1b4301458921978eb61f100be6a0d7021 Mon Sep 17 00:00:00 2001 From: fangyangci_qnn Date: Tue, 1 Sep 2026 09:24:49 +0200 Subject: [PATCH 1/2] fix: stabilize hf export mask handling for e2e --- src/winml/modelkit/export/io.py | 56 +++++++++++++- src/winml/modelkit/models/hf/blip.py | 8 +- .../modelkit/models/hf/decoder_wrapper.py | 22 +++++- src/winml/modelkit/transformers_compat.py | 74 +++++++++++++++++-- tests/unit/export/test_blip_onnx_config.py | 9 +-- tests/unit/test_transformers_compat.py | 48 ++++++++++++ 6 files changed, 198 insertions(+), 19 deletions(-) diff --git a/src/winml/modelkit/export/io.py b/src/winml/modelkit/export/io.py index 3bedca226..ad4b02881 100644 --- a/src/winml/modelkit/export/io.py +++ b/src/winml/modelkit/export/io.py @@ -360,6 +360,52 @@ def _populate_sequence_length_from_config( ) +def _coerce_binary_attention_masks_to_bool( + dummy_inputs: dict[str, torch.Tensor], +) -> dict[str, torch.Tensor]: + """Normalize binary attention masks to bool for export tracing. + + Some transformer attention paths reject integral masks under newer torch + SDPA checks, while accepting bool masks. To keep behavior architecture- + agnostic, coerce only tensors whose names indicate an attention mask and + whose runtime values are binary (0/1). + """ + import torch + + integer_dtypes = { + torch.int8, + torch.int16, + torch.int32, + torch.int64, + torch.uint8, + } + normalized = dict(dummy_inputs) + + for name, tensor in dummy_inputs.items(): + if "attention_mask" not in name: + continue + if tensor.dtype not in integer_dtypes: + continue + if tensor.numel() == 0: + continue + + min_value = int(tensor.min().item()) + max_value = int(tensor.max().item()) + if min_value < 0 or max_value > 1: + logger.debug( + "Skipping bool coercion for %s: non-binary range=(%d, %d)", + name, + min_value, + max_value, + ) + continue + + normalized[name] = tensor.to(dtype=torch.bool) + logger.debug("Coerced %s from %s to bool for export tracing", name, tensor.dtype) + + return normalized + + def generate_dummy_inputs( model_type: str, task: str, @@ -419,10 +465,11 @@ def generate_dummy_inputs( ) # Optimum's OnnxConfig is untyped; the dummy-inputs dict matches our return type. - return cast( + dummy_inputs = cast( "dict[str, torch.Tensor]", onnx_config.generate_dummy_inputs(framework="pt", **shape_kwargs), ) + return _coerce_binary_attention_masks_to_bool(dummy_inputs) def resolve_io_specs( @@ -482,7 +529,12 @@ def resolve_io_specs( # Generate dummy inputs for concrete shapes and dtypes, # intercepting value ranges from Optimum's tensor gen methods with intercept_value_ranges() as value_ranges: - dummy_inputs = onnx_config.generate_dummy_inputs(framework="pt", **shape_kwargs) + dummy_inputs = cast( + "dict[str, torch.Tensor]", + onnx_config.generate_dummy_inputs(framework="pt", **shape_kwargs), + ) + + dummy_inputs = _coerce_binary_attention_masks_to_bool(dummy_inputs) input_shapes = [tuple(t.shape) for t in dummy_inputs.values()] input_dtypes = [str(t.dtype).replace("torch.", "") for t in dummy_inputs.values()] diff --git a/src/winml/modelkit/models/hf/blip.py b/src/winml/modelkit/models/hf/blip.py index fb96ba539..cef262481 100644 --- a/src/winml/modelkit/models/hf/blip.py +++ b/src/winml/modelkit/models/hf/blip.py @@ -265,15 +265,13 @@ def _invoke_hf(self, cache: Any, inputs: dict[str, torch.Tensor]) -> torch.Tenso dtype=torch.long, device=encoder_hidden_states.device, ) - decoder_mask = (1 - inputs["decoder_attention_mask"]).to(dtype=encoder_hidden_states.dtype) - decoder_mask = decoder_mask.unsqueeze(1).unsqueeze(1) - decoder_mask = decoder_mask * torch.finfo(encoder_hidden_states.dtype).min + decoder_mask = inputs["decoder_attention_mask"] # self.model is nn.Module; torch's __getattr__ types text_decoder as # Tensor | Module, so narrow to a callable Module. outputs = cast("nn.Module", self.model.text_decoder)( input_ids=inputs["decoder_input_ids"], - # HF's causal-mask reconstruction traces as ops the NPU analyzer - # doesn't support; pass an additive 4-D mask to bypass reconstruction. + # BLIP text decoder expects a 2-D/3-D mask and internally expands + # it with cache-aware causal masking. attention_mask=decoder_mask, # Without explicit position_ids, BlipTextModel would derive them # from past_kv_len=0 (a frozen constant in the trace), giving every diff --git a/src/winml/modelkit/models/hf/decoder_wrapper.py b/src/winml/modelkit/models/hf/decoder_wrapper.py index f5396b9d6..752563629 100644 --- a/src/winml/modelkit/models/hf/decoder_wrapper.py +++ b/src/winml/modelkit/models/hf/decoder_wrapper.py @@ -95,8 +95,26 @@ def get_export_args(self, inputs: dict[str, torch.Tensor]) -> tuple[torch.Tensor """Order dict inputs positionally to match the IOConfig's input order.""" return tuple(inputs.values()) - def forward(self, *args: torch.Tensor) -> tuple[torch.Tensor, ...]: - """Execute the three-step adapter. Subclasses override ``_invoke_hf``.""" + def forward( + self, + *args: torch.Tensor, + **kwargs: torch.Tensor, + ) -> tuple[torch.Tensor, ...]: + """Execute the three-step adapter. Subclasses override ``_invoke_hf``. + + Accepts either positional args in ONNX input order or keyword args + keyed by ONNX input names. + """ + if args and kwargs: + raise TypeError("Provide either positional args or keyword args, not both") + + if kwargs: + input_order = list(self.onnx_config.inputs.keys()) + missing = [name for name in input_order if name not in kwargs] + if missing: + raise TypeError(f"Missing decoder input(s): {missing}") + args = tuple(kwargs[name] for name in input_order) + inputs = dict(zip(self.onnx_config.inputs.keys(), args, strict=True)) # 1. Create cache aliased to ONNX past-KV inputs. diff --git a/src/winml/modelkit/transformers_compat.py b/src/winml/modelkit/transformers_compat.py index 49d5b10b5..aabf848b5 100644 --- a/src/winml/modelkit/transformers_compat.py +++ b/src/winml/modelkit/transformers_compat.py @@ -48,11 +48,45 @@ @contextlib.contextmanager def use_eager_attention_for_export(model: nn.Module) -> Iterator[None]: - """Temporarily prefer eager attention on HF-style module configs.""" + """Temporarily prefer eager attention and normalize SDPA mask dtypes. + + Some exporter paths feed integer attention masks into SDPA, which newer + torch versions reject. During export-only tracing, coerce integral masks + to bool before dispatching to the currently-installed SDPA implementation. + """ + import torch + restored: list[tuple[int, Any, Any]] = [] configs: dict[int, Any] = {} children: dict[int, set[int]] = {} seen_configs: set[int] = set() + original_sdpa = torch.nn.functional.scaled_dot_product_attention + + def _coerce_integer_mask_sdpa(*args: Any, **kwargs: Any) -> Any: + attn_mask: Any | None = None + use_positional = len(args) >= 4 + + if use_positional: + attn_mask = args[3] + elif "attn_mask" in kwargs: + attn_mask = kwargs["attn_mask"] + + if isinstance(attn_mask, torch.Tensor) and attn_mask.dtype in { + torch.int8, + torch.int16, + torch.int32, + torch.int64, + torch.uint8, + }: + cast_mask = attn_mask.to(dtype=torch.bool) + if use_positional: + args_list = list(args) + args_list[3] = cast_mask + args = tuple(args_list) + else: + kwargs["attn_mask"] = cast_mask + + return original_sdpa(*args, **kwargs) for module in model.modules(): _collect_attention_configs(getattr(module, "config", None), configs, children, seen_configs) @@ -64,10 +98,12 @@ def use_eager_attention_for_export(model: nn.Module) -> Iterator[None]: for config in configs.values(): if config._attn_implementation != "eager": config._attn_implementation = "eager" + torch.nn.functional.scaled_dot_product_attention = _coerce_integer_mask_sdpa try: yield finally: + torch.nn.functional.scaled_dot_product_attention = original_sdpa for _config_id, config, previous in _parent_before_child(restored, children): config._attn_implementation = previous @@ -339,11 +375,39 @@ def _sdpa_mask_without_vmap_tf5( ) -> Any: if mask_function is None: mask_function = causal_mask_function - padding_mask = prepare_padding_mask(attention_mask, kv_length, kv_offset) + try: + # transformers >=5 uses _slice to preserve compile-friendly indexing. + padding_mask = prepare_padding_mask(attention_mask, kv_length, kv_offset, _slice=False) + except TypeError: + # Older helper signatures do not expose _slice. + padding_mask = prepare_padding_mask(attention_mask, kv_length, kv_offset) ignore_causal_mask_sdpa = cast("Any", _ignore_causal_mask_sdpa) - if allow_is_causal_skip and ignore_causal_mask_sdpa( - padding_mask, q_length, kv_length, q_offset, kv_offset, local_size - ): + should_skip = False + if allow_is_causal_skip: + try: + # transformers >=5.1 signature + should_skip = bool( + ignore_causal_mask_sdpa( + padding_mask, + q_length, + kv_length, + kv_offset, + local_size, + ) + ) + except TypeError: + # transformers 5.0 / late 4.x signature + should_skip = bool( + ignore_causal_mask_sdpa( + padding_mask, + q_length, + kv_length, + q_offset, + kv_offset, + local_size, + ) + ) + if should_skip: return None if padding_mask is not None: mask_function = and_masks(mask_function, padding_mask_function(padding_mask)) diff --git a/tests/unit/export/test_blip_onnx_config.py b/tests/unit/export/test_blip_onnx_config.py index 49dd3f040..2bb8219b1 100644 --- a/tests/unit/export/test_blip_onnx_config.py +++ b/tests/unit/export/test_blip_onnx_config.py @@ -136,7 +136,7 @@ def test_past_kv_is_full_buffer(self, blip_config) -> None: assert tuple(inputs["past_0_key"].shape) == expected assert tuple(inputs["past_0_value"].shape) == expected - def test_decoder_passes_four_dimensional_additive_attention_mask(self, monkeypatch) -> None: + def test_decoder_passes_two_dimensional_attention_mask(self, monkeypatch) -> None: from types import SimpleNamespace from unittest.mock import MagicMock @@ -165,10 +165,9 @@ def decode(**kwargs): wrapper._invoke_hf(object(), inputs) mask = captured["attention_mask"] - assert mask.shape == (1, 1, 1, 2) - assert mask.is_floating_point() - assert mask[0, 0, 0, 0] == 0 - assert mask[0, 0, 0, 1] == torch.finfo(mask.dtype).min + assert mask.shape == (1, 2) + assert mask.dtype == torch.int64 + assert torch.equal(mask, inputs["decoder_attention_mask"]) def test_decoder_cache_uses_position_input_when_model_omits_cache_kwargs( self, blip_config diff --git a/tests/unit/test_transformers_compat.py b/tests/unit/test_transformers_compat.py index f443d9236..d3a78a316 100644 --- a/tests/unit/test_transformers_compat.py +++ b/tests/unit/test_transformers_compat.py @@ -17,6 +17,7 @@ import inspect import sys +from types import SimpleNamespace import pytest @@ -115,3 +116,50 @@ def test_install_is_idempotent_and_reentrant_safe(): transformers_compat.install() assert _is_transformers5_patched(model_patcher) + + +def test_patched_sdpa_mask_without_vmap_accepts_current_signature(): + import optimum.exporters.onnx.model_patcher as model_patcher + import torch + + mask = torch.tensor([[1, 1, 1, 0]], dtype=torch.int32) + result = model_patcher.sdpa_mask_without_vmap( + batch_size=1, + q_length=4, + kv_length=4, + attention_mask=mask, + device="cpu", + ) + + assert result is None or isinstance(result, torch.Tensor) + + +def test_use_eager_attention_for_export_coerces_integral_sdpa_mask(monkeypatch): + import torch + import torch.nn as nn + + captured: dict[str, torch.dtype] = {} + + def fake_sdpa(*args, **kwargs): + attn_mask = kwargs.get("attn_mask", args[3] if len(args) >= 4 else None) + if isinstance(attn_mask, torch.Tensor): + captured["dtype"] = attn_mask.dtype + return kwargs.get("query", args[0]) + + monkeypatch.setattr(torch.nn.functional, "scaled_dot_product_attention", fake_sdpa) + + class DummyModel(nn.Module): + def __init__(self) -> None: + super().__init__() + self.config = SimpleNamespace(_attn_implementation="sdpa") + + model = DummyModel() + q = torch.zeros((1, 1, 1, 1), dtype=torch.float32) + mask = torch.ones((1, 1, 1, 1), dtype=torch.int32) + + with transformers_compat.use_eager_attention_for_export(model): + assert model.config._attn_implementation == "eager" + torch.nn.functional.scaled_dot_product_attention(q, q, q, attn_mask=mask) + + assert captured["dtype"] == torch.bool + assert model.config._attn_implementation == "sdpa" From 1e496972710fc2bc66be5b80a2ae40a9a92cb084 Mon Sep 17 00:00:00 2001 From: fangyangci <133664123+fangyangci@users.noreply.github.com> Date: Wed, 2 Sep 2026 03:19:44 +0800 Subject: [PATCH 2/2] Potential fix for pull request finding 'CodeQL / Unnecessary lambda' Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com> --- tests/unit/export/test_blip_onnx_config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/export/test_blip_onnx_config.py b/tests/unit/export/test_blip_onnx_config.py index 32e94a671..4e6bf6280 100644 --- a/tests/unit/export/test_blip_onnx_config.py +++ b/tests/unit/export/test_blip_onnx_config.py @@ -187,7 +187,7 @@ def decode(**kwargs): wrapper.model.text_decoder.side_effect = decode monkeypatch.setattr(blip_module, "EncoderDecoderCache", lambda *_args: object()) - monkeypatch.setattr(blip_module, "DynamicCache", lambda: object()) + monkeypatch.setattr(blip_module, "DynamicCache", object) inputs = { "decoder_input_ids": torch.zeros((1, 1), dtype=torch.int32), "decoder_attention_mask": torch.tensor([[[1, 0]]], dtype=torch.int64),