fix: stabilize HF export mask handling for e2e

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()]

src/winml/modelkit/models/hf/blip.py

@@ -20,10 +20,9 @@
 
 BLIP-specific trace-time adjustments:
 
-- **3-D ``decoder_attention_mask``** — ``BlipTextModel.get_extended_attention_mask``
-  has a ``dim == 3`` branch that broadcasts our mask without reconstructing a
-  causal triangle.  Passing a ``[1, 1, max_cache_len]`` mask routes through
-  that branch.
+- **Binary ``decoder_attention_mask`` passthrough** — the decoder accepts
+    2-D or 3-D binary masks and expands them internally. Preserve caller rank
+    so export paths can provide either shape.
 - **Explicit ``position_ids``** — ``BlipTextEmbeddings`` would otherwise
   derive positions from ``past_key_values_length`` (which traces as 0 for a
   static cache), baking the wrong position into the embedding lookup.
@@ -265,10 +264,9 @@ def _invoke_hf(self, cache: Any, inputs: dict[str, torch.Tensor]) -> torch.Tenso
             dtype=torch.long,
             device=encoder_hidden_states.device,
         )
-        # Transformers 4.57 accepts a 2-D or 3-D binary attention mask and
-        # expands it to additive form internally. Preserve the explicit query
-        # dimension required by this one-token static-cache decoder step.
-        decoder_mask = inputs["decoder_attention_mask"].unsqueeze(1)
+        decoder_mask = inputs["decoder_attention_mask"]
+        if decoder_mask.dim() not in (2, 3):
+            raise ValueError("decoder_attention_mask must be a 2-D or 3-D tensor")
         # 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)(

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.

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))

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_three_dimensional_binary_attention_mask(self, monkeypatch) -> None:
+    def test_decoder_preserves_two_dimensional_attention_mask(self, monkeypatch) -> None:
         from types import SimpleNamespace
         from unittest.mock import MagicMock
 
@@ -164,10 +164,43 @@ def decode(**kwargs):
 
         wrapper._invoke_hf(object(), inputs)
 
+        mask = captured["attention_mask"]
+        assert mask.shape == (1, 2)
+        assert mask.dtype == torch.int64
+        assert torch.equal(mask, inputs["decoder_attention_mask"])
+
+    def test_decoder_preserves_three_dimensional_attention_mask(self, monkeypatch) -> None:
+        from types import SimpleNamespace
+        from unittest.mock import MagicMock
+
+        import torch
+
+        from winml.modelkit.models.hf import blip as blip_module
+
+        wrapper = blip_module.BlipDecoderWrapper()
+        wrapper.model = MagicMock()
+        captured: dict[str, torch.Tensor] = {}
+
+        def decode(**kwargs):
+            captured["attention_mask"] = kwargs["attention_mask"]
+            return SimpleNamespace(logits=torch.zeros((1, 1, 4)))
+
+        wrapper.model.text_decoder.side_effect = decode
+        monkeypatch.setattr(blip_module, "EncoderDecoderCache", lambda *_args: 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),
+            "encoder_hidden_states": torch.zeros((1, 3, 4)),
+            "cache_position": torch.zeros((1,), dtype=torch.int64),
+        }
+
+        wrapper._invoke_hf(object(), inputs)
+
         mask = captured["attention_mask"]
         assert mask.shape == (1, 1, 2)
         assert mask.dtype == torch.int64
-        torch.testing.assert_close(mask, inputs["decoder_attention_mask"].unsqueeze(1))
+        torch.testing.assert_close(mask, inputs["decoder_attention_mask"])
 
     def test_decoder_cache_uses_position_input_when_model_omits_cache_kwargs(
         self, blip_config

tests/unit/test_transformers_compat.py

@@ -17,6 +17,7 @@
 
 import inspect
 import sys
+from types import SimpleNamespace
 
 import pytest
 import torch
@@ -118,6 +119,53 @@ def test_install_is_idempotent_and_reentrant_safe():
     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"
+
+
 def test_traceable_sdpa_accepts_integral_visibility_mask() -> None:
     """Optimum must normalize the integral masks emitted by Transformers 4.57."""
     import optimum.exporters.onnx.model_patcher as model_patcher