onnxruntime
diff --git a/‎src/mobius/_builder.py‎
Lines changed: 86 additions & 20 deletions b/‎src/mobius/_builder.py‎
Lines changed: 86 additions & 20 deletions
diff --git a/‎src/mobius/_builder_test.py‎
Lines changed: 170 additions & 0 deletions b/‎src/mobius/_builder_test.py‎
Lines changed: 170 additions & 0 deletions
@@ -213,29 +213,95 @@ def forward(self, op, input_ids, attention_mask,
             trace=trace_optimization,
         )
 
-    # Lower default-domain opset from 24 to 23 when the target EP doesn't
-    # register opset 24 kernels for standard ops (Reshape, RMSNormalization,
-    # etc.). Without this, those ops fall to CPU and produce ~280 memcpy
-    # nodes that destroy performance. The flag defaults to True; set
-    # MOBIUS_ORT_LOWER_OPSET_FOR_EP=0 to disable for EPs that support
-    # opset 24 natively.
-    if flags.ort_lower_opset_for_ep and execution_provider != "default":
-        for name, model in pkg.items():
-            if "" in model.graph.opset_imports:
-                original = model.graph.opset_imports[""]
-                model.graph.opset_imports[""] = 23
-                logger.warning(
-                    "Lowered opset %d→23 for '%s' (EP=%s). "
-                    "ORT does not yet register opset %d kernels for this EP. "
-                    "Track https://github.com/microsoft/onnxruntime/issues/27729",
-                    original,
-                    name,
-                    execution_provider,
-                    original,
-                )
+    _maybe_apply_opset_lowering(pkg, execution_provider)
     return pkg
 
 
+# Attention input index for the optional ``nonpad_kv_seqlen`` operand. This
+# operand (external/static KV cache length) and the TensorScatter op are
+# defined only in opset 24, so a graph using either must not declare opset 23.
+_ATTENTION_NONPAD_KV_SEQLEN_INPUT_INDEX = 6
+
+
+def _maybe_apply_opset_lowering(pkg: ModelPackage, execution_provider: str) -> None:
+    """Lower the default-domain opset from 24 to 23 where it is safe to do so.
+
+    Some EPs (older ORT builds) don't register opset 24 kernels for standard
+    ops (Reshape, RMSNormalization, etc.). Without lowering, those ops fall to
+    CPU and produce ~280 memcpy nodes that destroy performance. The
+    ``MOBIUS_ORT_LOWER_OPSET_FOR_EP`` flag (default False) opts a deployment
+    into the lowering; it is a no-op for the ``"default"`` and ``"cpu"`` EPs
+    (the CPU EP already has opset-24 kernels), matching the inference-side gate
+    ``ort_inference._should_lower_opset``.
+
+    Any sub-model that uses opset-24-only semantics (TensorScatter, or
+    Attention with a non-empty input #6 ``nonpad_kv_seqlen``) is left at opset
+    24: declaring opset 23 on such a graph is invalid and would strip the
+    static-cache Flash path. See :func:`_graph_requires_opset24`. The decision
+    is made per sub-model, so a mixed package lowers its standard sub-models
+    while preserving its static-cache sub-models.
+    """
+    if not flags.ort_lower_opset_for_ep:
+        return
+    # Mirror ort_inference._should_lower_opset: the "default" EP is a no-op, and
+    # the CPU EP already registers opset-24 kernels, so lowering there is both
+    # unnecessary and inconsistent with the inference-side gate.
+    if execution_provider in ("default", "cpu"):
+        return
+    for name, model in pkg.items():
+        if "" not in model.graph.opset_imports:
+            continue
+        if _graph_requires_opset24(model.graph):
+            logger.info(
+                "Skipped opset→23 lowering for '%s' (EP=%s): graph uses "
+                "opset-24-only ops (TensorScatter / Attention nonpad_kv_seqlen). "
+                "Preserving opset 24 to keep the static-cache Flash path valid.",
+                name,
+                execution_provider,
+            )
+            continue
+        original = model.graph.opset_imports[""]
+        model.graph.opset_imports[""] = 23
+        logger.warning(
+            "Lowered opset %d→23 for '%s' (EP=%s). "
+            "ORT does not yet register opset %d kernels for this EP. "
+            "Track https://github.com/microsoft/onnxruntime/issues/27729",
+            original,
+            name,
+            execution_provider,
+            original,
+        )
+
+
+def _graph_requires_opset24(graph: ir.Graph) -> bool:
+    """Return True if the graph uses opset-24-only default-domain semantics.
+
+    Lowering the default-domain opset import to 23 on such a graph is invalid
+    and would silently break the static-cache Flash-attention path. A graph
+    requires opset 24 when it contains:
+
+    - a ``TensorScatter`` node (default domain), or
+    - an ``Attention`` node consuming a non-empty input #6 (``nonpad_kv_seqlen``).
+
+    The scan is recursive: nodes nested inside ``If``/``Loop``/``Scan``
+    subgraphs are inspected too, so a future graph that buries one of these ops
+    in a control-flow body is still detected.
+    """
+    for node in ir.traversal.RecursiveGraphIterator(graph):
+        if node.domain not in ("", "ai.onnx"):
+            continue
+        if node.op_type == "TensorScatter":
+            return True
+        if node.op_type == "Attention":
+            inputs = node.inputs
+            if (
+                len(inputs) > _ATTENTION_NONPAD_KV_SEQLEN_INPUT_INDEX
+                and inputs[_ATTENTION_NONPAD_KV_SEQLEN_INPUT_INDEX] is not None
+            ):
+                return True
+    return False
+
+
 def build(
     model_id: str,
     task: str | ModelTask | None = None,
 
@@ -0,0 +1,170 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Unit tests for :mod:`mobius._builder` opset-lowering logic.
+
+These tests are CPU-authorable: they construct tiny ONNX-IR graphs directly
+(no weights, no network) and exercise the conditional opset 24→23 lowering
+that preserves the static-cache Flash-attention path.
+"""
+
+from __future__ import annotations
+
+import onnx_ir as ir
+import pytest
+
+from mobius._builder import _graph_requires_opset24, _maybe_apply_opset_lowering, flags
+from mobius._model_package import ModelPackage
+
+
+def _make_value(name: str) -> ir.Value:
+    return ir.Value(name=name)
+
+
+def _graph_with(nodes: list[ir.Node], *, opset: int = 24) -> ir.Graph:
+    return ir.Graph(inputs=[], outputs=[], nodes=nodes, opset_imports={"": opset})
+
+
+def _model_with(nodes: list[ir.Node], *, opset: int = 24) -> ir.Model:
+    return ir.Model(_graph_with(nodes, opset=opset), ir_version=10)
+
+
+def _attention_nonpad_inputs() -> list[ir.Value | None]:
+    # q, k, v, attn_mask, past_key, past_value, nonpad_kv_seqlen (input #6).
+    return [
+        _make_value("q"),
+        _make_value("k"),
+        _make_value("v"),
+        None,
+        None,
+        None,
+        _make_value("nonpad_kv_seqlen"),
+    ]
+
+
+def _static_cache_nodes() -> list[ir.Node]:
+    return [
+        ir.Node(
+            "",
+            "TensorScatter",
+            inputs=[_make_value("p"), _make_value("u"), _make_value("i")],
+        ),
+        ir.Node("", "Attention", inputs=_attention_nonpad_inputs(), num_outputs=1),
+    ]
+
+
+def _standard_nodes() -> list[ir.Node]:
+    return [ir.Node("", "Reshape", inputs=[_make_value("x"), _make_value("shape")])]
+
+
+def test_graph_requires_opset24_tensor_scatter() -> None:
+    # A TensorScatter node (opset-24-only) must force opset 24 retention.
+    node = ir.Node(
+        "",
+        "TensorScatter",
+        inputs=[_make_value("past"), _make_value("update"), _make_value("idx")],
+    )
+    assert _graph_requires_opset24(_graph_with([node])) is True
+
+
+def test_graph_requires_opset24_attention_nonpad_kv_seqlen() -> None:
+    # Attention consuming input #6 (nonpad_kv_seqlen) is opset-24-only.
+    node = ir.Node("", "Attention", inputs=_attention_nonpad_inputs())
+    assert _graph_requires_opset24(_graph_with([node])) is True
+
+
+def test_graph_requires_opset24_attention_without_nonpad() -> None:
+    # A plain Attention (no input #6) does not require opset 24.
+    inputs = [_make_value("q"), _make_value("k"), _make_value("v")]
+    node = ir.Node("", "Attention", inputs=inputs)
+    assert _graph_requires_opset24(_graph_with([node])) is False
+
+
+def test_graph_requires_opset24_standard_ops_only() -> None:
+    # Standard ops (Reshape) are safe to lower.
+    node = ir.Node("", "Reshape", inputs=[_make_value("x"), _make_value("shape")])
+    assert _graph_requires_opset24(_graph_with([node])) is False
+
+
+def test_graph_requires_opset24_ignores_custom_domain() -> None:
+    # Same op names in a non-default domain must not trigger retention.
+    node = ir.Node(
+        "com.example",
+        "TensorScatter",
+        inputs=[_make_value("a"), _make_value("b")],
+    )
+    assert _graph_requires_opset24(_graph_with([node])) is False
+
+
+def test_graph_requires_opset24_detects_nested_subgraph() -> None:
+    # An opset-24-only op buried inside an If/Loop/Scan body must still be
+    # detected: the scan is recursive (RecursiveGraphIterator).
+    then_branch = _graph_with([ir.Node("", "TensorScatter", inputs=[_make_value("a")])])
+    else_branch = _graph_with([ir.Node("", "Identity", inputs=[_make_value("b")])])
+    if_node = ir.Node(
+        "",
+        "If",
+        inputs=[_make_value("cond")],
+        attributes=[
+            ir.AttrGraph("then_branch", then_branch),
+            ir.AttrGraph("else_branch", else_branch),
+        ],
+    )
+    outer = _graph_with([if_node])
+    assert _graph_requires_opset24(outer) is True
+
+
+def test_maybe_apply_opset_lowering_mixed_package(monkeypatch: pytest.MonkeyPatch) -> None:
+    # Drive the REAL lowering branch in build_from_module via _maybe_apply_opset_lowering.
+    # A mixed package must be handled per sub-model: the static-cache sub-model
+    # keeps opset 24, the standard sub-model is lowered to 23.
+    monkeypatch.setattr(flags, "ort_lower_opset_for_ep", True)
+    pkg = ModelPackage(
+        {
+            "model": _model_with(_static_cache_nodes()),
+            "embedding": _model_with(_standard_nodes()),
+        }
+    )
+
+    _maybe_apply_opset_lowering(pkg, execution_provider="cuda")
+
+    assert pkg["model"].graph.opset_imports[""] == 24
+    assert pkg["embedding"].graph.opset_imports[""] == 23
+
+
+def test_maybe_apply_opset_lowering_skipped_for_default_ep(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    # The "default" EP gate: lowering never fires even with the flag enabled.
+    monkeypatch.setattr(flags, "ort_lower_opset_for_ep", True)
+    pkg = ModelPackage({"embedding": _model_with(_standard_nodes())})
+
+    _maybe_apply_opset_lowering(pkg, execution_provider="default")
+
+    assert pkg["embedding"].graph.opset_imports[""] == 24
+
+
+def test_maybe_apply_opset_lowering_skipped_for_cpu_ep(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    # The "cpu" EP gate: the CPU EP already registers opset-24 kernels, so
+    # lowering never fires for it even with the flag enabled (mirrors the
+    # inference-side ort_inference._should_lower_opset CPU skip).
+    monkeypatch.setattr(flags, "ort_lower_opset_for_ep", True)
+    pkg = ModelPackage({"embedding": _model_with(_standard_nodes())})
+
+    _maybe_apply_opset_lowering(pkg, execution_provider="cpu")
+
+    assert pkg["embedding"].graph.opset_imports[""] == 24
+
+
+def test_maybe_apply_opset_lowering_skipped_when_flag_disabled(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    # The flag gate: lowering never fires when the flag is off (the default).
+    monkeypatch.setattr(flags, "ort_lower_opset_for_ep", False)
+    pkg = ModelPackage({"embedding": _model_with(_standard_nodes())})
+
+    _maybe_apply_opset_lowering(pkg, execution_provider="cuda")
+
+    assert pkg["embedding"].graph.opset_imports[""] == 24