Split ONNX opsets for legacy and dynamo export

QEfficient/base/modeling_qeff.py

@@ -367,7 +367,7 @@ def _export_via_legacy(
                 output_names=output_names,
                 dynamic_axes=dynamic_axes,
                 dynamo=False,
-                opset_version=constants.ONNX_EXPORT_OPSET,
+                opset_version=constants.ONNX_LEGACY_EXPORT_OPSET,
                 **export_kwargs,
             )
 
@@ -416,7 +416,7 @@ def _export_via_dynamo(
                 output_names=output_names,
                 dynamic_axes=None,
                 dynamic_shapes=dynamic_shapes,
-                opset_version=constants.ONNX_EXPORT_OPSET,
+                opset_version=constants.ONNX_DYNAMO_EXPORT_OPSET,
                 **export_kwargs,
             )
             if onnx_program is None:
@@ -580,6 +580,7 @@ def _resolve_pkv_names(layer_idx, layer_state):
                 "onnx_base_dir": str(export_dir) if needs_external_tensor_data else None,
                 "model_name": self.model_name,
                 "dynamic_axes": None if dynamo else dynamic_axes,  # dynamo uses dynamic_shapes, not axes
+                "onnx_export_opset": constants.get_onnx_export_opset(dynamo),
             }
             if onnx_transform_kwargs is not None:
                 transform_kwargs.update(onnx_transform_kwargs)
@@ -880,7 +881,7 @@ def _resolve_pkv_names(layer_idx, layer_state):
                     input_names=input_names,
                     output_names=output_names,
                     dynamic_axes=dynamic_axes,
-                    opset_version=constants.ONNX_EXPORT_OPSET,
+                    opset_version=constants.ONNX_LEGACY_EXPORT_OPSET,
                     dynamo=False,
                     **export_kwargs,
                 )

QEfficient/base/onnx_transforms.py

@@ -46,8 +46,10 @@
 )
 
 # from QEfficient.customop.quantization_ops import CastToUInt4, CastToUInt4Func
+from QEfficient.customop.onnxscript_utils import get_onnxscript_func
 from QEfficient.customop.rms_norm import CustomRMSNorm, CustomRMSNormFunc
-from QEfficient.utils.constants import FILE_CHUNK_SIZE_DEFAULT, ONNX_EXPORT_OPSET, SIZE_THRESHOLD_DEFAULT
+from QEfficient.utils import constants
+from QEfficient.utils.constants import FILE_CHUNK_SIZE_DEFAULT, SIZE_THRESHOLD_DEFAULT
 
 logger = logging.getLogger(__name__)
 
@@ -114,13 +116,13 @@ class CustomOpTransform(BaseOnnxTransform):
     }
 
     @classmethod
-    def apply(cls, model: ModelProto) -> bool:
+    def apply(cls, model: ModelProto, onnx_export_opset: int = constants.ONNX_LEGACY_EXPORT_OPSET) -> bool:
         op_applied = False
 
         # Register with PyTorch ONNX exporter (for export time)
         for op_name, (func_class, _) in cls._custom_ops.items():
             if hasattr(func_class, "symbolic"):
-                torch.onnx.register_custom_op_symbolic(f"::{op_name}", func_class.symbolic, ONNX_EXPORT_OPSET)
+                torch.onnx.register_custom_op_symbolic(f"::{op_name}", func_class.symbolic, onnx_export_opset)
 
         used_op_types = {node.op_type for node in model.graph.node}
         for function_proto in model.functions:
@@ -130,7 +132,7 @@ def apply(cls, model: ModelProto) -> bool:
         existing = {f.name for f in model.functions}
 
         for func_name, onnxscript_func in cls._custom_ops.values():
-            proto = onnxscript_func.to_function_proto()
+            proto = get_onnxscript_func(onnxscript_func, onnx_export_opset).to_function_proto()
             if proto.name not in used_op_types:
                 continue
             if proto.name not in existing:
@@ -643,7 +645,9 @@ def _set_external_data(tensor, file_name):
 
         # Non-looping transforms
         if CustomOpTransform in requested:
-            applied[CustomOpTransform] = CustomOpTransform.apply(model)
+            applied[CustomOpTransform] = CustomOpTransform.apply(
+                model, onnx_export_opset=kwargs.get("onnx_export_opset", constants.ONNX_LEGACY_EXPORT_OPSET)
+            )
 
         if RenameFunctionOutputsTransform in requested:
             applied[RenameFunctionOutputsTransform] = RenameFunctionOutputsTransform.apply(

QEfficient/customop/ctx_scatter_gather.py

@@ -8,12 +8,13 @@
 import onnxscript
 import torch
 
+from QEfficient.customop.onnxscript_utils import qeff_custom_op
 from QEfficient.utils import constants
 
-ops = getattr(onnxscript, "opset" + str(constants.ONNX_EXPORT_OPSET))
+ops = getattr(onnxscript, "opset" + str(constants.ONNX_LEGACY_EXPORT_OPSET))
 
 
-@onnxscript.script(onnxscript.values.Opset("com.qualcomm.cloud", 1))
+@qeff_custom_op("com.qualcomm.cloud", 1)
 def CtxScatter(data: onnxscript.FLOAT, position_ids: onnxscript.INT32, updates: onnxscript.FLOAT) -> onnxscript.FLOAT:
     # Find dims
     batch_size = ops.Gather(ops.Shape(data), [0])
@@ -56,7 +57,7 @@ def symbolic(g: torch.Graph, data: torch.Value, position_ids: torch.Value, updat
         return g.onnxscript_op(CtxScatter, data, position_ids, updates).setTypeAs(data)
 
 
-@onnxscript.script(onnxscript.values.Opset("com.qualcomm.cloud", 1))
+@qeff_custom_op("com.qualcomm.cloud", 1)
 def CtxScatter3D(data: onnxscript.FLOAT, position_ids: onnxscript.INT32, updates: onnxscript.FLOAT) -> onnxscript.FLOAT:
     # Find dims
     batch_size = ops.Gather(ops.Shape(data), [0])
@@ -122,7 +123,7 @@ def symbolic(g: torch.Graph, data: torch.Value, position_ids: torch.Value, updat
         return g.onnxscript_op(CtxScatter3D, data, position_ids, updates).setTypeAs(data)
 
 
-@onnxscript.script(onnxscript.values.Opset("com.qualcomm.cloud", 1))
+@qeff_custom_op("com.qualcomm.cloud", 1)
 def CtxScatter3DInt(
     data: onnxscript.INT32, position_ids: onnxscript.INT32, updates: onnxscript.INT32
 ) -> onnxscript.INT32:
@@ -164,7 +165,7 @@ def symbolic(g: torch.Graph, data: torch.Value, position_ids: torch.Value, updat
         return g.onnxscript_op(CtxScatter3DInt, data, position_ids, updates).setTypeAs(data)
 
 
-@onnxscript.script(onnxscript.values.Opset("com.qualcomm.cloud", 1))
+@qeff_custom_op("com.qualcomm.cloud", 1)
 def CtxGather3D(data: onnxscript.FLOAT, ctx_indices: onnxscript.INT32) -> onnxscript.FLOAT:
     batch_size = ops.Slice(ops.Shape(data), starts=[0], ends=[1], axes=[0])
     idx_seq_len = ops.Slice(ops.Shape(ctx_indices), starts=[1], ends=[2], axes=[0])
@@ -215,7 +216,7 @@ def symbolic(g: torch.Graph, data: torch.Value, ctx_indices: torch.Value) -> tor
         return g.onnxscript_op(CtxGather3D, data, ctx_indices)
 
 
-@onnxscript.script(onnxscript.values.Opset("com.qualcomm.cloud", 1))
+@qeff_custom_op("com.qualcomm.cloud", 1)
 def CtxGather(
     data: onnxscript.FLOAT, ctx_indices: onnxscript.INT32, comp_ctx_len: onnxscript.INT32
 ) -> onnxscript.FLOAT:
@@ -249,7 +250,7 @@ def symbolic(g: torch.Graph, data: torch.Value, ctx_indices: torch.Value, comp_c
         return g.onnxscript_op(CtxGather, data, ctx_indices, comp_ctx_len).setTypeAs(data)
 
 
-@onnxscript.script(onnxscript.values.Opset("com.qualcomm.cloud", 1))
+@qeff_custom_op("com.qualcomm.cloud", 1)
 def CtxGatherBlockedKV(data: onnxscript.FLOAT, ctx_indices: onnxscript.INT32) -> onnxscript.FLOAT:
     ctx_indices = ops.Unsqueeze(ctx_indices, [-1])
     return ops.GatherND(data, ctx_indices, batch_dims=2)

QEfficient/customop/ctx_scatter_gather_cb.py

@@ -8,12 +8,13 @@
 import onnxscript
 import torch
 
+from QEfficient.customop.onnxscript_utils import qeff_custom_op
 from QEfficient.utils import constants
 
-ops = getattr(onnxscript, "opset" + str(constants.ONNX_EXPORT_OPSET))
+ops = getattr(onnxscript, "opset" + str(constants.ONNX_LEGACY_EXPORT_OPSET))
 
 
-@onnxscript.script(onnxscript.values.Opset("com.qualcomm.cloud", 1))
+@qeff_custom_op("com.qualcomm.cloud", 1)
 def CtxScatterCB(
     data: onnxscript.FLOAT, batch_index: onnxscript.INT32, position_ids: onnxscript.INT32, updates: onnxscript.FLOAT
 ) -> onnxscript.FLOAT:
@@ -58,7 +59,7 @@ def symbolic(
         return g.onnxscript_op(CtxScatterCB, data, batch_index, position_ids, updates).setTypeAs(data)
 
 
-@onnxscript.script(onnxscript.values.Opset("com.qualcomm.cloud", 1))
+@qeff_custom_op("com.qualcomm.cloud", 1)
 def CtxScatterCB3D(
     data: onnxscript.FLOAT, batch_index: onnxscript.INT32, position_ids: onnxscript.INT32, updates: onnxscript.FLOAT
 ) -> onnxscript.FLOAT:
@@ -99,7 +100,7 @@ def symbolic(
         return g.onnxscript_op(CtxScatterCB3D, data, batch_index, position_ids, updates).setTypeAs(data)
 
 
-@onnxscript.script(onnxscript.values.Opset("com.qualcomm.cloud", 1))
+@qeff_custom_op("com.qualcomm.cloud", 1)
 def CtxGatherCB(
     data: onnxscript.FLOAT, batch_index: onnxscript.INT32, ctx_indices: onnxscript.INT32, comp_ctx_len: onnxscript.INT32
 ) -> onnxscript.FLOAT:
@@ -146,7 +147,7 @@ def symbolic(
         return g.onnxscript_op(CtxGatherCB, data, batch_index, ctx_indices, comp_ctx_len).setTypeAs(data)
 
 
-@onnxscript.script(onnxscript.values.Opset("com.qualcomm.cloud", 1))
+@qeff_custom_op("com.qualcomm.cloud", 1)
 def CtxGatherBlockedKVCB(
     data: onnxscript.FLOAT, batch_index: onnxscript.INT32, ctx_indices: onnxscript.INT32
 ) -> onnxscript.FLOAT:
@@ -185,7 +186,7 @@ def symbolic(g: torch.Graph, data: torch.Value, batch_index: torch.Value, ctx_in
         return g.onnxscript_op(CtxGatherBlockedKVCB, data, batch_index, ctx_indices).setTypeAs(data)
 
 
-@onnxscript.script(onnxscript.values.Opset("com.qualcomm.cloud", 1))
+@qeff_custom_op("com.qualcomm.cloud", 1)
 def CtxGatherCB3D(
     data: onnxscript.FLOAT, batch_index: onnxscript.INT32, ctx_indices: onnxscript.INT32
 ) -> onnxscript.FLOAT:

QEfficient/customop/dynamo_ops.py

@@ -368,21 +368,22 @@ def _(data: torch.Tensor, position_ids: torch.Tensor, updates: torch.Tensor) ->
     CtxScatterCB,
     CtxScatterCB3D,
 )
+from QEfficient.customop.onnxscript_utils import get_dynamo_onnxscript_func  # noqa: E402
 from QEfficient.customop.rms_norm import CustomRMSNorm  # noqa: E402
 
 DYNAMO_CUSTOM_OP_TABLE = {
-    torch.ops.qefficient.rms_norm.default: CustomRMSNorm,
-    torch.ops.qefficient.ctx_scatter.default: CtxScatter,
-    torch.ops.qefficient.ctx_scatter_3d.default: CtxScatter3D,
-    torch.ops.qefficient.ctx_scatter_cb.default: CtxScatterCB,
-    torch.ops.qefficient.ctx_scatter_cb_3d.default: CtxScatterCB3D,
-    torch.ops.qefficient.ctx_scatter_3d_int.default: CtxScatter3DInt,
-    torch.ops.qefficient.ctx_scatter_3d_generalized.default: CtxScatter3D,
-    torch.ops.qefficient.ctx_gather.default: CtxGather,
-    torch.ops.qefficient.ctx_gather_3d.default: CtxGather3D,
-    torch.ops.qefficient.ctx_gather_cb.default: CtxGatherCB,
-    torch.ops.qefficient.ctx_gather_cb_3d.default: CtxGatherCB3D,
-    torch.ops.qefficient.ctx_gather_blocked_kv.default: CtxGatherBlockedKV,
-    torch.ops.qefficient.ctx_gather_blocked_kv_cb.default: CtxGatherBlockedKVCB,
-    torch.ops.qefficient.ctx_gather_3d_generalized.default: CtxGather3D,
+    torch.ops.qefficient.rms_norm.default: get_dynamo_onnxscript_func(CustomRMSNorm),
+    torch.ops.qefficient.ctx_scatter.default: get_dynamo_onnxscript_func(CtxScatter),
+    torch.ops.qefficient.ctx_scatter_3d.default: get_dynamo_onnxscript_func(CtxScatter3D),
+    torch.ops.qefficient.ctx_scatter_cb.default: get_dynamo_onnxscript_func(CtxScatterCB),
+    torch.ops.qefficient.ctx_scatter_cb_3d.default: get_dynamo_onnxscript_func(CtxScatterCB3D),
+    torch.ops.qefficient.ctx_scatter_3d_int.default: get_dynamo_onnxscript_func(CtxScatter3DInt),
+    torch.ops.qefficient.ctx_scatter_3d_generalized.default: get_dynamo_onnxscript_func(CtxScatter3D),
+    torch.ops.qefficient.ctx_gather.default: get_dynamo_onnxscript_func(CtxGather),
+    torch.ops.qefficient.ctx_gather_3d.default: get_dynamo_onnxscript_func(CtxGather3D),
+    torch.ops.qefficient.ctx_gather_cb.default: get_dynamo_onnxscript_func(CtxGatherCB),
+    torch.ops.qefficient.ctx_gather_cb_3d.default: get_dynamo_onnxscript_func(CtxGatherCB3D),
+    torch.ops.qefficient.ctx_gather_blocked_kv.default: get_dynamo_onnxscript_func(CtxGatherBlockedKV),
+    torch.ops.qefficient.ctx_gather_blocked_kv_cb.default: get_dynamo_onnxscript_func(CtxGatherBlockedKVCB),
+    torch.ops.qefficient.ctx_gather_3d_generalized.default: get_dynamo_onnxscript_func(CtxGather3D),
 }

QEfficient/customop/onnxscript_utils.py

@@ -0,0 +1,72 @@
+# -----------------------------------------------------------------------------
+#
+# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# -----------------------------------------------------------------------------
+
+import inspect
+from collections.abc import Callable
+
+import onnxscript
+
+from QEfficient.utils import constants
+
+_DYNAMO_FUNC_ATTR = "_qeff_dynamo_onnxscript_func"
+_MISSING = object()
+
+
+def _onnx_opset(opset_version: int) -> onnxscript.values.Opset:
+    return getattr(onnxscript, f"opset{opset_version}")
+
+
+def _compile_with_default_opset(
+    fn: Callable,
+    custom_opset: onnxscript.values.Opset,
+    default_opset_version: int,
+):
+    """Compile an ONNXScript function with the requested default-domain opset.
+
+    Existing custom-op definitions refer to a module-level ``ops`` symbol. The
+    ONNXScript compiler captures module globals while the decorator runs, so we
+    temporarily point that symbol at the requested default opset and immediately
+    restore it after compilation.
+    """
+    module = inspect.getmodule(fn)
+    if module is None:
+        raise RuntimeError(f"Unable to resolve module for ONNXScript function {fn.__name__}")
+
+    previous_ops = module.__dict__.get("ops", _MISSING)
+    module.__dict__["ops"] = _onnx_opset(default_opset_version)
+    try:
+        return onnxscript.script(custom_opset)(fn)
+    finally:
+        if previous_ops is _MISSING:
+            module.__dict__.pop("ops", None)
+        else:
+            module.__dict__["ops"] = previous_ops
+
+
+def qeff_custom_op(domain: str, version: int):
+    """Compile one custom op body into legacy and dynamo ONNXScript variants."""
+    custom_opset = onnxscript.values.Opset(domain, version)
+
+    def decorator(fn: Callable):
+        legacy_func = _compile_with_default_opset(fn, custom_opset, constants.ONNX_LEGACY_EXPORT_OPSET)
+        dynamo_func = _compile_with_default_opset(fn, custom_opset, constants.ONNX_DYNAMO_EXPORT_OPSET)
+        setattr(legacy_func, _DYNAMO_FUNC_ATTR, dynamo_func)
+        return legacy_func
+
+    return decorator
+
+
+def get_dynamo_onnxscript_func(onnxscript_func):
+    """Return the dynamo/opset18 variant attached by ``qeff_custom_op``."""
+    return getattr(onnxscript_func, _DYNAMO_FUNC_ATTR)
+
+
+def get_onnxscript_func(onnxscript_func, onnx_export_opset: int):
+    """Return the ONNXScript variant matching the requested export opset."""
+    if onnx_export_opset == constants.ONNX_DYNAMO_EXPORT_OPSET:
+        return get_dynamo_onnxscript_func(onnxscript_func)
+    return onnxscript_func

QEfficient/customop/rms_norm.py

@@ -9,13 +9,14 @@
 import torch
 from torch import nn
 
+from QEfficient.customop.onnxscript_utils import qeff_custom_op
 from QEfficient.customop.utils import select_interface
 from QEfficient.utils import constants
 
-ops = getattr(onnxscript, "opset" + str(constants.ONNX_EXPORT_OPSET))
+ops = getattr(onnxscript, "opset" + str(constants.ONNX_LEGACY_EXPORT_OPSET))
 
 
-@onnxscript.script(onnxscript.values.Opset(domain="com.qti.aisw.onnx", version=1))
+@qeff_custom_op("com.qti.aisw.onnx", 1)
 def CustomRMSNorm(hidden_states: onnxscript.FLOAT, weight: onnxscript.FLOAT, epsilon: float) -> onnxscript.FLOAT:
     weight = ops.Cast(weight, to=1)
     variance = ops.ReduceMean(ops.Pow(hidden_states, 2), axes=[-1], keepdims=1)

QEfficient/exporter/export_utils.py

@@ -100,7 +100,7 @@ def export_onnx(
                 input_names=input_names,
                 output_names=output_names,
                 dynamic_axes=dynamic_axes,
-                opset_version=constants.ONNX_EXPORT_OPSET,
+                opset_version=constants.ONNX_LEGACY_EXPORT_OPSET,
                 custom_opsets={"com.qti.aisw.onnx": 1},
             )
     except Exception as e:

QEfficient/utils/constants.py

@@ -98,11 +98,17 @@ def get_models_dir():
 ONNX_EXPORT_EXAMPLE_MAX_TOP_K_IDS = 512
 ONNX_EXPORT_EXAMPLE_TOP_PS = 0.80
 ONNX_EXPORT_EXAMPLE_MIN_PS = 0.99
-ONNX_EXPORT_OPSET = 18
+ONNX_LEGACY_EXPORT_OPSET = 17
+ONNX_DYNAMO_EXPORT_OPSET = 18
+ONNX_EXPORT_OPSET = ONNX_LEGACY_EXPORT_OPSET
 FILE_CHUNK_SIZE_DEFAULT = 10 * 2**30  # 10 GB
 SIZE_THRESHOLD_DEFAULT = 1024
 
 
+def get_onnx_export_opset(dynamo: bool = False) -> int:
+    return ONNX_DYNAMO_EXPORT_OPSET if dynamo else ONNX_LEGACY_EXPORT_OPSET
+
+
 COMPILER = ["/opt/qti-aic/exec/qaic-compile", "-aic-hw"]
 
 

tests/unit_test/transforms/test_onnx_transforms.py

@@ -81,6 +81,24 @@ def test_custom_op_transform_has_apply_method(self):
         assert hasattr(CustomOpTransform, "apply")
         assert callable(CustomOpTransform.apply)
 
+    def test_custom_op_variants_use_legacy_and_dynamo_opsets(self):
+        from QEfficient.customop.onnxscript_utils import get_onnxscript_func
+        from QEfficient.customop.rms_norm import CustomRMSNorm
+        from QEfficient.utils import constants
+
+        def default_domain_opset(function_proto):
+            return next(opset.version for opset in function_proto.opset_import if opset.domain == "")
+
+        legacy_proto = get_onnxscript_func(
+            CustomRMSNorm, constants.get_onnx_export_opset(dynamo=False)
+        ).to_function_proto()
+        dynamo_proto = get_onnxscript_func(
+            CustomRMSNorm, constants.get_onnx_export_opset(dynamo=True)
+        ).to_function_proto()
+
+        assert default_domain_opset(legacy_proto) == constants.ONNX_LEGACY_EXPORT_OPSET
+        assert default_domain_opset(dynamo_proto) == constants.ONNX_DYNAMO_EXPORT_OPSET
+
     def test_base_onnx_transform_importable(self):
         from QEfficient.base.onnx_transforms import BaseOnnxTransform